feat(llmobs): [MLOB-2662] [MLOB-3100] add agent manifest

ddtrace/contrib/internal/openai_agents/patch.py

@@ -7,6 +7,12 @@
 from ddtrace.contrib.internal.openai_agents.processor import LLMObsTraceProcessor
 from ddtrace.llmobs._integrations.openai_agents import OpenAIAgentsIntegration
 from ddtrace.trace import Pin
+from ddtrace.contrib.trace_utils import with_traced_module_async
+from ddtrace.contrib.trace_utils import wrap
+from ddtrace.contrib.trace_utils import unwrap
+from ddtrace.internal.utils import get_argument_value
+from ddtrace.contrib.internal.openai_agents.utils import create_agent_manifest
+from ddtrace.llmobs._constants import AGENT_MANIFEST
 
 
 config._add("openai_agents", {})
@@ -17,11 +23,22 @@ def get_version() -> str:
 
     return getattr(version, "__version__", "")
 
-
 def _supported_versions() -> Dict[str, str]:
     return {"agents": ">=0.0.2"}
 
 
+@with_traced_module_async
+async def patched_run_single_turn(agents, pin, func, instance, args, kwargs):
+    from ddtrace import tracer
+    s = tracer.current_span()
+    r = await func(*args, **kwargs)
+
+    agent = get_argument_value(args, kwargs, 0, "agent", None)
+    agent_manifest = create_agent_manifest(agent) if agent else None
+
+    s._set_ctx_item(AGENT_MANIFEST, agent_manifest)
+    return r
+
 def patch():
     """
     Patch the instrumented methods
@@ -34,6 +51,8 @@ def patch():
     Pin().onto(agents)
 
     add_trace_processor(LLMObsTraceProcessor(OpenAIAgentsIntegration(integration_config=config.openai_agents)))
+
+    wrap(agents.Runner, "_run_single_turn", patched_run_single_turn(agents))
 
 
 def unpatch():
@@ -44,3 +63,5 @@ def unpatch():
         return
 
     agents._datadog_patch = False
+
+    unwrap(agents.Runner, "_run_single_turn")

ddtrace/contrib/internal/openai_agents/utils.py

@@ -0,0 +1,145 @@
+from agents import Agent
+from agents import Handoff
+from agents import (
+    WebSearchTool,
+    FileSearchTool,
+    ComputerTool,
+)
+
+def create_agent_manifest(agent):
+    manifest = {}
+    manifest["framework"] = "OpenAI"
+    manifest["model_provider"] = "openai"
+
+    if hasattr(agent, "name"):
+        manifest["name"] = agent.name
+    if hasattr(agent, "instructions"):
+        manifest["instructions"] = agent.instructions
+    if hasattr(agent, "handoff_description"):
+        manifest["handoff_description"] = agent.handoff_description
+    if hasattr(agent, "model"):
+        manifest["model"] = agent.model
+
+    model_settings = extract_model_settings_from_agent(agent)
+    if model_settings:
+        manifest["model_settings"] = model_settings
+
+    tools = extract_tools_from_agent(agent)
+    if tools:
+        manifest["tools"] = tools
+
+    handoffs = extract_handoffs_from_agent(agent)
+    if handoffs:
+        manifest["handoffs"] = handoffs
+
+    guardrails = extract_guardrails_from_agent(agent)
+    if guardrails:
+        manifest["guardrails"] = guardrails
+
+    return manifest
+
+def extract_model_settings_from_agent(agent):
+    if not hasattr(agent, "model_settings"):
+        return None
+
+    # convert model_settings to dict if it's not already
+    model_settings = agent.model_settings
+    if type(model_settings) != dict:
+        if hasattr(model_settings, "__dict__"):
+            model_settings = model_settings.__dict__
+        else:
+            return None
+
+    return make_json_compatible(model_settings)
+
+def extract_tools_from_agent(agent):
+    if not hasattr(agent, "tools"):
+        return None
+
+    tools = []
+    for tool in agent.tools:
+        tool_dict = {}
+        if isinstance(tool, WebSearchTool):
+            if hasattr(tool, "user_location"):
+                tool_dict["user_location"] = tool.user_location
+            if hasattr(tool, "search_context_size"):
+                tool_dict["search_context_size"] = tool.search_context_size
+        elif isinstance(tool, FileSearchTool):
+            if hasattr(tool, "vector_store_ids"):
+                tool_dict["vector_store_ids"] = tool.vector_store_ids
+            if hasattr(tool, "max_num_results"):
+                tool_dict["max_num_results"] = tool.max_num_results
+            if hasattr(tool, "include_search_results"):
+                tool_dict["include_search_results"] = tool.include_search_results
+        elif isinstance(tool, ComputerTool):
+            if hasattr(tool, "name"):
+                tool_dict["name"] = tool.name
+        else:
+            if hasattr(tool, "name"):
+                tool_dict["name"] = tool.name
+            if hasattr(tool, "description"):
+                tool_dict["description"] = tool.description
+            if hasattr(tool, "strict_json_schema"):
+                tool_dict["strict_json_schema"] = tool.strict_json_schema
+            if hasattr(tool, "params_json_schema"):
+                parameter_schema = tool.params_json_schema
+                required_params = get_required_param_dict(parameter_schema.get("required", [])) 
+                parameters = {}
+                for param, schema in parameter_schema.get("properties", {}).items():
+                    param_dict = {}
+                    if "type" in schema:
+                        param_dict["type"] = schema["type"]
+                    if "title" in schema:
+                        param_dict["title"] = schema["title"]
+                    if param in required_params:
+                        param_dict["required"] = True
+                    parameters[param] = param_dict
+                tool_dict["parameters"] = parameters
+        tools.append(tool_dict)                     
+
+    return tools
+
+def extract_handoffs_from_agent(agent):
+    if not hasattr(agent, "handoffs"):
+        return None
+
+    handoffs = []
+    for handoff in agent.handoffs:
+        handoff_dict = {}
+        if isinstance(handoff, Agent):
+            if hasattr(handoff, "handoff_description"):
+                handoff_dict["handoff_description"] = handoff.handoff_description
+            if hasattr(handoff, "name"):
+                handoff_dict["agent_name"] = handoff.name
+        elif isinstance(handoff, Handoff):
+            if hasattr(handoff, "tool_name"):
+                handoff_dict["tool_name"] = handoff.tool_name
+            if hasattr(handoff, "tool_description"):
+                handoff_dict["handoff_description"] = handoff.tool_description
+            if hasattr(handoff, "agent_name"):
+                handoff_dict["agent_name"] = handoff.agent_name
+        if handoff_dict:
+            handoffs.append(handoff_dict)
+
+    return handoffs
+
+def extract_guardrails_from_agent(agent):
+    guardrails = []
+    if hasattr(agent, "input_guardrails"):
+        guardrails.extend([getattr(guardrail, "name", "") for guardrail in agent.input_guardrails])
+    if hasattr(agent, "output_guardrails"):
+        guardrails.extend([getattr(guardrail, "name", "") for guardrail in agent.output_guardrails])
+    return guardrails
+
+def get_required_param_dict(required_params):
+    return {param: True for param in required_params}
+
+def make_json_compatible(obj):
+    if isinstance(obj, dict):
+        return {str(k): make_json_compatible(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple, set)):
+        return [make_json_compatible(v) for v in obj]
+    elif isinstance(obj, (int, float, str, bool)) or obj is None:
+        return obj
+    else:
+        return str(obj)

ddtrace/contrib/internal/pydantic_ai/patch.py

@@ -1,7 +1,9 @@
+import sys
 from typing import Dict
 
 from ddtrace import config
 from ddtrace.contrib.internal.pydantic_ai.utils import TracedPydanticAsyncContextManager
+from ddtrace.contrib.internal.pydantic_ai.utils import TracedPydanticRunStream
 from ddtrace.contrib.internal.trace_utils import unwrap
 from ddtrace.contrib.internal.trace_utils import wrap
 from ddtrace.contrib.trace_utils import with_traced_module
@@ -22,22 +24,53 @@ def _supported_versions() -> Dict[str, str]:
     return {"pydantic_ai": "*"}
 
 
+@with_traced_module
+def traced_agent_run_stream(pydantic_ai, pin, func, instance, args, kwargs):
+    integration = pydantic_ai._datadog_integration
+    integration._run_stream_active = True
+    span = integration.trace(
+        pin, "Pydantic Agent", submit_to_llmobs=True, model=getattr(instance, "model", None), kind="agent"
+    )
+    span.name = getattr(instance, "name", None) or "Pydantic Agent"
+
+    result = func(*args, **kwargs)
+    kwargs["instance"] = instance
+    return TracedPydanticRunStream(result, span, integration, args, kwargs)
+
+
 @with_traced_module
 def traced_agent_iter(pydantic_ai, pin, func, instance, args, kwargs):
     integration = pydantic_ai._datadog_integration
-    span = integration.trace(pin, "Pydantic Agent", submit_to_llmobs=False, model=getattr(instance, "model", None))
+    # avoid double tracing if run_stream has already been called
+    if integration._run_stream_active:
+        integration._run_stream_active = False
+        return func(*args, **kwargs)
+    span = integration.trace(
+        pin, "Pydantic Agent", submit_to_llmobs=True, model=getattr(instance, "model", None), kind="agent"
+    )
     span.name = getattr(instance, "name", None) or "Pydantic Agent"
 
     result = func(*args, **kwargs)
-    return TracedPydanticAsyncContextManager(result, span)
+    kwargs["instance"] = instance
+    return TracedPydanticAsyncContextManager(result, span, instance, integration, args, kwargs)
 
 
 @with_traced_module
 async def traced_tool_run(pydantic_ai, pin, func, instance, args, kwargs):
     integration = pydantic_ai._datadog_integration
-    with integration.trace(pin, "Pydantic Tool", submit_to_llmobs=False) as span:
+    resp = None
+    try:
+        span = integration.trace(pin, "Pydantic Tool", submit_to_llmobs=True, kind="tool")
         span.name = getattr(instance, "name", None) or "Pydantic Tool"
-        return await func(*args, **kwargs)
+        resp = await func(*args, **kwargs)
+        return resp
+    except Exception:
+        span.set_exc_info(*sys.exc_info())
+        raise
+    finally:
+        kwargs["instance"] = instance
+        integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=resp)
+        span.finish()
 
 
 def patch():
@@ -53,6 +86,7 @@ def patch():
 
     wrap(pydantic_ai, "agent.Agent.iter", traced_agent_iter(pydantic_ai))
     wrap(pydantic_ai, "tools.Tool.run", traced_tool_run(pydantic_ai))
+    wrap(pydantic_ai, "agent.Agent.run_stream", traced_agent_run_stream(pydantic_ai))
 
 
 def unpatch():
@@ -65,6 +99,7 @@ def unpatch():
 
     unwrap(pydantic_ai.agent.Agent, "iter")
     unwrap(pydantic_ai.tools.Tool, "run")
+    unwrap(pydantic_ai.agent.Agent, "run_stream")
 
     delattr(pydantic_ai, "_datadog_integration")
     Pin().remove_from(pydantic_ai)