fix(openai-agents): fix broken traces with agents handoff on run_stream (#3143)

Author: galkleinman

package-lock.json

@@ -1,3 +1,5 @@
+import dataclasses
+import json
 import os
 from opentelemetry import context as context_api
 
@@ -17,3 +19,20 @@ def should_send_prompts():
     env_setting = os.getenv("TRACELOOP_TRACE_CONTENT", "true")
     override = context_api.get_value("override_enable_content_tracing")
     return _is_truthy(env_setting) or bool(override)
+
+
+class JSONEncoder(json.JSONEncoder):
+    def default(self, o):
+        if dataclasses.is_dataclass(o):
+            return dataclasses.asdict(o)
+
+        if hasattr(o, "to_json"):
+            return o.to_json()
+
+        if hasattr(o, "json"):
+            return o.json()
+
+        if hasattr(o, "__class__"):
+            return o.__class__.__name__
+
+        return super().default(o)

packages/opentelemetry-instrumentation-openai-agents/opentelemetry/instrumentation/openai_agents/__init__.py

@@ -18,6 +18,7 @@
 
 from agents import Agent, function_tool, ModelSettings, WebSearchTool
 from pydantic import BaseModel
+from typing import List, Dict, Union
 
 pytest_plugins = []
 
@@ -45,6 +46,13 @@ def environment():
 @pytest.fixture(autouse=True)
 def clear_exporter(exporter):
     exporter.clear()
+    from opentelemetry.instrumentation.openai_agents import (
+        _root_span_storage,
+        _instrumented_tools,
+    )
+
+    _root_span_storage.clear()
+    _instrumented_tools.clear()
 
 
 @pytest.fixture(scope="session")
@@ -88,9 +96,7 @@ async def get_weather(city: str) -> str:
 
     return Agent(
         name="WeatherAgent",
-        instructions=(
-            "You get the weather for a city using the get_weather tool."
-        ),
+        instructions=("You get the weather for a city using the get_weather tool."),
         model="gpt-4.1",
         tools=[get_weather],
     )
@@ -109,10 +115,12 @@ def web_search_tool_agent():
 @pytest.fixture(scope="session")
 def handoff_agent():
 
-    agent_a = Agent(name="AgentA", instructions="Agent A does something.",
-                    model="gpt-4.1")
-    agent_b = Agent(name="AgentB", instructions="Agent B does something else.",
-                    model="gpt-4.1")
+    agent_a = Agent(
+        name="AgentA", instructions="Agent A does something.", model="gpt-4.1"
+    )
+    agent_b = Agent(
+        name="AgentB", instructions="Agent B does something else.", model="gpt-4.1"
+    )
 
     class HandoffExample(BaseModel):
         message: str
@@ -131,11 +139,136 @@ class HandoffExample(BaseModel):
         instructions="You decide which agent to handoff to.",
         model="gpt-4.1",
         handoffs=[agent_a, agent_b],
-        tools=[handoff_tool_a, handoff_tool_b]
+        tools=[handoff_tool_a, handoff_tool_b],
     )
     return triage_agent
 
 
+@pytest.fixture(scope="session")
+def recipe_workflow_agents():
+    """Create Main Chat Agent and Recipe Editor Agent with function tools for recipe management."""
+
+    class Recipe(BaseModel):
+        id: str
+        name: str
+        ingredients: List[str]
+        instructions: List[str]
+        prep_time: str
+        cook_time: str
+        servings: int
+
+    class SearchResponse(BaseModel):
+        status: str
+        message: str
+        recipes: Union[Dict[str, Recipe], None] = None
+        recipe_count: Union[int, None] = None
+        query: Union[str, None] = None
+
+    class EditResponse(BaseModel):
+        status: str
+        message: str
+        modified_recipe: Union[Recipe, None] = None
+        changes_made: Union[List[str], None] = None
+        original_recipe: Union[Recipe, None] = None
+
+    # Mock recipe database
+    MOCK_RECIPES = {
+        "spaghetti_carbonara": {
+            "id": "spaghetti_carbonara",
+            "name": "Spaghetti Carbonara",
+            "ingredients": [
+                "400g spaghetti",
+                "200g pancetta",
+                "4 large eggs",
+                "100g Pecorino Romano cheese",
+            ],
+            "instructions": [
+                "Cook spaghetti",
+                "Dice pancetta",
+                "Whisk eggs with cheese",
+            ],
+            "prep_time": "10 minutes",
+            "cook_time": "15 minutes",
+            "servings": 4,
+        }
+    }
+
+    @function_tool
+    async def search_recipes(query: str = "") -> SearchResponse:
+        """Search and browse recipes in the database."""
+        if "carbonara" in query.lower():
+            recipe_data = MOCK_RECIPES["spaghetti_carbonara"]
+            recipes_dict = {"spaghetti_carbonara": Recipe(**recipe_data)}
+            return SearchResponse(
+                status="success",
+                message=f'Found 1 recipes matching "{query}"',
+                recipes=recipes_dict,
+                recipe_count=1,
+                query=query,
+            )
+        return SearchResponse(
+            status="success",
+            message="No recipes found",
+            recipes={},
+            recipe_count=0,
+            query=query,
+        )
+
+    @function_tool
+    async def plan_and_apply_recipe_modifications(
+        recipe: Recipe, modification_request: str
+    ) -> EditResponse:
+        """Plan modifications to a recipe based on user request and apply them."""
+
+        if (
+            "vegetarian" in modification_request.lower()
+            and "carbonara" in recipe.name.lower()
+        ):
+            modified_recipe = Recipe(
+                id=recipe.id,
+                name="Vegetarian Carbonara",
+                ingredients=[
+                    "400g spaghetti",
+                    "200g mushrooms",
+                    "4 large eggs",
+                    "100g Pecorino Romano cheese",
+                ],
+                instructions=[
+                    "Cook spaghetti",
+                    "Sauté mushrooms",
+                    "Whisk eggs with cheese",
+                ],
+                prep_time=recipe.prep_time,
+                cook_time=recipe.cook_time,
+                servings=recipe.servings,
+            )
+            return EditResponse(
+                status="success",
+                message="Successfully modified Spaghetti Carbonara to be vegetarian",
+                modified_recipe=modified_recipe,
+                changes_made=["Replaced pancetta with mushrooms"],
+                original_recipe=recipe,
+            )
+
+        return EditResponse(status="error", message="Could not modify recipe")
+
+    recipe_editor_agent = Agent(
+        name="Recipe Editor Agent",
+        instructions="You are a recipe editor specialist. Help users search and modify recipes using your tools.",
+        model="gpt-4o",
+        tools=[search_recipes, plan_and_apply_recipe_modifications],
+    )
+
+    main_chat_agent = Agent(
+        name="Main Chat Agent",
+        instructions="You handle general conversation and route recipe tasks to the recipe editor agent.",
+        model="gpt-4o",
+        handoffs=[recipe_editor_agent],
+    )
+
+    return main_chat_agent, recipe_editor_agent
+
+
 @pytest.fixture(scope="module")
 def vcr_config():
     return {"filter_headers": ["authorization", "api-key"]}