samples-python/tests/langsmith_tracing/test_chatbot.py at 1c563c13414cdd279a4c46a70f1a616f074759ee · temporalio/samples-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import json
import uuid
from unittest.mock import MagicMock

from temporalio import activity
from temporalio.client import Client
from temporalio.contrib.langsmith import LangSmithPlugin
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

from langsmith_tracing.chatbot.activities import ChatResponse, OpenAIRequest, ToolCall
from langsmith_tracing.chatbot.workflows import ChatbotWorkflow
from tests.langsmith_tracing.helpers import make_text_response


def _make_function_call_response(
    name: str, arguments: dict, call_id: str = "call_123"
) -> ChatResponse:
    return ChatResponse(
        id="resp_tool",
        tool_calls=[
            ToolCall(
                call_id=call_id,
                name=name,
                arguments=json.dumps(arguments),
            )
        ],
    )


async def test_chatbot_save_note(
    client: Client, env: WorkflowEnvironment, mock_ls_client: MagicMock
):
    """Test save_note tool call loop — save_note runs as a workflow method."""
    call_count = 0

    @activity.defn(name="call_openai")
    async def mock_call_openai(request: OpenAIRequest) -> ChatResponse:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return _make_function_call_response(
                name="save_note",
                arguments={"name": "greeting", "content": "Hello world"},
            )
        return make_text_response("Note saved successfully!")

    async with Worker(
        client,
        task_queue="test-langsmith-chatbot",
        workflows=[ChatbotWorkflow],
        activities=[mock_call_openai],
        plugins=[LangSmithPlugin(client=mock_ls_client)],
    ):
        wf_handle = await client.start_workflow(
            ChatbotWorkflow.run,
            id=f"test-chatbot-{uuid.uuid4().hex[:8]}",
            task_queue="test-langsmith-chatbot",
        )

        response = await wf_handle.execute_update(
            ChatbotWorkflow.message_from_user, "Save a note"
        )
        assert response == "Note saved successfully!"

        notes = await wf_handle.query(ChatbotWorkflow.notes)
        assert notes == {"greeting": "Hello world"}

        await wf_handle.signal(ChatbotWorkflow.exit)
        result = await wf_handle.result()
        assert result == "Session ended."


async def test_chatbot_read_note(
    client: Client, env: WorkflowEnvironment, mock_ls_client: MagicMock
):
    """Test read_note tool call loop — saves a note first, then reads it back."""
    call_count = 0

    @activity.defn(name="call_openai")
    async def mock_call_openai(request: OpenAIRequest) -> ChatResponse:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return _make_function_call_response(
                name="save_note",
                arguments={"name": "todo", "content": "Buy milk"},
                call_id="call_save",
            )
        if call_count == 2:
            return make_text_response("Saved your todo!")
        if call_count == 3:
            return _make_function_call_response(
                name="read_note",
                arguments={"name": "todo"},
                call_id="call_read",
            )
        return make_text_response("Your todo says: Buy milk")

    async with Worker(
        client,
        task_queue="test-langsmith-chatbot-read",
        workflows=[ChatbotWorkflow],
        activities=[mock_call_openai],
        plugins=[LangSmithPlugin(client=mock_ls_client)],
    ):
        wf_handle = await client.start_workflow(
            ChatbotWorkflow.run,
            id=f"test-chatbot-read-{uuid.uuid4().hex[:8]}",
            task_queue="test-langsmith-chatbot-read",
        )

        response = await wf_handle.execute_update(
            ChatbotWorkflow.message_from_user, "Save my todo"
        )
        assert response == "Saved your todo!"

        response = await wf_handle.execute_update(
            ChatbotWorkflow.message_from_user, "Read my todo"
        )
        assert response == "Your todo says: Buy milk"

        await wf_handle.signal(ChatbotWorkflow.exit)
        await wf_handle.result()