import json from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from mlflow.types.responses import ResponsesAgentStreamEvent, output_to_responses_items_stream def test_output_to_responses_items_stream_langchain(): """ Tests langchain message stream to responses items stream conversion. Accounts for: - AIMessage w/ and w/o tool calls - ToolMessage - Filtering out HumanMessage from the stream - Message """ messages = [ AIMessage( content="test text0", additional_kwargs={}, response_metadata={}, name="query_result", id="e0eafab0-f008-49d4-ac0d-f17a70096fe1", ), AIMessage( content="Transferring back to supervisor", additional_kwargs={}, response_metadata={"__is_handoff_back": True}, name="revenue-genie", id="5e88662b-29e7-4659-a521-f8175e7642ee", tool_calls=[ { "name": "transfer_back_to_supervisor", "args": {}, "id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "type": "tool_call", } ], ), ToolMessage( content="Successfully transferred back to supervisor", name="transfer_back_to_supervisor", id="6fd471d8-57d4-46ec-a21a-9bb20dfda4d3", tool_call_id="543a6b6b-dc73-463c-9b6e-5d5a941b7669", ), HumanMessage( content="Which companies do I have revenue data for", additional_kwargs={}, response_metadata={}, id="43d0cf0a-d687-4302-8562-4f2e09603473", ), AIMessage( content="I'll help you", additional_kwargs={ "tool_calls": [ { "id": "toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", "function": { "arguments": "{}", "name": "transfer_to_revenue-genie", }, "type": "function", } ] }, response_metadata={ "usage": { "prompt_tokens": 551, "completion_tokens": 73, "total_tokens": 624, }, "prompt_tokens": 551, "completion_tokens": 73, "total_tokens": 624, "model": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_name": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "finish_reason": "tool_calls", }, name="supervisor", id="run--e112332b-ed6d-4e10-b17c-adf637fb67eb-0", tool_calls=[ { "name": "transfer_to_revenue-genie", "args": {}, "id": "toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", "type": "tool_call", } ], ), ToolMessage( content="Successfully transferred to revenue-genie", name="transfer_to_revenue-genie", id="92acdb97-babc-4239-979e-b16880e7f58f", tool_call_id="toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", ), AIMessage( content="test text1", additional_kwargs={}, response_metadata={}, name="query_result", id="e0eafab0-f008-49d4-ac0d-f17a70096fe1", ), AIMessage( content="Transferring back to supervisor", additional_kwargs={}, response_metadata={"__is_handoff_back": True}, name="revenue-genie", id="5e88662b-29e7-4659-a521-f8175e7642ee", tool_calls=[ { "name": "transfer_back_to_supervisor", "args": {}, "id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "type": "tool_call", } ], ), ToolMessage( content="Successfully transferred back to supervisor", name="transfer_back_to_supervisor", id="6fd471d8-57d4-46ec-a21a-9bb20dfda4d3", tool_call_id="543a6b6b-dc73-463c-9b6e-5d5a941b7669", ), AIMessage( content="test text2", additional_kwargs={}, response_metadata={ "usage": { "prompt_tokens": 813, "completion_tokens": 108, "total_tokens": 921, }, "prompt_tokens": 813, "completion_tokens": 108, "total_tokens": 921, "model": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_name": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "finish_reason": "stop", }, name="supervisor", id="run--2622edf9-37b6-4e25-9e97-6351d145b198-0", ), ] expected = [ ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "e0eafab0-f008-49d4-ac0d-f17a70096fe1", "content": [{"text": "test text0", "type": "output_text", "annotations": []}], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "5e88662b-29e7-4659-a521-f8175e7642ee", "content": [ { "text": "Transferring back to supervisor", "type": "output_text", "annotations": [], } ], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call", "id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "call_id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "name": "transfer_back_to_supervisor", "arguments": "{}", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call_output", "call_id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "output": "Successfully transferred back to supervisor", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "run--e112332b-ed6d-4e10-b17c-adf637fb67eb-0", "content": [{"text": "I'll help you", "type": "output_text", "annotations": []}], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call", "id": "toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", "call_id": "toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", "name": "transfer_to_revenue-genie", "arguments": "{}", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call_output", "call_id": "toolu_bdrk_01FtmRmzFm89zDtwYu3xdFkh", "output": "Successfully transferred to revenue-genie", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "e0eafab0-f008-49d4-ac0d-f17a70096fe1", "content": [{"text": "test text1", "type": "output_text", "annotations": []}], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "5e88662b-29e7-4659-a521-f8175e7642ee", "content": [ { "text": "Transferring back to supervisor", "type": "output_text", "annotations": [], } ], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call", "id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "call_id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "name": "transfer_back_to_supervisor", "arguments": "{}", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call_output", "call_id": "543a6b6b-dc73-463c-9b6e-5d5a941b7669", "output": "Successfully transferred back to supervisor", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "run--2622edf9-37b6-4e25-9e97-6351d145b198-0", "content": [{"text": "test text2", "type": "output_text", "annotations": []}], "role": "assistant", "type": "message", }, ), ] result = list(output_to_responses_items_stream(messages)) assert result == expected def test_output_to_responses_items_stream_langchain_multiple_tool_calls_unique_ids(): messages = [ AIMessage( content="I'll look up both the current and historical stock prices for Apple.", id="lc_run--019c2b6f-4d8d-70e3-b620-b80545b4cb30", tool_calls=[ { "name": "get_current_stock_price", "args": {"ticker": "AAPL"}, "id": "toolu_bdrk_01X3zqC3kknbchSJB6XYmtHQ", "type": "tool_call", }, { "name": "get_historical_stock_prices", "args": { "ticker": "AAPL", "start_date": "2023-01-01", "end_date": "2024-01-01", }, "id": "toolu_bdrk_01Dte8KnRx9Tk7pMz9h83okn", "type": "tool_call", }, ], ), ToolMessage( content='{"ticker": "AAPL", "price": 276.49}', name="get_current_stock_price", id="tool-output-1", tool_call_id="toolu_bdrk_01X3zqC3kknbchSJB6XYmtHQ", ), ToolMessage( content="Error executing tool", name="get_historical_stock_prices", id="tool-output-2", tool_call_id="toolu_bdrk_01Dte8KnRx9Tk7pMz9h83okn", ), ] expected = [ ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "id": "lc_run--019c2b6f-4d8d-70e3-b620-b80545b4cb30", "content": [ { "text": "I'll look up both the current and historical " "stock prices for Apple.", "type": "output_text", "annotations": [], } ], "role": "assistant", "type": "message", }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call", "id": "toolu_bdrk_01X3zqC3kknbchSJB6XYmtHQ", "call_id": "toolu_bdrk_01X3zqC3kknbchSJB6XYmtHQ", "name": "get_current_stock_price", "arguments": '{"ticker": "AAPL"}', }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call", "id": "toolu_bdrk_01Dte8KnRx9Tk7pMz9h83okn", "call_id": "toolu_bdrk_01Dte8KnRx9Tk7pMz9h83okn", "name": "get_historical_stock_prices", "arguments": '{"ticker": "AAPL", "start_date": "2023-01-01", ' '"end_date": "2024-01-01"}', }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call_output", "call_id": "toolu_bdrk_01X3zqC3kknbchSJB6XYmtHQ", "output": '{"ticker": "AAPL", "price": 276.49}', }, ), ResponsesAgentStreamEvent( type="response.output_item.done", custom_outputs=None, item={ "type": "function_call_output", "call_id": "toolu_bdrk_01Dte8KnRx9Tk7pMz9h83okn", "output": "Error executing tool", }, ), ] result = list(output_to_responses_items_stream(messages)) assert result == expected def test_output_to_responses_items_stream_langchain_non_string_tool_content(): mcp_content_blocks = [{"type": "text", "text": "result from mcp tool"}] messages = [ AIMessage( content="Calling tools", id="ai-1", tool_calls=[ {"name": "mcp_tool", "args": {}, "id": "call-1", "type": "tool_call"}, ], ), ToolMessage( content=mcp_content_blocks, name="mcp_tool", id="tool-1", tool_call_id="call-1", ), ] result = list(output_to_responses_items_stream(messages)) tool_outputs = [r for r in result if r.item.get("type") == "function_call_output"] assert len(tool_outputs) == 1 assert tool_outputs[0].item["output"] == json.dumps(mcp_content_blocks)