litellm/tests/unit/integrations/test_lago.py
Mateo Wang 3f38730d2c
refactor(typing): clear about 1,100 basedpyright Any errors across 303 files (#49986)
* refactor(typing): clear basedpyright Any errors (wave 1)

* refactor(typing): clear basedpyright Any errors (wave 2)

* refactor(typing): clear basedpyright Any errors (wave 3)

* refactor(typing): clear basedpyright Any errors (wave 4)

* refactor(typing): clear basedpyright Any errors (wave 5)

* refactor(typing): clear basedpyright Any errors (wave 6)

* fix(typing): keep the cached-plan fallback's direct query_first calls so span naming resolves

* fix(logging): keep callbacks running when litellm_params is a params model

should_run_callback and the sync success and failure handlers read litellm_params through any object with a get method again, as they did before this branch, so a logging hook that returns a GenericLiteLLMParams no longer stops sync logging for that request

Adds behavior tests for the changed lines codecov listed as not run

* fix: read the regional uplift for every data residency again and reuse one dict TypeAdapter

The typed lookup only read the eu and us multipliers, so a deployment that set regional_processing_uplift_multiplier_fedramp stopped applying it. /model/update and the model write paths now share one module-level TypeAdapter instead of building one per call.

* fix: bill regional uplift multipliers float() accepts, bytes included

* fix: raise the in operator's own error from one shared json_contains, and name the bad uplift type
2026-11-06 13:31:15 -08:00

55 lines
1.9 KiB
Python

import json
from datetime import datetime, timezone
from typing import Final
import httpx
import pytest
from litellm.integrations.lago import LagoLogger
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.utils import ModelResponse, Usage
CALL_TIME: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
@pytest.fixture
def lago_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LAGO_API_KEY", "lago-key")
monkeypatch.setenv("LAGO_API_BASE", "https://lago.example")
monkeypatch.setenv("LAGO_API_EVENT_CODE", "llm-usage")
monkeypatch.delenv("LAGO_API_CHARGE_BY", raising=False)
def test_lago_event_bills_the_response_token_usage_to_the_end_user(lago_env: None) -> None:
sent: Final[list[httpx.Request]] = []
def lago(request: httpx.Request) -> httpx.Response:
sent.append(request)
return httpx.Response(200, json={})
logger: Final = LagoLogger()
logger.sync_http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lago)))
response: Final = ModelResponse(model="gpt-4o", usage=Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17))
logger.log_success_event(
kwargs={
"model": "gpt-4o",
"response_cost": 0.002,
"litellm_params": {"metadata": {}, "proxy_server_request": {"body": {"user": "end-user-1"}}},
},
response_obj=response,
start_time=CALL_TIME,
end_time=CALL_TIME,
)
assert [str(request.url) for request in sent] == ["https://lago.example/api/v1/events"]
event: Final = json.loads(sent[0].content)["event"]
assert event["external_subscription_id"] == "end-user-1"
assert event["code"] == "llm-usage"
assert event["properties"] == {
"model": "gpt-4o",
"response_cost": 0.002,
"prompt_tokens": 12,
"completion_tokens": 5,
"total_tokens": 17,
}