generated from kgod/ai-review-template
提交
This commit is contained in:
kgod
@@ -0,0 +1,8 @@
|
||||
"""LangSmith integration for OpenAI Agents SDK.
|
||||
|
||||
This module provides tracing support for the OpenAI Agents SDK.
|
||||
"""
|
||||
|
||||
from ._openai_agents import OpenAIAgentsTracingProcessor
|
||||
|
||||
__all__ = ["OpenAIAgentsTracingProcessor"]
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
+228
@@ -0,0 +1,228 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Literal
|
||||
|
||||
try:
|
||||
from agents import tracing # type: ignore[import]
|
||||
|
||||
HAVE_AGENTS = True
|
||||
except ImportError:
|
||||
HAVE_AGENTS = False
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RunTypeT = Literal["tool", "chain", "llm", "retriever", "embedding", "prompt", "parser"]
|
||||
|
||||
if HAVE_AGENTS:
|
||||
|
||||
def parse_io(data: Any, default_key: str = "output") -> dict:
|
||||
"""Parse inputs or outputs into a dictionary format.
|
||||
|
||||
Args:
|
||||
data: The data to parse (can be inputs or outputs)
|
||||
default_key: The default key to use if data is not a dict
|
||||
(`'input'` or `'output'`)
|
||||
|
||||
Returns:
|
||||
Dict: The parsed data as a dictionary
|
||||
"""
|
||||
if isinstance(data, list):
|
||||
if len(data) == 0:
|
||||
return {}
|
||||
# Check if this is a list of output blocks (reasoning, message, etc.)
|
||||
if len(data) > 0 and isinstance(data[0], dict):
|
||||
if "type" in data[0]:
|
||||
return {default_key: data}
|
||||
elif len(data) == 1:
|
||||
return data[0]
|
||||
return {default_key: data}
|
||||
elif isinstance(data, dict):
|
||||
data_ = data
|
||||
elif isinstance(data, str):
|
||||
try:
|
||||
parsed_json = json.loads(data)
|
||||
if isinstance(parsed_json, dict):
|
||||
data_ = parsed_json
|
||||
else:
|
||||
data_ = {default_key: data}
|
||||
except json.JSONDecodeError:
|
||||
data_ = {default_key: data}
|
||||
elif (
|
||||
data is not None
|
||||
and hasattr(data, "model_dump")
|
||||
and callable(data.model_dump)
|
||||
and not isinstance(data, type)
|
||||
):
|
||||
try:
|
||||
data_ = data.model_dump(exclude_none=True, mode="json")
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
f"Failed to use model_dump to serialize {type(data)} to JSON: {e}"
|
||||
)
|
||||
data_ = {default_key: data}
|
||||
else:
|
||||
data_ = {default_key: data}
|
||||
|
||||
return data_
|
||||
|
||||
def get_run_type(span: tracing.Span) -> RunTypeT:
|
||||
span_type = getattr(span.span_data, "type", None)
|
||||
if span_type in ["agent", "handoff", "custom"]:
|
||||
return "chain"
|
||||
elif span_type in ["function", "guardrail"]:
|
||||
return "tool"
|
||||
elif span_type in ["generation", "response"]:
|
||||
return "llm"
|
||||
else:
|
||||
return "chain"
|
||||
|
||||
def get_run_name(span: tracing.Span) -> str:
|
||||
if hasattr(span.span_data, "name") and span.span_data.name:
|
||||
return span.span_data.name
|
||||
span_type = getattr(span.span_data, "type", None)
|
||||
if span_type == "generation":
|
||||
return "Generation"
|
||||
elif span_type == "response":
|
||||
return "Response"
|
||||
elif span_type == "handoff":
|
||||
return "Handoff"
|
||||
else:
|
||||
return "Span"
|
||||
|
||||
def _extract_function_span_data(
|
||||
span_data: tracing.FunctionSpanData,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"inputs": parse_io(span_data.input, "input"),
|
||||
"outputs": parse_io(span_data.output, "output"),
|
||||
}
|
||||
|
||||
def _extract_generation_span_data(
|
||||
span_data: tracing.GenerationSpanData,
|
||||
) -> dict[str, Any]:
|
||||
data = {
|
||||
"inputs": parse_io(span_data.input, "input"),
|
||||
"outputs": parse_io(span_data.output, "output"),
|
||||
"invocation_params": {
|
||||
"model": span_data.model,
|
||||
"model_config": span_data.model_config,
|
||||
},
|
||||
}
|
||||
if span_data.usage:
|
||||
from langsmith.wrappers._openai import _create_usage_metadata
|
||||
|
||||
if "metadata" not in data:
|
||||
data["metadata"] = {}
|
||||
data["metadata"]["usage_metadata"] = _create_usage_metadata(span_data.usage)
|
||||
return data
|
||||
|
||||
def _extract_response_span_data(
|
||||
span_data: tracing.ResponseSpanData,
|
||||
) -> dict[str, Any]:
|
||||
data: dict[str, Any] = {}
|
||||
if span_data.input is not None:
|
||||
data["inputs"] = {
|
||||
"input": span_data.input,
|
||||
"instructions": (
|
||||
span_data.response.instructions
|
||||
if span_data.response is not None
|
||||
and span_data.response.instructions
|
||||
else ""
|
||||
),
|
||||
}
|
||||
if span_data.response is not None:
|
||||
response = span_data.response.model_dump(exclude_none=True, mode="json")
|
||||
output_data = response.pop("output", [])
|
||||
data["outputs"] = parse_io(output_data, "output")
|
||||
data["invocation_params"] = {
|
||||
k: v
|
||||
for k, v in response.items()
|
||||
if k
|
||||
in (
|
||||
"max_output_tokens",
|
||||
"model",
|
||||
"parallel_tool_calls",
|
||||
"reasoning",
|
||||
"temperature",
|
||||
"text",
|
||||
"tool_choice",
|
||||
"tools",
|
||||
"top_p",
|
||||
"truncation",
|
||||
)
|
||||
}
|
||||
metadata = {
|
||||
k: v
|
||||
for k, v in response.items()
|
||||
if k
|
||||
not in (
|
||||
{"output", "usage", "instructions"}.union(data["invocation_params"])
|
||||
)
|
||||
}
|
||||
metadata.update(
|
||||
{
|
||||
"ls_model_name": data["invocation_params"].get("model"),
|
||||
"ls_max_tokens": data["invocation_params"].get("max_output_tokens"),
|
||||
"ls_temperature": data["invocation_params"].get("temperature"),
|
||||
"ls_model_type": "chat",
|
||||
"ls_provider": "openai",
|
||||
}
|
||||
)
|
||||
if usage := response.pop("usage", None):
|
||||
from langsmith.wrappers._openai import _create_usage_metadata
|
||||
|
||||
metadata["usage_metadata"] = _create_usage_metadata(usage)
|
||||
data["metadata"] = metadata
|
||||
|
||||
return data
|
||||
|
||||
def _extract_agent_span_data(span_data: tracing.AgentSpanData) -> dict[str, Any]:
|
||||
return {
|
||||
"invocation_params": {
|
||||
"tools": span_data.tools,
|
||||
"handoffs": span_data.handoffs,
|
||||
},
|
||||
"metadata": {
|
||||
"output_type": span_data.output_type,
|
||||
},
|
||||
}
|
||||
|
||||
def _extract_handoff_span_data(
|
||||
span_data: tracing.HandoffSpanData,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"inputs": {
|
||||
"from_agent": span_data.from_agent,
|
||||
"to_agent": span_data.to_agent,
|
||||
}
|
||||
}
|
||||
|
||||
def _extract_guardrail_span_data(
|
||||
span_data: tracing.GuardrailSpanData,
|
||||
) -> dict[str, Any]:
|
||||
return {"metadata": {"triggered": span_data.triggered}}
|
||||
|
||||
def _extract_custom_span_data(span_data: tracing.CustomSpanData) -> dict[str, Any]:
|
||||
return {"metadata": span_data.data}
|
||||
|
||||
def extract_span_data(span: tracing.Span) -> dict[str, Any]:
|
||||
data: dict[str, Any] = {}
|
||||
|
||||
if isinstance(span.span_data, tracing.FunctionSpanData):
|
||||
data.update(_extract_function_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.GenerationSpanData):
|
||||
data.update(_extract_generation_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.ResponseSpanData):
|
||||
data.update(_extract_response_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.AgentSpanData):
|
||||
data.update(_extract_agent_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.HandoffSpanData):
|
||||
data.update(_extract_handoff_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.GuardrailSpanData):
|
||||
data.update(_extract_guardrail_span_data(span.span_data))
|
||||
elif isinstance(span.span_data, tracing.CustomSpanData):
|
||||
data.update(_extract_custom_span_data(span.span_data))
|
||||
else:
|
||||
return {}
|
||||
|
||||
return data
|
||||
@@ -0,0 +1,436 @@
|
||||
import logging
|
||||
import weakref
|
||||
from datetime import datetime
|
||||
from functools import cache
|
||||
from typing import Optional
|
||||
|
||||
from langsmith import run_trees as rt
|
||||
from langsmith._internal import _context
|
||||
from langsmith.run_helpers import get_current_run_tree
|
||||
|
||||
try:
|
||||
from agents import tracing # type: ignore[import]
|
||||
|
||||
required = (
|
||||
"TracingProcessor",
|
||||
"Trace",
|
||||
"Span",
|
||||
"ResponseSpanData",
|
||||
)
|
||||
if not all(hasattr(tracing, name) for name in required):
|
||||
raise ImportError("The `agents` package is not installed.")
|
||||
|
||||
from langsmith.integrations.openai_agents_sdk import (
|
||||
_openai_agent_utils as agent_utils,
|
||||
)
|
||||
|
||||
HAVE_AGENTS = True
|
||||
except ImportError:
|
||||
HAVE_AGENTS = False
|
||||
|
||||
class OpenAIAgentsTracingProcessor:
|
||||
"""Tracing processor for the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/).
|
||||
|
||||
Traces all intermediate steps of your OpenAI Agent to LangSmith.
|
||||
|
||||
Requirements: Make sure to install `pip install -U langsmith[openai-agents]`.
|
||||
|
||||
Args:
|
||||
client: An instance of `langsmith.client.Client`. If not provided, a default
|
||||
client is created.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from agents import (
|
||||
Agent,
|
||||
FileSearchTool,
|
||||
Runner,
|
||||
WebSearchTool,
|
||||
function_tool,
|
||||
set_trace_processors,
|
||||
)
|
||||
|
||||
from langsmith.wrappers import OpenAIAgentsTracingProcessor
|
||||
|
||||
set_trace_processors([OpenAIAgentsTracingProcessor()])
|
||||
|
||||
|
||||
@function_tool
|
||||
def get_weather(city: str) -> str:
|
||||
return f"The weather in {city} is sunny"
|
||||
|
||||
|
||||
haiku_agent = Agent(
|
||||
name="Haiku agent",
|
||||
instructions="Always respond in haiku form",
|
||||
model="o3-mini",
|
||||
tools=[get_weather],
|
||||
)
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
tools=[WebSearchTool()],
|
||||
instructions="speak in spanish. use Haiku agent if they ask for a haiku or for the weather",
|
||||
handoffs=[haiku_agent],
|
||||
)
|
||||
|
||||
result = await Runner.run(
|
||||
agent,
|
||||
"write a haiku about the weather today and tell me a recent news story about new york",
|
||||
)
|
||||
print(result.final_output)
|
||||
```
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise ImportError(
|
||||
"The `agents` package is not installed. "
|
||||
"Please install it with `pip install langsmith[openai-agents]`."
|
||||
)
|
||||
|
||||
|
||||
from langsmith import client as ls_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@cache
|
||||
def _get_package_version(package_name: str) -> str | None:
|
||||
try:
|
||||
from importlib.metadata import version
|
||||
|
||||
return version(package_name)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
if HAVE_AGENTS:
|
||||
|
||||
class OpenAIAgentsTracingProcessor(tracing.TracingProcessor): # type: ignore[no-redef]
|
||||
"""Tracing processor for the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/).
|
||||
|
||||
Traces all intermediate steps of your OpenAI Agent to LangSmith.
|
||||
|
||||
Requirements: Make sure to install `pip install -U langsmith[openai-agents]`.
|
||||
|
||||
Args:
|
||||
client: An instance of `langsmith.client.Client`. If not provided,
|
||||
a default client is created.
|
||||
metadata: Metadata to associate with all traces.
|
||||
tags: Tags to associate with all traces.
|
||||
project_name: LangSmith project to trace to.
|
||||
name: Name of the root trace.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from agents import (
|
||||
Agent,
|
||||
FileSearchTool,
|
||||
Runner,
|
||||
WebSearchTool,
|
||||
function_tool,
|
||||
set_trace_processors,
|
||||
)
|
||||
|
||||
from langsmith.wrappers import OpenAIAgentsTracingProcessor
|
||||
|
||||
set_trace_processors([OpenAIAgentsTracingProcessor()])
|
||||
|
||||
|
||||
@function_tool
|
||||
def get_weather(city: str) -> str:
|
||||
return f"The weather in {city} is sunny"
|
||||
|
||||
|
||||
haiku_agent = Agent(
|
||||
name="Haiku agent",
|
||||
instructions="Always respond in haiku form",
|
||||
model="o3-mini",
|
||||
tools=[get_weather],
|
||||
)
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
tools=[WebSearchTool()],
|
||||
instructions="speak in spanish. use Haiku agent if they ask for a haiku or for the weather",
|
||||
handoffs=[haiku_agent],
|
||||
)
|
||||
|
||||
result = await Runner.run(
|
||||
agent,
|
||||
"write a haiku about the weather today and tell me a recent news story about new york",
|
||||
)
|
||||
print(result.final_output)
|
||||
```
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: Optional[ls_client.Client] = None,
|
||||
*,
|
||||
metadata: Optional[dict] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
project_name: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
self.client = client or rt.get_cached_client()
|
||||
self._metadata = metadata
|
||||
self._tags = tags
|
||||
self._project_name = project_name
|
||||
self._name = name
|
||||
self._first_response_inputs: dict = {}
|
||||
self._last_response_outputs: dict = {}
|
||||
|
||||
self._runs: dict[str, rt.RunTree] = {}
|
||||
self._span_data_types: dict[
|
||||
str, type
|
||||
] = {} # Track span data types by span_id
|
||||
self._unposted_traces: set[str] = set()
|
||||
self._unposted_spans: set[str] = set()
|
||||
|
||||
def on_trace_start(self, trace: tracing.Trace) -> None:
|
||||
current_run_tree = get_current_run_tree()
|
||||
|
||||
# Determine run name
|
||||
if self._name:
|
||||
run_name = self._name
|
||||
elif trace.name:
|
||||
run_name = trace.name
|
||||
else:
|
||||
run_name = "Agent workflow"
|
||||
|
||||
# Build metadata
|
||||
run_extra = {
|
||||
"metadata": {
|
||||
**(self._metadata or {}),
|
||||
"ls_integration": "openai-agents-sdk",
|
||||
"ls_integration_version": _get_package_version("openai-agents"),
|
||||
"ls_agent_type": "root",
|
||||
}
|
||||
}
|
||||
trace_dict = trace.export() or {}
|
||||
if trace_dict.get("group_id") is not None:
|
||||
run_extra["metadata"]["thread_id"] = trace_dict["group_id"]
|
||||
|
||||
try:
|
||||
if current_run_tree is not None:
|
||||
# Nest under existing trace
|
||||
new_run = current_run_tree.create_child(
|
||||
name=run_name,
|
||||
run_type="chain",
|
||||
inputs={},
|
||||
extra=run_extra,
|
||||
tags=self._tags,
|
||||
)
|
||||
else:
|
||||
# Create new root trace
|
||||
run_kwargs = {
|
||||
"name": run_name,
|
||||
"run_type": "chain",
|
||||
"inputs": {},
|
||||
"extra": run_extra,
|
||||
"tags": self._tags,
|
||||
"client": self.client,
|
||||
}
|
||||
if self._project_name is not None:
|
||||
run_kwargs["project_name"] = self._project_name
|
||||
new_run = rt.RunTree(**run_kwargs) # type: ignore[arg-type]
|
||||
|
||||
# Delay posting until first response/generation span ends
|
||||
# so inputs can be included in the POST.
|
||||
self._unposted_traces.add(trace.trace_id)
|
||||
if new_run is not None:
|
||||
_context._PARENT_RUN_TREE_REF.set(weakref.ref(new_run))
|
||||
self._runs[trace.trace_id] = new_run
|
||||
except Exception as e:
|
||||
logger.exception(f"Error creating trace run: {e}")
|
||||
|
||||
def on_trace_end(self, trace: tracing.Trace) -> None:
|
||||
run = self._runs.pop(trace.trace_id, None)
|
||||
if not run:
|
||||
return
|
||||
|
||||
trace_dict = trace.export() or {}
|
||||
metadata = {**(trace_dict.get("metadata") or {}), **(self._metadata or {})}
|
||||
|
||||
try:
|
||||
# Update run with final inputs/outputs
|
||||
run.outputs = self._last_response_outputs.pop(trace.trace_id, {})
|
||||
|
||||
# Update metadata
|
||||
if "metadata" not in run.extra:
|
||||
run.extra["metadata"] = {}
|
||||
run.extra["metadata"].update(metadata)
|
||||
|
||||
# End and patch
|
||||
run.end()
|
||||
|
||||
if trace.trace_id in self._unposted_traces:
|
||||
# No response/generation spans ended, post now
|
||||
run.inputs = self._first_response_inputs.pop(trace.trace_id, {})
|
||||
self._unposted_traces.discard(trace.trace_id)
|
||||
run.post()
|
||||
else:
|
||||
self._first_response_inputs.pop(trace.trace_id, None)
|
||||
run.patch(exclude_inputs=True)
|
||||
|
||||
# Restore parent context
|
||||
if run.parent_run is not None:
|
||||
_context._PARENT_RUN_TREE_REF.set(weakref.ref(run.parent_run))
|
||||
else:
|
||||
_context._PARENT_RUN_TREE_REF.set(None)
|
||||
except Exception as e:
|
||||
logger.exception(f"Error updating trace run: {e}")
|
||||
|
||||
def on_span_start(self, span: tracing.Span) -> None:
|
||||
# Find parent run
|
||||
parent_run = (
|
||||
self._runs.get(span.parent_id)
|
||||
if span.parent_id
|
||||
else self._runs.get(span.trace_id)
|
||||
)
|
||||
|
||||
if parent_run is None:
|
||||
logger.warning(
|
||||
f"No trace info found for span, skipping: {span.span_id}"
|
||||
)
|
||||
return
|
||||
|
||||
# Extract span data
|
||||
run_name = agent_utils.get_run_name(span)
|
||||
if isinstance(span.span_data, tracing.ResponseSpanData):
|
||||
parent_name = parent_run.name
|
||||
raw_span_name = getattr(span, "name", None) or getattr(
|
||||
span.span_data, "name", None
|
||||
)
|
||||
span_name = str(raw_span_name) if raw_span_name else run_name
|
||||
if parent_name:
|
||||
run_name = f"{parent_name} {span_name}".strip()
|
||||
else:
|
||||
run_name = span_name
|
||||
|
||||
run_type = agent_utils.get_run_type(span)
|
||||
extracted = agent_utils.extract_span_data(span)
|
||||
|
||||
try:
|
||||
# Create child run
|
||||
child_run = parent_run.create_child(
|
||||
name=run_name,
|
||||
run_type=run_type,
|
||||
inputs=extracted.get("inputs", {}),
|
||||
extra=extracted,
|
||||
start_time=datetime.fromisoformat(span.started_at)
|
||||
if span.started_at
|
||||
else None,
|
||||
)
|
||||
|
||||
# Add ls_agent_type metadata for agent spans that are children of
|
||||
# function spans (i.e., agents used as tools via as_tool()).
|
||||
# Note: Handoff agents are considered root agents, not subagents,
|
||||
# since they take over the conversation rather than being called
|
||||
# as tools.
|
||||
if isinstance(span.span_data, tracing.AgentSpanData):
|
||||
# Check if parent span is a function span (agent used as tool)
|
||||
parent_span_data_type = (
|
||||
self._span_data_types.get(span.parent_id)
|
||||
if span.parent_id
|
||||
else None
|
||||
)
|
||||
if parent_span_data_type is tracing.FunctionSpanData:
|
||||
if "metadata" not in child_run.extra:
|
||||
child_run.extra["metadata"] = {}
|
||||
child_run.extra["metadata"]["ls_agent_type"] = "subagent"
|
||||
|
||||
# Track span data type for parent lookups
|
||||
self._span_data_types[span.span_id] = type(span.span_data)
|
||||
|
||||
# Delay posting for spans whose inputs aren't available at start
|
||||
if isinstance(
|
||||
span.span_data,
|
||||
(
|
||||
tracing.GenerationSpanData,
|
||||
tracing.ResponseSpanData,
|
||||
tracing.FunctionSpanData,
|
||||
),
|
||||
):
|
||||
self._unposted_spans.add(span.span_id)
|
||||
else:
|
||||
child_run.post()
|
||||
self._runs[span.span_id] = child_run
|
||||
except Exception as e:
|
||||
logger.exception(f"Error creating span run: {e}")
|
||||
|
||||
def on_span_end(self, span: tracing.Span) -> None:
|
||||
run = self._runs.pop(span.span_id, None)
|
||||
self._span_data_types.pop(
|
||||
span.span_id, None
|
||||
) # Clean up span data type tracking
|
||||
if not run:
|
||||
return
|
||||
|
||||
try:
|
||||
# Extract outputs and metadata
|
||||
extracted = agent_utils.extract_span_data(span)
|
||||
outputs = extracted.pop("outputs", {})
|
||||
inputs = extracted.pop("inputs", {})
|
||||
|
||||
# Update run
|
||||
run.outputs = outputs
|
||||
if inputs:
|
||||
run.inputs = inputs
|
||||
if error := span.error:
|
||||
run.error = str(error)
|
||||
|
||||
# Add OpenAI metadata
|
||||
if "metadata" not in run.extra:
|
||||
run.extra["metadata"] = {}
|
||||
run.extra["metadata"].update(
|
||||
{
|
||||
"openai_parent_id": span.parent_id,
|
||||
"openai_trace_id": span.trace_id,
|
||||
"openai_span_id": span.span_id,
|
||||
}
|
||||
)
|
||||
if metadata := extracted.get("metadata"):
|
||||
run.extra["metadata"].update(metadata)
|
||||
if invocation_params := extracted.get("invocation_params"):
|
||||
run.extra["invocation_params"] = invocation_params
|
||||
|
||||
if isinstance(span.span_data, tracing.ResponseSpanData):
|
||||
self._first_response_inputs[span.trace_id] = (
|
||||
self._first_response_inputs.get(span.trace_id) or inputs
|
||||
)
|
||||
self._last_response_outputs[span.trace_id] = outputs
|
||||
self._maybe_post_trace(span.trace_id, inputs)
|
||||
elif isinstance(span.span_data, tracing.GenerationSpanData):
|
||||
self._first_response_inputs[span.trace_id] = (
|
||||
self._first_response_inputs.get(span.trace_id) or inputs
|
||||
)
|
||||
self._last_response_outputs[span.trace_id] = outputs
|
||||
self._maybe_post_trace(span.trace_id, inputs)
|
||||
|
||||
if span.ended_at:
|
||||
run.end_time = datetime.fromisoformat(span.ended_at)
|
||||
else:
|
||||
run.end()
|
||||
|
||||
if span.span_id in self._unposted_spans:
|
||||
self._unposted_spans.discard(span.span_id)
|
||||
run.post()
|
||||
else:
|
||||
run.patch(exclude_inputs=True)
|
||||
except Exception as e:
|
||||
logger.exception(f"Error updating span run: {e}")
|
||||
|
||||
def _maybe_post_trace(self, trace_id: str, inputs: dict) -> None:
|
||||
"""Post the trace if it hasn't been posted yet."""
|
||||
if trace_id in self._unposted_traces:
|
||||
trace_run = self._runs.get(trace_id)
|
||||
if trace_run:
|
||||
trace_run.inputs = inputs
|
||||
trace_run.post()
|
||||
self._unposted_traces.discard(trace_id)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self.client.flush()
|
||||
|
||||
def force_flush(self) -> None:
|
||||
self.client.flush()
|
||||
Reference in New Issue
Block a user