diff --git a/src/uipath_langchain/runtime/messages.py b/src/uipath_langchain/runtime/messages.py index 6a19d8f8..d1fb426b 100644 --- a/src/uipath_langchain/runtime/messages.py +++ b/src/uipath_langchain/runtime/messages.py @@ -181,6 +181,17 @@ def map_event( tool_args = tool_chunk_block.get("args") if tool_call_id: + # Extract metadata including simulated flag from response_metadata + metadata: dict[str, Any] = {} + if ( + self.current_message.response_metadata + and "simulated" in self.current_message.response_metadata + ): + metadata["simulated"] = self.current_message.response_metadata["simulated"] + logger.info( + f"Tool call {tool_call_id} for {tool_name} is marked as simulated" + ) + tool_event = UiPathConversationMessageEvent( message_id=message.id, tool_call=UiPathConversationToolCallEvent( @@ -189,6 +200,7 @@ def map_event( tool_name=tool_name, timestamp=timestamp, input=UiPathInlineValue(inline=tool_args), + metadata=metadata if metadata else None, ), ), ) diff --git a/tests/runtime/test_simulated_tool_calls.py b/tests/runtime/test_simulated_tool_calls.py new file mode 100644 index 00000000..510baa5e --- /dev/null +++ b/tests/runtime/test_simulated_tool_calls.py @@ -0,0 +1,148 @@ +"""Tests for simulated tool call logging and tracing.""" + +import pytest +from langchain_core.messages import AIMessageChunk, ToolCallChunk +from uipath_langchain.runtime.messages import UiPathChatMessagesMapper + + +def test_simulated_tool_call_metadata_added(): + """Test that simulated flag from response_metadata is added to tool call event metadata.""" + mapper = UiPathChatMessagesMapper() + + # Create first chunk with tool call + chunk1 = AIMessageChunk( + content="", + id="test-message-id", + response_metadata={"simulated": True}, + content_blocks=[ + ToolCallChunk( + type="tool_call_chunk", + id="call_123", + name="test_tool", + args={"arg1": "value1"}, + ) + ], + ) + + # Process first chunk + events1 = mapper.map_event(chunk1) + assert events1 is not None + + # Create last chunk + chunk2 = AIMessageChunk( + content="", + id="test-message-id", + chunk_position="last", + ) + + # Process last chunk - this should emit tool call start event + events2 = mapper.map_event(chunk2) + assert events2 is not None + assert len(events2) > 0 + + # Find the tool call start event + tool_call_event = None + for event in events2: + if event.tool_call and event.tool_call.start: + tool_call_event = event + break + + assert tool_call_event is not None, "Tool call start event should be emitted" + assert tool_call_event.tool_call.start.metadata is not None + assert tool_call_event.tool_call.start.metadata.get("simulated") is True + + +def test_non_simulated_tool_call_no_metadata(): + """Test that tool calls without simulated flag don't have metadata set.""" + mapper = UiPathChatMessagesMapper() + + # Create first chunk with tool call (no simulated flag) + chunk1 = AIMessageChunk( + content="", + id="test-message-id-2", + content_blocks=[ + ToolCallChunk( + type="tool_call_chunk", + id="call_456", + name="test_tool", + args={"arg1": "value1"}, + ) + ], + ) + + # Process first chunk + events1 = mapper.map_event(chunk1) + assert events1 is not None + + # Create last chunk + chunk2 = AIMessageChunk( + content="", + id="test-message-id-2", + chunk_position="last", + ) + + # Process last chunk + events2 = mapper.map_event(chunk2) + assert events2 is not None + assert len(events2) > 0 + + # Find the tool call start event + tool_call_event = None + for event in events2: + if event.tool_call and event.tool_call.start: + tool_call_event = event + break + + assert tool_call_event is not None + # Metadata should be None or empty when simulated flag is not present + assert ( + tool_call_event.tool_call.start.metadata is None + or tool_call_event.tool_call.start.metadata == {} + ) + + +def test_simulated_false_not_added(): + """Test that simulated=False is still added to metadata.""" + mapper = UiPathChatMessagesMapper() + + # Create first chunk with tool call and simulated=False + chunk1 = AIMessageChunk( + content="", + id="test-message-id-3", + response_metadata={"simulated": False}, + content_blocks=[ + ToolCallChunk( + type="tool_call_chunk", + id="call_789", + name="test_tool", + args={"arg1": "value1"}, + ) + ], + ) + + # Process first chunk + events1 = mapper.map_event(chunk1) + assert events1 is not None + + # Create last chunk + chunk2 = AIMessageChunk( + content="", + id="test-message-id-3", + chunk_position="last", + ) + + # Process last chunk + events2 = mapper.map_event(chunk2) + assert events2 is not None + assert len(events2) > 0 + + # Find the tool call start event + tool_call_event = None + for event in events2: + if event.tool_call and event.tool_call.start: + tool_call_event = event + break + + assert tool_call_event is not None + assert tool_call_event.tool_call.start.metadata is not None + assert tool_call_event.tool_call.start.metadata.get("simulated") is False