diff --git a/litellm/tracing/store.py b/litellm/tracing/store.py index b3cefa8eac0..4cca5336099 100644 --- a/litellm/tracing/store.py +++ b/litellm/tracing/store.py @@ -117,6 +117,12 @@ ORDER BY o.Timestamp, abs(toUnixTimestamp64Milli(s.start_time) - toUnixTimestamp LIMIT 1 BY o.SpanId """ +TRACE_IO_SQL: Final = f""" +SELECT SpanId AS span_id, Input AS input, Output AS output +FROM {OTEL_TRACES_TABLE} +WHERE TraceId = {{trace_id:String}} AND {_SCOPE_OTEL} +""" + SPAN_DETAIL_SQL: Final = f""" SELECT SpanId AS span_id, Input AS input, Output AS output, SpanAttributes AS attributes FROM {OTEL_TRACES_TABLE} @@ -341,3 +347,8 @@ class ClickHouseTraceStore: output=rows[0]["output"], attributes=dict(rows[0]["attributes"]), ) + + async def get_span_io(self, trace_id: str, scope: TraceScope) -> dict[str, tuple[str, str]]: + """span_id -> (input, output) for every span in the trace, in one query (for exports).""" + rows = await self.client.query(TRACE_IO_SQL, {**scope, "trace_id": trace_id}) + return {r["span_id"]: (r["input"], r["output"]) for r in rows}