reformat tracing

deepeval/evaluate/execute.py

@@ -8,7 +8,6 @@
 
 from deepeval.tracing.tracing import (
     Observer,
-    get_current_trace,
     trace_manager,
     Trace,
     BaseSpan,
@@ -19,6 +18,7 @@
     perf_counter_to_datetime,
     to_zod_compatible_iso,
 )
+from deepeval.tracing.context import current_trace_context
 from deepeval.tracing.api import (
     TraceApi,
     BaseApiSpan,
@@ -818,7 +818,7 @@ def evaluate_test_cases(
                         loop.run_until_complete(observed_callback(golden.input))
                     else:
                         observed_callback(golden.input)
-                    current_trace: Trace = get_current_trace()
+                    current_trace: Trace = current_trace_context.get()
 
                 if pbar_callback is not None:
                     pbar_callback.update(1)
@@ -1117,7 +1117,7 @@ async def a_execute_agentic_test_case(
             await observed_callback(golden.input)
         else:
             observed_callback(golden.input)
-        current_trace: Trace = get_current_trace()
+        current_trace: Trace = current_trace_context.get()
 
     if pbar_callback is not None:
         pbar_callback.update(1)

deepeval/tracing/__init__.py

@@ -1,11 +1,9 @@
-from .tracing import (
-    observe,
-    update_current_span,
-    update_current_trace,
+from .context import update_current_span, update_current_trace
+from .attributes import (
     LlmAttributes,
     RetrieverAttributes,
     ToolAttributes,
     AgentAttributes,
-    get_current_trace,
-    trace_manager,
 )
+from .types import BaseSpan, Trace
+from .tracing import observe, trace_manager

deepeval/tracing/attributes.py

@@ -0,0 +1,55 @@
+from pydantic import BaseModel, Field
+from typing import Any, Dict, List, Optional, Union
+from deepeval.prompt import Prompt
+
+
+class AgentAttributes(BaseModel):
+    # input
+    input: Union[str, Dict, list]
+    # output
+    output: Union[str, Dict, list]
+
+
+class LlmAttributes(BaseModel):
+    # input
+    input: Union[str, List[Dict[str, Any]]]
+    # output
+    output: str
+    prompt: Optional[Prompt] = None
+
+    # Optional variables
+    input_token_count: Optional[int] = Field(
+        None, serialization_alias="inputTokenCount"
+    )
+    output_token_count: Optional[int] = Field(
+        None, serialization_alias="outputTokenCount"
+    )
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class RetrieverAttributes(BaseModel):
+    # input
+    embedding_input: str = Field(serialization_alias="embeddingInput")
+    # output
+    retrieval_context: List[str] = Field(serialization_alias="retrievalContext")
+
+    # Optional variables
+    top_k: Optional[int] = Field(None, serialization_alias="topK")
+    chunk_size: Optional[int] = Field(None, serialization_alias="chunkSize")
+
+
+# Don't have to call this manually, will be taken as input and output of function
+# Can be overridden by user
+class ToolAttributes(BaseModel):
+    # input
+    input_parameters: Optional[Dict[str, Any]] = Field(
+        None, serialization_alias="inputParameters"
+    )
+    # output
+    output: Optional[Any] = None
+
+
+Attributes = Union[
+    AgentAttributes, LlmAttributes, RetrieverAttributes, ToolAttributes
+]

deepeval/tracing/context.py

@@ -0,0 +1,55 @@
+from typing import Any, Dict, List, Optional
+from contextvars import ContextVar
+from deepeval.tracing.types import BaseSpan, Trace
+from deepeval.test_case import LLMTestCase
+from deepeval.tracing.attributes import Attributes
+
+
+current_span_context: ContextVar[Optional[BaseSpan]] = ContextVar(
+    "current_span", default=None
+)
+
+current_trace_context: ContextVar[Optional[Trace]] = ContextVar(
+    "current_trace", default=None
+)
+
+
+def update_current_span(
+    test_case: Optional[LLMTestCase] = None,
+    attributes: Optional[Attributes] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+):
+    current_span = current_span_context.get()
+    if not current_span:
+        return
+    if attributes:
+        current_span.set_attributes(attributes)
+    if test_case:
+        current_span.llm_test_case = test_case
+    if metadata:
+        current_span.metadata = metadata
+
+
+def update_current_trace(
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    thread_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+    input: Optional[Any] = None,
+    output: Optional[Any] = None,
+):
+    current_trace = current_trace_context.get()
+    if not current_trace:
+        return
+    if tags:
+        current_trace.tags = tags
+    if metadata:
+        current_trace.metadata = metadata
+    if thread_id:
+        current_trace.thread_id = thread_id
+    if user_id:
+        current_trace.user_id = user_id
+    if input:
+        current_trace.input = input
+    if output:
+        current_trace.output = output

deepeval/tracing/patchers.py

@@ -0,0 +1,89 @@
+from openai import OpenAI
+import functools
+
+from deepeval.tracing.attributes import LlmAttributes
+from deepeval.tracing.context import update_current_span
+from deepeval.tracing.context import current_span_context
+from deepeval.tracing.types import LlmSpan
+
+
+def patch_openai_client(client: OpenAI):
+
+    original_methods = {}
+
+    # patches these methods
+    methods_to_patch = [
+        "chat.completions.create",
+        "beta.chat.completions.parse",
+    ]
+
+    for method_path in methods_to_patch:
+        # Split the path into components
+        parts = method_path.split(".")
+        current_obj = client
+
+        # Navigate to the parent object
+        for part in parts[:-1]:
+            if not hasattr(current_obj, part):
+                print(f"Warning: Cannot find {part} in the path {method_path}")
+                continue
+            current_obj = getattr(current_obj, part)
+
+        method_name = parts[-1]
+        if not hasattr(current_obj, method_name):
+            print(
+                f"Warning: Cannot find method {method_name} in the path {method_path}"
+            )
+            continue
+
+        method = getattr(current_obj, method_name)
+
+        if callable(method) and not isinstance(method, type):
+            original_methods[method_path] = method
+
+            # Capture the current 'method' using a default argument
+            @functools.wraps(method)
+            def wrapped_method(*args, original_method=method, **kwargs):
+                current_span = current_span_context.get()
+                # call the original method using the captured default argument
+                response = original_method(*args, **kwargs)
+                if isinstance(current_span, LlmSpan):
+                    # extract model
+                    model = kwargs.get("model", None)
+                    if model is None:
+                        raise ValueError("model not found in client")
+
+                    # set model
+                    current_span.model = model
+
+                    # extract output message
+                    output = None
+                    try:
+                        output = response.choices[0].message.content
+                    except Exception as e:
+                        pass
+
+                    # extract input output token counts
+                    input_token_count = None
+                    output_token_count = None
+                    try:
+                        input_token_count = response.usage.prompt_tokens
+                        output_token_count = response.usage.completion_tokens
+                    except Exception as e:
+                        pass
+
+                    update_current_span(
+                        attributes=LlmAttributes(
+                            input=kwargs.get(
+                                "messages", "INPUT_MESSAGE_NOT_FOUND"
+                            ),
+                            output=(
+                                output if output else "OUTPUT_MESSAGE_NOT_FOUND"
+                            ),
+                            input_token_count=input_token_count,
+                            output_token_count=output_token_count,
+                        )
+                    )
+                return response
+
+            setattr(current_obj, method_name, wrapped_method)