wehub-resource-sync
701 lines
25 KiB
Python
701 lines
25 KiB
Python
from contextlib import contextmanager
|
|
from typing import Any, Dict, Iterator, List, Optional
|
|
from contextvars import ContextVar
|
|
|
|
from deepeval.tracing.types import (
|
|
AgentSpan,
|
|
BaseSpan,
|
|
LlmSpan,
|
|
RetrieverSpan,
|
|
ToolSpan,
|
|
Trace,
|
|
)
|
|
from deepeval.test_case.llm_test_case import ToolCall, LLMTestCase
|
|
from deepeval.prompt.prompt import Prompt
|
|
from deepeval.metrics import BaseMetric
|
|
|
|
|
|
class SpanContext:
|
|
def __init__(self):
|
|
self.current_span: ContextVar[Optional[BaseSpan]] = ContextVar(
|
|
"current_span", default=None
|
|
)
|
|
|
|
def get(self):
|
|
return self.current_span.get()
|
|
|
|
def set(self, value):
|
|
return self.current_span.set(value)
|
|
|
|
def reset(self, value):
|
|
return self.current_span.reset(value)
|
|
|
|
def drop(self):
|
|
span = self.current_span.get()
|
|
if span:
|
|
span.drop = True
|
|
|
|
|
|
class TraceContext:
|
|
def __init__(self):
|
|
self.current_trace: ContextVar[Optional[Trace]] = ContextVar(
|
|
"current_trace", default=None
|
|
)
|
|
|
|
def get(self):
|
|
return self.current_trace.get()
|
|
|
|
def set(self, value):
|
|
return self.current_trace.set(value)
|
|
|
|
def reset(self, value):
|
|
return self.current_trace.reset(value)
|
|
|
|
def drop(self):
|
|
trace = self.current_trace.get()
|
|
if trace:
|
|
trace.drop = True
|
|
|
|
|
|
current_span_context = SpanContext()
|
|
current_trace_context = TraceContext()
|
|
|
|
|
|
def update_current_span(
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
integration: Optional[str] = None,
|
|
provider: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
):
|
|
current_span = current_span_context.get()
|
|
if not current_span:
|
|
return
|
|
if test_case:
|
|
|
|
current_span.input = test_case.input
|
|
current_span.output = test_case.actual_output
|
|
current_span.expected_output = test_case.expected_output
|
|
current_span.retrieval_context = test_case.retrieval_context
|
|
current_span.context = test_case.context
|
|
current_span.tools_called = test_case.tools_called
|
|
current_span.expected_tools = test_case.expected_tools
|
|
if metadata:
|
|
current_span.metadata = metadata
|
|
if input:
|
|
current_span.input = input
|
|
if output:
|
|
current_span.output = output
|
|
if retrieval_context:
|
|
current_span.retrieval_context = retrieval_context
|
|
if context:
|
|
current_span.context = context
|
|
if expected_output:
|
|
current_span.expected_output = expected_output
|
|
if tools_called:
|
|
current_span.tools_called = tools_called
|
|
if expected_tools:
|
|
current_span.expected_tools = expected_tools
|
|
if name:
|
|
current_span.name = name
|
|
if integration is not None:
|
|
current_span.integration = integration
|
|
if provider is not None and hasattr(current_span, "provider"):
|
|
current_span.provider = provider
|
|
if metric_collection:
|
|
current_span.metric_collection = metric_collection
|
|
if metrics:
|
|
current_span.metrics = metrics
|
|
|
|
|
|
def update_current_trace(
|
|
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,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
confident_api_key: Optional[str] = None,
|
|
test_case_id: Optional[str] = None,
|
|
test_run_id: Optional[str] = None,
|
|
turn_id: Optional[str] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
):
|
|
current_trace = current_trace_context.get()
|
|
if not current_trace:
|
|
return
|
|
if test_case:
|
|
current_trace.input = test_case.input
|
|
current_trace.output = test_case.actual_output
|
|
current_trace.expected_output = test_case.expected_output
|
|
current_trace.retrieval_context = test_case.retrieval_context
|
|
current_trace.context = test_case.context
|
|
current_trace.tools_called = test_case.tools_called
|
|
current_trace.expected_tools = test_case.expected_tools
|
|
if name:
|
|
current_trace.name = name
|
|
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
|
|
if retrieval_context:
|
|
current_trace.retrieval_context = retrieval_context
|
|
if context:
|
|
current_trace.context = context
|
|
if expected_output:
|
|
current_trace.expected_output = expected_output
|
|
if tools_called:
|
|
current_trace.tools_called = tools_called
|
|
if expected_tools:
|
|
current_trace.expected_tools = expected_tools
|
|
if confident_api_key:
|
|
current_trace.confident_api_key = confident_api_key
|
|
if test_case_id:
|
|
current_trace.test_case_id = test_case_id
|
|
if test_run_id:
|
|
current_trace.test_run_id = test_run_id
|
|
if turn_id:
|
|
current_trace.turn_id = turn_id
|
|
if metric_collection:
|
|
current_trace.metric_collection = metric_collection
|
|
if metrics:
|
|
current_trace.metrics = metrics
|
|
|
|
|
|
def update_llm_span(
|
|
model: Optional[str] = None,
|
|
input_token_count: Optional[float] = None,
|
|
output_token_count: Optional[float] = None,
|
|
cost_per_input_token: Optional[float] = None,
|
|
cost_per_output_token: Optional[float] = None,
|
|
token_intervals: Optional[Dict[float, str]] = None,
|
|
prompt: Optional[Prompt] = None,
|
|
):
|
|
current_span = current_span_context.get()
|
|
if not current_span or not isinstance(current_span, LlmSpan):
|
|
return
|
|
if model:
|
|
current_span.model = model
|
|
if input_token_count:
|
|
current_span.input_token_count = input_token_count
|
|
if output_token_count:
|
|
current_span.output_token_count = output_token_count
|
|
if cost_per_input_token:
|
|
current_span.cost_per_input_token = cost_per_input_token
|
|
if cost_per_output_token:
|
|
current_span.cost_per_output_token = cost_per_output_token
|
|
if token_intervals:
|
|
current_span.token_intervals = token_intervals
|
|
if prompt:
|
|
current_span.prompt = prompt
|
|
# Updating on span as well
|
|
current_span.prompt_alias = prompt.alias
|
|
current_span.prompt_commit_hash = prompt.hash
|
|
current_span.prompt_label = prompt.label
|
|
current_span.prompt_version = prompt.version
|
|
|
|
|
|
def update_agent_span(
|
|
available_tools: Optional[List[str]] = None,
|
|
agent_handoffs: Optional[List[str]] = None,
|
|
):
|
|
"""Mutate the active ``AgentSpan`` with agent-specific fields.
|
|
|
|
Type-specific counterpart to ``update_current_span(...)``: only
|
|
handles fields unique to ``AgentSpan``. Generic fields (name,
|
|
metadata, metric_collection, input/output, ...) still go through
|
|
``update_current_span(...)``. No-op if the current span isn't an
|
|
``AgentSpan``.
|
|
"""
|
|
current_span = current_span_context.get()
|
|
if not current_span or not isinstance(current_span, AgentSpan):
|
|
return
|
|
if available_tools is not None:
|
|
current_span.available_tools = available_tools
|
|
if agent_handoffs is not None:
|
|
current_span.agent_handoffs = agent_handoffs
|
|
|
|
|
|
def update_tool_span(
|
|
description: Optional[str] = None,
|
|
):
|
|
"""Mutate the active ``ToolSpan`` with tool-specific fields.
|
|
|
|
Type-specific counterpart to ``update_current_span(...)``: only
|
|
handles fields unique to ``ToolSpan``. ``ToolSpan.name`` is set at
|
|
span creation; use ``update_current_span(name=...)`` to rename
|
|
after the fact. No-op if the current span isn't a ``ToolSpan``.
|
|
"""
|
|
current_span = current_span_context.get()
|
|
if not current_span or not isinstance(current_span, ToolSpan):
|
|
return
|
|
if description is not None:
|
|
current_span.description = description
|
|
|
|
|
|
def update_retriever_span(
|
|
embedder: Optional[str] = None,
|
|
top_k: Optional[int] = None,
|
|
chunk_size: Optional[int] = None,
|
|
):
|
|
current_span = current_span_context.get()
|
|
if not current_span or not isinstance(current_span, RetrieverSpan):
|
|
return
|
|
if embedder:
|
|
current_span.embedder = embedder
|
|
if top_k:
|
|
current_span.top_k = top_k
|
|
if chunk_size:
|
|
current_span.chunk_size = chunk_size
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# next_*_span: declarative defaults for the NEXT span of a given type.
|
|
#
|
|
# Counterpart to ``update_current_*_span(...)`` for spans without a
|
|
# user-code seam — i.e. spans the user never executes code inside, so
|
|
# ``update_current_*_span`` from "their" body isn't reachable. The
|
|
# canonical case is an integration-emitted agent / LLM span where the
|
|
# only callsite the user owns is the one wrapping the framework call.
|
|
#
|
|
# Semantics:
|
|
# - One-shot: the dict is consumed by the FIRST span of the matching
|
|
# type that the consumer (typically an integration's OTel processor)
|
|
# creates inside the active scope. Subsequent spans see an empty slot.
|
|
# - Per-type isolation: each ``next_*_span`` writes to its own
|
|
# ``ContextVar``, so stacking ``with next_agent_span(...),
|
|
# next_llm_span(...):`` is safe and unambiguous.
|
|
# - One-stop kwargs: each helper accepts BASE fields (everything
|
|
# ``update_current_span`` takes) AND its type-specific fields in a
|
|
# single call. Diverges intentionally from the
|
|
# ``update_*_span`` family (which is decomposed) — see commit msg.
|
|
# - Consumer responsibility: integrations call ``_pop_pending_*(...)``
|
|
# when classifying a fresh span and apply the dict to the placeholder
|
|
# they push onto ``current_span_context``. If no integration is
|
|
# listening the dict is silently discarded on ``with`` exit.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _PendingSlot:
|
|
"""Mutable wrapper around a pending-defaults dict.
|
|
|
|
Why a wrapper instead of putting the dict directly into the
|
|
``ContextVar``: APIs like ``Agent.run_sync(...)`` call
|
|
``asyncio.run(...)`` internally, which creates a NEW asyncio context
|
|
that inherits a *snapshot* of the parent's contextvars. A
|
|
``ContextVar.set(...)`` inside that snapshot does not propagate back
|
|
to the outer ``with`` block — so a naive design that does
|
|
``slot.set(None)`` from inside the consumer would let a second
|
|
``agent.run_sync(...)`` in the same ``with`` re-consume the
|
|
still-populated value.
|
|
|
|
Mutating ``self.payload`` instead works because ContextVar
|
|
inheritance copies the REFERENCE to this wrapper. Both the outer
|
|
``with`` block and the inner asyncio sub-context see the same
|
|
``_PendingSlot`` instance, so ``slot.payload = None`` is visible
|
|
everywhere.
|
|
"""
|
|
|
|
__slots__ = ("payload",)
|
|
|
|
def __init__(self, payload: Optional[Dict[str, Any]]):
|
|
self.payload: Optional[Dict[str, Any]] = payload
|
|
|
|
|
|
_pending_next_span: ContextVar[Optional[_PendingSlot]] = ContextVar(
|
|
"pending_next_span", default=None
|
|
)
|
|
_pending_next_agent_span: ContextVar[Optional[_PendingSlot]] = ContextVar(
|
|
"pending_next_agent_span", default=None
|
|
)
|
|
_pending_next_llm_span: ContextVar[Optional[_PendingSlot]] = ContextVar(
|
|
"pending_next_llm_span", default=None
|
|
)
|
|
_pending_next_tool_span: ContextVar[Optional[_PendingSlot]] = ContextVar(
|
|
"pending_next_tool_span", default=None
|
|
)
|
|
_pending_next_retriever_span: ContextVar[Optional[_PendingSlot]] = ContextVar(
|
|
"pending_next_retriever_span", default=None
|
|
)
|
|
|
|
|
|
def _drop_none(d: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Strip keys whose value is None — keeps the pending dict tight so
|
|
consumers don't have to re-check every kwarg they passed through."""
|
|
return {k: v for k, v in d.items() if v is not None}
|
|
|
|
|
|
# --- base: applies to the next span of ANY type ----------------------------
|
|
|
|
|
|
@contextmanager
|
|
def next_span(
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
) -> Iterator[None]:
|
|
"""Set base-span defaults for the next span of any type.
|
|
|
|
Mirrors ``update_current_span(...)`` kwargs. Use when the type of
|
|
the upcoming span doesn't matter or isn't known. For a typed match,
|
|
use ``next_agent_span`` / ``next_llm_span`` / ``next_tool_span`` /
|
|
``next_retriever_span``.
|
|
"""
|
|
payload = _drop_none(
|
|
{
|
|
"input": input,
|
|
"output": output,
|
|
"retrieval_context": retrieval_context,
|
|
"context": context,
|
|
"expected_output": expected_output,
|
|
"tools_called": tools_called,
|
|
"expected_tools": expected_tools,
|
|
"metadata": metadata,
|
|
"name": name,
|
|
"test_case": test_case,
|
|
"metric_collection": metric_collection,
|
|
"metrics": metrics,
|
|
}
|
|
)
|
|
token = _pending_next_span.set(_PendingSlot(payload))
|
|
try:
|
|
yield
|
|
finally:
|
|
_pending_next_span.reset(token)
|
|
|
|
|
|
# --- agent: base + agent-specific (one-stop) -------------------------------
|
|
|
|
|
|
@contextmanager
|
|
def next_agent_span(
|
|
available_tools: Optional[List[str]] = None,
|
|
agent_handoffs: Optional[List[str]] = None,
|
|
# base fields (mirror update_current_span)
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
) -> Iterator[None]:
|
|
"""Set defaults for the next ``AgentSpan``. One-stop: accepts
|
|
agent-specific fields (``available_tools``, ``agent_handoffs``) AND
|
|
the same base fields ``update_current_span(...)`` takes."""
|
|
payload = _drop_none(
|
|
{
|
|
"available_tools": available_tools,
|
|
"agent_handoffs": agent_handoffs,
|
|
"input": input,
|
|
"output": output,
|
|
"retrieval_context": retrieval_context,
|
|
"context": context,
|
|
"expected_output": expected_output,
|
|
"tools_called": tools_called,
|
|
"expected_tools": expected_tools,
|
|
"metadata": metadata,
|
|
"name": name,
|
|
"test_case": test_case,
|
|
"metric_collection": metric_collection,
|
|
"metrics": metrics,
|
|
}
|
|
)
|
|
token = _pending_next_agent_span.set(_PendingSlot(payload))
|
|
try:
|
|
yield
|
|
finally:
|
|
_pending_next_agent_span.reset(token)
|
|
|
|
|
|
# --- llm: base + llm-specific (one-stop) -----------------------------------
|
|
|
|
|
|
@contextmanager
|
|
def next_llm_span(
|
|
model: Optional[str] = None,
|
|
input_token_count: Optional[float] = None,
|
|
output_token_count: Optional[float] = None,
|
|
cost_per_input_token: Optional[float] = None,
|
|
cost_per_output_token: Optional[float] = None,
|
|
token_intervals: Optional[Dict[float, str]] = None,
|
|
prompt: Optional[Prompt] = None,
|
|
# base fields
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
) -> Iterator[None]:
|
|
"""Set defaults for the next ``LlmSpan``. One-stop: accepts
|
|
LLM-specific fields (``model``, token counts, ``prompt``, ...) AND
|
|
the same base fields ``update_current_span(...)`` takes."""
|
|
payload = _drop_none(
|
|
{
|
|
"model": model,
|
|
"input_token_count": input_token_count,
|
|
"output_token_count": output_token_count,
|
|
"cost_per_input_token": cost_per_input_token,
|
|
"cost_per_output_token": cost_per_output_token,
|
|
"token_intervals": token_intervals,
|
|
"prompt": prompt,
|
|
"input": input,
|
|
"output": output,
|
|
"retrieval_context": retrieval_context,
|
|
"context": context,
|
|
"expected_output": expected_output,
|
|
"tools_called": tools_called,
|
|
"expected_tools": expected_tools,
|
|
"metadata": metadata,
|
|
"name": name,
|
|
"test_case": test_case,
|
|
"metric_collection": metric_collection,
|
|
"metrics": metrics,
|
|
}
|
|
)
|
|
token = _pending_next_llm_span.set(_PendingSlot(payload))
|
|
try:
|
|
yield
|
|
finally:
|
|
_pending_next_llm_span.reset(token)
|
|
|
|
|
|
# --- tool: base + tool-specific (one-stop) ---------------------------------
|
|
|
|
|
|
@contextmanager
|
|
def next_tool_span(
|
|
description: Optional[str] = None,
|
|
# base fields
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
) -> Iterator[None]:
|
|
"""Set defaults for the next ``ToolSpan``. One-stop: accepts
|
|
tool-specific fields (``description``) AND the same base fields
|
|
``update_current_span(...)`` takes."""
|
|
payload = _drop_none(
|
|
{
|
|
"description": description,
|
|
"input": input,
|
|
"output": output,
|
|
"retrieval_context": retrieval_context,
|
|
"context": context,
|
|
"expected_output": expected_output,
|
|
"tools_called": tools_called,
|
|
"expected_tools": expected_tools,
|
|
"metadata": metadata,
|
|
"name": name,
|
|
"test_case": test_case,
|
|
"metric_collection": metric_collection,
|
|
"metrics": metrics,
|
|
}
|
|
)
|
|
token = _pending_next_tool_span.set(_PendingSlot(payload))
|
|
try:
|
|
yield
|
|
finally:
|
|
_pending_next_tool_span.reset(token)
|
|
|
|
|
|
# --- retriever: base + retriever-specific (one-stop) -----------------------
|
|
|
|
|
|
@contextmanager
|
|
def next_retriever_span(
|
|
embedder: Optional[str] = None,
|
|
top_k: Optional[int] = None,
|
|
chunk_size: Optional[int] = None,
|
|
# base fields
|
|
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,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
name: Optional[str] = None,
|
|
test_case: Optional[LLMTestCase] = None,
|
|
metric_collection: Optional[str] = None,
|
|
metrics: Optional[List[BaseMetric]] = None,
|
|
) -> Iterator[None]:
|
|
"""Set defaults for the next ``RetrieverSpan``. One-stop: accepts
|
|
retriever-specific fields (``embedder``, ``top_k``, ``chunk_size``)
|
|
AND the same base fields ``update_current_span(...)`` takes."""
|
|
payload = _drop_none(
|
|
{
|
|
"embedder": embedder,
|
|
"top_k": top_k,
|
|
"chunk_size": chunk_size,
|
|
"input": input,
|
|
"output": output,
|
|
"retrieval_context": retrieval_context,
|
|
"context": context,
|
|
"expected_output": expected_output,
|
|
"tools_called": tools_called,
|
|
"expected_tools": expected_tools,
|
|
"metadata": metadata,
|
|
"name": name,
|
|
"test_case": test_case,
|
|
"metric_collection": metric_collection,
|
|
"metrics": metrics,
|
|
}
|
|
)
|
|
token = _pending_next_retriever_span.set(_PendingSlot(payload))
|
|
try:
|
|
yield
|
|
finally:
|
|
_pending_next_retriever_span.reset(token)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Consumer-facing pop helpers.
|
|
#
|
|
# Integrations (e.g. ``deepeval.integrations.pydantic_ai.SpanInterceptor``)
|
|
# call these the moment they classify a fresh span and BEFORE they push the
|
|
# placeholder onto ``current_span_context``. The pop is one-shot: the slot
|
|
# is reset to None for the rest of the active ``with`` scope.
|
|
#
|
|
# ``pop_pending_for(span_type)`` returns the merged dict of base + typed
|
|
# defaults — base values are overwritten by the typed slot's values when
|
|
# both are present, matching "more specific wins".
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_TYPED_SLOTS = {
|
|
"agent": _pending_next_agent_span,
|
|
"llm": _pending_next_llm_span,
|
|
"tool": _pending_next_tool_span,
|
|
"retriever": _pending_next_retriever_span,
|
|
}
|
|
|
|
|
|
def pop_pending_for(span_type: Optional[str]) -> Dict[str, Any]:
|
|
"""One-shot consume the pending-defaults dict for ``span_type``.
|
|
|
|
Returns a merged dict {**base_slot, **typed_slot}. Typed values win
|
|
on overlap. Drained slots have ``payload`` mutated to ``None`` —
|
|
NOT reassigned via ``ContextVar.set(...)``, because consumers often
|
|
run inside a sub-context (e.g. ``asyncio.run`` started by
|
|
``Agent.run_sync``) where a ``set`` would not propagate back.
|
|
Mutating ``_PendingSlot.payload`` is visible in BOTH the consumer's
|
|
sub-context and the outer ``with`` block, since both inherit the
|
|
same wrapper reference.
|
|
|
|
``span_type`` may be one of ``"agent" | "llm" | "tool" |
|
|
"retriever"`` or ``None`` to consume only the base slot.
|
|
"""
|
|
merged: Dict[str, Any] = {}
|
|
|
|
base_slot = _pending_next_span.get()
|
|
if base_slot is not None and base_slot.payload:
|
|
merged.update(base_slot.payload)
|
|
base_slot.payload = None
|
|
|
|
if span_type and span_type in _TYPED_SLOTS:
|
|
typed_slot = _TYPED_SLOTS[span_type].get()
|
|
if typed_slot is not None and typed_slot.payload:
|
|
merged.update(typed_slot.payload)
|
|
typed_slot.payload = None
|
|
|
|
return merged
|
|
|
|
|
|
def apply_pending_to_span(span: BaseSpan, payload: Dict[str, Any]) -> None:
|
|
"""Apply a popped pending-defaults dict to ``span`` in-place.
|
|
|
|
Mirrors ``update_current_span(...)`` semantics for the BASE keys —
|
|
notably the ``test_case`` unpacking path, which writes the
|
|
LLMTestCase's fields onto the span and overrides any individual
|
|
field set in the same payload. Typed kwargs (``available_tools``,
|
|
``model``, ``embedder``, ``description``, etc.) are setattr'd
|
|
directly when the span is the matching subclass; mismatches are
|
|
silently dropped (e.g. ``model`` on a ``ToolSpan``).
|
|
|
|
Used by integrations after pushing a fresh placeholder onto
|
|
``current_span_context`` so that ``next_*_span(...)`` defaults land
|
|
on the placeholder before user code or downstream serialization sees
|
|
it.
|
|
"""
|
|
if not payload:
|
|
return
|
|
|
|
test_case = payload.get("test_case")
|
|
if test_case is not None:
|
|
span.input = test_case.input
|
|
span.output = test_case.actual_output
|
|
span.expected_output = test_case.expected_output
|
|
span.retrieval_context = test_case.retrieval_context
|
|
span.context = test_case.context
|
|
span.tools_called = test_case.tools_called
|
|
span.expected_tools = test_case.expected_tools
|
|
|
|
for key, value in payload.items():
|
|
if key == "test_case" or value is None:
|
|
continue
|
|
# Only setattr keys the span actually declares — guards against
|
|
# cross-type leakage (e.g. ``embedder`` landing on an LlmSpan).
|
|
if not hasattr(span, key):
|
|
continue
|
|
try:
|
|
setattr(span, key, value)
|
|
except Exception:
|
|
# Pydantic validation errors / locked fields → skip silently.
|
|
continue
|