CopilotKit/sdk-python/tests/test_langchain_messages_to_copilotkit.py at main · samuelson-chen/CopilotKit · 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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
"""Tests for langchain_messages_to_copilotkit assistant message emission.

Covers the parentMessageId orphan bug where an `if content:` guard skipped
emitting the assistant message when content was empty. Tool call entries
reference their parent assistant message via parentMessageId, so the
assistant message must always be emitted — even when content is empty
(standard OpenAI behavior for tool-call-only responses).
"""

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from copilotkit.langgraph import langchain_messages_to_copilotkit


def _convert_and_split(messages):
    """Convert messages and split result into assistant vs tool-call entries."""
    result = langchain_messages_to_copilotkit(messages)
    assistant_msgs = [m for m in result if m.get("role") == "assistant"]
    tool_call_msgs = [m for m in result if "parentMessageId" in m]
    return result, assistant_msgs, tool_call_msgs


class TestAssistantMessageAlwaysEmitted:
    """The assistant message must always be present so tool call entries can
    reference it via parentMessageId. Without it, tool calls are orphaned
    and the frontend cannot reconstruct tool call rendering on reconnect."""

    def test_ai_message_with_content_and_tool_calls(self):
        """AIMessage with both content and tool_calls emits assistant + tool calls."""
        messages = [
            AIMessage(
                id="ai-1",
                content="Let me help with that.",
                tool_calls=[
                    {"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}}
                ],
            ),
        ]
        _, assistant_msgs, tool_call_msgs = _convert_and_split(messages)

        assert len(assistant_msgs) == 1
        assert assistant_msgs[0]["id"] == "ai-1"
        assert assistant_msgs[0]["content"] == "Let me help with that."

        assert len(tool_call_msgs) == 1
        assert tool_call_msgs[0]["parentMessageId"] == "ai-1"

    def test_ai_message_with_empty_content_and_tool_calls(self):
        """AIMessage with empty content (OpenAI-style) still emits the assistant message."""
        messages = [
            AIMessage(
                id="ai-1",
                content="",
                tool_calls=[
                    {"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}}
                ],
            ),
        ]
        _, assistant_msgs, tool_call_msgs = _convert_and_split(messages)

        assert len(assistant_msgs) == 1, (
            "Assistant message must be emitted even with empty content"
        )
        assert assistant_msgs[0]["id"] == "ai-1"
        assert assistant_msgs[0]["content"] == ""

        assert len(tool_call_msgs) == 1
        assert tool_call_msgs[0]["parentMessageId"] == "ai-1"

    def test_ai_message_with_none_content_and_tool_calls(self):
        """AIMessage with None content still emits the assistant message."""
        msg = AIMessage(
            id="ai-1",
            content="",
            tool_calls=[{"id": "tc-1", "name": "get_help", "args": {}}],
        )
        # Simulate None content (some models/edge cases)
        msg.content = None  # type: ignore[assignment]

        _, assistant_msgs, _ = _convert_and_split([msg])

        assert len(assistant_msgs) == 1, (
            "Assistant message must be emitted even with None content"
        )
        assert assistant_msgs[0]["content"] == ""

    def test_no_orphaned_parent_message_ids(self):
        """Every parentMessageId must reference an existing assistant message."""
        messages = [
            HumanMessage(id="h-1", content="help me"),
            AIMessage(
                id="ai-1",
                content="",
                tool_calls=[
                    {"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}},
                    {"id": "tc-2", "name": "search", "args": {"query": "docs"}},
                ],
            ),
            ToolMessage(id="tm-1", content="done", tool_call_id="tc-1"),
            ToolMessage(id="tm-2", content="found", tool_call_id="tc-2"),
        ]
        result, _, tool_call_msgs = _convert_and_split(messages)

        message_ids = {m["id"] for m in result if "role" in m}

        for tc in tool_call_msgs:
            assert tc["parentMessageId"] in message_ids, (
                f"Tool call {tc['id']} has orphaned parentMessageId {tc['parentMessageId']}"
            )

    def test_ai_message_with_list_content_and_tool_calls(self):
        """AIMessage with empty list content (Anthropic-style) still emits the assistant message."""
        msg = AIMessage(
            id="ai-1",
            content="",
            tool_calls=[{"id": "tc-1", "name": "get_help", "args": {}}],
        )
        # Anthropic models can return content as a list; empty list is falsy
        msg.content = []  # type: ignore[assignment]

        _, assistant_msgs, tool_call_msgs = _convert_and_split([msg])

        assert len(assistant_msgs) == 1, (
            "Assistant message must be emitted even with empty list content"
        )
        assert assistant_msgs[0]["content"] == ""

        assert len(tool_call_msgs) == 1
        assert tool_call_msgs[0]["parentMessageId"] == "ai-1"

    def test_ai_message_without_tool_calls(self):
        """Plain AIMessage (no tool calls) emits just the assistant message."""
        messages = [AIMessage(id="ai-1", content="Hello!")]
        result, _, _ = _convert_and_split(messages)

        assert len(result) == 1
        assert result[0]["role"] == "assistant"
        assert result[0]["content"] == "Hello!"