mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
* test: align integration fixtures with current behavior Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for the spend flush before asserting its trace placement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): key response-cache entries by model group from litellm_metadata /v1/responses routes through the router with model_group in litellm_metadata, which the cache key ignored, so identical requests to different model groups sharing one underlying model hit each other's cached responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
72 lines
3.3 KiB
Python
72 lines
3.3 KiB
Python
import uuid
|
|
from collections.abc import Callable, Iterator, Mapping
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
|
|
from integration._support.otlp_sink import Span, SpanSinks, recorded_spans
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import JsonValue
|
|
|
|
pytestmark: Final = pytest.mark.timeout(180)
|
|
|
|
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def gateway(audit_sinks: SpanSinks) -> Iterator[Gateway]:
|
|
with gateway_from_environment() as base:
|
|
yield base
|
|
|
|
|
|
def _traces(spans: tuple[Span, ...]) -> dict[str, frozenset[str]]:
|
|
trace_ids: Final = {span["trace_id"] for span in spans}
|
|
return {trace: frozenset(span["name"] for span in spans if span["trace_id"] == trace) for trace in trace_ids}
|
|
|
|
|
|
def test_default_otel_logger_keeps_spend_flush_outside_the_request_trace(
|
|
gateway: Gateway, audit_sinks: SpanSinks, otel_audit_config: AuditConfigWriter, tmp_path: Path
|
|
) -> None:
|
|
config: Final = otel_audit_config(tmp_path, {})
|
|
overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator}
|
|
with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate, candidate.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
key: Final = scenario.key(models=[model])
|
|
start, _ = recorded_spans(audit_sinks.operator)
|
|
response: Final = candidate.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"otel v1 {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
expected: Final = frozenset({"postgres", "redis", "raw_gen_ai_request"})
|
|
traces: Final = eventually(
|
|
lambda: _traces(recorded_spans(audit_sinks.operator, start)[1]),
|
|
lambda grouped: any(expected <= names for names in grouped.values()),
|
|
seconds=60,
|
|
return_last_on_timeout=True,
|
|
)
|
|
assert any(expected <= names for names in traces.values()), {
|
|
trace: sorted(names) for trace, names in traces.items()
|
|
}
|
|
request_trace: Final = next(trace for trace, names in traces.items() if expected <= names)
|
|
key_info: Final = eventually(
|
|
lambda: candidate.request("GET", "/key/info", key=key, params={"key": key}),
|
|
lambda response: response.status_code == 200
|
|
and float(str(object_value(response.json()["info"])["spend"])) > 0,
|
|
seconds=60,
|
|
)
|
|
assert key_info.status_code == 200, key_info.text
|
|
assert float(str(object_value(key_info.json()["info"])["spend"])) > 0
|
|
request_spans: Final = tuple(
|
|
span for span in recorded_spans(audit_sinks.operator, start)[1] if span["trace_id"] == request_trace
|
|
)
|
|
assert not any(span["name"] == "batch_write_to_db" for span in request_spans), request_spans
|
|
assert not any(
|
|
span["name"] == "postgres"
|
|
and span["attributes"].get("call_type") == "commit_spend_updates"
|
|
and span["attributes"].get("table_name") == "LiteLLM_VerificationToken"
|
|
for span in request_spans
|
|
), request_spans
|