Merge pull request #2084 from confident-ai/mayank/openai_agents_types

OpenAI agents types

Author: penguine-ip

deepeval/openai_agents/agent.py

@@ -1,37 +1,40 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field, replace
-from typing import Any, Optional, Awaitable, Callable, Generic, TypeVar
+from typing import Any, Optional, Awaitable, Callable, Generic, TypeVar, List
 
 from deepeval.tracing import observe
 from deepeval.prompt import Prompt
+from deepeval.tracing.tracing import Observer
+from deepeval.metrics import BaseMetric
+from deepeval.tracing.utils import make_json_serializable
+from deepeval.tracing.types import LlmSpan
+from deepeval.tracing.context import current_span_context
 
 try:
     from agents.agent import Agent as BaseAgent
     from agents.models.interface import Model, ModelProvider
+    from openai.types.responses import ResponseCompletedEvent
 except Exception as e:
     raise RuntimeError(
         "openai-agents is required for this integration. Please install it."
     ) from e
 
 TContext = TypeVar("TContext")
 
-
 class _ObservedModel(Model):
     def __init__(
         self,
         inner: Model,
-        *,
-        metrics: Optional[list[Any]] = None,
-        metric_collection: Optional[str] = None,
-        deepeval_prompt: Optional[Any] = None,
+        llm_metric_collection: str = None,
+        llm_metrics: List[BaseMetric] = None,
+        confident_prompt: Prompt = None,
     ) -> None:
         self._inner = inner
-        self._metrics = metrics
-        self._metric_collection = metric_collection
-        self._deepeval_prompt = deepeval_prompt
+        self._llm_metric_collection = llm_metric_collection
+        self._llm_metrics = llm_metrics
+        self._confident_prompt = confident_prompt
 
-    # Delegate attributes not overridden
     def __getattr__(self, name: str) -> Any:
         return getattr(self._inner, name)
 
@@ -59,29 +62,48 @@ async def get_response(
         previous_response_id,
         conversation_id,
         prompt,
+        **kwargs,
     ):
         model_name = self._get_model_name()
-
-        wrapped = observe(
-            metrics=self._metrics,
-            metric_collection=self._metric_collection,
-            type="llm",
-            model=model_name,
-            prompt=self._deepeval_prompt,
-        )(self._inner.get_response)
-
-        return await wrapped(
-            system_instructions,
-            input,
-            model_settings,
-            tools,
-            output_schema,
-            handoffs,
-            tracing,
-            previous_response_id=previous_response_id,
-            conversation_id=conversation_id,
-            prompt=prompt,
-        )
+        with Observer(
+            span_type="llm",
+            func_name="LLM",
+            function_kwargs={
+                "system_instructions": system_instructions, 
+                "input": input, 
+                "model_settings": model_settings, 
+                "tools": tools, 
+                "output_schema": output_schema, 
+                "handoffs": handoffs, 
+                # "tracing": tracing, # not important for llm spans
+                # "previous_response_id": previous_response_id, # not important for llm spans
+                # "conversation_id": conversation_id, # not important for llm spans
+                "prompt": prompt,
+                **kwargs,
+            },
+            observe_kwargs={"model": model_name},
+            metrics=self._llm_metrics,
+            metric_collection=self._llm_metric_collection,
+        ) as observer:
+            result = await self._inner.get_response(
+                system_instructions,
+                input,
+                model_settings,
+                tools,
+                output_schema,
+                handoffs,
+                tracing,
+                previous_response_id=previous_response_id,
+                conversation_id=conversation_id,
+                prompt=prompt,
+                **kwargs,
+            )
+            llm_span: LlmSpan = current_span_context.get()
+            llm_span.prompt = self._confident_prompt
+
+            observer.result = make_json_serializable(result.output)
+        
+        return result
 
     def stream_response(
         self,
@@ -96,91 +118,74 @@ def stream_response(
         previous_response_id,
         conversation_id,
         prompt,
+        **kwargs,
     ):
-        # Optional: if you also want to observe streaming, uncomment and wrap similarly.
-        # wrapped = observe(
-        #     metrics=self._metrics,
-        #     metric_collection=self._metric_collection,
-        #     type="llm",
-        #     model=model_name,
-        # )(self._inner.stream_response)
-        # return wrapped(
-        #     system_instructions,
-        #     input,
-        #     model_settings,
-        #     tools,
-        #     output_schema,
-        #     handoffs,
-        #     tracing,
-        #     previous_response_id=previous_response_id,
-        #     conversation_id=conversation_id,
-        #     prompt=prompt,
-        # )
-        return self._inner.stream_response(
-            system_instructions,
-            input,
-            model_settings,
-            tools,
-            output_schema,
-            handoffs,
-            tracing,
-            previous_response_id=previous_response_id,
-            conversation_id=conversation_id,
-            prompt=prompt,
-        )
-
-
-class _ObservedProvider(ModelProvider):
-    def __init__(
-        self,
-        base: ModelProvider,
-        *,
-        metrics: Optional[list[Any]] = None,
-        metric_collection: Optional[str] = None,
-        deepeval_prompt: Optional[Any] = None,
-    ) -> None:
-        self._base = base
-        self._metrics = metrics
-        self._metric_collection = metric_collection
-        self._deepeval_prompt = deepeval_prompt
+        model_name = self._get_model_name()
 
-    def get_model(self, model_name: str | None) -> Model:
-        model = self._base.get_model(model_name)
-        return _ObservedModel(
-            model,
-            metrics=self._metrics,
-            metric_collection=self._metric_collection,
-            deepeval_prompt=self._deepeval_prompt,
-        )
+        async def _gen():
+            observer = Observer(
+                span_type="llm",
+                func_name="LLM",
+                function_kwargs={
+                    "system_instructions": system_instructions,
+                    "input": input,
+                    "model_settings": model_settings,
+                    "tools": tools,
+                    "output_schema": output_schema,
+                    "handoffs": handoffs,
+                    # "tracing": tracing,
+                    # "previous_response_id": previous_response_id,
+                    # "conversation_id": conversation_id,
+                    "prompt": prompt,
+                    **kwargs,
+                },
+                observe_kwargs={"model": model_name},
+                metrics=self._llm_metrics,
+                metric_collection=self._llm_metric_collection,
+            )
+            observer.__enter__()
+
+            llm_span: LlmSpan = current_span_context.get()
+            llm_span.prompt = self._confident_prompt
 
+            try:
+                async for event in self._inner.stream_response(
+                    system_instructions,
+                    input,
+                    model_settings,
+                    tools,
+                    output_schema,
+                    handoffs,
+                    tracing,
+                    previous_response_id=previous_response_id,
+                    conversation_id=conversation_id,
+                    prompt=prompt,
+                ):
+
+                    if isinstance(event, ResponseCompletedEvent):
+                        observer.result = event.response.output_text #TODO: support other response types
+
+                    yield event
+
+                observer.__exit__(None, None, None)
+            except Exception as e:
+                observer.__exit__(type(e), e, e.__traceback__)
+                raise
+            finally:
+
+                observer.__exit__(None, None, None)
+
+        return _gen()
 
 @dataclass
 class DeepEvalAgent(BaseAgent[TContext], Generic[TContext]):
     """
-    A subclass of agents.Agent that accepts `metrics` and `metric_collection`
-    and ensures the underlying model's `get_response` is wrapped with deepeval.observe.
+    A subclass of agents.Agent.
     """
-
-    metrics: list[Any] | None = field(default=None)
-    metric_collection: str | None = field(default=None)
-    deepeval_prompt: Prompt | None = field(default=None)
+    llm_metric_collection: str = None
+    llm_metrics: List[BaseMetric] = None
+    confident_prompt: Prompt = None
 
     def __post_init__(self):
         super().__post_init__()
-        # If a direct Model instance is set on the agent, wrap it here.
-        if self.model is not None and not isinstance(self.model, str):
-            try:
-                from agents.models.interface import (
-                    Model as _Model,
-                )  # local import for safety
-
-                if isinstance(self.model, _Model):
-                    self.model = _ObservedModel(
-                        self.model,
-                        metrics=self.metrics,
-                        metric_collection=self.metric_collection,
-                        deepeval_prompt=self.deepeval_prompt,
-                    )
-            except Exception:
-                # If we can't import or wrap, silently skip.
-                pass
+

deepeval/openai_agents/callback_handler.py

@@ -46,17 +46,7 @@ def on_span_start(self, span: "Span") -> None:
         if not span.started_at:
             return
         span_type = self.get_span_kind(span.span_data)
-        if span_type == "agent":
-            if isinstance(span.span_data, AgentSpanData):
-                current_trace = current_trace_context.get()
-                if current_trace:
-                    current_trace.name = span.span_data.name
-
-        if span_type == "tool":
-            return
-        elif span_type == "llm":
-            return
-        else:
+        if span_type and span_type == "agent":
             observer = Observer(span_type=span_type, func_name="NA")
             observer.update_span_properties = (
                 lambda base_span: update_span_properties(
@@ -68,13 +58,13 @@ def on_span_start(self, span: "Span") -> None:
 
     def on_span_end(self, span: "Span") -> None:
         span_type = self.get_span_kind(span.span_data)
-        if span_type == "llm":
+        if span_type and span_type == "agent":
             current_span = current_span_context.get()
             if current_span:
                 update_span_properties(current_span, span.span_data)
-        observer = self.span_observers.pop(span.span_id, None)
-        if observer:
-            observer.__exit__(None, None, None)
+            observer = self.span_observers.pop(span.span_id, None)
+            if observer:
+                observer.__exit__(None, None, None)
 
     def force_flush(self) -> None:
         pass
@@ -85,18 +75,19 @@ def shutdown(self) -> None:
     def get_span_kind(self, span_data: "SpanData") -> str:
         if isinstance(span_data, AgentSpanData):
             return "agent"
-        if isinstance(span_data, FunctionSpanData):
-            return "tool"
-        if isinstance(span_data, MCPListToolsSpanData):
-            return "tool"
-        if isinstance(span_data, GenerationSpanData):
-            return "llm"
-        if isinstance(span_data, ResponseSpanData):
-            return "llm"
-        if isinstance(span_data, HandoffSpanData):
-            return "custom"
-        if isinstance(span_data, CustomSpanData):
-            return "base"
-        if isinstance(span_data, GuardrailSpanData):
-            return "base"
-        return "base"
+        # if isinstance(span_data, FunctionSpanData):
+        #     return "tool"
+        # if isinstance(span_data, MCPListToolsSpanData):
+        #     return "tool"
+        # if isinstance(span_data, GenerationSpanData):
+        #     return "llm"
+        # if isinstance(span_data, ResponseSpanData):
+        #     return "llm"
+        # if isinstance(span_data, HandoffSpanData):
+        #     return "custom"
+        # if isinstance(span_data, CustomSpanData):
+        #     return "base"
+        # if isinstance(span_data, GuardrailSpanData):
+        #     return "base"
+        # return "base"
+        return None