.

Author: spike-spiegel-21

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,44 @@ 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=_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=[],
+                )
+            )
 
     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)
 
     def on_span_start(self, span: "Span") -> None:
         if not span.started_at:

tests/test_integrations/test_openai_agents/agents_app.py

@@ -98,24 +98,40 @@ 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
 
-
-# Usage example
-async def main():
+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_2 = await Runner.run(weather_agent, user_query, metric_collection="test_collection_1")
+        response_1 = await Runner.run(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
+
+# without trace (group_id and metadata not present)
+async def main2():
     user_query = "What's the weather like in London today?"
-    response = await run_weather_agent(user_query)
-    print(f"Agent Response: {response}")
+    response_2 = await Runner.run(weather_agent, user_query, metric_collection="test_collection_1")
+    response_1 = await Runner.run(triage_agent, "Hola, ¿cómo estás?", metric_collection="test_collection_1", thread_id="test")
 
+async def main3():
+    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"}):
+        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")
 
 def execute_agent():
-    return asyncio.run(main())
+    asyncio.run(main1())
+    asyncio.run(main2())
+    asyncio.run(main3())
 
 
-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")