"""Payload builders for Opik traces and spans."""
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from litellm import _logging
from litellm.integrations.opik import utils
from . import types
def build_trace_payload(
project_name: str,
trace_id: str,
response_obj: Dict[str, Any],
start_time: datetime,
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: Dict[str, Any],
tags: List[str],
thread_id: Optional[str],
) -> types.TracePayload:
"""Build a complete trace payload."""
trace_name = response_obj.get("object", "unknown type")
return types.TracePayload(
project_name=project_name,
id=trace_id,
name=trace_name,
start_time=(
start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
),
end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
input=input_data,
output=output_data,
metadata=metadata,
tags=tags,
thread_id=thread_id,
)
def build_span_payload(
project_name: str,
trace_id: str,
parent_span_id: Optional[str],
response_obj: Dict[str, Any],
start_time: datetime,
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: Dict[str, Any],
tags: List[str],
usage: Dict[str, int],
provider: Optional[str] = None,
cost: Optional[float] = None,
) -> types.SpanPayload:
"""Build a complete span payload."""
span_id = utils.create_uuid7()
model = response_obj.get("model", "unknown-model")
obj_type = response_obj.get("object", "unknown-object")
created = response_obj.get("created", 0)
span_name = f"{model}_{obj_type}_{created}"
_logging.verbose_logger.debug(
f"OpikLogger creating span with id {span_id} for trace {trace_id}"
)
return types.SpanPayload(
id=span_id,
project_name=project_name,
trace_id=trace_id,
parent_span_id=parent_span_id,
name=span_name,
type="llm",
model=model,
start_time=(
start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
),
end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
input=input_data,
output=output_data,
metadata=metadata,
tags=tags,
usage=usage,
provider=provider,
total_cost=cost,
)