fix(langchain): preserve tool calls in generation traces

backend/chainlit/langchain/callbacks.py

@@ -167,6 +167,8 @@ def _convert_message_dict(
         )
         if name := kwargs.get("name"):
             msg["name"] = name
+        if tool_call_id := kwargs.get("tool_call_id"):
+            msg["tool_call_id"] = tool_call_id
         if function_call:
             msg["function_call"] = function_call
         else:
@@ -196,8 +198,25 @@ def _convert_message_dict(
             else:
                 msg["content"] = content  # type: ignore
 
+        if normalized_tool_calls := kwargs.get("tool_calls"):
+            msg["tool_calls"] = self._convert_tool_calls(normalized_tool_calls)
+
         return msg
 
+    def _convert_tool_calls(self, tool_calls: List[Dict]) -> List[Dict]:
+        """Convert LangChain's normalized tool calls to the generation format."""
+        return [
+            {
+                "id": tool_call.get("id"),
+                "type": "function",
+                "function": {
+                    "name": tool_call["name"],
+                    "arguments": tool_call["args"],
+                },
+            }
+            for tool_call in tool_calls
+        ]
+
     def _convert_message(
         self,
         message: Union[Dict, BaseMessage],
@@ -220,6 +239,8 @@ def _convert_message(
 
         if name := getattr(message, "name", None):
             msg["name"] = name
+        if tool_call_id := getattr(message, "tool_call_id", None):
+            msg["tool_call_id"] = tool_call_id
 
         if function_call:
             msg["function_call"] = function_call
@@ -251,6 +272,9 @@ def _convert_message(
             else:
                 msg["content"] = message.content  # type: ignore
 
+        if normalized_tool_calls := getattr(message, "tool_calls", None):
+            msg["tool_calls"] = self._convert_tool_calls(normalized_tool_calls)
+
         return msg
 
     def _build_llm_settings(

backend/tests/langchain/test_async_callback.py

@@ -5,7 +5,9 @@
 from uuid import uuid4
 
 import pytest
-from langchain_core.outputs import GenerationChunk
+from langchain_core.callbacks import AsyncCallbackManager
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.outputs import ChatGeneration, GenerationChunk, LLMResult
 
 from chainlit.langchain.callbacks import LangchainTracer
 from chainlit.step import Step
@@ -166,3 +168,34 @@ async def test_error_handling(mock_chainlit_context):
 
         assert step.is_error is True
         assert step.output == "Test error"
+
+
+async def test_generation_trace_preserves_tool_exchange(mock_chainlit_context):
+    """Tool calls and results stay correlated in the persisted generation payload."""
+    call = AIMessage(
+        content="",
+        tool_calls=[{"name": "weather", "args": {"city": "Paris"}, "id": "call_1"}],
+    )
+    result = ToolMessage(content="Sunny", tool_call_id="call_1")
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+        manager = AsyncCallbackManager([tracer])
+        runs = await manager.on_chat_model_start(
+            serialized={"name": "test_llm"},
+            messages=[[HumanMessage(content="Weather?"), call, result]],
+            invocation_params={"_type": "test", "model": "test-model"},
+        )
+        step = tracer.steps[str(runs[0].run_id)]
+
+        with patch.object(step, "update", new_callable=AsyncMock):
+            await runs[0].on_llm_end(
+                LLMResult(generations=[[ChatGeneration(message=call)]])
+            )
+
+        generation = step.generation
+        assert generation.messages[1]["tool_calls"][0]["id"] == "call_1"
+        assert generation.messages[2]["tool_call_id"] == "call_1"
+        assert generation.message_completion["tool_calls"][0]["function"] == {
+            "name": "weather",
+            "arguments": {"city": "Paris"},
+        }

backend/tests/langchain/test_sync_callback.py

@@ -4,7 +4,8 @@
 from unittest.mock import Mock
 from uuid import uuid4
 
-from langchain_core.messages import AIMessage, HumanMessage
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
 from chainlit.langchain.callbacks import (
     FinalStreamHelper,
@@ -111,6 +112,102 @@ def test_build_llm_settings(self):
         assert model == "gpt-4"
         assert settings["temperature"] == 0.7
 
+    @pytest.mark.parametrize("serialized", [False, True])
+    @pytest.mark.parametrize("content", ["Checking the weather.", []])
+    def test_convert_standard_tool_calls(self, serialized, content):
+        message = AIMessage(
+            content=content,
+            tool_calls=[
+                {"name": "weather", "args": {"city": "上海"}, "id": "call_weather"},
+                {"name": "clock", "args": {}, "id": "call_clock"},
+            ],
+        )
+        original = message.model_dump()
+        result = GenerationHelper()._convert_message(
+            message.to_json() if serialized else message
+        )
+
+        assert result["tool_calls"] == [
+            {
+                "id": "call_weather",
+                "type": "function",
+                "function": {"name": "weather", "arguments": {"city": "上海"}},
+            },
+            {
+                "id": "call_clock",
+                "type": "function",
+                "function": {"name": "clock", "arguments": {}},
+            },
+        ]
+        assert result["content"] == (content or "")
+        assert message.model_dump() == original
+
+    @pytest.mark.parametrize("serialized", [False, True])
+    def test_convert_tool_result_preserves_call_id(self, serialized):
+        message = ToolMessage(
+            content="Sunny", tool_call_id="call_weather", name="weather"
+        )
+        result = GenerationHelper()._convert_message(
+            message.to_json() if serialized else message
+        )
+
+        assert result == {
+            "role": "tool",
+            "content": "Sunny",
+            "name": "weather",
+            "tool_call_id": "call_weather",
+        }
+
+    @pytest.mark.parametrize("serialized", [False, True])
+    @pytest.mark.parametrize("normalized", [False, True])
+    def test_convert_anthropic_tool_calls_without_duplicates(
+        self, serialized, normalized
+    ):
+        message = AIMessage(
+            content=[
+                {"type": "text", "text": "Checking."},
+                {
+                    "type": "tool_use",
+                    "id": "call_weather",
+                    "name": "weather",
+                    "input": {"city": "Paris"},
+                },
+            ],
+            tool_calls=(
+                [{"name": "weather", "args": {"city": "Paris"}, "id": "call_weather"}]
+                if normalized
+                else []
+            ),
+        )
+        result = GenerationHelper()._convert_message(
+            message.to_json() if serialized else message
+        )
+
+        assert result["tool_calls"] == [
+            {
+                "id": "call_weather",
+                "type": "function",
+                "function": {"name": "weather", "arguments": {"city": "Paris"}},
+            }
+        ]
+        assert result["content"] == [{"type": "text", "text": "Checking."}]
+
+    @pytest.mark.parametrize("serialized", [False, True])
+    def test_convert_legacy_function_call(self, serialized):
+        function_call = {"name": "weather", "arguments": '{"city":"Paris"}'}
+        message = AIMessage(
+            content="", additional_kwargs={"function_call": function_call}
+        )
+        result = GenerationHelper()._convert_message(
+            message.to_json() if serialized else message
+        )
+
+        assert result == {
+            "role": "assistant",
+            "content": "",
+            "function_call": function_call,
+        }
+
 
 async def test_should_ignore_run(mock_chainlit_context):
     """Test _should_ignore_run method."""