fix(agno): fix streaming tool-call arguments and add structured output message attributes

Author: satyadevai

python/instrumentation/openinference-instrumentation-agno/examples/streaming_tool_calls.py

@@ -1,5 +1,3 @@
-import json
-
 from agno.agent import Agent
 from agno.models.anthropic import Claude
 from agno.models.google import Gemini
@@ -21,20 +19,27 @@ def get_weather(city: str) -> str:
     """Return current weather for a city."""
     return f"Sunny, 22°C in {city}."
 
+
 if __name__ == "__main__":
-    for event in Agent(name="OpenAI Agent", model=OpenAIChat(id="gpt-4o-mini"), tools=[get_weather]).run(
-        "What is the weather in Paris and New Delhi?", stream=True
-    ):
+    for event in Agent(
+        name="OpenAI Agent", model=OpenAIChat(id="gpt-4o-mini"), tools=[get_weather]
+    ).run("What is the weather in Paris and New Delhi?", stream=True):
         print(event)
 
+    # print(
+    #     Agent(name="OpenAI Agent", model=OpenAIChat(id="gpt-4o-mini"), tools=[get_weather]).run(
+    #         "What is the weather in Paris and New Delhi?", stream=False
+    #     )
+    # )
+
     print("\n=== Anthropic ===")
-    for event in Agent(name="Anthropic Agent", model=Claude(id="claude-haiku-4-5-20251001"), tools=[get_weather]).run(
-        "What is the weather in Paris and New Delhi?", stream=True
-    ):
+    for event in Agent(
+        name="Anthropic Agent", model=Claude(id="claude-haiku-4-5-20251001"), tools=[get_weather]
+    ).run("What is the weather in Paris and New Delhi?", stream=True):
         print(event)
 
     print("\n=== Gemini ===")
-    for event in Agent(name="Gemini Agent", model=Gemini(id="gemini-2.0-flash"), tools=[get_weather]).run(
-        "What is the weather in Paris and New Delhi?", stream=True
-    ):
+    for event in Agent(
+        name="Gemini Agent", model=Gemini(id="gemini-2.0-flash"), tools=[get_weather]
+    ).run("What is the weather in Paris and New Delhi?", stream=True):
         print(event)

python/instrumentation/openinference-instrumentation-agno/pyproject.toml

@@ -78,6 +78,7 @@ testpaths = [
 [tool.mypy]
 strict = true
 explicit_package_bases = true
+warn_unused_ignores = false
 exclude = [
   "examples",
   "dist",
@@ -89,6 +90,7 @@ ignore_missing_imports = true
 module = [
   "agno",
   "wrapt",
+  "wrapt.*",
 ]
 
 [tool.ruff]

python/instrumentation/openinference-instrumentation-agno/src/openinference/instrumentation/agno/_model_wrapper.py

@@ -72,6 +72,9 @@ def _bind_arguments(method: Callable[..., Any], *args: Any, **kwargs: Any) -> Di
 def _llm_input_messages(arguments: Mapping[str, Any]) -> Iterator[Tuple[str, Any]]:
     def process_message(idx: int, message: Any) -> Iterator[Tuple[str, Any]]:
         yield f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_ROLE}", message.role
+        # tool_call_id links a role=tool result message back to the original tool call request
+        if message.tool_call_id:
+            yield f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_TOOL_CALL_ID}", message.tool_call_id
         if message.content:
             yield f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_CONTENT}", message.get_content_string()
         if message.tool_calls:
@@ -85,9 +88,10 @@ def process_message(idx: int, message: Any) -> Iterator[Tuple[str, Any]]:
                     f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
                     _get_attr(function_obj, "name"),
                 )
+                arguments = _get_attr(function_obj, "arguments")
                 yield (
                     f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                    safe_json_dumps(_get_attr(function_obj, "arguments", {})),
+                    json.dumps(arguments) if isinstance(arguments, dict) else (arguments or "{}"),
                 )
 
     messages = arguments.get("messages", [])
@@ -199,9 +203,12 @@ def _output_value_and_mime_type(output: str) -> Iterator[Tuple[str, Any]]:
                             f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
                             _get_attr(function_obj, "name"),
                         )
+                        arguments = _get_attr(function_obj, "arguments")
                         yield (
                             f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                            safe_json_dumps(_get_attr(function_obj, "arguments", {})),
+                            json.dumps(arguments)
+                            if isinstance(arguments, dict)
+                            else (arguments or "{}"),
                         )
 
             yield OUTPUT_VALUE, safe_json_dumps(messages)
@@ -246,8 +253,8 @@ def _parse_model_output_stream(output: Any) -> Dict[str, Any]:
     non_indexed_tool_calls: list[Dict[str, Any]] = []
 
     for chunk in output:
-        if chunk.content:
-            accumulated_content += chunk.content
+        if chunk.content is not None and chunk.content != "":
+            accumulated_content += str(chunk.content)
 
         if chunk.tool_calls:
             for tool_call in chunk.tool_calls:
@@ -306,6 +313,33 @@ def _parse_model_output_stream(output: Any) -> Dict[str, Any]:
     return {"messages": messages}
 
 
+def _stream_output_messages(output: Dict[str, Any]) -> Iterator[Tuple[str, Any]]:
+    for i, message in enumerate(output.get("messages", [])):
+        if role := _get_attr(message, "role"):
+            yield f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_ROLE}", role
+        if content := _get_attr(message, "content"):
+            yield f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_CONTENT}", content
+        if tool_calls := _get_attr(message, "tool_calls"):
+            for j, tool_call in enumerate(tool_calls):
+                if tc_id := _get_attr(tool_call, "id"):
+                    yield (
+                        f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{j}.{TOOL_CALL_ID}",
+                        tc_id,
+                    )
+                function_obj = _get_attr(tool_call, "function", {})
+                if fn_name := _get_attr(function_obj, "name"):
+                    yield (
+                        f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{j}.{TOOL_CALL_FUNCTION_NAME}",
+                        fn_name,
+                    )
+                fn_args = _get_attr(function_obj, "arguments")
+                if fn_args is not None:
+                    yield (
+                        f"{LLM_OUTPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{j}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
+                        fn_args,
+                    )
+
+
 class _ModelWrapper:
     def __init__(self, tracer: trace_api.Tracer) -> None:
         self._tracer = tracer
@@ -407,6 +441,7 @@ def run_stream(
             output_message = json.dumps(output_message_dict)
             span.set_attribute(OUTPUT_MIME_TYPE, JSON)
             span.set_attribute(OUTPUT_VALUE, output_message)
+            span.set_attributes(dict(_stream_output_messages(output_message_dict)))
 
             # Find the final response with complete metrics (last one with response_usage)
             final_response_with_metrics = None
@@ -540,6 +575,7 @@ async def arun_stream(
             output_message = json.dumps(output_message_dict)
             span.set_attribute(OUTPUT_MIME_TYPE, JSON)
             span.set_attribute(OUTPUT_VALUE, output_message)
+            span.set_attributes(dict(_stream_output_messages(output_message_dict)))
 
             # Find the final response with complete metrics (last one with response_usage)
             final_response_with_metrics = None
@@ -597,6 +633,7 @@ async def arun_stream(
 MESSAGE_FUNCTION_CALL_NAME = MessageAttributes.MESSAGE_FUNCTION_CALL_NAME
 MESSAGE_NAME = MessageAttributes.MESSAGE_NAME
 MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
+MESSAGE_TOOL_CALL_ID = MessageAttributes.MESSAGE_TOOL_CALL_ID
 MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
 
 # mime types