diff --git a/litellm-rust/crates/traces-clickhouse/src/query.rs b/litellm-rust/crates/traces-clickhouse/src/query.rs index d8f7041386c..e564dcd31d2 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query.rs @@ -502,7 +502,7 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result litellm_traces::schema::request_schemas(), + Some("--responses") => litellm_traces::schema::response_schemas(), + _ => litellm_traces::schema::schemas(), }; println!("{}", serde_json::to_string_pretty(&schemas).unwrap()); } diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index 7a72183de2c..8ac0d2756e0 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -17,6 +17,7 @@ pub mod query; mod query_access; pub mod request; mod resolve; +pub mod response; #[cfg(feature = "schema")] pub mod schema; mod shared; diff --git a/litellm-rust/crates/traces/src/response.rs b/litellm-rust/crates/traces/src/response.rs new file mode 100644 index 00000000000..cd8538ed602 --- /dev/null +++ b/litellm-rust/crates/traces/src/response.rs @@ -0,0 +1,10 @@ +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug)] +#[serde(deny_unknown_fields)] +pub struct TraceSQLResponse { + #[cfg_attr( + feature = "schema", + schemars(extend("x-python-normalized" = {"type": "tuple[Mapping[str, JsonValue], ...]"})) + )] + pub data: Vec>, +} diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index 56e166019d7..7b368a966bf 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -90,3 +90,10 @@ pub fn request_schemas() -> BTreeMap<&'static str, Schema> { ), ]) } + +pub fn response_schemas() -> BTreeMap<&'static str, Schema> { + BTreeMap::from([( + "TraceSQLResponse", + emitted::(), + )]) +} diff --git a/litellm-rust/crates/traces/tests/response_schema.rs b/litellm-rust/crates/traces/tests/response_schema.rs new file mode 100644 index 00000000000..4386acbb454 --- /dev/null +++ b/litellm-rust/crates/traces/tests/response_schema.rs @@ -0,0 +1,46 @@ +#![cfg(feature = "schema")] + +use litellm_traces::response::TraceSQLResponse; +use litellm_traces::schema::response_schemas; +use rstest::rstest; +use serde_json::{Value, json}; + +#[rstest] +fn sql_response_serializes_only_data() { + let response = TraceSQLResponse { + data: vec![ + serde_json::from_value(json!({ + "answer": 42, + "nested": {"items": [true, null, "9007199254740993"]} + })) + .unwrap(), + ], + }; + + assert_eq!( + serde_json::to_value(response).unwrap(), + json!({"data": [{"answer": 42, "nested": {"items": [true, null, "9007199254740993"]}}]}) + ); +} + +#[rstest] +fn sql_response_schema_requires_data_and_leaves_rows_open() { + let schemas = response_schemas(); + let schema: Value = serde_json::to_value(&schemas["TraceSQLResponse"]).unwrap(); + + assert_eq!(schema["additionalProperties"], false); + assert_eq!(schema["required"], json!(["data"])); + assert_eq!(schema["properties"].as_object().unwrap().len(), 1); + assert!( + schema["properties"] + .as_object() + .unwrap() + .contains_key("data") + ); + assert_eq!(schema["properties"]["data"]["type"], "array"); + assert_eq!(schema["properties"]["data"]["items"]["type"], "object"); + assert_ne!( + schema["properties"]["data"]["items"]["additionalProperties"], + false + ); +} diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 6ec0c6da3e7..6536ccd8493 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -36,6 +36,7 @@ from litellm.rust_bridge.trace.generated.requests import ( TraceQueryRequest, TraceSpanRequest, ) +from litellm.rust_bridge.trace.generated.responses import TraceSQLResponse from litellm.rust_bridge.trace.generated.types import ( AllQueryScope, OwnedQueryScope, @@ -46,7 +47,6 @@ from litellm.rust_bridge.trace.generated.types import ( TracePage, TraceScope, ) -from litellm.rust_bridge.trace.queries import TraceSQLResponse from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError from litellm.tracing.otlp_http import InvalidOTLPPayloadError, encode_otlp_response diff --git a/litellm/rust_bridge/trace/generated/responses.py b/litellm/rust_bridge/trace/generated/responses.py new file mode 100644 index 00000000000..b6de931c7ad --- /dev/null +++ b/litellm/rust_bridge/trace/generated/responses.py @@ -0,0 +1,20 @@ +# @generated by scripts/generate_trace_types.py, do not edit + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Annotated, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + + +class TraceSQLResponse(BaseModel): + model_config = ConfigDict( + extra="forbid", + frozen=True, + ) + + data: tuple[Mapping[str, JsonValue], ...] + + +TraceWireResponses: TypeAlias = Annotated[TraceSQLResponse, Field(..., title="TraceWireResponses")] diff --git a/litellm/rust_bridge/trace/queries.py b/litellm/rust_bridge/trace/queries.py index f40e1f19944..681b6af9f7e 100644 --- a/litellm/rust_bridge/trace/queries.py +++ b/litellm/rust_bridge/trace/queries.py @@ -28,7 +28,7 @@ class TraceQueryStatistics(BaseModel): bytes_read: int | str -class TraceSQLResponse(BaseModel): +class ClickHouseSQLEnvelope(BaseModel): model_config = _RESPONSE_CONFIG meta: tuple[TraceQueryColumn, ...] data: tuple[Mapping[str, JsonValue], ...] diff --git a/litellm/rust_bridge/trace/storage.py b/litellm/rust_bridge/trace/storage.py index a80e41a5aa3..67899cfe931 100644 --- a/litellm/rust_bridge/trace/storage.py +++ b/litellm/rust_bridge/trace/storage.py @@ -17,6 +17,7 @@ from litellm.rust_bridge.trace.generated.models import ( LensSampleParams, PartRow, ) +from litellm.rust_bridge.trace.generated.responses import TraceSQLResponse from litellm.rust_bridge.trace.generated.types import ReadQueryName from litellm.rust_bridge.trace.queries import ( LENS_AGENTS, @@ -24,6 +25,7 @@ from litellm.rust_bridge.trace.queries import ( LENS_CONTENT, LENS_EVIDENCE, LENS_SAMPLE, + ClickHouseSQLEnvelope, ParamsT, ReadQuery, RowT, @@ -38,7 +40,6 @@ from .generated.types import ( TracePage, TraceScope, ) -from .queries import TraceSQLResponse @dataclass(frozen=True, slots=True) @@ -99,7 +100,7 @@ class NativeTraces(Protocol): QUERY_PARAMETERS: Final = TypeAdapter(dict[str, str | int | float | list[str]]) -_SQL_RESPONSE: Final = TypeAdapter(TraceSQLResponse) +_SQL_ENVELOPE: Final = TypeAdapter(ClickHouseSQLEnvelope) _HELP_RESPONSE: Final = TypeAdapter(TraceQueryHelp) _TRACE_PAGE: Final = TypeAdapter(TracePage) _TRACE: Final[TypeAdapter[Trace | None]] = TypeAdapter(Trace | None) @@ -219,7 +220,8 @@ class ClickHouseStorage: async def query_sql(self, sql: str, scope: QueryScope, secret: str) -> TraceSQLResponse: result: Final = await self._native.query_sql(sql, scope, secret) - return _decode_query_response(_SQL_RESPONSE, result) + envelope: Final = _decode_query_response(_SQL_ENVELOPE, result) + return TraceSQLResponse(data=envelope.data) async def query_help(self, scope: QueryScope, secret: str) -> TraceQueryHelp: result: Final = await self._native.query_help(scope, secret) diff --git a/scripts/generate_trace_types.py b/scripts/generate_trace_types.py index 52a61452711..a8f4054a646 100644 --- a/scripts/generate_trace_types.py +++ b/scripts/generate_trace_types.py @@ -101,7 +101,7 @@ def generate( "pydantic_v2.BaseModel", "--enable-faux-immutability", "--additional-imports", - "collections.abc.Mapping,typing.TypeAlias", + "collections.abc.Mapping,typing.TypeAlias,pydantic.JsonValue", ) ) subprocess.run( @@ -165,8 +165,9 @@ def main() -> int: return 1 domain: Final = export("traces") requests: Final = export("traces", ("--requests",)) + responses: Final = export("traces", ("--responses",)) clickhouse: Final = export("traces-clickhouse") - exported: Final = tuple(schema_files(domain, clickhouse, requests)) + exported: Final = tuple(schema_files(domain, clickhouse, requests, responses)) schema_results: Final = tuple(publish(path, content, args.check) for path, content in exported) schema_set_matches: Final = reconcile_schemas(frozenset(path for path, _ in exported), args.check) with TemporaryDirectory(prefix="trace-codegen-") as temporary: @@ -179,10 +180,12 @@ def main() -> int: config, ) request_models: Final = generate(requests, "requests", directory, config) + response_models: Final = generate(responses, "responses", directory, config) python_results: Final = ( publish(GENERATED / "types.py", types.read_text(), args.check), publish(GENERATED / "models.py", models.read_text(), args.check), publish(GENERATED / "requests.py", request_models.read_text(), args.check), + publish(GENERATED / "responses.py", response_models.read_text(), args.check), ) return 0 if all((schema_set_matches, *schema_results, *python_results)) else 1 @@ -191,8 +194,14 @@ def schema_files( domain: Mapping[str, Mapping[str, JsonValue]], clickhouse: Mapping[str, Mapping[str, JsonValue]], requests: Mapping[str, Mapping[str, JsonValue]], + responses: Mapping[str, Mapping[str, JsonValue]], ) -> Iterator[tuple[Path, str]]: - for crate, schemas in (("traces", domain), ("traces-clickhouse", clickhouse), ("traces", requests)): + for crate, schemas in ( + ("traces", domain), + ("traces-clickhouse", clickhouse), + ("traces", requests), + ("traces", responses), + ): for name, schema in schemas.items(): yield TOOLING / "schemas" / crate / f"{name}.json", json.dumps(schema, indent=2, sort_keys=True) + "\n" diff --git a/scripts/trace_codegen/README.md b/scripts/trace_codegen/README.md index bb4f511a569..a784a51dc89 100644 --- a/scripts/trace_codegen/README.md +++ b/scripts/trace_codegen/README.md @@ -10,4 +10,4 @@ The templates preserve tuple conversion, immutable tuple defaults, and bounded ` Changing the public request contract requires a separate behavior-change PR -Edit the owning Rust contract, schema annotation, or generation configuration, then regenerate. Never edit `litellm/rust_bridge/trace/generated/` manually. The SQL response envelope remains handwritten in `queries.py` +Edit the owning Rust contract, schema annotation, or generation configuration, then regenerate. Never edit `litellm/rust_bridge/trace/generated/` manually. `responses.py` holds the generated SQL response, and `queries.py` keeps the ClickHouse envelope as an internal validator diff --git a/scripts/trace_codegen/schemas/traces/TraceSQLResponse.json b/scripts/trace_codegen/schemas/traces/TraceSQLResponse.json new file mode 100644 index 00000000000..da943b40cfd --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceSQLResponse.json @@ -0,0 +1,21 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "properties": { + "data": { + "items": { + "additionalProperties": true, + "type": "object" + }, + "type": "array", + "x-python-normalized": { + "type": "tuple[Mapping[str, JsonValue], ...]" + } + } + }, + "required": [ + "data" + ], + "title": "TraceSQLResponse", + "type": "object" +} diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 43950a2f011..50982f9774e 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -300,7 +300,7 @@ async def test_insert_validates_values_without_pydantic_copy(recording_server: R ("internal_user", None, 403), ), ) -def test_trace_sql_endpoint_enforces_ownership_and_preserves_clickhouse_envelope( +def test_trace_sql_endpoint_returns_data_only_and_enforces_ownership( recording_server: RecordingServer, role: str, user_id: str | None, expected_status: int ) -> None: from fastapi import FastAPI @@ -316,6 +316,7 @@ def test_trace_sql_endpoint_enforces_ownership_and_preserves_clickhouse_envelope "data": [{"answer": 42}], "rows": 1, "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 1}, + "rows_before_limit_at_least": 1, } recording_server.expected_requests = 12 if expected_status == 200 else 0 if expected_status == 200: @@ -339,7 +340,7 @@ def test_trace_sql_endpoint_enforces_ownership_and_preserves_clickhouse_envelope if expected_status == 403: assert result.json() == {"detail": "Not allowed to view logs"} return - assert result.json() == envelope + assert result.json() == {"data": envelope["data"]} assert recording_server.requests[-1].raw_body == b"SELECT 42 AS answer" assert client.post("/v1/traces/query", json={"sql": " "}).status_code == 400 assert client.post("/v1/traces/query", json={}).status_code == 422 @@ -428,6 +429,7 @@ def test_trace_sql_endpoint_distinguishes_query_errors_from_reader_failures( "data": [{"answer": 42}], "rows": 1, "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 1}, + "rows_before_limit_at_least": 1, } recording_server.enqueue(ResponseSpec(body=envelope)) storage: Final = ClickHouseStorage(TraceStorageConfig(recording_server.base_url, "trace_test")) @@ -441,7 +443,7 @@ def test_trace_sql_endpoint_distinguishes_query_errors_from_reader_failures( assert failed.status_code == expected_status, failed.text recovered: Final = client.post("/v1/traces/query", json={"sql": "SELECT 42 AS answer"}) assert recovered.status_code == 200, recovered.text - assert recovered.json() == envelope + assert recovered.json() == {"data": envelope["data"]} assert recording_server.requests[-2].raw_body == b"SELEC 42" diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 11d63ad545c..bd6350ad94e 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -12,7 +12,7 @@ import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from httpx import Response -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter from typing_extensions import ReadOnly from litellm.constants import TRACE_READ_RETRY_AFTER_SECONDS @@ -25,22 +25,27 @@ from litellm.proxy.tracing_runtime import manage_tracing, provide_storage from litellm.rust_bridge import loader from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp +from litellm.rust_bridge.trace.generated.responses import TraceSQLResponse from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope -from litellm.rust_bridge.trace.queries import TraceSQLResponse from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig from litellm.tracing import Tenant, TraceReceiver, TracingPayloadTooLargeError -SQL_ENVELOPE: Final = { - "meta": [{"name": "value", "type": "UInt64"}], - "data": [{"value": "9007199254740993"}], - "rows": 1, - "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 8}, - "rows_before_limit_at_least": 1, -} +SQL_ROWS: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "value": "9007199254740993", + "count": 42, + "fraction": 2.5, + "nested": {"values": [True, None, "text"]}, + }, +) +SQL_RESPONSE: Final = TraceSQLResponse(data=SQL_ROWS) QUERY_HELP: Final[Mapping[str, object]] = { "dialect": "test SQL", "access": "authenticated scope", - "response": "JSON envelope", + "response": ( + 'JSON object {"data": [rows]}; each row maps selected columns to values; ' + "64-bit integers may be strings" + ), "tables": [{"name": "otel_traces", "columns": [{"name": "value", "type": "String", "comment": "label"}]}], "normalized_fields": [], "metadata": { @@ -783,13 +788,14 @@ def test_sql_and_help_use_authenticated_scope( ) -> None: client.app.dependency_overrides[user_api_key_auth] = lambda: auth client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" - receiver.storage.query_sql = AsyncMock(return_value=TraceSQLResponse.model_validate(SQL_ENVELOPE)) + receiver.storage.query_sql = AsyncMock(return_value=SQL_RESPONSE) receiver.storage.query_help = AsyncMock(return_value=TraceQueryHelp.model_validate(QUERY_HELP)) result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) assert result.status_code == 200, result.text - assert result.json() == SQL_ENVELOPE + assert result.json() == {"data": list(SQL_ROWS)} assert type(result.json()["data"][0]["value"]) is str - assert type(result.json()["statistics"]["elapsed"]) is float + assert type(result.json()["data"][0]["count"]) is int + assert type(result.json()["data"][0]["fraction"]) is float receiver.storage.query_sql.assert_awaited_once_with("SELECT * FROM otel_traces", expected_scope, "test-secret") help_result: Final = client.get("/v1/traces/query/help") assert help_result.status_code == 200, help_result.text @@ -800,6 +806,28 @@ def test_sql_and_help_use_authenticated_scope( assert receiver.storage.query_sql.await_count == 1 +def test_sql_query_returns_empty_data(client: TestClient, receiver: MagicMock) -> None: + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + receiver.storage.query_sql = AsyncMock(return_value=TraceSQLResponse(data=())) + + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) + + assert result.status_code == 200, result.text + assert result.json() == {"data": []} + + +def test_sql_query_openapi_declares_a_closed_response_object(client: TestClient) -> None: + openapi: Final = client.app.openapi() + response: Final = openapi["paths"]["/v1/traces/query"]["post"]["responses"]["200"]["content"][ + "application/json" + ]["schema"] + component_name: Final = response["$ref"].rsplit("/", 1)[-1] + component: Final = openapi["components"]["schemas"][component_name] + + assert set(component["properties"]) == {"data"} + assert component["additionalProperties"] is False + + @pytest.mark.parametrize( ("body", "error_type", "location"), ( @@ -868,7 +896,7 @@ def test_queries_require_a_proxy_secret( from litellm.proxy import proxy_server monkeypatch.setattr(proxy_server, "master_key", secret) - receiver.storage.query_sql = AsyncMock(return_value=TraceSQLResponse.model_validate(SQL_ENVELOPE)) + receiver.storage.query_sql = AsyncMock(return_value=SQL_RESPONSE) result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"}) if secret is None: assert result.status_code == 503, result.text @@ -904,7 +932,7 @@ def test_shared_trace_permissions_reach_read_and_sql_boundaries( team_lookup: Final = AsyncMock(side_effect=lookup) storage: Final = MagicMock(spec=ClickHouseStorage) storage.get_span = AsyncMock(return_value=SPAN_DETAIL_RESPONSE) - storage.query_sql = AsyncMock(return_value=TraceSQLResponse.model_validate(SQL_ENVELOPE)) + storage.query_sql = AsyncMock(return_value=SQL_RESPONSE) storage.query_help = AsyncMock(return_value=TraceQueryHelp.model_validate(QUERY_HELP)) client.app.dependency_overrides[user_api_key_auth] = lambda: auth client.app.dependency_overrides[get_log_team_lookup] = lambda: team_lookup @@ -919,7 +947,7 @@ def test_shared_trace_permissions_reach_read_and_sql_boundaries( ) sql_response: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) assert sql_response.status_code == 200, sql_response.text - assert sql_response.json() == SQL_ENVELOPE + assert sql_response.json() == {"data": list(SQL_ROWS)} assert client.get("/v1/traces/query/help").json() == QUERY_HELP query_scope: Final = ( {"kind": "all"} diff --git a/tests/unit/rust_bridge/trace/test_queries.py b/tests/unit/rust_bridge/trace/test_queries.py index 1556cb5f46f..650f6bbafe9 100644 --- a/tests/unit/rust_bridge/trace/test_queries.py +++ b/tests/unit/rust_bridge/trace/test_queries.py @@ -1,11 +1,10 @@ -from collections.abc import Mapping from typing import Final import pytest -from pydantic import JsonValue, ValidationError +from pydantic import ValidationError from litellm.rust_bridge.trace.generated.models import LensContentParams -from litellm.rust_bridge.trace.queries import LENS_CONTENT, LENS_EVIDENCE, TraceSQLResponse +from litellm.rust_bridge.trace.queries import LENS_CONTENT, LENS_EVIDENCE @pytest.mark.parametrize("offset", (-1, 2**32)) @@ -50,18 +49,6 @@ def test_named_query_rejects_rows_missing_required_result_fields() -> None: assert error.value.error_count() == 4 -def test_sql_envelope_preserves_nested_data_large_integer_strings_and_extra_fields() -> None: - envelope: Final[Mapping[str, JsonValue]] = { - "meta": [{"name": "count", "type": "UInt64", "comment": "label"}], - "data": [{"count": "9007199254740993", "nested": [True, None, {"value": 2}]}], - "rows": "1", - "statistics": {"elapsed": 0.01, "rows_read": "1", "bytes_read": "8", "extra_stat": 4}, - "totals": {"count": "9007199254740993"}, - } - result: Final = TraceSQLResponse.model_validate(envelope) - assert result.model_dump(mode="json", exclude_unset=True) == envelope - - @pytest.mark.parametrize("count", (0, "9007199254740993", 2**64 - 1)) def test_clickhouse_rows_normalize_numbers_and_preserve_tuples(count: int | str) -> None: from litellm.rust_bridge.trace.queries import LENS_SAMPLE diff --git a/tests/unit/rust_bridge/trace/test_storage.py b/tests/unit/rust_bridge/trace/test_storage.py new file mode 100644 index 00000000000..19e0efa4496 --- /dev/null +++ b/tests/unit/rust_bridge/trace/test_storage.py @@ -0,0 +1,116 @@ +import json +from collections.abc import Mapping +from types import ModuleType +from typing import Final + +import pytest +from pydantic import JsonValue + +from litellm.rust_bridge import loader +from litellm.rust_bridge.trace.generated.types import QueryScope +from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig + + +class _NativeConfig: + def __init__(self, database: str, url: str, retention_days: int, max_attribute_value_bytes: int) -> None: + pass + + +class _NativeBridge(ModuleType): + def __init__(self, response: str) -> None: + super().__init__("native_traces") + + class Storage: + def __init__(self, config: _NativeConfig) -> None: + pass + + async def query_sql(self, sql: str, scope: QueryScope, secret: str) -> str: + return response + + self.NativeTraceConfig: Final = _NativeConfig + self.NativeTraceStorage: Final = Storage + + def trace_encode_error(self, message: str) -> bytes: + return b"" + + def trace_span_rows( + self, body: bytes, content_type: str | None, tenant: Mapping[str, str], max_attribute_value_bytes: int + ) -> list[dict[str, JsonValue]]: + return [] + + +@pytest.mark.parametrize( + ("body", "rows"), + ( + ( + json.dumps( + { + "meta": [{"name": "value", "type": "UInt64"}], + "data": [ + { + "integer": 9007199254740993, + "integer_string": "9007199254740993", + "fraction": 2.5, + "nested": {"values": [True, None, "text", {"count": 2}]}, + } + ], + "rows": 1, + "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 8}, + "rows_before_limit_at_least": 1, + "unknown_extra": {"kept_by_internal_validator": True}, + } + ), + [ + { + "integer": 9007199254740993, + "integer_string": "9007199254740993", + "fraction": 2.5, + "nested": {"values": [True, None, "text", {"count": 2}]}, + } + ], + ), + ( + json.dumps( + { + "meta": [], + "data": [], + "rows": 0, + "statistics": {"elapsed": 0, "rows_read": 0, "bytes_read": 0}, + } + ), + [], + ), + ), +) +async def test_sql_query_returns_only_rows_without_normalizing_json_values( + monkeypatch: pytest.MonkeyPatch, body: str, rows: list[dict[str, JsonValue]] +) -> None: + monkeypatch.setattr(loader, "_cached_bridge", _NativeBridge(body)) + storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) + + result: Final = await storage.query_sql("SELECT 1", {"kind": "all"}, "secret") + + assert result.model_dump(mode="json") == {"data": rows} + if rows: + row: Final = result.data[0] + assert type(row["integer"]) is int + assert type(row["integer_string"]) is str + assert type(row["fraction"]) is float + assert result.data == tuple(rows) + + +@pytest.mark.parametrize( + "body", + ( + '{"meta":[],"data":[],"rows":0}', + '{"meta":[],"data":{},"rows":0,"statistics":{"elapsed":0,"rows_read":0,"bytes_read":0}}', + ), +) +async def test_sql_query_rejects_malformed_clickhouse_envelopes( + monkeypatch: pytest.MonkeyPatch, body: str +) -> None: + monkeypatch.setattr(loader, "_cached_bridge", _NativeBridge(body)) + storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) + + with pytest.raises(RuntimeError, match="Native trace query returned an invalid response"): + await storage.query_sql("SELECT 1", {"kind": "all"}, "secret") diff --git a/tests/unit/test_seed_tracing_fixtures.py b/tests/unit/test_seed_tracing_fixtures.py index 8dac769b3c1..2608612ff31 100644 --- a/tests/unit/test_seed_tracing_fixtures.py +++ b/tests/unit/test_seed_tracing_fixtures.py @@ -11,7 +11,7 @@ import pytest from prisma import Json, Prisma from pydantic import InstanceOf, TypeAdapter -from litellm.rust_bridge.trace.queries import TraceSQLResponse +from litellm.rust_bridge.trace.generated.responses import TraceSQLResponse from litellm.rust_bridge.trace.storage import Tenant, span_rows from litellm.tracing.types import SpendLogRecord from scripts.seed_tracing_fixtures import ( @@ -205,13 +205,8 @@ async def test_first_copy_stamps_the_authenticated_tenant_and_writes_both_stores pattern: Final = response_pattern(tuple(chain.from_iterable(rows for _, rows in fixtures))) replays: Final = fixture_replays(TRACE_FIXTURES, 1_800_000_000_000, "first", pattern) storage: Final = AsyncMock(spec=ClickHouseStorage) - storage.query_sql.return_value = TraceSQLResponse.model_validate( - { - "meta": (), - "data": [{"team_id": "local-team", "api_key": "local-hash", "user": "admin"}], - "rows": 1, - "statistics": {"elapsed": 0, "rows_read": 1, "bytes_read": 1}, - } + storage.query_sql.return_value = TraceSQLResponse( + data=({"team_id": "local-team", "api_key": "local-hash", "user": "admin"},) ) database: Final = AsyncMock(spec=Prisma, litellm_spendlogs=AsyncMock()) client: Final = AsyncMock(spec=httpx.AsyncClient) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 73336ca1823..bc407c30c70 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46915,17 +46915,6 @@ export interface components { /** Sql */ sql: string; }; - /** TraceQueryStatistics */ - TraceQueryStatistics: { - /** Bytes Read */ - bytes_read: number | string; - /** Elapsed */ - elapsed: number; - /** Rows Read */ - rows_read: number | string; - } & { - [key: string]: unknown; - }; /** TraceQueryTable */ TraceQueryTable: { /** Columns */ @@ -46942,13 +46931,6 @@ export interface components { data: { [key: string]: components["schemas"]["JsonValue"]; }[]; - /** Meta */ - meta: components["schemas"]["TraceQueryColumn"][]; - /** Rows */ - rows: number | string; - statistics: components["schemas"]["TraceQueryStatistics"]; - } & { - [key: string]: unknown; }; /** TraceSummary */ TraceSummary: { diff --git a/ui/litellm-dashboard/src/lib/http/traceQueryTypes.test-d.ts b/ui/litellm-dashboard/src/lib/http/traceQueryTypes.test-d.ts index d977a628012..0d5275f5657 100644 --- a/ui/litellm-dashboard/src/lib/http/traceQueryTypes.test-d.ts +++ b/ui/litellm-dashboard/src/lib/http/traceQueryTypes.test-d.ts @@ -5,11 +5,9 @@ import type { components, paths } from "@/lib/http/schema"; type SQLResponse = paths["/v1/traces/query"]["post"]["responses"][200]["content"]["application/json"]; type QueryHelp = paths["/v1/traces/query/help"]["get"]["responses"][200]["content"]["application/json"]; -test("SQL results expose column metadata, JSON values and row counts", () => { - expectTypeOf().toEqualTypeOf(); +test("SQL results expose only row data and JSON values", () => { + expectTypeOf().toEqualTypeOf<"data">(); expectTypeOf().toEqualTypeOf(); - expectTypeOf().toEqualTypeOf(); - expectTypeOf().toEqualTypeOf(); }); test("query help exposes schema discovery, examples and discovery failures", () => {