Merge pull request #2087 from confident-ai/mayank/agents_trace_support_1

open ai agents trace support

Author: penguine-ip

deepeval/openai_agents/callback_handler.py

@@ -1,9 +1,13 @@
 from deepeval.tracing.tracing import (
     Observer,
     current_span_context,
+    trace_manager,
 )
 from deepeval.openai_agents.extractors import *
 from deepeval.tracing.context import current_trace_context
+from deepeval.tracing.utils import make_json_serializable
+from time import perf_counter
+from deepeval.tracing.types import TraceSpanStatus
 
 try:
     from agents.tracing import Span, Trace, TracingProcessor
@@ -33,14 +37,49 @@ def _check_openai_agents_available():
 class DeepEvalTracingProcessor(TracingProcessor):
     def __init__(self) -> None:
         _check_openai_agents_available()
-        self.root_span_observers: dict[str, Observer] = {}
         self.span_observers: dict[str, Observer] = {}
 
     def on_trace_start(self, trace: "Trace") -> None:
-        pass
+        trace_dict = trace.export()
+        _trace_uuid = trace_dict.get("id")
+        _thread_id = trace_dict.get("group_id")
+        _trace_name = trace_dict.get("workflow_name")
+        _trace_metadata = trace_dict.get("metadata")
+
+        if _thread_id or _trace_metadata:
+            _trace = trace_manager.start_new_trace(trace_uuid=str(_trace_uuid))
+            _trace.thread_id = str(_thread_id)
+            _trace.name = str(_trace_name)
+            _trace.metadata = make_json_serializable(_trace_metadata)
+            current_trace_context.set(_trace)
+
+            trace_manager.add_span( # adds a dummy root span
+                BaseSpan(
+                    uuid=_trace_uuid,
+                    trace_uuid=_trace_uuid,
+                    parent_uuid=None,
+                    start_time=perf_counter(),
+                    name=_trace_name,
+                    status=TraceSpanStatus.IN_PROGRESS,
+                    children=[],
+                )
+            )
+        else:
+            current_trace = current_trace_context.get()
+            if current_trace:
+                current_trace.name = str(_trace_name)
 
     def on_trace_end(self, trace: "Trace") -> None:
-        pass
+        trace_dict = trace.export()
+        _trace_uuid = trace_dict.get("id")
+        _thread_id = trace_dict.get("group_id")
+        _trace_name = trace_dict.get("workflow_name")
+        _trace_metadata = trace_dict.get("metadata")
+
+        if _thread_id or _trace_metadata:
+            trace_manager.remove_span(_trace_uuid) # removing the dummy root span
+            trace_manager.end_trace(_trace_uuid)
+            current_trace_context.set(None)
 
     def on_span_start(self, span: "Span") -> None:
         if not span.started_at:

deepeval/openai_agents/runner.py

@@ -109,10 +109,9 @@ async def run(
             metric_collection=metric_collection,
             metrics=metrics,
             func_name="run",
-            function_kwargs={"input": input},
+            function_kwargs={"input": input}, # also set below
         ) as observer:
             update_trace_attributes(
-                input=input,
                 name=name,
                 tags=tags,
                 metadata=metadata,
@@ -123,7 +122,8 @@ async def run(
             )
             current_span = current_span_context.get()
             current_trace = current_trace_context.get()
-            current_trace.input = input
+            if not current_trace.input:
+                current_trace.input = input
             if current_span:
                 current_span.input = input
             res = await super().run(
@@ -138,8 +138,9 @@ async def run(
                 session=session,
                 **kwargs,  # backwards compatibility
             )
+            current_trace_thread_id = current_trace_context.get().thread_id
             _output = None
-            if thread_id:
+            if current_trace_thread_id:
                 _output = res.final_output
             else:
                 _output = str(res)
@@ -170,30 +171,30 @@ def run_sync(
         **kwargs,
     ) -> RunResult:
         is_agents_available()
-        input_val = input
-
-        update_trace_attributes(
-            input=input_val,
-            name=name,
-            tags=tags,
-            metadata=metadata,
-            thread_id=thread_id,
-            user_id=user_id,
-            metric_collection=metric_collection,
-            metrics=metrics,
-        )
 
         with Observer(
             span_type="custom",
             metric_collection=metric_collection,
             metrics=metrics,
             func_name="run_sync",
-            function_kwargs={"input": input_val},
+            function_kwargs={"input": input}, # also set below
         ) as observer:
+            update_trace_attributes(
+                name=name,
+                tags=tags,
+                metadata=metadata,
+                thread_id=thread_id,
+                user_id=user_id,
+                metric_collection=metric_collection,
+                metrics=metrics,
+            )
+
             current_span = current_span_context.get()
             current_trace = current_trace_context.get()
+            if not current_trace.input:
+                current_trace.input = input
             if current_span:
-                current_span.input = input_val
+                current_span.input = input
             res = super().run_sync(
                 starting_agent,
                 input,
@@ -206,8 +207,9 @@ def run_sync(
                 session=session,
                 **kwargs,  # backwards compatibility
             )
+            current_trace_thread_id = current_trace_context.get().thread_id
             _output = None
-            if thread_id:
+            if current_trace_thread_id:
                 _output = res.final_output
             else:
                 _output = str(res)
@@ -250,7 +252,6 @@ def run_streamed(
         observer.__enter__()
 
         update_trace_attributes(
-            input=input,
             name=name,
             tags=tags,
             metadata=metadata,
@@ -259,7 +260,10 @@ def run_streamed(
             metric_collection=metric_collection,
             metrics=metrics,
         )
-
+        current_trace = current_trace_context.get()
+        if not current_trace.input:
+            current_trace.input = input
+            
         current_span = current_span_context.get()
         if current_span:
             current_span.input = input

tests/test_integrations/test_openai_agents/agents_app.py

@@ -7,6 +7,7 @@
     DeepEvalTracingProcessor,
 )
 
+from deepeval.tracing.context import update_current_trace
 from deepeval.prompt import Prompt
 
 add_trace_processor(DeepEvalTracingProcessor())
@@ -98,24 +99,55 @@ async def run_weather_agent(user_input: str):
         weather_agent,
         user_input,
         metric_collection="test_collection_1",
-        name="test_name_1",
-        user_id="test_user_id_1",
-        thread_id="test_thread_id_1",
-        tags=["test_tag_1"],
-        metadata={"test_metadata_1": "test_metadata_1"},
+        # name="test_name_1",
+        # user_id="test_user_id_1",
+        # thread_id="test_thread_id_1",
+        # tags=["test_tag_1"],
+        # metadata={"test_metadata_1": "test_metadata_1"},
     )
     return result.final_output
 
+from agents import trace
+from multi_agents import triage_agent
+# with trace (group_id and metadata)
+async def main1():
+    with trace(workflow_name="test_workflow_1", group_id="test_group_id_1", metadata={"test_metadata_1": "test_metadata_1"}):
+        user_query = "What's the weather like in London today?"
+        response_1 = await Runner.run(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
+        response_2 = await Runner.run(weather_agent, user_query, metric_collection="test_collection_1")
+        update_current_trace(input="initial input", output="final output")
+
+# without trace (group_id and metadata not present)
+async def main2():
+    user_query = "What's the weather like in London today?"
+    response_1 = await Runner.run(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
+    response_2 = await Runner.run(weather_agent, user_query, metric_collection="test_collection_1")
+    
 
-# Usage example
-async def main():
+async def main3():
     user_query = "What's the weather like in London today?"
-    response = await run_weather_agent(user_query)
-    print(f"Agent Response: {response}")
+    with trace(workflow_name="test_workflow_1", group_id="test_group_id_1", metadata={"test_metadata_1": "test_metadata_1"}):
+        response_2 = await Runner.run(weather_agent, user_query, metric_collection="test_collection_1")
+    with trace(workflow_name="test_workflow_2", group_id="test_group_id_2", metadata={"test_metadata_2": "test_metadata_2"}):
+        response_1 = await Runner.run(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
 
+async def main4():
+    user_query = "What's the weather like in London today?"
+    with trace(workflow_name="test_workflow_1", group_id="test_group_id_1", metadata={"test_metadata_1": "test_metadata_1"}):
+        run_streamed_1 = Runner.run_streamed(weather_agent, user_query, metric_collection="test_collection_1")
+        async for chunk in run_streamed_1.stream_events():
+            print(chunk, end="", flush=True)
+            print("=" * 50)
+        run_streamed_2 = Runner.run_streamed(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
+        async for chunk in run_streamed_2.stream_events():
+            print(chunk, end="", flush=True)
+            print("=" * 50)
 
 def execute_agent():
-    return asyncio.run(main())
+    asyncio.run(main1())
+    # asyncio.run(main2())
+    # asyncio.run(main3())
+    # asyncio.run(main4())
 
 
-execute_agent()
+# execute_agent()

tests/test_integrations/test_openai_agents/multi_agents.py

@@ -1,11 +1,11 @@
 import asyncio
 from deepeval.openai_agents import Agent, Runner
 from deepeval.prompt import Prompt
-from deepeval.openai_agents import DeepEvalTracingProcessor
+# from deepeval.openai_agents import DeepEvalTracingProcessor
 
-from agents import add_trace_processor
+# from agents import add_trace_processor
 
-add_trace_processor(DeepEvalTracingProcessor())
+# add_trace_processor(DeepEvalTracingProcessor())
 
 prompt = Prompt(alias="asd")
 prompt.pull(version="00.00.01")