diff --git a/backend/chainlit/langchain/callbacks.py b/backend/chainlit/langchain/callbacks.py index ef0b9b7182..3348748522 100644 --- a/backend/chainlit/langchain/callbacks.py +++ b/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( diff --git a/backend/tests/langchain/test_async_callback.py b/backend/tests/langchain/test_async_callback.py index 910a885f62..38e3b999b8 100644 --- a/backend/tests/langchain/test_async_callback.py +++ b/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"}, + } diff --git a/backend/tests/langchain/test_sync_callback.py b/backend/tests/langchain/test_sync_callback.py index 906cac9779..22d5ece152 100644 --- a/backend/tests/langchain/test_sync_callback.py +++ b/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."""