From 77289793c5b9c0299d23011ce8de36be9b453bc4 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 26 Nov 2024 18:43:46 -0300 Subject: [PATCH] fix: handle different output formats in agent message processing (#4867) * fix: handle different output formats in agent message processing * fix: correct condition for handling Anthropic message output * refactor: extract text output handling into separate function --- .../base/langflow/base/agents/events.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/base/agents/events.py b/src/backend/base/langflow/base/agents/events.py index 44f9a9d25..2ea4299a6 100644 --- a/src/backend/base/langflow/base/agents/events.py +++ b/src/backend/base/langflow/base/agents/events.py @@ -80,12 +80,25 @@ def handle_on_chain_start( return agent_message, start_time +def _extract_output_text(output: str | list) -> str: + if isinstance(output, str): + text = output + elif isinstance(output, list) and len(output) == 1 and isinstance(output[0], dict) and "text" in output[0]: + text = output[0]["text"] + else: + msg = f"Output is not a string or list of dictionaries with 'text' key: {output}" + raise ValueError(msg) + return text + + def handle_on_chain_end( event: dict[str, Any], agent_message: Message, send_message_method: SendMessageFunctionType, start_time: float ) -> tuple[Message, float]: data_output = event["data"].get("output") if data_output and isinstance(data_output, AgentFinish) and data_output.return_values.get("output"): - agent_message.text = data_output.return_values.get("output") + output = data_output.return_values.get("output") + + agent_message.text = _extract_output_text(output) agent_message.properties.state = "complete" # Add duration to the last content if it exists if agent_message.content_blocks: @@ -194,7 +207,9 @@ def handle_on_chain_stream( ) -> tuple[Message, float]: data_chunk = event["data"].get("chunk", {}) if isinstance(data_chunk, dict) and data_chunk.get("output"): - agent_message.text = data_chunk.get("output") + output = data_chunk.get("output") + if output and isinstance(output, str | list): + agent_message.text = _extract_output_text(output) agent_message.properties.state = "complete" agent_message = send_message_method(message=agent_message) start_time = perf_counter()