wehub-resource-sync
115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
from contextvars import ContextVar
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
from deepeval.metrics import BaseMetric
|
|
from deepeval.prompt import Prompt
|
|
from deepeval.test_case.llm_test_case import ToolCall
|
|
from deepeval.tracing.context import current_trace_context, update_current_trace
|
|
from deepeval.tracing.tracing import trace_manager
|
|
from deepeval.tracing.types import TraceWorkerStatus
|
|
from deepeval.tracing.utils import is_async_context
|
|
|
|
|
|
@dataclass
|
|
class LlmSpanContext:
|
|
prompt: Optional[Prompt] = None
|
|
metrics: Optional[List[BaseMetric]] = None
|
|
metric_collection: Optional[str] = None
|
|
expected_output: Optional[str] = None
|
|
expected_tools: Optional[List[ToolCall]] = None
|
|
context: Optional[List[str]] = None
|
|
retrieval_context: Optional[List[str]] = None
|
|
|
|
|
|
@dataclass
|
|
class AgentSpanContext:
|
|
metrics: Optional[List[BaseMetric]] = None
|
|
metric_collection: Optional[str] = None
|
|
expected_output: Optional[str] = None
|
|
expected_tools: Optional[List[ToolCall]] = None
|
|
context: Optional[List[str]] = None
|
|
retrieval_context: Optional[List[str]] = None
|
|
|
|
|
|
current_llm_context: ContextVar[Optional[LlmSpanContext]] = ContextVar(
|
|
"current_llm_context", default=LlmSpanContext()
|
|
)
|
|
|
|
current_agent_context: ContextVar[Optional[AgentSpanContext]] = ContextVar(
|
|
"current_agent_context", default=AgentSpanContext()
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def trace(
|
|
llm_span_context: Optional[LlmSpanContext] = None,
|
|
agent_span_context: Optional[AgentSpanContext] = None,
|
|
name: Optional[str] = None,
|
|
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,
|
|
retrieval_context: Optional[List[str]] = None,
|
|
context: Optional[List[str]] = None,
|
|
expected_output: Optional[str] = None,
|
|
tools_called: Optional[List[ToolCall]] = None,
|
|
expected_tools: Optional[List[ToolCall]] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
metric_collection: Optional[str] = None,
|
|
):
|
|
if is_async_context():
|
|
trace_manager._print_trace_status(
|
|
message="Warning: Detected use of the synchronous 'trace' context manager within an async method",
|
|
trace_worker_status=TraceWorkerStatus.WARNING,
|
|
description="Wrapping an async method with the synchronous 'trace' context manager may lead to unexpected behavior.",
|
|
)
|
|
|
|
current_trace = current_trace_context.get()
|
|
started_new_trace = False
|
|
|
|
if not current_trace:
|
|
current_trace = trace_manager.start_new_trace()
|
|
started_new_trace = True
|
|
|
|
if metrics:
|
|
current_trace.metrics = metrics
|
|
|
|
if metric_collection:
|
|
current_trace.metric_collection = metric_collection
|
|
|
|
trace_ctx_token = current_trace_context.set(current_trace)
|
|
|
|
update_current_trace(
|
|
name=name,
|
|
tags=tags,
|
|
metadata=metadata,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
input=input,
|
|
output=output,
|
|
retrieval_context=retrieval_context,
|
|
context=context,
|
|
expected_output=expected_output,
|
|
tools_called=tools_called,
|
|
expected_tools=expected_tools,
|
|
)
|
|
|
|
if llm_span_context:
|
|
current_llm_context.set(llm_span_context)
|
|
if agent_span_context:
|
|
current_agent_context.set(agent_span_context)
|
|
try:
|
|
yield current_trace
|
|
finally:
|
|
if started_new_trace:
|
|
trace_manager.end_trace(current_trace.uuid)
|
|
|
|
current_trace_context.reset(trace_ctx_token)
|
|
|
|
current_llm_context.set(LlmSpanContext())
|
|
current_agent_context.set(AgentSpanContext())
|