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
This commit is contained in:
parent
11286778dc
commit
77289793c5
1 changed files with 17 additions and 2 deletions
|
|
@ -80,12 +80,25 @@ def handle_on_chain_start(
|
||||||
return agent_message, start_time
|
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(
|
def handle_on_chain_end(
|
||||||
event: dict[str, Any], agent_message: Message, send_message_method: SendMessageFunctionType, start_time: float
|
event: dict[str, Any], agent_message: Message, send_message_method: SendMessageFunctionType, start_time: float
|
||||||
) -> tuple[Message, float]:
|
) -> tuple[Message, float]:
|
||||||
data_output = event["data"].get("output")
|
data_output = event["data"].get("output")
|
||||||
if data_output and isinstance(data_output, AgentFinish) and data_output.return_values.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"
|
agent_message.properties.state = "complete"
|
||||||
# Add duration to the last content if it exists
|
# Add duration to the last content if it exists
|
||||||
if agent_message.content_blocks:
|
if agent_message.content_blocks:
|
||||||
|
|
@ -194,7 +207,9 @@ def handle_on_chain_stream(
|
||||||
) -> tuple[Message, float]:
|
) -> tuple[Message, float]:
|
||||||
data_chunk = event["data"].get("chunk", {})
|
data_chunk = event["data"].get("chunk", {})
|
||||||
if isinstance(data_chunk, dict) and data_chunk.get("output"):
|
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.properties.state = "complete"
|
||||||
agent_message = send_message_method(message=agent_message)
|
agent_message = send_message_method(message=agent_message)
|
||||||
start_time = perf_counter()
|
start_time = perf_counter()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue