From 0313dcea61f8df29ce1438d8f371291cf8e10c69 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 13:38:28 -0700 Subject: [PATCH 01/35] ci: test Python 3.10 through 3.14 compatibility --- .github/workflows/_test-unit-base.yml | 25 +++-- .../workflows/test-python-compatibility.yml | 105 ++++++++++++++++++ 2 files changed, 118 insertions(+), 12 deletions(-) create mode 100644 .github/workflows/test-python-compatibility.yml diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index c4045a08ffb..46baf2ac26f 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -57,9 +57,15 @@ permissions: jobs: run: - name: Run tests + name: ${{ matrix.python-version == '3.12' && 'Run tests' || format('Run tests (Python {0})', matrix.python-version) }} runs-on: ubuntu-latest timeout-minutes: ${{ inputs.job-timeout-minutes }} + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + env: + UV_PYTHON: ${{ matrix.python-version }} permissions: contents: read pull-requests: read @@ -82,7 +88,7 @@ jobs: timeout-minutes: 3 uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: - python-version: "3.12" + python-version: ${{ matrix.python-version }} - name: Set up uv if: steps.changes.outputs.decision != 'skip' @@ -99,9 +105,9 @@ jobs: path: | ~/.cache/uv .venv - key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }} + key: ${{ runner.os }}-uv-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} restore-keys: | - ${{ runner.os }}-uv- + ${{ runner.os }}-uv-py${{ matrix.python-version }}- - name: Cache the Rust build if: steps.changes.outputs.decision != 'skip' @@ -113,6 +119,7 @@ jobs: timeout-minutes: 8 run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml + uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -134,13 +141,7 @@ jobs: WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} DIST: ${{ inputs.dist }} - # coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has. - # It is only the default from Python 3.14, and these shards run 3.12, so it - # has to be asked for. Coverage refuses it when branch measurement is on - # (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with - # a `no-sysmon` warning, so turning on `branch = true` here means giving this - # back until the runners move to 3.14. - COVERAGE_CORE: sysmon + COVERAGE_CORE: ${{ contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) && 'ctrace' || 'sysmon' }} run: | if [ "${WORKERS}" = "0" ]; then uv run --no-sync pytest ${TEST_PATH:?} \ @@ -167,7 +168,7 @@ jobs: fi - name: Save coverage report - if: always() && steps.changes.outputs.decision != 'skip' + if: always() && matrix.python-version == '3.12' && steps.changes.outputs.decision != 'skip' uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 with: name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }} diff --git a/.github/workflows/test-python-compatibility.yml b/.github/workflows/test-python-compatibility.yml new file mode 100644 index 00000000000..d3a89d7141c --- /dev/null +++ b/.github/workflows/test-python-compatibility.yml @@ -0,0 +1,105 @@ +name: Python Compatibility + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + push: + branches: + - main + - litellm_internal_staging + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + compatibility: + name: Python ${{ matrix.python-version }} + runs-on: ubuntu-latest + timeout-minutes: 45 + permissions: + contents: read + pull-requests: read + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + env: + UV_PYTHON: ${{ matrix.python-version }} + LITELLM_LOCAL_MODEL_COST_MAP: "True" + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + timeout-minutes: 3 + with: + persist-credentials: false + + - name: Detect relevant changes + id: changes + timeout-minutes: 2 + uses: ./.github/actions/detect-changes + + - name: Set up Python + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Set up uv + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache uv dependencies + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: ~/.cache/uv + key: ${{ runner.os }}-uv-compat-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} + restore-keys: | + ${{ runner.os }}-uv-compat-py${{ matrix.python-version }}- + + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/cache-cargo-build + + - name: Install SDK and proxy dependencies + id: install + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 8 + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --extra proxy + uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' + + - name: Compile package + if: ${{ !cancelled() && steps.install.outcome == 'success' }} + timeout-minutes: 2 + run: uv run --no-sync python -m compileall -q litellm + + - name: Import SDK + if: ${{ !cancelled() && steps.install.outcome == 'success' }} + timeout-minutes: 2 + run: uv run --no-sync python -c 'from litellm import *; print("SDK import passed")' + + - name: Import proxy + if: ${{ !cancelled() && steps.install.outcome == 'success' }} + timeout-minutes: 2 + run: uv run --no-sync python -c 'import litellm.proxy.proxy_server; print("Proxy import passed")' + + - name: Construct a completion response + if: ${{ !cancelled() && steps.install.outcome == 'success' }} + timeout-minutes: 2 + run: uv run --no-sync python -c 'from litellm.types.utils import ModelResponse; ModelResponse(); print("Response construction passed")' From 8010f267e2ce01ad67d2d4479486def1d5288cd4 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 13:42:22 -0700 Subject: [PATCH 02/35] ci: keep Python compatibility coverage in unit workflow --- .../workflows/test-python-compatibility.yml | 105 ------------------ 1 file changed, 105 deletions(-) delete mode 100644 .github/workflows/test-python-compatibility.yml diff --git a/.github/workflows/test-python-compatibility.yml b/.github/workflows/test-python-compatibility.yml deleted file mode 100644 index d3a89d7141c..00000000000 --- a/.github/workflows/test-python-compatibility.yml +++ /dev/null @@ -1,105 +0,0 @@ -name: Python Compatibility - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_staging - - "litellm_**" - push: - branches: - - main - - litellm_internal_staging - workflow_dispatch: - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - compatibility: - name: Python ${{ matrix.python-version }} - runs-on: ubuntu-latest - timeout-minutes: 45 - permissions: - contents: read - pull-requests: read - strategy: - fail-fast: false - matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - env: - UV_PYTHON: ${{ matrix.python-version }} - LITELLM_LOCAL_MODEL_COST_MAP: "True" - - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - timeout-minutes: 3 - with: - persist-credentials: false - - - name: Detect relevant changes - id: changes - timeout-minutes: 2 - uses: ./.github/actions/detect-changes - - - name: Set up Python - if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 3 - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: ${{ matrix.python-version }} - - - name: Set up uv - if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 3 - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache uv dependencies - if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 5 - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: ~/.cache/uv - key: ${{ runner.os }}-uv-compat-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} - restore-keys: | - ${{ runner.os }}-uv-compat-py${{ matrix.python-version }}- - - - name: Cache the Rust build - if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 5 - uses: ./.github/actions/cache-cargo-build - - - name: Install SDK and proxy dependencies - id: install - if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 8 - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --extra proxy - uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - - - name: Compile package - if: ${{ !cancelled() && steps.install.outcome == 'success' }} - timeout-minutes: 2 - run: uv run --no-sync python -m compileall -q litellm - - - name: Import SDK - if: ${{ !cancelled() && steps.install.outcome == 'success' }} - timeout-minutes: 2 - run: uv run --no-sync python -c 'from litellm import *; print("SDK import passed")' - - - name: Import proxy - if: ${{ !cancelled() && steps.install.outcome == 'success' }} - timeout-minutes: 2 - run: uv run --no-sync python -c 'import litellm.proxy.proxy_server; print("Proxy import passed")' - - - name: Construct a completion response - if: ${{ !cancelled() && steps.install.outcome == 'success' }} - timeout-minutes: 2 - run: uv run --no-sync python -c 'from litellm.types.utils import ModelResponse; ModelResponse(); print("Response construction passed")' From c8f6531be2d908750fcb487062738869784ddb48 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 13:52:35 -0700 Subject: [PATCH 03/35] fix: restore Python 3.10 imports and response construction --- .../websearch_interception/handler.py | 4 +-- .../_experimental/mcp_server/tool_search.py | 4 +-- litellm/proxy/agent_endpoints/endpoints.py | 4 +-- .../proxy/common_utils/reset_budget_job.py | 4 ++- .../batch_file_validation.py | 4 ++- .../types/llms/gemini_audio_transcription.py | 4 +-- litellm/types/llms/openai.py | 2 +- .../types/llms/test_types_llms_openai.py | 35 +++++++++++++++++++ type-discipline-budget.json | 2 +- 9 files changed, 51 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index dc61ee38a8c..900ab6408ba 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,9 +10,9 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly import litellm from litellm._logging import verbose_logger diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index f79765f6d01..f59c7c2f45f 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -3,9 +3,9 @@ from __future__ import annotations import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from typing import TYPE_CHECKING, Any, Final, TypedDict -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index b6c41a17503..147f2ba0fcf 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,10 +13,10 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict, assert_never +from typing import Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..47f69732e95 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,9 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar, assert_never +from typing import Final, Literal, Protocol, TypeVar + +from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 0aee5e8cc54..a41bd36d510 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -2,7 +2,9 @@ import json from collections.abc import Iterator from dataclasses import dataclass from itertools import chain -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException diff --git a/litellm/types/llms/gemini_audio_transcription.py b/litellm/types/llms/gemini_audio_transcription.py index cb12e0f45b8..f7e74ba4bf8 100644 --- a/litellm/types/llms/gemini_audio_transcription.py +++ b/litellm/types/llms/gemini_audio_transcription.py @@ -1,7 +1,7 @@ -from typing import Literal, Required +from typing import Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict class GeminiTranscriptionAudioInput(TypedDict): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 32d88da0085..943696326d0 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -631,7 +631,7 @@ class ChatCompletionReasoningItem(TypedDict, total=False): type: Required[Literal["reasoning"]] id: str encrypted_content: str | None - summary: list["ChatCompletionReasoningSummaryTextBlock"] + summary: ReadOnly[list[ChatCompletionReasoningSummaryTextBlock]] class WebSearchOptionsUserLocationApproximate(TypedDict, total=False): diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 42719ce838b..64ec09838e8 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -10,6 +10,41 @@ import litellm from litellm.types.llms.openai import HttpxBinaryResponseContent +@pytest.mark.parametrize("stream", (False, True)) +def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: + from typing import Final + + from litellm.types.llms.openai import ( + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ) + from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + ) + + reasoning_item: Final = ChatCompletionReasoningItem( + type="reasoning", + id="rs_123", + encrypted_content="encrypted", + summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], + ) + response: Final = ( + ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) + if stream + else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) + ) + message_key: Final = "delta" if stream else "message" + assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + restored: Final = type(response).model_validate_json(response.model_dump_json()) + assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + def test_generic_event(): from litellm.types.llms.openai import GenericEvent diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..5b43c0f6778 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -33,6 +33,6 @@ "limit": 5535 }, "LIT012": { - "limit": 4495 + "limit": 4494 } } From 77d6aedf0ad192a61ce0c424d0e7720ae3446bf1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 14:17:19 -0700 Subject: [PATCH 04/35] fix: address cross-version CI failures --- .github/workflows/_test-unit-base.yml | 8 +-- basedpyright-code-budget.json | 2 +- litellm/proxy/schema.prisma | 1 + litellm/utils.py | 10 +-- pyproject.toml | 1 + ruff-strict-budget.json | 4 +- scripts/budget_ratchet_check.py | 7 +- scripts/mutation_report.py | 6 +- test-quality-budget.json | 2 +- tests/code_coverage_tests/check_licenses.py | 6 +- .../test_litellm/caching/test_redis_cache.py | 12 ++-- tests/test_litellm/caching/test_s3_cache.py | 18 ++--- .../test_max_streaming_duration.py | 13 ++-- .../messages/test_mcp_handler.py | 13 ++-- .../llms/test_file_search_responses.py | 65 ++++++++++--------- .../guardrail_hooks/test_straiker.py | 2 +- tests/test_litellm/proxy/test_proxy_server.py | 35 ++++++---- .../test_responses_api_request_body.py | 5 +- .../test_streaming_iterator_error_events.py | 17 ++--- .../router_strategy/test_litellm_encoder.py | 4 ++ tests/test_litellm/test_gpt_realtime_mode.py | 6 +- tests/test_litellm/test_main.py | 5 +- tests/test_litellm/test_ruff_strict_gate.py | 6 +- tests/test_litellm/test_utils.py | 10 +++ uv.lock | 2 + 25 files changed, 152 insertions(+), 108 deletions(-) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 46baf2ac26f..d2cc0aa6d8d 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -102,12 +102,10 @@ jobs: timeout-minutes: 5 uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: - path: | - ~/.cache/uv - .venv - key: ${{ runner.os }}-uv-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} + path: ${{ env.UV_CACHE_DIR }} + key: ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} restore-keys: | - ${{ runner.os }}-uv-py${{ matrix.python-version }}- + ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}- - name: Cache the Rust build if: steps.changes.outputs.decision != 'skip' diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..609391d02ab 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -18,7 +18,7 @@ "limit": 40 }, "reportDeprecated": { - "limit": 211 + "limit": 210 }, "reportDuplicateImport": { "limit": 19 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7604ceadf7a..2a2665f9731 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -5,6 +5,7 @@ datasource client { generator client { provider = "prisma-client-py" + recursive_type_depth = -1 binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"] } diff --git a/litellm/utils.py b/litellm/utils.py index 252b6756937..c894f441687 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5149,14 +5149,8 @@ def get_api_key(llm_provider: str, dynamic_api_key: str | None): return api_key -def get_utc_datetime(): - import datetime as dt - from datetime import datetime - - if hasattr(dt, "UTC"): - return datetime.now(dt.UTC) - else: - return datetime.utcnow() +def get_utc_datetime() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc) def get_max_tokens(model: str) -> int | None: diff --git a/pyproject.toml b/pyproject.toml index 60162544612..57bce6b470f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,7 @@ dev = [ "basedpyright==1.39.7", "keyring==25.7.0", "pytest==9.0.3", + "tomli==2.4.1; python_version < '3.11'", "pytest-mock==3.15.1", "pytest-asyncio==1.3.0", "pytest-postgresql==7.0.2", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..37b4486a7c9 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 809 }, "ANN201": { - "limit": 2001 + "limit": 2000 }, "ANN202": { "limit": 835 @@ -87,7 +87,7 @@ "limit": 2 }, "DTZ003": { - "limit": 26 + "limit": 25 }, "DTZ005": { "limit": 233 diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 34dd234477a..adc4c0664be 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -24,7 +24,6 @@ seen the red and accepted it. Usage: python scripts/budget_ratchet_check.py [--base REF] [budget.json ...] -Stdlib only. """ from __future__ import annotations @@ -33,11 +32,15 @@ import argparse import json import subprocess import sys -import tomllib from pathlib import Path from types import MappingProxyType from typing import NamedTuple +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + REPO_ROOT = Path(__file__).resolve().parent.parent DEFAULT_BASE = "origin/litellm_internal_staging" DEFAULT_BUDGETS: tuple[str, ...] = ( diff --git a/scripts/mutation_report.py b/scripts/mutation_report.py index e0d4d569484..d0f9ddf0491 100644 --- a/scripts/mutation_report.py +++ b/scripts/mutation_report.py @@ -18,13 +18,17 @@ import json import re import subprocess import sys -import tomllib from collections import defaultdict from difflib import SequenceMatcher from pathlib import Path from typing import Final, NamedTuple from textwrap import dedent +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + ROOT = Path(__file__).resolve().parent.parent MUTMUT_INVOCATION = ["uv", "run", "--no-sync", "--with", "mutmut==3.5.0", "mutmut"] diff --git a/test-quality-budget.json b/test-quality-budget.json index d834c581609..74a1cb307df 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11135 + "limit": 11103 } } diff --git a/tests/code_coverage_tests/check_licenses.py b/tests/code_coverage_tests/check_licenses.py index 158e25180e1..a9eddc3fabb 100644 --- a/tests/code_coverage_tests/check_licenses.py +++ b/tests/code_coverage_tests/check_licenses.py @@ -6,12 +6,16 @@ from pathlib import Path import re import sys import time -import tomllib from typing import Callable, Dict, Final, List, Optional, Protocol, Set, Tuple from packaging.requirements import Requirement import requests +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + DEFAULT_TRANSITIVE_PIN_PACKAGES = ( "aiofiles", "anyio", diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 487a64797d1..0743824f85b 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -484,7 +484,7 @@ def _closed_port() -> int: pytest.param(lambda c: c.async_get_ttl("lit4930"), id="async_get_ttl"), ], ) -async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no_ping, call_method): +async def test_circuit_breaker_opens_when_method_swallows_redis_failure(call_method): """A guarded method that swallows its own Redis error must still count as a failure. These methods catch connection errors and return a default so callers degrade instead @@ -495,7 +495,7 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no """ from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + cache = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): await call_method(cache) @@ -505,7 +505,7 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no @pytest.mark.asyncio -async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ping): +async def test_circuit_breaker_success_still_resets_the_failure_streak(): """A reachable Redis must keep the breaker closed, however many earlier calls failed. The guard now records success only when nothing failed while the method ran, so this @@ -514,7 +514,7 @@ async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ """ from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + cache = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1): await cache.async_get_cache("lit4930") @@ -532,7 +532,7 @@ async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ @pytest.mark.asyncio -async def test_circuit_breaker_covers_lua_script_execution(redis_no_ping): +async def test_circuit_breaker_covers_lua_script_execution(): """Lua script execution must feed the breaker like every other Redis call. The v3 rate limiter issues all of its Redis traffic through async_register_script, so @@ -544,7 +544,7 @@ async def test_circuit_breaker_covers_lua_script_execution(redis_no_ping): from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + cache = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) run_script = cache.async_register_script("return 1") for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index f9a0b165e12..f86f2da30ef 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -258,11 +258,9 @@ async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies): # Verify each call calls = cache.s3_client.put_object.call_args_list - for i, (key, value) in enumerate(cache_list): - call_args = calls[i][1] - assert call_args["Bucket"] == "test-bucket" - assert call_args["Key"] == key - assert call_args["Body"] == json.dumps(value) + assert {(call.kwargs["Bucket"], call.kwargs["Key"], call.kwargs["Body"]) for call in calls} == { + ("test-bucket", key, json.dumps(value)) for key, value in cache_list + } @pytest.mark.asyncio @@ -285,10 +283,12 @@ async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies): # Verify each call had correct parameters calls = cache.s3_client.put_object.call_args_list - for i, call in enumerate(calls): - call_args = call[1] - assert call_args["Bucket"] == "test-bucket" - assert f"concurrent_key_{i}" == call_args["Key"] + assert {call.kwargs["Key"] for call in calls} == {f"concurrent_key_{i}" for i in range(5)} + for call in calls: + assert call.kwargs["Bucket"] == "test-bucket" + payload = json.loads(call.kwargs["Body"]) + assert call.kwargs["Key"] == f"concurrent_key_{payload['id']}" + assert payload["data"] == f"test_data_{payload['id']}" @pytest.mark.asyncio diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py index c768be22a9e..fc8daab3899 100644 --- a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py @@ -7,6 +7,7 @@ Covers: """ import time +from importlib import import_module from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -107,16 +108,16 @@ class TestResponsesStreamingIteratorMaxDuration: def test_should_not_raise_when_duration_is_none(self): it = self._make_base_iterator() - with patch( - "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", + with patch.object( + import_module("litellm.responses.streaming_iterator"), "LITELLM_MAX_STREAMING_DURATION_SECONDS", None, ): it._check_max_streaming_duration() def test_should_not_raise_when_under_limit(self): it = self._make_base_iterator() - with patch( - "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", + with patch.object( + import_module("litellm.responses.streaming_iterator"), "LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0, ): it._check_max_streaming_duration() @@ -124,8 +125,8 @@ class TestResponsesStreamingIteratorMaxDuration: def test_should_raise_timeout_when_exceeded(self): it = self._make_base_iterator() it._stream_created_time = time.time() - 20 - with patch( - "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", + with patch.object( + import_module("litellm.responses.streaming_iterator"), "LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0, ): with pytest.raises(litellm.Timeout, match="max streaming duration"): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py index b6914809263..f8c48e46b2f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -1,3 +1,4 @@ +from importlib import import_module from unittest.mock import AsyncMock, patch import pytest @@ -163,8 +164,8 @@ async def test_anthropic_messages_with_mcp_forwards_the_callers_mcp_credentials( ).LiteLLM_Proxy_MCP_Handler, "_process_mcp_tools_without_openai_transform", new=process, - ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._execute_tool_calls", + ), patch.object( + import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "_execute_tool_calls", new=execute, ), patch( "litellm.anthropic_messages", new=AsyncMock(side_effect=responses) @@ -218,11 +219,11 @@ async def test_anthropic_messages_with_mcp_stops_when_every_tool_call_is_skipped with patch.object( MCPRequestContext, "resolve", return_value=MCPRequestContext(user_api_key_auth="auth") - ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform", + ), patch.object( + import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "_process_mcp_tools_without_openai_transform", new=AsyncMock(return_value=([], {})), - ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._execute_tool_calls", + ), patch.object( + import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "_execute_tool_calls", new=AsyncMock(return_value=[]), ), patch( "litellm.anthropic_messages", new=anthropic_messages_mock diff --git a/tests/test_litellm/llms/test_file_search_responses.py b/tests/test_litellm/llms/test_file_search_responses.py index 2f7ad3874fa..887f14ce80e 100644 --- a/tests/test_litellm/llms/test_file_search_responses.py +++ b/tests/test_litellm/llms/test_file_search_responses.py @@ -13,6 +13,7 @@ Coverage: import base64 from typing import Any, Dict, List, Optional +from importlib import import_module from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -223,28 +224,28 @@ class TestFileSearchGuardInResponsesMain: expected = {"ok": True} with ( - patch( - "litellm.responses.main.litellm.get_llm_provider", + patch.object( + import_module("litellm.responses.main").litellm, "get_llm_provider", return_value=("claude-sonnet-4-5", "anthropic", None, None), ), - patch( - "litellm.responses.main.update_responses_input_with_model_file_ids", + patch.object( + import_module("litellm.responses.main"), "update_responses_input_with_model_file_ids", return_value="hello", ), - patch( - "litellm.responses.main.update_responses_tools_with_model_file_ids", + patch.object( + import_module("litellm.responses.main"), "update_responses_tools_with_model_file_ids", return_value=tools, ), - patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", + patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, ), - patch( - "litellm.responses.main.ResponsesAPIRequestUtils.get_requested_response_api_optional_param", + patch.object( + import_module("litellm.responses.main").ResponsesAPIRequestUtils, "get_requested_response_api_optional_param", return_value={}, ), - patch( - "litellm.responses.main.run_async_function", return_value=expected + patch.object( + import_module("litellm.responses.main"), "run_async_function", return_value=expected ) as run_async_mock, ): result = responses( @@ -274,28 +275,28 @@ class TestFileSearchGuardInResponsesMain: mock_config.supports_native_file_search.return_value = False with ( - patch( - "litellm.responses.main.litellm.get_llm_provider", + patch.object( + import_module("litellm.responses.main").litellm, "get_llm_provider", return_value=("claude-sonnet-4-5", "anthropic", None, None), ), - patch( - "litellm.responses.main.update_responses_input_with_model_file_ids", + patch.object( + import_module("litellm.responses.main"), "update_responses_input_with_model_file_ids", return_value="hello", ), - patch( - "litellm.responses.main.update_responses_tools_with_model_file_ids", + patch.object( + import_module("litellm.responses.main"), "update_responses_tools_with_model_file_ids", return_value=tools, ), - patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", + patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=mock_config, ), - patch( - "litellm.responses.main.ResponsesAPIRequestUtils.get_requested_response_api_optional_param", + patch.object( + import_module("litellm.responses.main").ResponsesAPIRequestUtils, "get_requested_response_api_optional_param", return_value={}, ), - patch( - "litellm.responses.main.run_async_function", return_value=expected + patch.object( + import_module("litellm.responses.main"), "run_async_function", return_value=expected ) as run_async_mock, ): result = responses( @@ -758,8 +759,8 @@ class TestEmulatedFileSearchHandler: mock_search_response.data = [search_result] with ( - patch( - "litellm.responses.file_search.emulated_handler._call_aresponses", + patch.object( + import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses", new=AsyncMock(side_effect=[first_resp, final_resp]), ), patch( @@ -821,8 +822,8 @@ class TestEmulatedFileSearchHandler: mock_search_response.data = [search_result] with ( - patch( - "litellm.responses.file_search.emulated_handler._call_aresponses", + patch.object( + import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses", new=AsyncMock(side_effect=[first_resp_plural, final_resp]), ), patch( @@ -855,8 +856,8 @@ class TestEmulatedFileSearchHandler: text="I already know the answer." ) - with patch( - "litellm.responses.file_search.emulated_handler._call_aresponses", + with patch.object( + import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses", new=AsyncMock(return_value=direct_resp), ): result = await aresponses_with_emulated_file_search( @@ -905,8 +906,8 @@ class TestEmulatedFileSearchHandler: mock_search_response.data = [search_result] with ( - patch( - "litellm.responses.file_search.emulated_handler._call_aresponses", + patch.object( + import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses", new=AsyncMock(side_effect=[first_resp, final_resp]), ) as mock_call, patch( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index 81604e22c87..d5d1c9bf176 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -109,7 +109,7 @@ def test_supported_hooks_limited_to_pre_and_post(): def test_during_call_mode_rejected_at_init(): - with pytest.raises(ValueError, match='Event hook GuardrailEventHooks\\.during_call is not in the'): + with pytest.raises(ValueError, match="during_call is not in the supported event hooks"): StraikerGuardrail(api_key="k", event_hook="during_call") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 91fca8f1e27..0ae51508655 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9,8 +9,9 @@ import subprocess import types from datetime import datetime, timedelta, timezone from pathlib import Path +from typing import Final from unittest import mock -from unittest.mock import AsyncMock, MagicMock, mock_open, patch +from unittest.mock import AsyncMock, MagicMock, create_autospec, mock_open, patch import click import httpx @@ -696,6 +697,18 @@ def test_restructure_always_happens(monkeypatch): assert ui_path == packaged_ui_path +def _mock_scheduled_proxy_config() -> MagicMock: + config: Final = proxy_server_module.ProxyConfig() + return MagicMock( + spec=proxy_server_module.ProxyConfig, + check_periodic_reloads=create_autospec(config.check_periodic_reloads), + get_credentials=create_autospec(config.get_credentials), + add_deployment=create_autospec(config.add_deployment), + reload_search_tools_from_db=create_autospec(config.reload_search_tools_from_db), + reload_mcp_servers_from_db=create_autospec(config.reload_mcp_servers_from_db), + ) + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_credentials(monkeypatch): """ @@ -711,7 +724,7 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), @@ -771,7 +784,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() scheduler = AsyncIOScheduler() try: @@ -812,7 +825,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() mock_scheduler = MagicMock() configured_interval = 47 @@ -861,7 +874,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() mock_scheduler = MagicMock() with ( @@ -908,7 +921,7 @@ async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_fal mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), @@ -7258,7 +7271,7 @@ async def test_batch_cost_poller_is_confirmed_before_serving(monkeypatch): mock_proxy_logging.db_spend_update_writer = MagicMock() with ( - patch("litellm.proxy.proxy_server.proxy_config", AsyncMock()), + patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()), patch("litellm.proxy.proxy_server.store_model_in_db", False), patch("litellm.proxy.proxy_server.llm_router", MagicMock()), patch("litellm.proxy.proxy_server.PROXY_BATCH_POLLING_ENABLED", True), @@ -7300,7 +7313,7 @@ async def test_store_model_in_db_db_override_when_config_false(): mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), @@ -7343,7 +7356,7 @@ async def test_store_model_in_db_db_check_skipped_when_already_true(monkeypatch) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), @@ -7386,7 +7399,7 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), @@ -11640,7 +11653,7 @@ async def _run_scheduled_background_jobs(): mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_logging.db_spend_update_writer = MagicMock() - mock_proxy_config = AsyncMock() + mock_proxy_config = _mock_scheduled_proxy_config() with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 5fd53fda01b..3e60906ec6d 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -7,6 +7,7 @@ in expected_responses_api_request/. import copy import json from pathlib import Path +from importlib import import_module from unittest.mock import AsyncMock, patch import httpx @@ -405,8 +406,8 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_ from litellm.responses.main import _aresponses_websocket - with patch( - "litellm.responses.main.base_llm_http_handler.async_responses_websocket", + with patch.object( + import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", new_callable=AsyncMock, ) as mock_ws: await _aresponses_websocket( diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 9c344fc6894..ad74861c096 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -15,6 +15,7 @@ Pydantic ValidationError (previously typed as Optional[str]). """ import json +from importlib import import_module from unittest.mock import Mock, patch import pytest @@ -259,8 +260,8 @@ def test_handle_logging_failed_response_maps_rate_limit_to_429(): {"type": "tokens", "code": "rate_limit_exceeded", "message": "throttled"} ) with ( - patch("litellm.responses.streaming_iterator.run_async_function") as mock_run_async, - patch("litellm.responses.streaming_iterator.executor"), + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async, + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] @@ -276,8 +277,8 @@ def test_handle_logging_failed_response_maps_type_field_to_400(): {"type": "invalid_request_error", "code": "invalid_prompt", "message": "bad prompt"} ) with ( - patch("litellm.responses.streaming_iterator.run_async_function") as mock_run_async, - patch("litellm.responses.streaming_iterator.executor"), + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async, + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] @@ -296,8 +297,8 @@ def test_handle_logging_failed_response_records_usage_and_cost(): iterator.completed_response = chunk iterator.logging_obj._response_cost_calculator.return_value = 0.0042 with ( - patch("litellm.responses.streaming_iterator.run_async_function"), - patch("litellm.responses.streaming_iterator.executor"), + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function"), + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), ): iterator._handle_logging_failed_response() combined_usage = iterator.logging_obj.model_call_details["combined_usage_object"] @@ -315,8 +316,8 @@ def test_handle_logging_failed_response_without_usage_skips_recording(): {"type": "server_error", "code": "server_error", "message": "boom"} ) with ( - patch("litellm.responses.streaming_iterator.run_async_function"), - patch("litellm.responses.streaming_iterator.executor"), + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function"), + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), ): iterator._handle_logging_failed_response() assert "combined_usage_object" not in iterator.logging_obj.model_call_details diff --git a/tests/test_litellm/router_strategy/test_litellm_encoder.py b/tests/test_litellm/router_strategy/test_litellm_encoder.py index ebd6efe309c..46187f52adb 100644 --- a/tests/test_litellm/router_strategy/test_litellm_encoder.py +++ b/tests/test_litellm/router_strategy/test_litellm_encoder.py @@ -1,5 +1,6 @@ """Tests for litellm/router_strategy/auto_router/litellm_encoder.py""" +import sys from typing import Any, Final import pytest @@ -7,6 +8,9 @@ import pytest import litellm from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS +if sys.version_info >= (3, 14): + pytest.skip("The semantic-router extra excludes Python 3.14", allow_module_level=True) + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py index ed593228621..314fd63c4cc 100644 --- a/tests/test_litellm/test_gpt_realtime_mode.py +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -1,8 +1,8 @@ import json -import typing from pathlib import Path import pytest +from typing_extensions import get_args, get_type_hints import litellm from litellm.types.utils import ModelInfoBase @@ -50,8 +50,8 @@ def _load_cost_map() -> dict: def test_realtime_is_a_valid_mode_literal(): - hints = typing.get_type_hints(ModelInfoBase, include_extras=False) - assert "realtime" in typing.get_args(hints["mode"]) + hints = get_type_hints(ModelInfoBase, include_extras=False) + assert "realtime" in get_args(hints["mode"]) @pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 7c2b9d0be05..5ac72462a34 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -16,6 +16,7 @@ from fastapi.testclient import TestClient import urllib.parse +from importlib import import_module from unittest.mock import MagicMock, patch import litellm @@ -2515,8 +2516,8 @@ def test_completion_forwards_store_and_prompt_cache_key_to_mcp_gateway(): prompt_cache_key are named params, so they no longer travel via **kwargs and must be forwarded explicitly like safety_identifier and service_tier. """ - with patch( - "litellm.responses.mcp.chat_completions_handler.acompletion_with_mcp" + with patch.object( + import_module("litellm.responses.mcp.chat_completions_handler"), "acompletion_with_mcp" ) as mock_mcp: result = litellm.completion( model="openai/gpt-4o", diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index 206207acb09..8fa9a18cf53 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -4,11 +4,15 @@ import re import shutil import subprocess import sys -import tomllib from pathlib import Path import pytest +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py" _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 0790b41c349..61d809f2ba0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2,6 +2,7 @@ import asyncio import json import logging import os +from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -53,6 +54,15 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_get_utc_datetime_returns_current_aware_utc_time() -> None: + before: Final = datetime.now(timezone.utc) + result: Final = litellm.utils.get_utc_datetime() + after: Final = datetime.now(timezone.utc) + + assert result.utcoffset() == timedelta(0) + assert before <= result <= after + + def test_usage_openai_cache_write_tokens_populates_both_names(): """OpenAI reports cache-write tokens as prompt_tokens_details.cache_write_tokens. The Usage constructor must expose it under both cache_write_tokens (canonical, diff --git a/uv.lock b/uv.lock index aa59ff7b229..67aa1014dd3 100644 --- a/uv.lock +++ b/uv.lock @@ -4453,6 +4453,7 @@ dev = [ { name = "responses" }, { name = "respx" }, { name = "ruff" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "types-boto3", extra = ["bedrock", "bedrock-agent", "bedrock-runtime", "kms", "s3", "sagemaker-runtime", "sts"] }, { name = "types-pyyaml" }, { name = "types-redis" }, @@ -4638,6 +4639,7 @@ dev = [ { name = "responses", specifier = "==0.26.0" }, { name = "respx", specifier = "==0.22.0" }, { name = "ruff", specifier = "==0.15.3" }, + { name = "tomli", marker = "python_full_version < '3.11'", specifier = "==2.4.1" }, { name = "types-boto3", extras = ["bedrock", "bedrock-agent", "bedrock-runtime", "kms", "s3", "sagemaker-runtime", "sts"], specifier = "==1.43.30" }, { name = "types-pyyaml", specifier = "==6.0.12.20250915" }, { name = "types-redis", specifier = "==4.6.0.20241004" }, From 729ec4c7b80cde9dd3527a55c49e9996581e1e75 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 14:18:43 -0700 Subject: [PATCH 05/35] fix: synchronize Prisma generator configuration --- litellm-proxy-extras/litellm_proxy_extras/schema.prisma | 1 + schema.prisma | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..2a2665f9731 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -5,6 +5,7 @@ datasource client { generator client { provider = "prisma-client-py" + recursive_type_depth = -1 binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"] } diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..2a2665f9731 100644 --- a/schema.prisma +++ b/schema.prisma @@ -5,6 +5,7 @@ datasource client { generator client { provider = "prisma-client-py" + recursive_type_depth = -1 binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"] } From cfcaaa03d6a16832d5346786ade10d1caf2bcf6a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 14:35:38 -0700 Subject: [PATCH 06/35] fix: resolve Python 3.14 OCR annotations and remaining matrix failures --- basedpyright-code-budget.json | 8 +-- litellm/llms/base_llm/ocr/transformation.py | 9 +-- test-quality-budget.json | 2 +- .../test_reducto_ocr_route.py | 4 ++ tests/test_litellm/caching/test_gcs_cache.py | 10 ++- .../caching/test_redis_cluster_cache.py | 5 +- .../caching/test_redis_connection_pool.py | 3 +- .../caching/test_redis_semantic_cache.py | 3 +- .../test_mcp_oauth_passthrough_tools.py | 7 ++- .../mcp_server/test_semantic_tool_filter.py | 26 ++++++++ .../test_responses_api_bridge_flag.py | 63 ++++++++++--------- .../test_responses_prompt_management.py | 28 ++++----- .../responses/test_responses_utils.py | 17 ++--- .../responses/test_text_format_conversion.py | 5 +- .../router_strategy/test_complexity_router.py | 20 ++++++ 15 files changed, 134 insertions(+), 76 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 609391d02ab..45dbb836552 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -45,7 +45,7 @@ "limit": 25 }, "reportInvalidTypeForm": { - "limit": 34 + "limit": 32 }, "reportInvalidTypeVarUse": { "limit": 2 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38350 + "limit": 38348 }, "reportUnknownParameterType": { - "limit": 19626 + "limit": 19624 }, "reportUnknownVariableType": { - "limit": 29890 + "limit": 29889 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 3b302837032..75306cd572a 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -2,6 +2,7 @@ Base OCR transformation configuration. """ +import builtins from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal @@ -93,8 +94,8 @@ class OCRResponse(LiteLLMPydanticObjectBase): document_annotation: Any | None = None usage_info: OCRUsageInfo | None = None content: str | None = None - tables: list[dict[str, object]] | None = None - keyValuePairs: list[dict[str, object]] | None = None + tables: list[dict[str, builtins.object]] | None = None + keyValuePairs: list[dict[str, builtins.object]] | None = None object: str = "ocr" model_config = {"extra": "allow"} @@ -102,11 +103,11 @@ class OCRResponse(LiteLLMPydanticObjectBase): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) - def set_provider_native_response(self, native_response: Mapping[str, object]) -> None: + def set_provider_native_response(self, native_response: Mapping[str, builtins.object]) -> None: """Keep the provider's own response payload alongside the normalized one.""" self._hidden_params[PROVIDER_NATIVE_RESPONSE_KEY] = native_response - def get_provider_native_response(self) -> Mapping[str, object] | None: + def get_provider_native_response(self) -> Mapping[str, builtins.object] | None: """The provider's own response payload, when `req_format=native` was requested.""" native_response: Final = self._hidden_params.get(PROVIDER_NATIVE_RESPONSE_KEY) return native_response if isinstance(native_response, dict) else None diff --git a/test-quality-budget.json b/test-quality-budget.json index 74a1cb307df..281c17748ee 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11103 + "limit": 11069 } } diff --git a/tests/proxy_unit_tests/test_reducto_ocr_route.py b/tests/proxy_unit_tests/test_reducto_ocr_route.py index dc658a74ee8..de0b4f55616 100644 --- a/tests/proxy_unit_tests/test_reducto_ocr_route.py +++ b/tests/proxy_unit_tests/test_reducto_ocr_route.py @@ -100,6 +100,8 @@ def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth): pages=[OCRPage(index=0, markdown="Proxy OCR")], model="parse-v3", usage_info=OCRUsageInfo(pages_processed=1, credits=1), + tables=[{"cells": [["Total", 42]], "page": 1}], + keyValuePairs=[{"key": "approved", "value": True, "confidence": 0.9}], ) data_uri = "data:application/pdf;base64,JVBERi0xLjQK" @@ -135,3 +137,5 @@ def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth): assert response_body["object"] == "ocr" assert response_body["usage_info"]["credits"] == 1 assert response_body["pages"][0]["markdown"] == "Proxy OCR" + assert response_body["tables"] == [{"cells": [["Total", 42]], "page": 1}] + assert response_body["keyValuePairs"] == [{"key": "approved", "value": True, "confidence": 0.9}] diff --git a/tests/test_litellm/caching/test_gcs_cache.py b/tests/test_litellm/caching/test_gcs_cache.py index 6222cf4760a..4dba0e76a57 100644 --- a/tests/test_litellm/caching/test_gcs_cache.py +++ b/tests/test_litellm/caching/test_gcs_cache.py @@ -1,3 +1,4 @@ +from importlib import import_module from unittest.mock import MagicMock, AsyncMock, patch import pytest @@ -13,15 +14,12 @@ def mock_gcs_dependencies(): mock_async_client = AsyncMock() with ( - patch( - "litellm.caching.gcs_cache._get_httpx_client", return_value=mock_sync_client + patch.object(import_module("litellm.caching.gcs_cache"), "_get_httpx_client", return_value=mock_sync_client ), - patch( - "litellm.caching.gcs_cache.get_async_httpx_client", + patch.object(import_module("litellm.caching.gcs_cache"), "get_async_httpx_client", return_value=mock_async_client, ), - patch( - "litellm.caching.gcs_cache.GCSBucketBase.sync_construct_request_headers", + patch.object(import_module("litellm.caching.gcs_cache").GCSBucketBase, "sync_construct_request_headers", return_value={}, ), ): diff --git a/tests/test_litellm/caching/test_redis_cluster_cache.py b/tests/test_litellm/caching/test_redis_cluster_cache.py index 372425aa9fa..0763b5110d5 100644 --- a/tests/test_litellm/caching/test_redis_cluster_cache.py +++ b/tests/test_litellm/caching/test_redis_cluster_cache.py @@ -1,3 +1,4 @@ +from importlib import import_module import json from unittest.mock import MagicMock, patch @@ -64,7 +65,7 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): @patch("litellm._redis.get_redis_connection_pool") @patch("litellm._redis.get_redis_client") -@patch("litellm.caching.redis_cache.RedisCache._setup_health_pings") +@patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings") def test_cache_init_creates_cluster_cache_from_env_var( mock_health, mock_get_client, mock_get_pool, monkeypatch ): @@ -91,7 +92,7 @@ def test_cache_init_creates_cluster_cache_from_env_var( @patch("litellm._redis.get_redis_connection_pool") @patch("litellm._redis.get_redis_client") -@patch("litellm.caching.redis_cache.RedisCache._setup_health_pings") +@patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings") def test_cache_init_creates_redis_cache_without_cluster_config( mock_health, mock_get_client, mock_get_pool, monkeypatch ): diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/test_litellm/caching/test_redis_connection_pool.py index 54dbe5361d7..74f7901cb7b 100644 --- a/tests/test_litellm/caching/test_redis_connection_pool.py +++ b/tests/test_litellm/caching/test_redis_connection_pool.py @@ -1,3 +1,4 @@ +from importlib import import_module from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -92,7 +93,7 @@ def _make_redis_cache(): patches = [ patch("litellm._redis.get_redis_client", return_value=mock_sync_client), patch("litellm._redis.get_redis_connection_pool", return_value=mock_async_pool), - patch("litellm.caching.redis_cache.RedisCache._setup_health_pings"), + patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings"), ] for p in patches: p.start() diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index be4367fd8bd..df990c43530 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -1,3 +1,4 @@ +from importlib import import_module import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -1453,7 +1454,7 @@ def test_cache_forwards_semantic_cache_embedding_timeout(): from litellm.caching.caching import Cache from litellm.types.caching import LiteLLMCacheType - with patch("litellm.caching.caching.RedisSemanticCache") as backend: + with patch.object(import_module("litellm.caching.caching"), "RedisSemanticCache") as backend: Cache( type=LiteLLMCacheType.REDIS_SEMANTIC, similarity_threshold=0.8, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 6d66748bf3f..9667224de98 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,11 +1,16 @@ """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import sys from unittest.mock import AsyncMock, MagicMock import httpx import pytest +if sys.version_info < (3, 11): + from exceptiongroup import ExceptionGroup + + from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, @@ -37,7 +42,7 @@ def test_extract_upstream_auth_failure_walks_exception_group(): inner = httpx.HTTPStatusError("401", request=response.request, response=response) try: - raise ExceptionGroup("wrapped", [inner]) # noqa: F821 (PEP 654, py3.11+) + raise ExceptionGroup("wrapped", [inner]) except Exception as group: result = _extract_upstream_auth_failure(group) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 054146d474d..bf0df17fafb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -18,6 +18,12 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 from mcp.types import Tool as MCPTool +requires_semantic_router = pytest.mark.skipif( + sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" +) + + +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_basic_filtering(): """ @@ -145,6 +151,7 @@ async def test_semantic_filter_basic_filtering(): print(f" Filter respects top_k parameter correctly") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_top_k_limiting(): """ @@ -328,6 +335,7 @@ async def test_semantic_filter_extract_user_query(): assert query3 == "" +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_triggers_on_completion(): """ @@ -453,6 +461,7 @@ async def test_semantic_filter_hook_skips_no_tools(): print("✅ Hook correctly skips requests without tools") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_preserves_native_tools(): """ @@ -584,6 +593,7 @@ async def test_semantic_filter_hook_preserves_native_tools(): ) +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_all_native_tools(): """ @@ -684,6 +694,7 @@ async def test_semantic_filter_hook_all_native_tools(): ) +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_responses_api_name_collision(): """ @@ -774,6 +785,7 @@ async def test_semantic_filter_hook_responses_api_name_collision(): print("✅ Responses API tool with MCP-matching name correctly classified as native") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): """ @@ -889,6 +901,7 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(allowed_tools)}, stats={stats}") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_narrows_mcp_reference_for_chat_completions(): """ @@ -1008,6 +1021,7 @@ async def test_semantic_filter_hook_narrows_mcp_reference_for_chat_completions() print(f"✅ chat completions: MCP reference preserved, narrowed to {allowed_tools}") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_zero_matches_exposes_all_tools_on_both_paths(): """ @@ -1126,6 +1140,7 @@ async def test_semantic_filter_hook_zero_matches_exposes_all_tools_on_both_paths print("✅ zero matches: both the MCP reference path and the plain tool path expose every tool") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_filters_expanded_tools_with_string_input(): """ @@ -1266,6 +1281,7 @@ async def test_semantic_filter_hook_expansion_skips_filter_when_disabled(): print("✅ Disabled filter: MCP reference untouched, no spurious stats") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_preserves_tool_order(): """ @@ -1651,6 +1667,7 @@ def _make_context_window_filter(state, top_k: int = 3): ) +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_fails_closed_on_query_time_context_window_error(): """ @@ -1682,6 +1699,7 @@ async def test_semantic_filter_fails_closed_on_query_time_context_window_error() print("✅ Query-time context window overflow fails closed") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_records_build_time_context_window_error(): """ @@ -1715,6 +1733,7 @@ async def test_semantic_filter_records_build_time_context_window_error(): print("✅ Build-time context window overflow is recorded and fails closed") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_fails_closed_on_context_window_error(): """ @@ -1762,6 +1781,7 @@ async def test_semantic_filter_hook_fails_closed_on_context_window_error(): print("✅ Hook fails closed with actionable 400 on context window overflow") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_fails_closed_on_expanded_tools_context_window_error(): """ @@ -1828,6 +1848,7 @@ async def test_semantic_filter_hook_fails_closed_on_expanded_tools_context_windo print("✅ Expansion path fails closed with actionable 400 on context window overflow") +@requires_semantic_router @pytest.mark.asyncio async def test_semantic_filter_hook_ignores_build_error_for_native_only_tools(): """ @@ -2018,6 +2039,7 @@ def _weather_tool(): ) +@requires_semantic_router @pytest.mark.asyncio async def test_filter_indexes_request_tools_when_startup_index_is_empty(): """ @@ -2042,6 +2064,7 @@ async def test_filter_indexes_request_tools_when_startup_index_is_empty(): print("✅ Empty startup index is built from authed request-time tools") +@requires_semantic_router @pytest.mark.asyncio async def test_filter_indexes_tools_missing_from_partial_index(): """ @@ -2070,6 +2093,7 @@ async def test_filter_indexes_tools_missing_from_partial_index(): print("✅ Partial startup index is completed from request-time tools, embedding each tool once") +@requires_semantic_router @pytest.mark.asyncio async def test_filter_fails_open_when_matches_are_not_in_available_tools(): """ @@ -2093,6 +2117,7 @@ async def test_filter_fails_open_when_matches_are_not_in_available_tools(): print("✅ Matches outside available_tools fail open instead of dropping every tool") +@requires_semantic_router @pytest.mark.asyncio async def test_request_time_context_window_error_is_request_scoped(): """ @@ -2129,6 +2154,7 @@ async def test_request_time_context_window_error_is_request_scoped(): print("✅ Request-time context window overflow is scoped to the request, not the worker") +@requires_semantic_router @pytest.mark.asyncio async def test_foreign_index_routes_cannot_displace_available_tools(): """ diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index d76fa59a888..57aa2a6baa2 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -6,6 +6,7 @@ Includes file_search emulation: the flag must be forwarded on inner aresponses calls so routed requests do not hit a custom api_base /v1/responses endpoint. """ +from importlib import import_module from unittest.mock import MagicMock, patch @@ -17,11 +18,11 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage class TestUseResponsesApiBridgeFlag: """Test that bridge opt-in forces the chat completions path.""" - @patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler" + @patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) def test_bridge_used_when_use_chat_completions_api_true( self, mock_get_config, mock_bridge_handler @@ -39,11 +40,11 @@ class TestUseResponsesApiBridgeFlag: mock_bridge_handler.assert_called_once() - @patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler" + @patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) def test_bridge_used_when_model_uses_chat_completions_prefix( self, mock_get_config, mock_bridge_handler @@ -62,9 +63,9 @@ class TestUseResponsesApiBridgeFlag: # Model string is provider-normalized after resolution; prefix only forces the bridge. assert mock_bridge_handler.call_args.kwargs["model"].endswith("my-custom-model") - @patch("litellm.responses.main.base_llm_http_handler.response_api_handler") - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object(import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler") + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) def test_native_forwarding_when_flag_absent( self, mock_get_config, mock_native_handler @@ -82,11 +83,11 @@ class TestUseResponsesApiBridgeFlag: mock_native_handler.assert_called_once() - @patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler" + @patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) def test_flag_does_not_leak_into_kwargs(self, mock_get_config, mock_bridge_handler): """use_chat_completions_api should be popped and not passed to the bridge handler.""" @@ -104,11 +105,11 @@ class TestUseResponsesApiBridgeFlag: all_kwargs = call_kwargs.kwargs if call_kwargs.kwargs else {} assert "use_chat_completions_api" not in all_kwargs - @patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler" + @patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) def test_bridge_used_when_provider_config_none( self, mock_get_config, mock_bridge_handler @@ -127,8 +128,8 @@ class TestUseResponsesApiBridgeFlag: mock_bridge_handler.assert_called_once() @patch("litellm.acompletion") - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) async def test_allowed_openai_params_forwarded_through_bridge( self, mock_get_config, mock_acompletion @@ -164,9 +165,9 @@ class TestUseResponsesApiBridgeFlag: "reasoning_effort" ] - @patch("litellm.responses.file_search.emulated_handler._call_aresponses") - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object(import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses") + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) async def test_bridge_flag_forwarded_to_file_search_emulation( self, mock_get_config, mock_call_aresponses @@ -206,12 +207,12 @@ class TestUseResponsesApiBridgeFlag: call_kwargs.get("use_chat_completions_api") is True ), "use_chat_completions_api should be forwarded to inner aresponses call" - @patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler" + @patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) @patch("litellm.vector_stores.main.asearch") - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) async def test_bridge_flag_prevents_native_responses_endpoint_call( self, mock_get_config, mock_asearch, mock_bridge_handler @@ -280,10 +281,10 @@ class TestUseResponsesApiBridgeFlag: assert result is not None assert result.id is not None - @patch("litellm.responses.main.base_llm_http_handler.response_api_handler") + @patch.object(import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler") @patch("litellm.vector_stores.main.asearch") - @patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + @patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" ) async def test_without_bridge_flag_uses_native_endpoint( self, mock_get_config, mock_asearch, mock_native_handler diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index 204b4d00f01..530afbd856b 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -13,6 +13,7 @@ Covers: I) async path propagates optional params to downstream handler """ +from importlib import import_module import asyncio from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -62,23 +63,20 @@ def _provider_by_model(model: str, **_: object) -> tuple[str, str, None, None]: def _patch_responses_dispatch(): """Patch everything after the prompt management block so tests stay unit-level.""" return [ - patch( - "litellm.responses.main.litellm.get_llm_provider", + patch.object( + import_module("litellm.responses.main").litellm, "get_llm_provider", side_effect=_provider_by_model, ), - patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler." - "LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", + patch.object( + import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "_should_use_litellm_mcp_gateway", return_value=False, ), - patch( - "litellm.responses.main.ProviderConfigManager" - ".get_provider_responses_api_config", + patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, ), - patch( - "litellm.responses.main.litellm_completion_transformation_handler" - ".response_api_handler", + patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler", return_value=MagicMock(), ), ] @@ -393,8 +391,8 @@ class TestResponsesAPIPromptManagement: patches = _patch_responses_dispatch() with ( - patch( - "litellm.responses.main.litellm.get_llm_provider", + patch.object( + import_module("litellm.responses.main").litellm, "get_llm_provider", side_effect=_provider_by_model, ), patches[1], @@ -599,8 +597,8 @@ def test_sync_prompt_swap_resolves_credentials_for_swapped_provider(monkeypatch: monkeypatch.setenv("XAI_API_KEY", "sk-xai-test") logging_obj = _make_logging_obj("gpt-4o-mini", [{"role": "user", "content": "hi"}]) - with patch( # test-quality-ok: handler boundary stub proves creds resolve for the swapped provider without network - "litellm.responses.main.base_llm_http_handler.response_api_handler", return_value=MagicMock() + with patch.object( # test-quality-ok: handler boundary stub proves creds resolve for the swapped provider without network + import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler", return_value=MagicMock() ) as mock_handler: litellm.responses(input="hi", model="xai/grok-4", prompt_id="p1", litellm_logging_obj=logging_obj) diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 6918ce0af13..cb6efa21036 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -1,3 +1,4 @@ +from importlib import import_module import base64 from unittest.mock import MagicMock, patch @@ -580,12 +581,12 @@ def test_responses_extra_body_forwarded_to_completion_transformation_handler(): so it was silently dropped. """ with ( - patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", + patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, ), - patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler", + patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler", ) as mock_handler, ): mock_handler.return_value = MagicMock() @@ -611,12 +612,12 @@ def test_responses_maps_reasoning_effort_from_litellm_params_to_reasoning(): that cannot set extra_body. """ with ( - patch( - "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", + patch.object( + import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, ), - patch( - "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler", + patch.object( + import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler", ) as mock_handler, ): mock_handler.return_value = MagicMock() diff --git a/tests/test_litellm/responses/test_text_format_conversion.py b/tests/test_litellm/responses/test_text_format_conversion.py index cca7748fd3a..c68ad16c4af 100644 --- a/tests/test_litellm/responses/test_text_format_conversion.py +++ b/tests/test_litellm/responses/test_text_format_conversion.py @@ -1,3 +1,4 @@ +from importlib import import_module import json import pytest @@ -148,8 +149,8 @@ class TestTextFormatConversion: incomplete_details=None, ) - with patch( - "litellm.responses.main.base_llm_http_handler.response_api_handler", + with patch.object( + import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler", new=mock_handler, ): litellm._turn_on_debug() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..c25b7f92270 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -6,6 +6,7 @@ Tests the rule-based complexity scoring and tier assignment logic. import asyncio import logging +import sys from typing import Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -46,6 +47,11 @@ from litellm.types.router import ( ) +requires_semantic_router = pytest.mark.skipif( + sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" +) + + @pytest.fixture def mock_router_instance(): """Create a mock LiteLLM Router instance.""" @@ -3413,6 +3419,7 @@ class FakeEmbeddingRouter: class TestSemanticKeywordTierRules: """Test embedding-based keyword_tier_rules matching.""" + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_match_routes_to_rule_tier(self, basic_config): """A paraphrase (no literal keyword) still routes via embedding similarity.""" @@ -3441,6 +3448,7 @@ class TestSemanticKeywordTierRules: assert result.model == "o1-preview" # REASONING via semantic match assert fake_router.async_embedding_calls, "expected an embedding call for the prompt" + @requires_semantic_router @pytest.mark.asyncio async def test_tier_matches_on_best_utterance_not_diluted_by_others(self, basic_config): """A tier with several keywords must match if the query is close to ANY of them, @@ -3475,6 +3483,7 @@ class TestSemanticKeywordTierRules: assert result is not None assert result.model == "o1-preview" # REASONING via best-utterance semantic match + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_embedding_call_carries_caller_metadata(self, basic_config): """The query embedding call must carry the caller's metadata/litellm_metadata @@ -3507,6 +3516,7 @@ class TestSemanticKeywordTierRules: assert fake_router.async_embedding_kwargs[0]["metadata"] == {**caller_metadata, **origin} assert fake_router.async_embedding_kwargs[0]["litellm_metadata"] == {**caller_litellm_metadata, **origin} + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_embedding_call_captures_request_body_in_proxy_server_request(self, basic_config): """The query embedding call must supply proxy_server_request so its request is logged. @@ -3540,6 +3550,7 @@ class TestSemanticKeywordTierRules: assert body["model"] == "fake-embed" assert body["input"] == ["roll out my k8s cluster"] + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_embedding_call_propagates_turn_off_message_logging(self, basic_config): """A caller's turn_off_message_logging must reach the query embedding call. @@ -3570,6 +3581,7 @@ class TestSemanticKeywordTierRules: assert fake_router.async_embedding_kwargs, "expected an embedding call for the prompt" assert fake_router.async_embedding_kwargs[0]["turn_off_message_logging"] is True + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_embedding_call_strips_budget_reservation(self, basic_config): """The embedding call must not carry the parent request's budget reservation. @@ -3623,6 +3635,7 @@ class TestSemanticKeywordTierRules: "budget_reservation": {"reserved_cost": 1.0}, } + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_routelayer_build_runs_off_event_loop(self, basic_config): """Building the SemanticRouter embeds route utterances via a synchronous provider @@ -3654,6 +3667,7 @@ class TestSemanticKeywordTierRules: # ...and none of it ran on the event-loop thread. assert all(tid != loop_thread_id for tid in fake_router.sync_embedding_thread_ids) + @requires_semantic_router @pytest.mark.asyncio async def test_concurrent_cold_start_builds_routelayer_once(self, basic_config): """Concurrent first requests must not each construct the route index (which would @@ -3717,6 +3731,7 @@ class TestSemanticKeywordTierRules: assert result is not None assert result.model == "gpt-4o-mini" # SIMPLE via scoring fallback + @requires_semantic_router @pytest.mark.asyncio async def test_route_embeddings_cached_across_requests(self, basic_config): """The route layer is built once and reused on subsequent requests.""" @@ -3932,6 +3947,7 @@ class TestKeywordOverrideEdgeCases: ) assert router._lexical_tier_override("deploy to k8s and reason step by step") is None + @requires_semantic_router def test_semantic_routelayer_requires_embedding_model(self, mock_router_instance, basic_config): """Building the route layer without an embedding model raises (defensive invariant).""" config = {**basic_config, "keyword_tier_rules": [{"keywords": ["k8s"], "tier": "REASONING"}]} @@ -3944,6 +3960,7 @@ class TestKeywordOverrideEdgeCases: with pytest.raises(ValueError, match="embedding_model is required"): router._get_or_create_semantic_routelayer() + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_override_maps_first_of_list(self, mock_router_instance, basic_config): """A list RouteChoice result maps to the first entry's tier.""" @@ -3953,6 +3970,7 @@ class TestKeywordOverrideEdgeCases: router._semantic_routelayer = _StubRouteLayer([RouteChoice(name="COMPLEX"), RouteChoice(name="SIMPLE")]) assert await router._semantic_tier_override("anything", {}) == ComplexityTier.COMPLEX + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_override_empty_list_returns_none(self, mock_router_instance, basic_config): """An empty list result falls through to scoring.""" @@ -3960,6 +3978,7 @@ class TestKeywordOverrideEdgeCases: router._semantic_routelayer = _StubRouteLayer([]) assert await router._semantic_tier_override("anything", {}) is None + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_override_unknown_route_name_returns_none(self, mock_router_instance, basic_config): """A matched route whose name is not a ComplexityTier is ignored.""" @@ -4027,6 +4046,7 @@ class TestRoutingDecisionCauseLogging: # A literal match must not be mislabelled as semantic. assert "cause=semantic_keyword_match" not in router_log_capture.text + @requires_semantic_router @pytest.mark.asyncio async def test_semantic_keyword_match_logs_its_cause(self, basic_config, router_log_capture): fake_router = FakeEmbeddingRouter() From cfb0b59e34045e5fe5ffe75680244543538cea9d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 14:44:13 -0700 Subject: [PATCH 07/35] fix: use typing backport in general upload validation --- .../proxy/openai_files_endpoints/general_upload_validation.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 9d450cb5b8d..8c59a520272 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose. from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.path_utils import safe_filename From 1988dfc4cc64b7b991f6dd696c99a74f9b4c1a07 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 15:02:04 -0700 Subject: [PATCH 08/35] fix: qualify list annotations for Python 3.14 runtime inspection --- litellm/proxy/client/models.py | 2 +- litellm/proxy/client/teams.py | 2 +- litellm/vector_stores/main.py | 10 +++++----- tests/test_litellm/vector_stores/test_main.py | 8 ++++++-- 4 files changed, 13 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 4b16087e15b..603597cc117 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -32,7 +32,7 @@ class ModelsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def list(self, return_request: bool = False) -> list[dict[str, Any]] | requests.Request: + def list(self, return_request: bool = False) -> builtins.list[dict[str, Any]] | requests.Request: """ Get the list of models supported by the server. diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index 105060e5ca9..54a6e869fef 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -40,7 +40,7 @@ class TeamsManagementClient: self, user_id: str | None = None, organization_id: str | None = None, - ) -> list[dict[str, Any]]: + ) -> builtins.list[dict[str, Any]]: """ List teams that the user belongs to. diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index cd576755f5f..662b4981304 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -39,7 +39,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() def mock_vector_store_search_response( - mock_results: list[VectorStoreSearchResult] | None = None, + mock_results: builtins.list[VectorStoreSearchResult] | None = None, ): """Mock response for vector store search""" if mock_results is None: @@ -93,7 +93,7 @@ def mock_vector_store_create_response( @client async def acreate( name: str | None = None, - file_ids: list[str] | None = None, + file_ids: builtins.list[str] | None = None, expires_after: dict | None = None, chunking_strategy: dict | None = None, metadata: dict[str, str] | None = None, @@ -157,7 +157,7 @@ async def acreate( @client def create( name: str | None = None, - file_ids: list[str] | None = None, + file_ids: builtins.list[str] | None = None, expires_after: dict | None = None, chunking_strategy: dict | None = None, metadata: dict[str, str] | None = None, @@ -270,7 +270,7 @@ def create( @client async def asearch( vector_store_id: str, - query: str | list[str], + query: str | builtins.list[str], filters: dict | None = None, max_num_results: int | None = None, ranking_options: dict | None = None, @@ -339,7 +339,7 @@ async def asearch( @client def search( vector_store_id: str, - query: str | list[str], + query: str | builtins.list[str], filters: dict | None = None, max_num_results: int | None = None, ranking_options: dict | None = None, diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py index d01e696906a..21f4165bd26 100644 --- a/tests/test_litellm/vector_stores/test_main.py +++ b/tests/test_litellm/vector_stores/test_main.py @@ -9,6 +9,8 @@ serialization trap). from unittest.mock import MagicMock, patch +import pytest + import litellm.vector_stores.main as vector_stores_main from litellm.vector_stores.main import search @@ -19,7 +21,8 @@ MOCK_SEARCH_RESPONSE = { } -def test_search_threads_router_to_handler(): +@pytest.mark.parametrize("query", ["q", ["q", "another question"]]) +def test_search_threads_router_to_handler(query: str | list[str]): """search() must pass its router param through to the HTTP handler""" mock_router = MagicMock() logger = MagicMock() @@ -37,7 +40,7 @@ def test_search_threads_router_to_handler(): ): response = search( vector_store_id="bkt:idx", - query="q", + query=query, custom_llm_provider="s3_vectors", router=mock_router, litellm_logging_obj=logger, @@ -46,6 +49,7 @@ def test_search_threads_router_to_handler(): assert response == MOCK_SEARCH_RESPONSE mock_handler.assert_called_once() assert mock_handler.call_args.kwargs["router"] is mock_router + assert mock_handler.call_args.kwargs["query"] == query def test_search_router_not_in_litellm_params(): From e03eb6961c990b716daa7d9b74814fd295ee8201 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:53:56 -0700 Subject: [PATCH 09/35] fix(deps): bump uvloop to 0.22.1 so the proxy boots on Python 3.14 --- pyproject.toml | 2 +- tests/test_litellm/proxy/test_proxy_cli.py | 8 +++ uv.lock | 66 +++++++++++++--------- 3 files changed, 48 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4b727a782c2..24ea1d3836b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", "granian>=2.7.4,<3.0", - "uvloop>=0.21.0,<1.0; sys_platform != 'win32'", + "uvloop>=0.22.1,<1.0; sys_platform != 'win32'", "fastapi>=0.136.3,<1.0", "starlette>=1.0.1,<2.0", "backoff>=2.2.1,<3.0", diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 7b3528f3a68..5f4cbcc5775 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -15,6 +15,8 @@ import urllib.parse as urlparse import uvicorn import yaml +from uvicorn.config import LOOP_FACTORIES +from uvicorn.importer import import_from_string from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server @@ -462,6 +464,12 @@ class TestProxyInitializationHelpers: with patch("sys.platform", "linux"): assert ProxyInitializationHelpers._get_loop_type() == "uvloop" + def test_selected_loop_factory_imports_on_this_interpreter(self): + loop_type = ProxyInitializationHelpers._get_loop_type() + if loop_type is None: + pytest.skip("uvicorn picks the loop itself on this platform") + assert callable(import_from_string(LOOP_FACTORIES[loop_type])) + @patch.dict(os.environ, {}, clear=True) def test_database_url_construction_with_special_characters(self): # Setup environment variables with special characters that need escaping diff --git a/uv.lock b/uv.lock index 67aa1014dd3..28cdd7fa711 100644 --- a/uv.lock +++ b/uv.lock @@ -4575,7 +4575,7 @@ requires-dist = [ { name = "tiktoken", specifier = ">=0.8.0,<1.0" }, { name = "tokenizers", specifier = ">=0.21.0,<1.0" }, { name = "uvicorn", marker = "extra == 'proxy'", specifier = ">=0.33.0,<1.0" }, - { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.21.0,<1.0" }, + { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] @@ -9947,34 +9947,46 @@ wheels = [ [[package]] name = "uvloop" -version = "0.21.0" +version = "0.22.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/af/c0/854216d09d33c543f12a44b393c402e89a920b1a0a7dc634c42de91b9cf6/uvloop-0.21.0.tar.gz", hash = "sha256:3bf12b0fda68447806a7ad847bfa591613177275d35b6724b1ee573faa3704e3", size = 2492741, upload-time = "2024-10-14T23:38:35.489Z" } +sdist = { url = "https://files.pythonhosted.org/packages/06/f0/18d39dbd1971d6d62c4629cc7fa67f74821b0dc1f5a77af43719de7936a7/uvloop-0.22.1.tar.gz", hash = "sha256:6c84bae345b9147082b17371e3dd5d42775bddce91f885499017f4607fdaf39f", size = 2443250, upload-time = "2025-10-16T22:17:19.342Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/76/44a55515e8c9505aa1420aebacf4dd82552e5e15691654894e90d0bd051a/uvloop-0.21.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:ec7e6b09a6fdded42403182ab6b832b71f4edaf7f37a9a0e371a01db5f0cb45f", size = 1442019, upload-time = "2024-10-14T23:37:20.068Z" }, - { url = "https://files.pythonhosted.org/packages/35/5a/62d5800358a78cc25c8a6c72ef8b10851bdb8cca22e14d9c74167b7f86da/uvloop-0.21.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:196274f2adb9689a289ad7d65700d37df0c0930fd8e4e743fa4834e850d7719d", size = 801898, upload-time = "2024-10-14T23:37:22.663Z" }, - { url = "https://files.pythonhosted.org/packages/f3/96/63695e0ebd7da6c741ccd4489b5947394435e198a1382349c17b1146bb97/uvloop-0.21.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f38b2e090258d051d68a5b14d1da7203a3c3677321cf32a95a6f4db4dd8b6f26", size = 3827735, upload-time = "2024-10-14T23:37:25.129Z" }, - { url = "https://files.pythonhosted.org/packages/61/e0/f0f8ec84979068ffae132c58c79af1de9cceeb664076beea86d941af1a30/uvloop-0.21.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:87c43e0f13022b998eb9b973b5e97200c8b90823454d4bc06ab33829e09fb9bb", size = 3825126, upload-time = "2024-10-14T23:37:27.59Z" }, - { url = "https://files.pythonhosted.org/packages/bf/fe/5e94a977d058a54a19df95f12f7161ab6e323ad49f4dabc28822eb2df7ea/uvloop-0.21.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:10d66943def5fcb6e7b37310eb6b5639fd2ccbc38df1177262b0640c3ca68c1f", size = 3705789, upload-time = "2024-10-14T23:37:29.385Z" }, - { url = "https://files.pythonhosted.org/packages/26/dd/c7179618e46092a77e036650c1f056041a028a35c4d76945089fcfc38af8/uvloop-0.21.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:67dd654b8ca23aed0a8e99010b4c34aca62f4b7fce88f39d452ed7622c94845c", size = 3800523, upload-time = "2024-10-14T23:37:32.048Z" }, - { url = "https://files.pythonhosted.org/packages/57/a7/4cf0334105c1160dd6819f3297f8700fda7fc30ab4f61fbf3e725acbc7cc/uvloop-0.21.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c0f3fa6200b3108919f8bdabb9a7f87f20e7097ea3c543754cabc7d717d95cf8", size = 1447410, upload-time = "2024-10-14T23:37:33.612Z" }, - { url = "https://files.pythonhosted.org/packages/8c/7c/1517b0bbc2dbe784b563d6ab54f2ef88c890fdad77232c98ed490aa07132/uvloop-0.21.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0878c2640cf341b269b7e128b1a5fed890adc4455513ca710d77d5e93aa6d6a0", size = 805476, upload-time = "2024-10-14T23:37:36.11Z" }, - { url = "https://files.pythonhosted.org/packages/ee/ea/0bfae1aceb82a503f358d8d2fa126ca9dbdb2ba9c7866974faec1cb5875c/uvloop-0.21.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b9fb766bb57b7388745d8bcc53a359b116b8a04c83a2288069809d2b3466c37e", size = 3960855, upload-time = "2024-10-14T23:37:37.683Z" }, - { url = "https://files.pythonhosted.org/packages/8a/ca/0864176a649838b838f36d44bf31c451597ab363b60dc9e09c9630619d41/uvloop-0.21.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8a375441696e2eda1c43c44ccb66e04d61ceeffcd76e4929e527b7fa401b90fb", size = 3973185, upload-time = "2024-10-14T23:37:40.226Z" }, - { url = "https://files.pythonhosted.org/packages/30/bf/08ad29979a936d63787ba47a540de2132169f140d54aa25bc8c3df3e67f4/uvloop-0.21.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:baa0e6291d91649c6ba4ed4b2f982f9fa165b5bbd50a9e203c416a2797bab3c6", size = 3820256, upload-time = "2024-10-14T23:37:42.839Z" }, - { url = "https://files.pythonhosted.org/packages/da/e2/5cf6ef37e3daf2f06e651aae5ea108ad30df3cb269102678b61ebf1fdf42/uvloop-0.21.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:4509360fcc4c3bd2c70d87573ad472de40c13387f5fda8cb58350a1d7475e58d", size = 3937323, upload-time = "2024-10-14T23:37:45.337Z" }, - { url = "https://files.pythonhosted.org/packages/8c/4c/03f93178830dc7ce8b4cdee1d36770d2f5ebb6f3d37d354e061eefc73545/uvloop-0.21.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:359ec2c888397b9e592a889c4d72ba3d6befba8b2bb01743f72fffbde663b59c", size = 1471284, upload-time = "2024-10-14T23:37:47.833Z" }, - { url = "https://files.pythonhosted.org/packages/43/3e/92c03f4d05e50f09251bd8b2b2b584a2a7f8fe600008bcc4523337abe676/uvloop-0.21.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7089d2dc73179ce5ac255bdf37c236a9f914b264825fdaacaded6990a7fb4c2", size = 821349, upload-time = "2024-10-14T23:37:50.149Z" }, - { url = "https://files.pythonhosted.org/packages/a6/ef/a02ec5da49909dbbfb1fd205a9a1ac4e88ea92dcae885e7c961847cd51e2/uvloop-0.21.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:baa4dcdbd9ae0a372f2167a207cd98c9f9a1ea1188a8a526431eef2f8116cc8d", size = 4580089, upload-time = "2024-10-14T23:37:51.703Z" }, - { url = "https://files.pythonhosted.org/packages/06/a7/b4e6a19925c900be9f98bec0a75e6e8f79bb53bdeb891916609ab3958967/uvloop-0.21.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:86975dca1c773a2c9864f4c52c5a55631038e387b47eaf56210f873887b6c8dc", size = 4693770, upload-time = "2024-10-14T23:37:54.122Z" }, - { url = "https://files.pythonhosted.org/packages/ce/0c/f07435a18a4b94ce6bd0677d8319cd3de61f3a9eeb1e5f8ab4e8b5edfcb3/uvloop-0.21.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:461d9ae6660fbbafedd07559c6a2e57cd553b34b0065b6550685f6653a98c1cb", size = 4451321, upload-time = "2024-10-14T23:37:55.766Z" }, - { url = "https://files.pythonhosted.org/packages/8f/eb/f7032be105877bcf924709c97b1bf3b90255b4ec251f9340cef912559f28/uvloop-0.21.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:183aef7c8730e54c9a3ee3227464daed66e37ba13040bb3f350bc2ddc040f22f", size = 4659022, upload-time = "2024-10-14T23:37:58.195Z" }, - { url = "https://files.pythonhosted.org/packages/3f/8d/2cbef610ca21539f0f36e2b34da49302029e7c9f09acef0b1c3b5839412b/uvloop-0.21.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:bfd55dfcc2a512316e65f16e503e9e450cab148ef11df4e4e679b5e8253a5281", size = 1468123, upload-time = "2024-10-14T23:38:00.688Z" }, - { url = "https://files.pythonhosted.org/packages/93/0d/b0038d5a469f94ed8f2b2fce2434a18396d8fbfb5da85a0a9781ebbdec14/uvloop-0.21.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:787ae31ad8a2856fc4e7c095341cccc7209bd657d0e71ad0dc2ea83c4a6fa8af", size = 819325, upload-time = "2024-10-14T23:38:02.309Z" }, - { url = "https://files.pythonhosted.org/packages/50/94/0a687f39e78c4c1e02e3272c6b2ccdb4e0085fda3b8352fecd0410ccf915/uvloop-0.21.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5ee4d4ef48036ff6e5cfffb09dd192c7a5027153948d85b8da7ff705065bacc6", size = 4582806, upload-time = "2024-10-14T23:38:04.711Z" }, - { url = "https://files.pythonhosted.org/packages/d2/19/f5b78616566ea68edd42aacaf645adbf71fbd83fc52281fba555dc27e3f1/uvloop-0.21.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f3df876acd7ec037a3d005b3ab85a7e4110422e4d9c1571d4fc89b0fc41b6816", size = 4701068, upload-time = "2024-10-14T23:38:06.385Z" }, - { url = "https://files.pythonhosted.org/packages/47/57/66f061ee118f413cd22a656de622925097170b9380b30091b78ea0c6ea75/uvloop-0.21.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bd53ecc9a0f3d87ab847503c2e1552b690362e005ab54e8a48ba97da3924c0dc", size = 4454428, upload-time = "2024-10-14T23:38:08.416Z" }, - { url = "https://files.pythonhosted.org/packages/63/9a/0962b05b308494e3202d3f794a6e85abe471fe3cafdbcf95c2e8c713aabd/uvloop-0.21.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a5c39f217ab3c663dc699c04cbd50c13813e31d917642d459fdcec07555cc553", size = 4660018, upload-time = "2024-10-14T23:38:10.888Z" }, + { url = "https://files.pythonhosted.org/packages/eb/14/ecceb239b65adaaf7fde510aa8bd534075695d1e5f8dadfa32b5723d9cfb/uvloop-0.22.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:ef6f0d4cc8a9fa1f6a910230cd53545d9a14479311e87e3cb225495952eb672c", size = 1343335, upload-time = "2025-10-16T22:16:11.43Z" }, + { url = "https://files.pythonhosted.org/packages/ba/ae/6f6f9af7f590b319c94532b9567409ba11f4fa71af1148cab1bf48a07048/uvloop-0.22.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:7cd375a12b71d33d46af85a3343b35d98e8116134ba404bd657b3b1d15988792", size = 742903, upload-time = "2025-10-16T22:16:12.979Z" }, + { url = "https://files.pythonhosted.org/packages/09/bd/3667151ad0702282a1f4d5d29288fce8a13c8b6858bf0978c219cd52b231/uvloop-0.22.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ac33ed96229b7790eb729702751c0e93ac5bc3bcf52ae9eccbff30da09194b86", size = 3648499, upload-time = "2025-10-16T22:16:14.451Z" }, + { url = "https://files.pythonhosted.org/packages/b3/f6/21657bb3beb5f8c57ce8be3b83f653dd7933c2fd00545ed1b092d464799a/uvloop-0.22.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:481c990a7abe2c6f4fc3d98781cc9426ebd7f03a9aaa7eb03d3bfc68ac2a46bd", size = 3700133, upload-time = "2025-10-16T22:16:16.272Z" }, + { url = "https://files.pythonhosted.org/packages/09/e0/604f61d004ded805f24974c87ddd8374ef675644f476f01f1df90e4cdf72/uvloop-0.22.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a592b043a47ad17911add5fbd087c76716d7c9ccc1d64ec9249ceafd735f03c2", size = 3512681, upload-time = "2025-10-16T22:16:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/bb/ce/8491fd370b0230deb5eac69c7aae35b3be527e25a911c0acdffb922dc1cd/uvloop-0.22.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:1489cf791aa7b6e8c8be1c5a080bae3a672791fcb4e9e12249b05862a2ca9cec", size = 3615261, upload-time = "2025-10-16T22:16:19.596Z" }, + { url = "https://files.pythonhosted.org/packages/c7/d5/69900f7883235562f1f50d8184bb7dd84a2fb61e9ec63f3782546fdbd057/uvloop-0.22.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c60ebcd36f7b240b30788554b6f0782454826a0ed765d8430652621b5de674b9", size = 1352420, upload-time = "2025-10-16T22:16:21.187Z" }, + { url = "https://files.pythonhosted.org/packages/a8/73/c4e271b3bce59724e291465cc936c37758886a4868787da0278b3b56b905/uvloop-0.22.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3b7f102bf3cb1995cfeaee9321105e8f5da76fdb104cdad8986f85461a1b7b77", size = 748677, upload-time = "2025-10-16T22:16:22.558Z" }, + { url = "https://files.pythonhosted.org/packages/86/94/9fb7fad2f824d25f8ecac0d70b94d0d48107ad5ece03769a9c543444f78a/uvloop-0.22.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:53c85520781d84a4b8b230e24a5af5b0778efdb39142b424990ff1ef7c48ba21", size = 3753819, upload-time = "2025-10-16T22:16:23.903Z" }, + { url = "https://files.pythonhosted.org/packages/74/4f/256aca690709e9b008b7108bc85fba619a2bc37c6d80743d18abad16ee09/uvloop-0.22.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56a2d1fae65fd82197cb8c53c367310b3eabe1bbb9fb5a04d28e3e3520e4f702", size = 3804529, upload-time = "2025-10-16T22:16:25.246Z" }, + { url = "https://files.pythonhosted.org/packages/7f/74/03c05ae4737e871923d21a76fe28b6aad57f5c03b6e6bfcfa5ad616013e4/uvloop-0.22.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:40631b049d5972c6755b06d0bfe8233b1bd9a8a6392d9d1c45c10b6f9e9b2733", size = 3621267, upload-time = "2025-10-16T22:16:26.819Z" }, + { url = "https://files.pythonhosted.org/packages/75/be/f8e590fe61d18b4a92070905497aec4c0e64ae1761498cad09023f3f4b3e/uvloop-0.22.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:535cc37b3a04f6cd2c1ef65fa1d370c9a35b6695df735fcff5427323f2cd5473", size = 3723105, upload-time = "2025-10-16T22:16:28.252Z" }, + { url = "https://files.pythonhosted.org/packages/3d/ff/7f72e8170be527b4977b033239a83a68d5c881cc4775fca255c677f7ac5d/uvloop-0.22.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fe94b4564e865d968414598eea1a6de60adba0c040ba4ed05ac1300de402cd42", size = 1359936, upload-time = "2025-10-16T22:16:29.436Z" }, + { url = "https://files.pythonhosted.org/packages/c3/c6/e5d433f88fd54d81ef4be58b2b7b0cea13c442454a1db703a1eea0db1a59/uvloop-0.22.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:51eb9bd88391483410daad430813d982010f9c9c89512321f5b60e2cddbdddd6", size = 752769, upload-time = "2025-10-16T22:16:30.493Z" }, + { url = "https://files.pythonhosted.org/packages/24/68/a6ac446820273e71aa762fa21cdcc09861edd3536ff47c5cd3b7afb10eeb/uvloop-0.22.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:700e674a166ca5778255e0e1dc4e9d79ab2acc57b9171b79e65feba7184b3370", size = 4317413, upload-time = "2025-10-16T22:16:31.644Z" }, + { url = "https://files.pythonhosted.org/packages/5f/6f/e62b4dfc7ad6518e7eff2516f680d02a0f6eb62c0c212e152ca708a0085e/uvloop-0.22.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b5b1ac819a3f946d3b2ee07f09149578ae76066d70b44df3fa990add49a82e4", size = 4426307, upload-time = "2025-10-16T22:16:32.917Z" }, + { url = "https://files.pythonhosted.org/packages/90/60/97362554ac21e20e81bcef1150cb2a7e4ffdaf8ea1e5b2e8bf7a053caa18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e047cc068570bac9866237739607d1313b9253c3051ad84738cbb095be0537b2", size = 4131970, upload-time = "2025-10-16T22:16:34.015Z" }, + { url = "https://files.pythonhosted.org/packages/99/39/6b3f7d234ba3964c428a6e40006340f53ba37993f46ed6e111c6e9141d18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:512fec6815e2dd45161054592441ef76c830eddaad55c8aa30952e6fe1ed07c0", size = 4296343, upload-time = "2025-10-16T22:16:35.149Z" }, + { url = "https://files.pythonhosted.org/packages/89/8c/182a2a593195bfd39842ea68ebc084e20c850806117213f5a299dfc513d9/uvloop-0.22.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:561577354eb94200d75aca23fbde86ee11be36b00e52a4eaf8f50fb0c86b7705", size = 1358611, upload-time = "2025-10-16T22:16:36.833Z" }, + { url = "https://files.pythonhosted.org/packages/d2/14/e301ee96a6dc95224b6f1162cd3312f6d1217be3907b79173b06785f2fe7/uvloop-0.22.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cdf5192ab3e674ca26da2eada35b288d2fa49fdd0f357a19f0e7c4e7d5077c8", size = 751811, upload-time = "2025-10-16T22:16:38.275Z" }, + { url = "https://files.pythonhosted.org/packages/b7/02/654426ce265ac19e2980bfd9ea6590ca96a56f10c76e63801a2df01c0486/uvloop-0.22.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6e2ea3d6190a2968f4a14a23019d3b16870dd2190cd69c8180f7c632d21de68d", size = 4288562, upload-time = "2025-10-16T22:16:39.375Z" }, + { url = "https://files.pythonhosted.org/packages/15/c0/0be24758891ef825f2065cd5db8741aaddabe3e248ee6acc5e8a80f04005/uvloop-0.22.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0530a5fbad9c9e4ee3f2b33b148c6a64d47bbad8000ea63704fa8260f4cf728e", size = 4366890, upload-time = "2025-10-16T22:16:40.547Z" }, + { url = "https://files.pythonhosted.org/packages/d2/53/8369e5219a5855869bcee5f4d317f6da0e2c669aecf0ef7d371e3d084449/uvloop-0.22.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bc5ef13bbc10b5335792360623cc378d52d7e62c2de64660616478c32cd0598e", size = 4119472, upload-time = "2025-10-16T22:16:41.694Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ba/d69adbe699b768f6b29a5eec7b47dd610bd17a69de51b251126a801369ea/uvloop-0.22.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1f38ec5e3f18c8a10ded09742f7fb8de0108796eb673f30ce7762ce1b8550cad", size = 4239051, upload-time = "2025-10-16T22:16:43.224Z" }, + { url = "https://files.pythonhosted.org/packages/90/cd/b62bdeaa429758aee8de8b00ac0dd26593a9de93d302bff3d21439e9791d/uvloop-0.22.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3879b88423ec7e97cd4eba2a443aa26ed4e59b45e6b76aabf13fe2f27023a142", size = 1362067, upload-time = "2025-10-16T22:16:44.503Z" }, + { url = "https://files.pythonhosted.org/packages/0d/f8/a132124dfda0777e489ca86732e85e69afcd1ff7686647000050ba670689/uvloop-0.22.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:4baa86acedf1d62115c1dc6ad1e17134476688f08c6efd8a2ab076e815665c74", size = 752423, upload-time = "2025-10-16T22:16:45.968Z" }, + { url = "https://files.pythonhosted.org/packages/a3/94/94af78c156f88da4b3a733773ad5ba0b164393e357cc4bd0ab2e2677a7d6/uvloop-0.22.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:297c27d8003520596236bdb2335e6b3f649480bd09e00d1e3a99144b691d2a35", size = 4272437, upload-time = "2025-10-16T22:16:47.451Z" }, + { url = "https://files.pythonhosted.org/packages/b5/35/60249e9fd07b32c665192cec7af29e06c7cd96fa1d08b84f012a56a0b38e/uvloop-0.22.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1955d5a1dd43198244d47664a5858082a3239766a839b2102a269aaff7a4e25", size = 4292101, upload-time = "2025-10-16T22:16:49.318Z" }, + { url = "https://files.pythonhosted.org/packages/02/62/67d382dfcb25d0a98ce73c11ed1a6fba5037a1a1d533dcbb7cab033a2636/uvloop-0.22.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b31dc2fccbd42adc73bc4e7cdbae4fc5086cf378979e53ca5d0301838c5682c6", size = 4114158, upload-time = "2025-10-16T22:16:50.517Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/f1171b4a882a5d13c8b7576f348acfe6074d72eaf52cccef752f748d4a9f/uvloop-0.22.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:93f617675b2d03af4e72a5333ef89450dfaa5321303ede6e67ba9c9d26878079", size = 4177360, upload-time = "2025-10-16T22:16:52.646Z" }, + { url = "https://files.pythonhosted.org/packages/79/7b/b01414f31546caf0919da80ad57cbfe24c56b151d12af68cee1b04922ca8/uvloop-0.22.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:37554f70528f60cad66945b885eb01f1bb514f132d92b6eeed1c90fd54ed6289", size = 1454790, upload-time = "2025-10-16T22:16:54.355Z" }, + { url = "https://files.pythonhosted.org/packages/d4/31/0bb232318dd838cad3fa8fb0c68c8b40e1145b32025581975e18b11fab40/uvloop-0.22.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b76324e2dc033a0b2f435f33eb88ff9913c156ef78e153fb210e03c13da746b3", size = 796783, upload-time = "2025-10-16T22:16:55.906Z" }, + { url = "https://files.pythonhosted.org/packages/42/38/c9b09f3271a7a723a5de69f8e237ab8e7803183131bc57c890db0b6bb872/uvloop-0.22.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:badb4d8e58ee08dad957002027830d5c3b06aea446a6a3744483c2b3b745345c", size = 4647548, upload-time = "2025-10-16T22:16:57.008Z" }, + { url = "https://files.pythonhosted.org/packages/c1/37/945b4ca0ac27e3dc4952642d4c900edd030b3da6c9634875af6e13ae80e5/uvloop-0.22.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b91328c72635f6f9e0282e4a57da7470c7350ab1c9f48546c0f2866205349d21", size = 4467065, upload-time = "2025-10-16T22:16:58.206Z" }, + { url = "https://files.pythonhosted.org/packages/97/cc/48d232f33d60e2e2e0b42f4e73455b146b76ebe216487e862700457fbf3c/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:daf620c2995d193449393d6c62131b3fbd40a63bf7b307a1527856ace637fe88", size = 4328384, upload-time = "2025-10-16T22:16:59.36Z" }, + { url = "https://files.pythonhosted.org/packages/e4/16/c1fd27e9549f3c4baf1dc9c20c456cd2f822dbf8de9f463824b0c0357e06/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6cde23eeda1a25c75b2e07d39970f3374105d5eafbaab2a4482be82f272d5a5e", size = 4296730, upload-time = "2025-10-16T22:17:00.744Z" }, ] [[package]] From f9f32b49a655b31b72014a5bf3e382065221bd38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:13:15 -0700 Subject: [PATCH 10/35] fix(tests): qualify traced function names on Python 3.10 --- tests/sdk_function_trace/profiler.py | 57 ++++++++++++++++++++--- tests/sdk_function_trace/test_profiler.py | 31 +++++++++++- 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py index c71c74ab0d3..44697e7b9f6 100644 --- a/tests/sdk_function_trace/profiler.py +++ b/tests/sdk_function_trace/profiler.py @@ -2,11 +2,12 @@ from __future__ import annotations import sys import threading -from collections.abc import Generator, Sequence +from collections.abc import Generator, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass +from functools import lru_cache from pathlib import Path -from types import CodeType, FrameType, FunctionType +from types import CodeType, FrameType, FunctionType, MappingProxyType from typing import Final @@ -27,11 +28,11 @@ class PythonProfiler: def __call__(self, frame: FrameType, event: str, _arg: object) -> None: if event != "call" or frame in self._seen_frames: return - function_name: Final = self.function_name(frame.f_code) + function_name: Final = self.function_name(frame) if function_name is None: return ancestors: Final = tuple( - name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None + name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor)) is not None ) self._seen_frames.add(frame) self.events.append( @@ -42,13 +43,57 @@ class PythonProfiler: ) ) - def function_name(self, code: CodeType) -> str | None: + def function_name(self, frame: FrameType) -> str | None: + code: Final = frame.f_code if self._source_root is None: return self._names_by_code.get(code) if not code.co_filename.startswith(self._source_root): return None relative: Final = code.co_filename.removeprefix(self._source_root) - return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + return f"{relative}:{code.co_firstlineno} {_qualified_name(frame)}" + + +def _qualified_name(frame: FrameType) -> str: + code: Final = frame.f_code + native: Final = getattr(code, "co_qualname", None) + if isinstance(native, str): + return native + module_name: Final = frame.f_globals.get("__name__") + if not isinstance(module_name, str): + return code.co_name + return _module_qualnames(module_name).get(code, code.co_name) + + +@lru_cache(maxsize=None) +def _module_qualnames(module_name: str) -> Mapping[CodeType, str]: + module: Final = sys.modules.get(module_name) + if module is None: + return MappingProxyType({}) + return MappingProxyType(dict(_declared_functions(vars(module), frozenset()))) + + +def _declared_functions(namespace: Mapping[str, object], visited: frozenset[int]) -> Iterator[tuple[CodeType, str]]: + for attribute in tuple(namespace.values()): + for value in _accessors(attribute): + if isinstance(value, FunctionType): + yield from ((wrapped.__code__, wrapped.__qualname__) for wrapped in _unwrapped(value)) + elif isinstance(value, type) and id(value) not in visited: + yield from _declared_functions(dict(vars(value)), visited | {id(value)}) + + +def _unwrapped(function: FunctionType) -> Iterator[FunctionType]: + yield function + inner: Final = getattr(function, "__wrapped__", None) + if isinstance(inner, FunctionType): + yield from _unwrapped(inner) + + +def _accessors(value: object) -> tuple[object, ...]: + if isinstance(value, (staticmethod, classmethod)): + return (value.__func__,) + if isinstance(value, property): + return tuple(accessor for accessor in (value.fget, value.fset, value.fdel) if accessor is not None) + return (value,) def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: diff --git a/tests/sdk_function_trace/test_profiler.py b/tests/sdk_function_trace/test_profiler.py index 10a266fb1e8..21a1ab134b2 100644 --- a/tests/sdk_function_trace/test_profiler.py +++ b/tests/sdk_function_trace/test_profiler.py @@ -2,9 +2,11 @@ from __future__ import annotations import asyncio import sys +from collections.abc import Callable +from functools import wraps from pathlib import Path from types import FunctionType -from typing import Final, cast +from typing import Final, ParamSpec, TypeVar, cast import pytest @@ -14,7 +16,18 @@ from tests.sdk_function_trace import ( TraceStep, assert_function_trace_parity, ) -from tests.sdk_function_trace.profiler import profile_python +from tests.sdk_function_trace.profiler import _module_qualnames, profile_python + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +def _passthrough(function: Callable[_P, _T]) -> Callable[_P, _T]: + @wraps(function) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T: + return function(*args, **kwargs) + + return wrapper class First: @@ -29,6 +42,20 @@ class Second: return None +class Decorated: + @_passthrough + def call(self) -> None: + return None + + +def test_source_profiler_qualifies_decorated_methods_by_class() -> None: + with profile_python(source_root=Path(__file__).parent) as profiler: + Decorated().call() + + assert any(event.function.endswith(" Decorated.call") for event in profiler.events) + assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call" + + def test_profiler_matches_code_objects_and_keeps_repeated_calls() -> None: with profile_python((First.run,)) as profiler: Second.run() From 63579f1e355dabb4df2b31f6c654d1d6f7b0ad44 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:12:57 +0000 Subject: [PATCH 11/35] fix(spend-tracking): keep batch spend keys joinable after v1.99 provenance gate Batch cost attribution and the legacy queue endpoint already store the VerificationToken hash in user_api_key, but omitted user_api_key_hash. Since v1.99 the spend-log writer re-hashes any key without that provenance flag, so DailyUserSpend.api_key no longer joins VerificationToken and Usage shows key-hash-... rows with null api_key_alias / user_email. Co-authored-by: Mateo Wang --- .../proxy/common_utils/check_batch_cost.py | 1 + litellm/proxy/proxy_server.py | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 43 +++++++++++++++++++ .../test_spend_tracking_utils.py | 36 ++++++++++++++++ 4 files changed, 81 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 354a6ed2fd0..e6f00877a26 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -161,6 +161,7 @@ class CheckBatchCost: metadata: dict[str, object] = { "user_api_key_user_id": job.created_by, "user_api_key": api_key, + "user_api_key_hash": api_key, "user_api_key_team_id": team_id, **(await self._get_user_info(batch_id, job.created_by)), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..222e79f7ca7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15197,6 +15197,7 @@ async def async_queue_request( # extra_body); see above for the same guard upstream. data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) _headers: Final = _safe_get_request_headers(request).copy() _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff5e8f89d64..9a6ab08e9b6 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2445,6 +2445,7 @@ class TestBatchCostAttribution: metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") assert metadata["user_api_key"] == "hash-alice" + assert metadata["user_api_key_hash"] == "hash-alice" assert metadata["user_api_key_user_id"] == "alice" assert metadata["user_api_key_team_id"] == "team-alpha" assert metadata["user_api_key_alias"] == "prod-key" @@ -2553,6 +2554,48 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self): + """ + CheckBatchCost stores the VerificationToken hash on the managed object. The + spend-log writer must receive matching user_api_key_hash provenance so it + does not re-hash that value; otherwise DailyUserSpend.api_key no longer joins + VerificationToken and Usage shows key-hash-... with a null alias/email. + """ + from datetime import datetime, timezone + from types import SimpleNamespace + + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.proxy.utils import hash_token + + token_hash = hash_token("sk-batch-creator-key") + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key"), + user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None), + ) + metadata = await instance._build_creator_attribution_metadata( + self._job(api_key=token_hash), "batch-1" + ) + + assert metadata["user_api_key"] == token_hash + assert metadata["user_api_key_hash"] == token_hash + + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o", + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": metadata}, + }, + response_obj={ + "id": "batch_123", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + assert payload["api_key"] == token_hash + assert payload["api_key"] != hash_token(token_hash) + class TestPollPageStarvation: """LIT-5462 regression: a row that can never be costed used to keep its slot in the diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9e5917637a8..bea1f8e2d6c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2755,6 +2755,42 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): assert meta["user_api_key"] == hash_token(already_hashed) +def test_get_logging_payload_batch_attribution_keeps_verification_token_hash(): + """ + Batch cost rebuilds metadata with the managed object's already-hashed api_key. + That hash must land in SpendLogs.api_key unchanged so Usage/CloudZero can join + LiteLLM_VerificationToken for api_key_alias and user_email. Regression: without + user_api_key_hash provenance, v1.99+ re-hashed the token and broke the join. + """ + token_hash = hash_token("sk-batch-creator-key") + kwargs = { + "model": "gpt-4o", + "call_type": "aretrieve_batch", + "litellm_params": { + "metadata": { + "user_api_key": token_hash, + "user_api_key_hash": token_hash, + "user_api_key_alias": "batch-creator", + "user_api_key_user_id": "alice", + "user_api_key_user_email": "alice@example.com", + "user_api_key_team_id": "team-1", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj={"id": "batch_123", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == token_hash + assert payload["api_key"] != hash_token(token_hash) + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] == token_hash + assert parsed_meta["user_api_key_alias"] == "batch-creator" + + def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): already_hashed = hash_token("sk-some-key") different_hash = hash_token("sk-other-key") From e07d58a4efd09e5dc6e044bdc089ba7e626489df Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:21:25 +0000 Subject: [PATCH 12/35] fix(usage): recover aliases for v1.99 double-hashed spend keys Callback log replay also omitted user_api_key_hash, so it could double-hash spend rows the same way batch costing did. On the read path, Usage key metadata now reverse-hashes orphaned DailyUserSpend.api_key values against VerificationToken and falls back to SpendLogs metadata so historical dirty rows show their api_key_alias again instead of key-hash-... Co-authored-by: Mateo Wang --- .../callback_logs_endpoints.py | 4 +- .../common_daily_activity.py | 163 +++++++++++++++++- .../test_callback_logs_endpoints.py | 1 + .../test_common_daily_activity.py | 62 +++++++ 4 files changed, 226 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index 66057f0dc16..cecadc03d71 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -88,8 +88,10 @@ class CallbackLogsReplayer: ) metadata: Final[dict[str, Any]] = payload.get("metadata") or {} + user_api_key_hash: Final = metadata.get("user_api_key_hash") litellm_metadata: Final[dict[str, Any]] = { - "user_api_key": metadata.get("user_api_key_hash"), + "user_api_key": user_api_key_hash, + "user_api_key_hash": user_api_key_hash, "user_api_key_alias": metadata.get("user_api_key_alias"), "user_api_key_user_id": metadata.get("user_api_key_user_id"), "user_api_key_team_id": metadata.get("user_api_key_team_id"), diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 91cd80b3c81..c1fec88d43d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -10,9 +10,10 @@ from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, hash_token from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, @@ -115,6 +116,25 @@ class _KeyMetadataDict(TypedDict, total=False): team_id: str | None +# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key +# values cannot pull an unbounded VerificationToken table into memory. +_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 + +_SPEND_LOGS_KEY_METADATA_SQL: Final = """ +SELECT DISTINCT ON (api_key) + api_key, + metadata->>'user_api_key_alias' AS key_alias, + metadata->>'user_api_key_team_id' AS team_id +FROM "LiteLLM_SpendLogs" +WHERE api_key = ANY($1::text[]) + AND ( + NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL + OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL + ) +ORDER BY api_key, "startTime" DESC NULLS LAST +""" + + _WhereValue = str | dict[str, object] @@ -439,6 +459,136 @@ def update_breakdown_metrics( return breakdown +class _TokenAliasRecord(Protocol): + @property + def token(self) -> str: ... + + @property + def key_alias(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + +def _token_digest_metadata( + records: Sequence[_TokenAliasRecord], + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + return { + digested: {"key_alias": record.key_alias, "team_id": record.team_id} + for record in records + for digested in (hash_token(record.token),) + if digested in wanted + } + + +async def _reverse_hash_active_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository( + prisma_client + ).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against active keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(active_records, wanted) + + +async def _reverse_hash_deleted_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( + take=_MAX_DOUBLE_HASH_TOKEN_SCAN, + order={"deleted_at": "desc"}, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(deleted_records, wanted) + + +async def _reverse_hash_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) + still_wanted: Final = wanted - frozenset(from_active) + if not still_wanted: + return from_active + return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))} + + +async def _spend_logs_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + spend_log_rows: Final = await prisma_client.db.query_raw( + _SPEND_LOGS_KEY_METADATA_SQL, + list(wanted), + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed SpendLogs metadata recovery for %d missing keys: %s", + len(wanted), + e, + ) + return {} + + if not isinstance(spend_log_rows, list): + return {} + + return { + row["api_key"]: { + "key_alias": row.get("key_alias"), + "team_id": row.get("team_id"), + } + for row in spend_log_rows + if isinstance(row, dict) + and isinstance(row.get("api_key"), str) + and row["api_key"] in wanted + } + + +async def _recover_double_hashed_key_metadata( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + """ + Recover key_alias/team_id for DailyUserSpend.api_key values that were + double-hashed by the v1.99 spend-log provenance gate. + + Those rows store hash(VerificationToken.token) instead of the token, so the + exact join misses. Prefer a bounded reverse-hash against active/deleted + tokens; fall back to the alias/team stamped into SpendLogs metadata (which + stayed correct even when api_key did not). + """ + sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) + if not sha_missing: + return {} + + from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) + still_missing: Final = sha_missing - frozenset(from_tokens) + if not still_missing: + return from_tokens + + return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + + async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], @@ -446,7 +596,8 @@ async def get_api_key_metadata( """Get api key metadata, falling back to deleted keys table for keys not found in active table. This ensures that key_alias and team_id are preserved in historical activity logs - even after a key is deleted or regenerated. + even after a key is deleted or regenerated. Also recovers aliases for api_key + values that were double-hashed by the v1.99 spend-log provenance gate. """ key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} @@ -479,7 +630,13 @@ async def get_api_key_metadata( e, ) - return result + still_missing: Final = api_keys - set(result.keys()) + if not still_missing: + return result + return { + **result, + **(await _recover_double_hashed_key_metadata(prisma_client, still_missing)), + } def _adjust_dates_for_timezone( diff --git a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py index 40e89329b8d..590d63fd868 100644 --- a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py +++ b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py @@ -61,6 +61,7 @@ def test_build_logging_obj_seeds_model_call_details(): # Metadata is mapped to the keys the cost-tracking callback reads. md = details["litellm_params"]["metadata"] assert md["user_api_key"] == "rust-gateway-test-key" + assert md["user_api_key_hash"] == "rust-gateway-test-key" assert md["user_api_key_user_id"] == "user-cb-logs-test" assert md["user_api_key_team_id"] == "team-cb-logs-test" diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index a258127acff..78752495f25 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -454,6 +454,68 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec assert result["old-key-hash"]["team_id"] == "latest-team" +@pytest.mark.asyncio +async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(): + """ + v1.99 spend logging re-hashed already-hashed api_key values when provenance was + missing. Usage joins DailyUserSpend.api_key to VerificationToken.token, so those + rows looked like key-hash-... with a null alias. Reverse-hash recovery must map + hash(token) back to the key's alias for historical dirty spend. + """ + from litellm.proxy.utils import hash_token + + token = "a" * 64 + double_hashed = hash_token(token) + mock_prisma = MagicMock() + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + side_effect=[ + [], # exact join miss + [SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")], + ] + ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={double_hashed}, + ) + + assert result[double_hashed]["key_alias"] == "batch-worker" + assert result[double_hashed]["team_id"] == "team-1" + mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): + """When the token tables cannot reverse-hash the dirty key, use SpendLogs metadata.""" + from litellm.proxy.utils import hash_token + + double_hashed = hash_token("b" * 64) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "api_key": double_hashed, + "key_alias": "from-spend-log", + "team_id": "team-spend", + } + ] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={double_hashed}, + ) + + assert result[double_hashed]["key_alias"] == "from-spend-log" + assert result[double_hashed]["team_id"] == "team-spend" + mock_prisma.db.query_raw.assert_called_once() + + @pytest.mark.asyncio async def test_tag_daily_activity_metadata_totals_not_zero(): """Test that tag daily activity returns correct metadata totals. From 9dae07175cd4c3c8a795d9e198903259cb05b78b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:24:08 +0000 Subject: [PATCH 13/35] fix(spend): share double-hash key alias recovery with CloudZero and Focus Extract the Usage reverse-hash / SpendLogs alias recovery into a shared helper and apply it when CloudZero and Focus export DailyUserSpend rows, so BI pulls get api_key_alias back for historical v1.99 double-hashed keys instead of null. Co-authored-by: Mateo Wang --- litellm/integrations/cloudzero/database.py | 11 +- litellm/integrations/focus/database.py | 9 +- .../common_daily_activity.py | 157 +----------- .../spend_tracking/key_metadata_recovery.py | 223 ++++++++++++++++++ .../test_key_metadata_recovery.py | 64 +++++ 5 files changed, 308 insertions(+), 156 deletions(-) create mode 100644 litellm/proxy/spend_tracking/key_metadata_recovery.py create mode 100644 tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index b050ee8e1ed..8fedd4edac4 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -94,8 +94,13 @@ class LiteLLMDatabase: try: db_response: Final = await client.db.query_raw(query, *params) - # Convert the response to polars DataFrame with full schema inference - # This prevents schema mismatch errors when data types vary across rows - return pl.DataFrame(db_response, infer_schema_length=None) + from litellm.proxy.spend_tracking.key_metadata_recovery import ( + fill_missing_api_key_aliases, + ) + + # v1.99 double-hashed DailyUserSpend.api_key values miss the + # VerificationToken join above; recover alias/team for those rows. + recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + return pl.DataFrame(list(recovered_rows), infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 815c38b9e9c..db9849bbbc9 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -96,7 +96,14 @@ class FocusLiteLLMDatabase: try: db_response: Final = await client.db.query_raw(query, *query_params) - return pl.DataFrame(db_response, infer_schema_length=None) + from litellm.proxy.spend_tracking.key_metadata_recovery import ( + fill_missing_api_key_aliases, + ) + + # v1.99 double-hashed DailyUserSpend.api_key values miss the + # VerificationToken join above; recover alias/team for those rows. + recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + return pl.DataFrame(list(recovered_rows), infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index c1fec88d43d..f687cedeee9 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -10,10 +10,12 @@ from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY -from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.spend_tracking.key_metadata_recovery import ( + recover_double_hashed_key_metadata, +) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled -from litellm.proxy.utils import PrismaClient, hash_token +from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, @@ -116,25 +118,6 @@ class _KeyMetadataDict(TypedDict, total=False): team_id: str | None -# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key -# values cannot pull an unbounded VerificationToken table into memory. -_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 - -_SPEND_LOGS_KEY_METADATA_SQL: Final = """ -SELECT DISTINCT ON (api_key) - api_key, - metadata->>'user_api_key_alias' AS key_alias, - metadata->>'user_api_key_team_id' AS team_id -FROM "LiteLLM_SpendLogs" -WHERE api_key = ANY($1::text[]) - AND ( - NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL - OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL - ) -ORDER BY api_key, "startTime" DESC NULLS LAST -""" - - _WhereValue = str | dict[str, object] @@ -459,136 +442,6 @@ def update_breakdown_metrics( return breakdown -class _TokenAliasRecord(Protocol): - @property - def token(self) -> str: ... - - @property - def key_alias(self) -> str | None: ... - - @property - def team_id(self) -> str | None: ... - - -def _token_digest_metadata( - records: Sequence[_TokenAliasRecord], - wanted: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - return { - digested: {"key_alias": record.key_alias, "team_id": record.team_id} - for record in records - for digested in (hash_token(record.token),) - if digested in wanted - } - - -async def _reverse_hash_active_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - try: - active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN) - except Exception as e: - verbose_proxy_logger.warning( - "Failed reverse-hash recovery against active keys for %d missing keys: %s", - len(wanted), - e, - ) - return {} - return _token_digest_metadata(active_records, wanted) - - -async def _reverse_hash_deleted_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - try: - deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository( - prisma_client - ).table.find_many( - take=_MAX_DOUBLE_HASH_TOKEN_SCAN, - order={"deleted_at": "desc"}, - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", - len(wanted), - e, - ) - return {} - return _token_digest_metadata(deleted_records, wanted) - - -async def _reverse_hash_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) - still_wanted: Final = wanted - frozenset(from_active) - if not still_wanted: - return from_active - return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))} - - -async def _spend_logs_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - try: - spend_log_rows: Final = await prisma_client.db.query_raw( - _SPEND_LOGS_KEY_METADATA_SQL, - list(wanted), - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed SpendLogs metadata recovery for %d missing keys: %s", - len(wanted), - e, - ) - return {} - - if not isinstance(spend_log_rows, list): - return {} - - return { - row["api_key"]: { - "key_alias": row.get("key_alias"), - "team_id": row.get("team_id"), - } - for row in spend_log_rows - if isinstance(row, dict) - and isinstance(row.get("api_key"), str) - and row["api_key"] in wanted - } - - -async def _recover_double_hashed_key_metadata( - prisma_client: PrismaClient, - missing_keys: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: - """ - Recover key_alias/team_id for DailyUserSpend.api_key values that were - double-hashed by the v1.99 spend-log provenance gate. - - Those rows store hash(VerificationToken.token) instead of the token, so the - exact join misses. Prefer a bounded reverse-hash against active/deleted - tokens; fall back to the alias/team stamped into SpendLogs metadata (which - stayed correct even when api_key did not). - """ - sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) - if not sha_missing: - return {} - - from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) - still_missing: Final = sha_missing - frozenset(from_tokens) - if not still_missing: - return from_tokens - - return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} - - async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], @@ -635,7 +488,7 @@ async def get_api_key_metadata( return result return { **result, - **(await _recover_double_hashed_key_metadata(prisma_client, still_missing)), + **(await recover_double_hashed_key_metadata(prisma_client, still_missing)), } diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py new file mode 100644 index 00000000000..4d85357b7a9 --- /dev/null +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -0,0 +1,223 @@ +from collections.abc import Mapping, Sequence, Set as AbstractSet +from typing import Final, Protocol + +from typing_extensions import TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash +from litellm.proxy.utils import PrismaClient, hash_token +from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) + +# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key +# values cannot pull an unbounded VerificationToken table into memory. +_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 + +_SPEND_LOGS_KEY_METADATA_SQL: Final = """ +SELECT DISTINCT ON (api_key) + api_key, + metadata->>'user_api_key_alias' AS key_alias, + metadata->>'user_api_key_team_id' AS team_id +FROM "LiteLLM_SpendLogs" +WHERE api_key = ANY($1::text[]) + AND ( + NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL + OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL + ) +ORDER BY api_key, "startTime" DESC NULLS LAST +""" + + +class KeyMetadataDict(TypedDict, total=False): + key_alias: str | None + team_id: str | None + + +class _TokenAliasRecord(Protocol): + @property + def token(self) -> str: ... + + @property + def key_alias(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + +def _token_digest_metadata( + records: Sequence[_TokenAliasRecord], + wanted: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + return { + digested: {"key_alias": record.key_alias, "team_id": record.team_id} + for record in records + for digested in (hash_token(record.token),) + if digested in wanted + } + + +async def _reverse_hash_active_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + try: + active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository( + prisma_client + ).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against active keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(active_records, wanted) + + +async def _reverse_hash_deleted_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + try: + deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( + take=_MAX_DOUBLE_HASH_TOKEN_SCAN, + order={"deleted_at": "desc"}, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(deleted_records, wanted) + + +async def _reverse_hash_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) + still_wanted: Final = wanted - frozenset(from_active) + if not still_wanted: + return from_active + return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))} + + +async def _spend_logs_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + try: + spend_log_rows: Final = await prisma_client.db.query_raw( + _SPEND_LOGS_KEY_METADATA_SQL, + list(wanted), + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed SpendLogs metadata recovery for %d missing keys: %s", + len(wanted), + e, + ) + return {} + + if not isinstance(spend_log_rows, list): + return {} + + return { + row["api_key"]: { + "key_alias": row.get("key_alias"), + "team_id": row.get("team_id"), + } + for row in spend_log_rows + if isinstance(row, dict) + and isinstance(row.get("api_key"), str) + and row["api_key"] in wanted + } + + +async def recover_double_hashed_key_metadata( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], +) -> dict[str, KeyMetadataDict]: + """ + Recover key_alias/team_id for DailyUserSpend.api_key values that were + double-hashed by the v1.99 spend-log provenance gate. + + Those rows store hash(VerificationToken.token) instead of the token, so the + exact join misses. Prefer a bounded reverse-hash against active/deleted + tokens; fall back to the alias/team stamped into SpendLogs metadata (which + stayed correct even when api_key did not). + """ + sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) + if not sha_missing: + return {} + + from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) + still_missing: Final = sha_missing - frozenset(from_tokens) + if not still_missing: + return from_tokens + + return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + + +def _row_with_recovered_alias( + row: Mapping[str, object], + recovered: Mapping[str, KeyMetadataDict], + *, + api_key_field: str, + alias_field: str, + team_id_field: str, +) -> Mapping[str, object]: + api_key: Final = row.get(api_key_field) + if not isinstance(api_key, str) or api_key not in recovered: + return row + meta: Final = recovered[api_key] + return { + **row, + alias_field: meta.get("key_alias") or row.get(alias_field), + team_id_field: meta.get("team_id") or row.get(team_id_field), + } + + +async def fill_missing_api_key_aliases( + prisma_client: PrismaClient, + rows: Sequence[Mapping[str, object]], + *, + api_key_field: str = "api_key", + alias_field: str = "api_key_alias", + team_id_field: str = "team_id", +) -> tuple[Mapping[str, object], ...]: + """ + Fill null api_key_alias / team_id on export rows whose api_key was double-hashed. + + Used by CloudZero and Focus, which join DailyUserSpend.api_key to + VerificationToken.token and otherwise export null aliases for those rows. + """ + missing_keys: Final = frozenset( + key + for row in rows + for key in (row.get(api_key_field),) + if isinstance(key, str) and key and row.get(alias_field) in (None, "") + ) + if not missing_keys: + return tuple(rows) + + recovered: Final = await recover_double_hashed_key_metadata(prisma_client, missing_keys) + if not recovered: + return tuple(rows) + + return tuple( + _row_with_recovered_alias( + row, + recovered, + api_key_field=api_key_field, + alias_field=alias_field, + team_id_field=team_id_field, + ) + for row in rows + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py new file mode 100644 index 00000000000..6b12676e82b --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -0,0 +1,64 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.spend_tracking.key_metadata_recovery import ( + fill_missing_api_key_aliases, + recover_double_hashed_key_metadata, +) +from litellm.proxy.utils import hash_token + + +@pytest.mark.asyncio +async def test_recover_double_hashed_key_metadata_via_reverse_hash(): + token = "a" * 64 + double_hashed = hash_token(token) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")] + ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) + + assert result[double_hashed]["key_alias"] == "batch-worker" + assert result[double_hashed]["team_id"] == "team-1" + mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_fill_missing_api_key_aliases_updates_null_alias_rows(): + token = "c" * 64 + double_hashed = hash_token(token) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token=token, key_alias="recovered-alias", team_id="team-9")] + ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + rows = ( + { + "api_key": double_hashed, + "api_key_alias": None, + "team_id": None, + "user_email": "owner@example.com", + "spend": 12.5, + }, + { + "api_key": "already-joined-token", + "api_key_alias": "named-key", + "team_id": "team-ok", + "user_email": "other@example.com", + "spend": 1.0, + }, + ) + + filled = await fill_missing_api_key_aliases(mock_prisma, rows) + + assert filled[0]["api_key_alias"] == "recovered-alias" + assert filled[0]["team_id"] == "team-9" + assert filled[0]["user_email"] == "owner@example.com" + assert filled[1]["api_key_alias"] == "named-key" From 6f0f2fcc8d00158fa7eb55b0b5781af6fd279acc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:27:12 +0000 Subject: [PATCH 14/35] fix(spend): restore user_email for double-hashed keys and persist it in spend logs Recovery now resolves the key owner's email from UserTable via the recovered token user_id, and SpendLogsMetadata keeps user_api_key_user_email so new batch/export consumers see email without a separate user join. Co-authored-by: Mateo Wang --- litellm/proxy/_types.py | 1 + .../spend_tracking/key_metadata_recovery.py | 101 +++++++++++++++--- .../spend_tracking/spend_tracking_utils.py | 1 + .../test_common_daily_activity.py | 13 ++- .../test_key_metadata_recovery.py | 31 +++++- 5 files changed, 127 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 849e54c65aa..54f1bfcad50 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3645,6 +3645,7 @@ class SpendLogsMetadata(TypedDict): user_api_key_project_alias: str | None user_api_key_org_id: str | None user_api_key_user_id: str | None + user_api_key_user_email: str | None user_api_key_team_alias: str | None spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call requester_ip_address: str | None diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 4d85357b7a9..7ad1c77ff57 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -7,6 +7,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient, hash_token from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) @@ -19,12 +20,15 @@ _SPEND_LOGS_KEY_METADATA_SQL: Final = """ SELECT DISTINCT ON (api_key) api_key, metadata->>'user_api_key_alias' AS key_alias, - metadata->>'user_api_key_team_id' AS team_id + metadata->>'user_api_key_team_id' AS team_id, + metadata->>'user_api_key_user_id' AS user_id, + metadata->>'user_api_key_user_email' AS user_email FROM "LiteLLM_SpendLogs" WHERE api_key = ANY($1::text[]) AND ( NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL + OR NULLIF(metadata->>'user_api_key_user_email', '') IS NOT NULL ) ORDER BY api_key, "startTime" DESC NULLS LAST """ @@ -33,6 +37,8 @@ ORDER BY api_key, "startTime" DESC NULLS LAST class KeyMetadataDict(TypedDict, total=False): key_alias: str | None team_id: str | None + user_id: str | None + user_email: str | None class _TokenAliasRecord(Protocol): @@ -45,13 +51,20 @@ class _TokenAliasRecord(Protocol): @property def team_id(self) -> str | None: ... + @property + def user_id(self) -> str | None: ... + def _token_digest_metadata( records: Sequence[_TokenAliasRecord], wanted: AbstractSet[str], ) -> dict[str, KeyMetadataDict]: return { - digested: {"key_alias": record.key_alias, "team_id": record.team_id} + digested: { + "key_alias": record.key_alias, + "team_id": record.team_id, + "user_id": getattr(record, "user_id", None), + } for record in records for digested in (hash_token(record.token),) if digested in wanted @@ -132,6 +145,8 @@ async def _spend_logs_key_metadata( row["api_key"]: { "key_alias": row.get("key_alias"), "team_id": row.get("team_id"), + "user_id": row.get("user_id"), + "user_email": row.get("user_email"), } for row in spend_log_rows if isinstance(row, dict) @@ -140,18 +155,67 @@ async def _spend_logs_key_metadata( } +async def _emails_for_user_ids( + prisma_client: PrismaClient, + user_ids: AbstractSet[str], +) -> Mapping[str, str]: + if not user_ids: + return {} + try: + users: Final = await UserRepository(prisma_client).table.find_many( + where={"user_id": {"in": list(user_ids)}} + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed user_email recovery for %d user ids: %s", + len(user_ids), + e, + ) + return {} + return { + user.user_id: user.user_email + for user in users + if getattr(user, "user_id", None) and getattr(user, "user_email", None) + } + + +def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict: + if meta.get("user_email"): + return meta + user_id: Final = meta.get("user_id") + if not isinstance(user_id, str) or user_id not in emails: + return meta + return {**meta, "user_email": emails[user_id]} + + +async def _with_user_emails( + prisma_client: PrismaClient, + recovered: Mapping[str, KeyMetadataDict], +) -> dict[str, KeyMetadataDict]: + needing_email: Final = frozenset( + user_id + for meta in recovered.values() + for user_id in (meta.get("user_id"),) + if isinstance(user_id, str) and user_id and not meta.get("user_email") + ) + emails: Final = await _emails_for_user_ids(prisma_client, needing_email) + if not emails: + return dict(recovered) + return {api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()} + + async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], ) -> dict[str, KeyMetadataDict]: """ - Recover key_alias/team_id for DailyUserSpend.api_key values that were - double-hashed by the v1.99 spend-log provenance gate. + Recover key_alias/team_id/user_email for DailyUserSpend.api_key values that + were double-hashed by the v1.99 spend-log provenance gate. Those rows store hash(VerificationToken.token) instead of the token, so the exact join misses. Prefer a bounded reverse-hash against active/deleted - tokens; fall back to the alias/team stamped into SpendLogs metadata (which - stayed correct even when api_key did not). + tokens; fall back to SpendLogs metadata. Emails come from SpendLogs when + present, otherwise from UserTable via the recovered key's user_id. """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: @@ -159,19 +223,22 @@ async def recover_double_hashed_key_metadata( from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) still_missing: Final = sha_missing - frozenset(from_tokens) - if not still_missing: - return from_tokens - - return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + recovered: Final = ( + from_tokens + if not still_missing + else {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + ) + return await _with_user_emails(prisma_client, recovered) -def _row_with_recovered_alias( +def _row_with_recovered_fields( row: Mapping[str, object], recovered: Mapping[str, KeyMetadataDict], *, api_key_field: str, alias_field: str, team_id_field: str, + user_email_field: str, ) -> Mapping[str, object]: api_key: Final = row.get(api_key_field) if not isinstance(api_key, str) or api_key not in recovered: @@ -181,6 +248,7 @@ def _row_with_recovered_alias( **row, alias_field: meta.get("key_alias") or row.get(alias_field), team_id_field: meta.get("team_id") or row.get(team_id_field), + user_email_field: meta.get("user_email") or row.get(user_email_field), } @@ -191,9 +259,11 @@ async def fill_missing_api_key_aliases( api_key_field: str = "api_key", alias_field: str = "api_key_alias", team_id_field: str = "team_id", + user_email_field: str = "user_email", ) -> tuple[Mapping[str, object], ...]: """ - Fill null api_key_alias / team_id on export rows whose api_key was double-hashed. + Fill null api_key_alias / team_id / user_email on export rows whose api_key + was double-hashed. Used by CloudZero and Focus, which join DailyUserSpend.api_key to VerificationToken.token and otherwise export null aliases for those rows. @@ -202,7 +272,9 @@ async def fill_missing_api_key_aliases( key for row in rows for key in (row.get(api_key_field),) - if isinstance(key, str) and key and row.get(alias_field) in (None, "") + if isinstance(key, str) + and key + and (row.get(alias_field) in (None, "") or row.get(user_email_field) in (None, "")) ) if not missing_keys: return tuple(rows) @@ -212,12 +284,13 @@ async def fill_missing_api_key_aliases( return tuple(rows) return tuple( - _row_with_recovered_alias( + _row_with_recovered_fields( row, recovered, api_key_field=api_key_field, alias_field=alias_field, team_id_field=team_id_field, + user_email_field=user_email_field, ) for row in rows ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 7442d71bd96..93d5dd09b0c 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -138,6 +138,7 @@ def _get_spend_logs_metadata( user_api_key_project_alias=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 78752495f25..2418a1cf245 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -471,10 +471,20 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( side_effect=[ [], # exact join miss - [SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")], + [ + SimpleNamespace( + token=token, + key_alias="batch-worker", + team_id="team-1", + user_id="alice", + ) + ], ] ) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) result = await get_api_key_metadata( @@ -484,6 +494,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( assert result[double_hashed]["key_alias"] == "batch-worker" assert result[double_hashed]["team_id"] == "team-1" + assert result[double_hashed]["user_email"] == "alice@example.com" mock_prisma.db.query_raw.assert_not_called() diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 6b12676e82b..3d08b2909c4 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -16,27 +16,48 @@ async def test_recover_double_hashed_key_metadata_via_reverse_hash(): double_hashed = hash_token(token) mock_prisma = MagicMock() mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")] + return_value=[ + SimpleNamespace( + token=token, + key_alias="batch-worker", + team_id="team-1", + user_id="alice", + ) + ] ) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) assert result[double_hashed]["key_alias"] == "batch-worker" assert result[double_hashed]["team_id"] == "team-1" + assert result[double_hashed]["user_email"] == "alice@example.com" mock_prisma.db.query_raw.assert_not_called() @pytest.mark.asyncio -async def test_fill_missing_api_key_aliases_updates_null_alias_rows(): +async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows(): token = "c" * 64 double_hashed = hash_token(token) mock_prisma = MagicMock() mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[SimpleNamespace(token=token, key_alias="recovered-alias", team_id="team-9")] + return_value=[ + SimpleNamespace( + token=token, + key_alias="recovered-alias", + team_id="team-9", + user_id="bob", + ) + ] ) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")] + ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) rows = ( @@ -44,7 +65,7 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_rows(): "api_key": double_hashed, "api_key_alias": None, "team_id": None, - "user_email": "owner@example.com", + "user_email": None, "spend": 12.5, }, { @@ -60,5 +81,5 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_rows(): assert filled[0]["api_key_alias"] == "recovered-alias" assert filled[0]["team_id"] == "team-9" - assert filled[0]["user_email"] == "owner@example.com" + assert filled[0]["user_email"] == "bob@example.com" assert filled[1]["api_key_alias"] == "named-key" From fd520e53abba19b35eaaa212515764022a253b05 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:34:44 +0000 Subject: [PATCH 15/35] style(spend): satisfy ruff format on key metadata recovery Lint CI failed because ruff format splits the Set alias import and collapses a couple of long lines. Co-authored-by: Mateo Wang --- litellm/proxy/spend_tracking/key_metadata_recovery.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 7ad1c77ff57..33339f16d52 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping, Sequence, Set as AbstractSet +from collections.abc import Mapping, Sequence +from collections.abc import Set as AbstractSet from typing import Final, Protocol from typing_extensions import TypedDict @@ -149,9 +150,7 @@ async def _spend_logs_key_metadata( "user_email": row.get("user_email"), } for row in spend_log_rows - if isinstance(row, dict) - and isinstance(row.get("api_key"), str) - and row["api_key"] in wanted + if isinstance(row, dict) and isinstance(row.get("api_key"), str) and row["api_key"] in wanted } @@ -162,9 +161,7 @@ async def _emails_for_user_ids( if not user_ids: return {} try: - users: Final = await UserRepository(prisma_client).table.find_many( - where={"user_id": {"in": list(user_ids)}} - ) + users: Final = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}) except Exception as e: verbose_proxy_logger.warning( "Failed user_email recovery for %d user ids: %s", From bbf4d1dc304827538e84ba66415466719136dcca Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:44:02 +0000 Subject: [PATCH 16/35] fix(spend): catch PrismaError instead of bare Exception in key recovery The Usage recovery path was adding four BLE001 hits and failing the strict-rule budget. Soft-fail only on PrismaError so a down token table still falls through to SpendLogs. Co-authored-by: Mateo Wang --- .../spend_tracking/key_metadata_recovery.py | 84 +++++++++---------- .../test_key_metadata_recovery.py | 28 +++++++ 2 files changed, 70 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 33339f16d52..78169283dff 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,7 +1,8 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from typing import Final, Protocol +from typing import Final, Protocol, TypeVar +from prisma.errors import PrismaError from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger @@ -13,6 +14,8 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) +_T = TypeVar("_T") + # Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key # values cannot pull an unbounded VerificationToken table into memory. _MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 @@ -56,6 +59,18 @@ class _TokenAliasRecord(Protocol): def user_id(self) -> str | None: ... +async def _db_or_empty( + load: Callable[[], Awaitable[_T]], + warning: str, + count: int, +) -> _T | None: + try: + return await load() + except PrismaError as e: + verbose_proxy_logger.warning(warning, count, e) + return None + + def _token_digest_metadata( records: Sequence[_TokenAliasRecord], wanted: AbstractSet[str], @@ -76,16 +91,12 @@ async def _reverse_hash_active_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], ) -> dict[str, KeyMetadataDict]: - try: - active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN) - except Exception as e: - verbose_proxy_logger.warning( - "Failed reverse-hash recovery against active keys for %d missing keys: %s", - len(wanted), - e, - ) + active_records: Final = await _db_or_empty( + lambda: VerificationTokenRepository(prisma_client).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN), + "Failed reverse-hash recovery against active keys for %d missing keys: %s", + len(wanted), + ) + if active_records is None: return {} return _token_digest_metadata(active_records, wanted) @@ -94,19 +105,15 @@ async def _reverse_hash_deleted_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], ) -> dict[str, KeyMetadataDict]: - try: - deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository( - prisma_client - ).table.find_many( + deleted_records: Final = await _db_or_empty( + lambda: DeletedVerificationTokenRepository(prisma_client).table.find_many( take=_MAX_DOUBLE_HASH_TOKEN_SCAN, order={"deleted_at": "desc"}, - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", - len(wanted), - e, - ) + ), + "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + len(wanted), + ) + if deleted_records is None: return {} return _token_digest_metadata(deleted_records, wanted) @@ -126,19 +133,14 @@ async def _spend_logs_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], ) -> dict[str, KeyMetadataDict]: - try: - spend_log_rows: Final = await prisma_client.db.query_raw( + spend_log_rows: Final = await _db_or_empty( + lambda: prisma_client.db.query_raw( _SPEND_LOGS_KEY_METADATA_SQL, list(wanted), - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed SpendLogs metadata recovery for %d missing keys: %s", - len(wanted), - e, - ) - return {} - + ), + "Failed SpendLogs metadata recovery for %d missing keys: %s", + len(wanted), + ) if not isinstance(spend_log_rows, list): return {} @@ -160,14 +162,12 @@ async def _emails_for_user_ids( ) -> Mapping[str, str]: if not user_ids: return {} - try: - users: Final = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}) - except Exception as e: - verbose_proxy_logger.warning( - "Failed user_email recovery for %d user ids: %s", - len(user_ids), - e, - ) + users: Final = await _db_or_empty( + lambda: UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}), + "Failed user_email recovery for %d user ids: %s", + len(user_ids), + ) + if users is None: return {} return { user.user_id: user.user_email diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 3d08b2909c4..6f0e912c5b1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -2,6 +2,7 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest +from prisma.errors import PrismaError from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, @@ -83,3 +84,30 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows(): assert filled[0]["team_id"] == "team-9" assert filled[0]["user_email"] == "bob@example.com" assert filled[1]["api_key_alias"] == "named-key" + + +@pytest.mark.asyncio +async def test_recover_falls_back_to_spend_logs_when_token_scan_raises_prisma_error(): + token = "b" * 64 + double_hashed = hash_token(token) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=PrismaError("db down")) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(side_effect=PrismaError("db down")) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "api_key": double_hashed, + "key_alias": "from-spend-logs", + "team_id": "team-sl", + "user_id": "carol", + "user_email": "carol@example.com", + } + ] + ) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) + + assert result[double_hashed]["key_alias"] == "from-spend-logs" + assert result[double_hashed]["team_id"] == "team-sl" + assert result[double_hashed]["user_email"] == "carol@example.com" From f610d205431a35f0bce1e19e6f2ebd8775a4af39 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:52:01 +0000 Subject: [PATCH 17/35] fix(spend): keep key recovery import-safe and LIT-clean The Python 3.10 smoke check imports the proxy without prisma, so PrismaError is loaded only inside the DB helper. Recovery now returns frozen mappings and ReadOnly TypedDict fields so the type-discipline budget stays put. Co-authored-by: Mateo Wang --- litellm/integrations/cloudzero/database.py | 2 +- litellm/integrations/focus/database.py | 2 +- litellm/proxy/_types.py | 2 +- .../common_daily_activity.py | 12 +- .../spend_tracking/key_metadata_recovery.py | 142 +++++++++++------- 5 files changed, 92 insertions(+), 68 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 8fedd4edac4..87f0c8bd160 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -101,6 +101,6 @@ class LiteLLMDatabase: # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) - return pl.DataFrame(list(recovered_rows), infer_schema_length=None) + return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index db9849bbbc9..96a32046e81 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -103,7 +103,7 @@ class FocusLiteLLMDatabase: # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) - return pl.DataFrame(list(recovered_rows), infer_schema_length=None) + return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 54f1bfcad50..a4519e175ee 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3645,7 +3645,7 @@ class SpendLogsMetadata(TypedDict): user_api_key_project_alias: str | None user_api_key_org_id: str | None user_api_key_user_id: str | None - user_api_key_user_email: str | None + user_api_key_user_email: ReadOnly[str | None] user_api_key_team_alias: str | None spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call requester_ip_address: str | None diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f687cedeee9..9853f05a068 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Protocol from fastapi import HTTPException, status @@ -445,7 +445,7 @@ def update_breakdown_metrics( async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], -) -> dict[str, _KeyMetadataDict]: +) -> Mapping[str, _KeyMetadataDict]: """Get api key metadata, falling back to deleted keys table for keys not found in active table. This ensures that key_alias and team_id are preserved in historical activity logs @@ -483,13 +483,11 @@ async def get_api_key_metadata( e, ) - still_missing: Final = api_keys - set(result.keys()) + still_missing: Final = api_keys - frozenset(result) if not still_missing: return result - return { - **result, - **(await recover_double_hashed_key_metadata(prisma_client, still_missing)), - } + recovered: Final = await recover_double_hashed_key_metadata(prisma_client, still_missing) + return MappingProxyType({**result, **recovered}) def _adjust_dates_for_timezone( diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 78169283dff..10ce4548214 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,9 +1,9 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet +from types import MappingProxyType from typing import Final, Protocol, TypeVar -from prisma.errors import PrismaError -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash @@ -16,8 +16,6 @@ from litellm.repositories.verification_token_repository import ( _T = TypeVar("_T") -# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key -# values cannot pull an unbounded VerificationToken table into memory. _MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 _SPEND_LOGS_KEY_METADATA_SQL: Final = """ @@ -39,10 +37,14 @@ ORDER BY api_key, "startTime" DESC NULLS LAST class KeyMetadataDict(TypedDict, total=False): - key_alias: str | None - team_id: str | None - user_id: str | None - user_email: str | None + key_alias: ReadOnly[str | None] + team_id: ReadOnly[str | None] + user_id: ReadOnly[str | None] + user_email: ReadOnly[str | None] + + +_EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) +_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) class _TokenAliasRecord(Protocol): @@ -64,6 +66,8 @@ async def _db_or_empty( warning: str, count: int, ) -> _T | None: + from prisma.errors import PrismaError + try: return await load() except PrismaError as e: @@ -71,89 +75,104 @@ async def _db_or_empty( return None +def _record_metadata(record: _TokenAliasRecord) -> KeyMetadataDict: + meta: Final[KeyMetadataDict] = { + "key_alias": record.key_alias, + "team_id": record.team_id, + "user_id": getattr(record, "user_id", None), + } + return meta + + +def _spend_log_row_metadata(row: Mapping[str, object]) -> KeyMetadataDict: + meta: Final[KeyMetadataDict] = { + "key_alias": row.get("key_alias") if isinstance(row.get("key_alias"), str) else None, + "team_id": row.get("team_id") if isinstance(row.get("team_id"), str) else None, + "user_id": row.get("user_id") if isinstance(row.get("user_id"), str) else None, + "user_email": row.get("user_email") if isinstance(row.get("user_email"), str) else None, + } + return meta + + def _token_digest_metadata( records: Sequence[_TokenAliasRecord], wanted: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: - return { - digested: { - "key_alias": record.key_alias, - "team_id": record.team_id, - "user_id": getattr(record, "user_id", None), +) -> Mapping[str, KeyMetadataDict]: + return MappingProxyType( + { + digested: _record_metadata(record) + for record in records + for digested in (hash_token(record.token),) + if digested in wanted } - for record in records - for digested in (hash_token(record.token),) - if digested in wanted - } + ) async def _reverse_hash_active_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: active_records: Final = await _db_or_empty( lambda: VerificationTokenRepository(prisma_client).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN), "Failed reverse-hash recovery against active keys for %d missing keys: %s", len(wanted), ) if active_records is None: - return {} + return _EMPTY_KEY_METADATA return _token_digest_metadata(active_records, wanted) async def _reverse_hash_deleted_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: deleted_records: Final = await _db_or_empty( lambda: DeletedVerificationTokenRepository(prisma_client).table.find_many( take=_MAX_DOUBLE_HASH_TOKEN_SCAN, - order={"deleted_at": "desc"}, + order={"deleted_at": "desc"}, # mutable-ok: Prisma find_many order= is a dict ), "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", len(wanted), ) if deleted_records is None: - return {} + return _EMPTY_KEY_METADATA return _token_digest_metadata(deleted_records, wanted) async def _reverse_hash_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) still_wanted: Final = wanted - frozenset(from_active) if not still_wanted: return from_active - return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))} + from_deleted: Final = await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted) + return MappingProxyType({**from_active, **from_deleted}) async def _spend_logs_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: spend_log_rows: Final = await _db_or_empty( lambda: prisma_client.db.query_raw( _SPEND_LOGS_KEY_METADATA_SQL, - list(wanted), + tuple(wanted), ), "Failed SpendLogs metadata recovery for %d missing keys: %s", len(wanted), ) if not isinstance(spend_log_rows, list): - return {} + return _EMPTY_KEY_METADATA - return { - row["api_key"]: { - "key_alias": row.get("key_alias"), - "team_id": row.get("team_id"), - "user_id": row.get("user_id"), - "user_email": row.get("user_email"), + return MappingProxyType( + { + row["api_key"]: _spend_log_row_metadata(row) + for row in spend_log_rows + if isinstance(row, dict) and isinstance(row.get("api_key"), str) and row["api_key"] in wanted } - for row in spend_log_rows - if isinstance(row, dict) and isinstance(row.get("api_key"), str) and row["api_key"] in wanted - } + ) async def _emails_for_user_ids( @@ -161,19 +180,23 @@ async def _emails_for_user_ids( user_ids: AbstractSet[str], ) -> Mapping[str, str]: if not user_ids: - return {} + return _EMPTY_EMAILS users: Final = await _db_or_empty( - lambda: UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}), + lambda: UserRepository(prisma_client).table.find_many( + where={"user_id": {"in": tuple(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict + ), "Failed user_email recovery for %d user ids: %s", len(user_ids), ) if users is None: - return {} - return { - user.user_id: user.user_email - for user in users - if getattr(user, "user_id", None) and getattr(user, "user_email", None) - } + return _EMPTY_EMAILS + return MappingProxyType( + { + user.user_id: user.user_email + for user in users + if getattr(user, "user_id", None) and getattr(user, "user_email", None) + } + ) def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict: @@ -182,13 +205,14 @@ def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMet user_id: Final = meta.get("user_id") if not isinstance(user_id, str) or user_id not in emails: return meta - return {**meta, "user_email": emails[user_id]} + updated: Final[KeyMetadataDict] = {**meta, "user_email": emails[user_id]} + return updated async def _with_user_emails( prisma_client: PrismaClient, recovered: Mapping[str, KeyMetadataDict], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: needing_email: Final = frozenset( user_id for meta in recovered.values() @@ -197,14 +221,14 @@ async def _with_user_emails( ) emails: Final = await _emails_for_user_ids(prisma_client, needing_email) if not emails: - return dict(recovered) - return {api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()} + return recovered + return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()}) async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], -) -> dict[str, KeyMetadataDict]: +) -> Mapping[str, KeyMetadataDict]: """ Recover key_alias/team_id/user_email for DailyUserSpend.api_key values that were double-hashed by the v1.99 spend-log provenance gate. @@ -216,14 +240,14 @@ async def recover_double_hashed_key_metadata( """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: - return {} + return _EMPTY_KEY_METADATA from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) still_missing: Final = sha_missing - frozenset(from_tokens) recovered: Final = ( from_tokens if not still_missing - else {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + else MappingProxyType({**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}) ) return await _with_user_emails(prisma_client, recovered) @@ -241,12 +265,14 @@ def _row_with_recovered_fields( if not isinstance(api_key, str) or api_key not in recovered: return row meta: Final = recovered[api_key] - return { - **row, - alias_field: meta.get("key_alias") or row.get(alias_field), - team_id_field: meta.get("team_id") or row.get(team_id_field), - user_email_field: meta.get("user_email") or row.get(user_email_field), - } + return MappingProxyType( + { + **row, + alias_field: meta.get("key_alias") or row.get(alias_field), + team_id_field: meta.get("team_id") or row.get(team_id_field), + user_email_field: meta.get("user_email") or row.get(user_email_field), + } + ) async def fill_missing_api_key_aliases( From e76a18ca3251489af9954e7a8ea30b3b763d6c01 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:56:26 +0000 Subject: [PATCH 18/35] fix(usage): return user_email on key activity metadata Recovery already resolved the owner email for double-hashed spend keys, then Key Activity dropped it. The Usage payload now carries user_email and the key label falls back to that email before key-hash-... Co-authored-by: Mateo Wang --- .../common_daily_activity.py | 29 ++++++++++++++----- .../spend_tracking/key_metadata_recovery.py | 4 +-- .../common_daily_activity.py | 1 + .../test_common_daily_activity.py | 19 ++++++++++++ .../src/components/UsagePage/types.ts | 1 + .../src/components/activity_metrics.test.tsx | 13 ++++++++- .../src/components/activity_metrics.tsx | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 ++ 8 files changed, 60 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 9853f05a068..833f463621a 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -6,12 +6,13 @@ from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Protocol from fastapi import HTTPException, status -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.key_metadata_recovery import ( + attach_user_emails, recover_double_hashed_key_metadata, ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled @@ -116,6 +117,8 @@ class DailySpendRecord(Protocol): class _KeyMetadataDict(TypedDict, total=False): key_alias: str | None team_id: str | None + user_id: ReadOnly[str | None] + user_email: ReadOnly[str | None] _WhereValue = str | dict[str, object] @@ -456,7 +459,12 @@ async def get_api_key_metadata( where={"token": {"in": list(api_keys)}} ) result: Final[dict[str, _KeyMetadataDict]] = { - k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records + k.token: { + "key_alias": k.key_alias, + "team_id": k.team_id, + "user_id": getattr(k, "user_id", None), + } + for k in key_records } # For any keys not found in the active table, check the deleted keys table @@ -475,6 +483,7 @@ async def get_api_key_metadata( result[k.token] = { "key_alias": k.key_alias, "team_id": k.team_id, + "user_id": getattr(k, "user_id", None), } except Exception as e: verbose_proxy_logger.warning( @@ -484,10 +493,12 @@ async def get_api_key_metadata( ) still_missing: Final = api_keys - frozenset(result) - if not still_missing: - return result - recovered: Final = await recover_double_hashed_key_metadata(prisma_client, still_missing) - return MappingProxyType({**result, **recovered}) + combined: Final = ( + result + if not still_missing + else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))}) + ) + return await attach_user_emails(prisma_client, combined) def _adjust_dates_for_timezone( @@ -961,7 +972,11 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata: meta: Final = api_key_metadata.get(api_key, {}) - return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id")) + return KeyMetadata( + key_alias=meta.get("key_alias"), + team_id=meta.get("team_id"), + user_email=meta.get("user_email"), + ) def _aggregate_grouping_sets_records_sync( diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 10ce4548214..fae091d4948 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -209,7 +209,7 @@ def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMet return updated -async def _with_user_emails( +async def attach_user_emails( prisma_client: PrismaClient, recovered: Mapping[str, KeyMetadataDict], ) -> Mapping[str, KeyMetadataDict]: @@ -249,7 +249,7 @@ async def recover_double_hashed_key_metadata( if not still_missing else MappingProxyType({**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}) ) - return await _with_user_emails(prisma_client, recovered) + return await attach_user_emails(prisma_client, recovered) def _row_with_recovered_fields( diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 101405abf50..2b39c5dbb9b 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -43,6 +43,7 @@ class KeyMetadata(BaseModel): key_alias: str | None = None team_id: str | None = None + user_email: str | None = None class KeyMetricWithMetadata(MetricBase): diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 2418a1cf245..f7ad01c3bba 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -516,6 +516,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): } ] ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -527,6 +528,24 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): mock_prisma.db.query_raw.assert_called_once() +def test_key_metadata_includes_recovered_user_email(): + from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata + + meta = _key_metadata( + { + "dirty-key": { + "key_alias": "batch-worker", + "team_id": "team-1", + "user_email": "alice@example.com", + } + }, + "dirty-key", + ) + + assert meta.key_alias == "batch-worker" + assert meta.user_email == "alice@example.com" + + @pytest.mark.asyncio async def test_tag_daily_activity_metadata_totals_not_zero(): """Test that tag daily activity returns correct metadata totals. diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index 8e7c1869df2..a10e9e68c4d 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -46,6 +46,7 @@ export interface KeyMetricWithMetadata { export interface KeyMetadata { key_alias: string | null; team_id: string | null; + user_email?: string | null; tags?: { tag: string; usage: number }[]; } diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index 74d258e2bd0..b0fc8dc7866 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -101,7 +101,7 @@ const createMockDailyData = ( }); const createMockKeyMetricWithMetadata = ( - metadata: { key_alias: string | null; team_id: string | null }, + metadata: { key_alias: string | null; team_id: string | null; user_email?: string | null }, metrics: typeof EMPTY_SPEND_METRICS = EMPTY_SPEND_METRICS, ): KeyMetricWithMetadata => ({ metrics, @@ -1450,6 +1450,17 @@ describe("formatKeyLabel", () => { expect(result).toBe("key-hash-actual-key (team: Test Team 1)"); }); + it("should use user_email when key_alias is null", () => { + const modelData = createMockKeyMetricWithMetadata({ + key_alias: null, + team_id: "team1", + user_email: "alice@example.com", + }); + + const result = formatKeyLabel(modelData, "actual-key", MOCK_TEAMS); + expect(result).toBe("alice@example.com (team: Test Team 1)"); + }); + it("should return key_alias with team_id when teams array is empty", () => { const modelData = createMockKeyMetricWithMetadata({ key_alias: "my-key", diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index a3fff08faae..e17263ab078 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -433,7 +433,7 @@ export const ActivityMetrics: React.FC = ({ modelMetrics, // Helper function to format key label export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => { - const keyAlias = modelData.metadata.key_alias || `key-hash-${model}`; + const keyAlias = modelData.metadata.key_alias || modelData.metadata.user_email || `key-hash-${model}`; const teamId = modelData.metadata.team_id; if (teamId) { const teamAlias = resolveTeamAliasFromTeamID(teamId, teams); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b0f8e618645..8085967cf6a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27994,6 +27994,8 @@ export interface components { key_alias?: string | null; /** Team Id */ team_id?: string | null; + /** User Email */ + user_email?: string | null; }; /** * KeyMetricWithMetadata From 6f4ea2d296311c93559f2f8d60d8690f35386e39 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:58:05 +0000 Subject: [PATCH 19/35] fix(usage): label key activity tables with email when alias is missing Key Activity charts already fell back to user_email. The top-keys tables and usage export still printed '-' or a truncated hash. They now use the same alias-then-email label. Co-authored-by: Mateo Wang --- .../EntityUsage/entityUsageAggregations.ts | 4 +++- .../_components/components/UsagePageView.tsx | 4 +++- .../src/components/EntityUsageExport/utils.ts | 3 ++- .../components/UsagePage/keyActivityLabel.test.ts | 15 +++++++++++++++ .../src/components/UsagePage/keyActivityLabel.ts | 8 ++++++++ .../src/components/activity_metrics.tsx | 5 +++-- 6 files changed, 34 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts create mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts index a53b1d2827b..d482a5576ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts @@ -1,3 +1,4 @@ +import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types"; export type ExtendedDailyData = DailyData & { @@ -118,6 +119,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number metadata: { key_alias: metrics.metadata.key_alias, team_id: metrics.metadata.team_id || null, + user_email: metrics.metadata.user_email, tags: tagDictionary[key] || [], }, }; @@ -137,7 +139,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, - key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias + key_alias: keyActivityLabel(metrics.metadata), tags: metrics.metadata.tags || "-", spend: metrics.metrics.spend, })) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index cbdfc8f39e6..29a81e1ae3f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -44,6 +44,7 @@ import { Tag } from "@/components/tag_management/types"; import UserAgentActivity from "@/components/user_agent_activity"; import ViewUserSpend from "@/components/view_user_spend"; import { usePaginatedDailyActivity } from "../hooks/usePaginatedDailyActivity"; +import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; import { DailyData, KeyMetricWithMetadata, MetricWithMetadata } from "@/components/UsagePage/types"; import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatters"; import { @@ -426,6 +427,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { metadata: { key_alias: metrics.metadata.key_alias, team_id: null, + user_email: metrics.metadata.user_email, tags: metrics.metadata.tags || [], }, }; @@ -445,7 +447,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, - key_alias: metrics.metadata.key_alias || "-", + key_alias: keyActivityLabel(metrics.metadata), tags: metrics.metadata.tags || [], spend: metrics.metrics.spend, })) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index de637d5d627..8fd75134bcc 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -1,6 +1,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import Papa from "papaparse"; +import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; const resolveEntityDisplay = ( @@ -186,7 +187,7 @@ export const generateDailyWithKeysData = ( // Iterate through each API key in the breakdown Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => { - const keyAlias = keyData?.metadata?.key_alias || null; + const keyAlias = keyActivityLabel(keyData?.metadata, "") || null; // Create unique key for aggregation: Date_EntityID_KeyID const uniqueKey = `${day.date}_${entityId}_${keyId}`; diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts new file mode 100644 index 00000000000..eaf1985c5fa --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts @@ -0,0 +1,15 @@ +import { keyActivityLabel } from "./keyActivityLabel"; + +describe("keyActivityLabel", () => { + it("prefers key_alias", () => { + expect(keyActivityLabel({ key_alias: "batch-worker", user_email: "alice@example.com" })).toBe("batch-worker"); + }); + + it("falls back to user_email when alias is missing", () => { + expect(keyActivityLabel({ key_alias: null, user_email: "alice@example.com" })).toBe("alice@example.com"); + }); + + it("uses the fallback when both alias and email are missing", () => { + expect(keyActivityLabel({ key_alias: null, user_email: null }, "key-hash-abc")).toBe("key-hash-abc"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts new file mode 100644 index 00000000000..8b3a7eec916 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts @@ -0,0 +1,8 @@ +import type { KeyMetadata } from "./types"; + +export function keyActivityLabel( + metadata: Pick | null | undefined, + fallback = "-", +): string { + return metadata?.key_alias || metadata?.user_email || fallback; +} diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index e17263ab078..f4348fb65ae 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -7,6 +7,7 @@ import { ChevronDown } from "lucide-react"; import React, { useState } from "react"; import { Team } from "./key_team_helpers/key_list"; import KeyModelUsageView from "./UsagePage/components/KeyModelUsageView"; +import { keyActivityLabel } from "./UsagePage/keyActivityLabel"; import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData, TopModelData } from "./UsagePage/types"; import { valueFormatter } from "./UsagePage/utils/value_formatters"; @@ -433,7 +434,7 @@ export const ActivityMetrics: React.FC = ({ modelMetrics, // Helper function to format key label export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => { - const keyAlias = modelData.metadata.key_alias || modelData.metadata.user_email || `key-hash-${model}`; + const keyAlias = keyActivityLabel(modelData.metadata, `key-hash-${model}`); const teamId = modelData.metadata.team_id; if (teamId) { const teamAlias = resolveTeamAliasFromTeamID(teamId, teams); @@ -516,7 +517,7 @@ export const processActivityData = ( if (!apiKeyBreakdown[apiKey]) { apiKeyBreakdown[apiKey] = { api_key: apiKey, - key_alias: keyData.metadata.key_alias, + key_alias: keyActivityLabel(keyData.metadata, "") || null, team_id: keyData.metadata.team_id, spend: 0, requests: 0, From eeb8b4f7f39afdd9f2148881d89d6c5a6fa951bb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:00:31 +0000 Subject: [PATCH 20/35] test(spend): assert batch spend metadata keeps user email Co-authored-by: Mateo Wang --- .../proxy/spend_tracking/test_spend_tracking_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index bea1f8e2d6c..6cdf9e3fe7e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2789,6 +2789,7 @@ def test_get_logging_payload_batch_attribution_keeps_verification_token_hash(): parsed_meta = json.loads(payload["metadata"]) assert parsed_meta["user_api_key"] == token_hash assert parsed_meta["user_api_key_alias"] == "batch-creator" + assert parsed_meta["user_api_key_user_email"] == "alice@example.com" def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): From e109d89c2041373a1744805e0397a181f49f427b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:11:45 +0000 Subject: [PATCH 21/35] fix(usage): include user_email on daily activity key breakdowns Co-authored-by: Mateo Wang --- .../common_daily_activity.py | 57 ++++++----------- .../test_common_daily_activity.py | 63 +++++++++++++++++++ 2 files changed, 81 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 833f463621a..ce6a97708ab 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -115,12 +115,21 @@ class DailySpendRecord(Protocol): class _KeyMetadataDict(TypedDict, total=False): - key_alias: str | None - team_id: str | None + key_alias: ReadOnly[str | None] + team_id: ReadOnly[str | None] user_id: ReadOnly[str | None] user_email: ReadOnly[str | None] +def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata: + meta: Final = api_key_metadata.get(api_key, {}) + return KeyMetadata( + key_alias=meta.get("key_alias"), + team_id=meta.get("team_id"), + user_email=meta.get("user_email"), + ) + + _WhereValue = str | dict[str, object] @@ -289,10 +298,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.models[model_key].api_key_breakdown: breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.models[model_key].api_key_breakdown[record.api_key].metrics, @@ -316,10 +322,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown: breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics, @@ -341,10 +344,7 @@ def update_breakdown_metrics( breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = ( KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) ) @@ -369,10 +369,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.providers[provider].api_key_breakdown: breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, @@ -394,10 +391,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown: breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics, @@ -409,10 +403,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.api_keys: breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), # Add any api_key-specific metadata here + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) @@ -432,10 +423,7 @@ def update_breakdown_metrics( if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + metadata=_key_metadata(api_key_metadata, record.api_key), ) breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, @@ -970,15 +958,6 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: ) -def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata: - meta: Final = api_key_metadata.get(api_key, {}) - return KeyMetadata( - key_alias=meta.get("key_alias"), - team_id=meta.get("team_id"), - user_email=meta.get("user_email"), - ) - - def _aggregate_grouping_sets_records_sync( *, records: Sequence[_GroupingSetsRow], diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index f7ad01c3bba..b9c0d953086 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -546,6 +546,69 @@ def test_key_metadata_includes_recovered_user_email(): assert meta.user_email == "alice@example.com" +def test_update_breakdown_metrics_includes_user_email(): + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + record = SimpleNamespace( + api_key="dirty-key", + model="gpt-4o-mini", + model_group="grp", + mcp_namespaced_tool_name="srv/tool", + custom_llm_provider="openai", + endpoint="/v1/chat/completions", + spend=1.23, + prompt_tokens=1, + completion_tokens=1, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0, + prompt_caching_savings_spend=0, + gateway_injected_caching_savings_spend=0, + autorouter_savings_spend=0, + total_tokens=2, + api_requests=1, + successful_requests=1, + failed_requests=0, + ptu_flat_cost=0.0, + user_id="alice", + ) + api_key_metadata = { + "dirty-key": { + "key_alias": "batch-worker", + "team_id": "team-1", + "user_email": "alice@example.com", + } + } + + update_breakdown_metrics( + breakdown, + record, + {}, + {}, + api_key_metadata, + entity_id_field="user_id", + ) + + expected = ("batch-worker", "alice@example.com") + top = breakdown.api_keys["dirty-key"].metadata + assert (top.key_alias, top.user_email) == expected + assert ( + breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.key_alias, + breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.user_email, + ) == expected + assert ( + breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.key_alias, + breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.user_email, + ) == expected + assert ( + breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.key_alias, + breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.user_email, + ) == expected + + @pytest.mark.asyncio async def test_tag_daily_activity_metadata_totals_not_zero(): """Test that tag daily activity returns correct metadata totals. From 7e4032cfcc319cf646537ee288fc3291e332af98 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:16:37 +0000 Subject: [PATCH 22/35] chore(openapi): sync lazy snapshot with KeyMetadata.user_email Co-authored-by: Mateo Wang --- litellm/proxy/_lazy_openapi_snapshot.json | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 13c7a4c7cfa..e8e1f53b473 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3151,6 +3151,17 @@ } ], "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" } }, "title": "KeyMetadata", From d3c839147edf831efc6ccc83ce8d927f751d4077 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:21:35 +0000 Subject: [PATCH 23/35] fix(spend): keep CloudZero export and spend-log snapshots compatible with email recovery Co-authored-by: Mateo Wang --- litellm/integrations/cloudzero/database.py | 7 ++++++- litellm/integrations/focus/database.py | 7 ++++++- .../test_litellm/integrations/cloudzero/test_cloudzero.py | 5 +++++ .../spend_tracking/test_spend_management_endpoints.py | 1 + 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 87f0c8bd160..6a63868a08b 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -98,9 +98,14 @@ class LiteLLMDatabase: fill_missing_api_key_aliases, ) + usage_rows: Final = ( + db_response.to_dicts() + if isinstance(db_response, pl.DataFrame) + else db_response if isinstance(db_response, list) else [] + ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. - recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 96a32046e81..da6458c9369 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -100,9 +100,14 @@ class FocusLiteLLMDatabase: fill_missing_api_key_aliases, ) + usage_rows: Final = ( + db_response.to_dicts() + if isinstance(db_response, pl.DataFrame) + else db_response if isinstance(db_response, list) else [] + ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. - recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response) + recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 2d51eeb9944..c543156eedd 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -74,6 +74,8 @@ class TestCloudZeroHourlyExport: fake_db = MagicMock() async def query_raw_mock(query: str, *params): + if "LiteLLM_SpendLogs" in query: + return [] start_time_utc = params[0] if len(params) > 0 else None end_time_utc = params[1] if len(params) > 1 else None limit = params[2] if len(params) > 2 else None @@ -146,6 +148,9 @@ class TestCloudZeroHourlyExport: return joined fake_db.query_raw = AsyncMock(side_effect=query_raw_mock) + fake_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + fake_db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + fake_db.litellm_usertable.find_many = AsyncMock(return_value=[]) fake_client.db = fake_db mock_prisma_client_getter.return_value = fake_client diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 30b086bab61..b62c1c076c8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -483,6 +483,7 @@ ignored_keys = [ "metadata.user_api_key_project_alias", "metadata.user_api_key_org_id", "metadata.user_api_key_user_id", + "metadata.user_api_key_user_email", "metadata.user_api_key_team_alias", "metadata.spend_logs_metadata", "metadata.requester_ip_address", From 4db370e85194bded3df0ec73a71a777ff7f15ce5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:25:46 +0000 Subject: [PATCH 24/35] style(spend): format CloudZero and Focus recovery row coercion Required lint failed ruff format on the ternary that unwraps a polars DataFrame or list before alias recovery Co-authored-by: Mateo Wang --- litellm/integrations/cloudzero/database.py | 4 +++- litellm/integrations/focus/database.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 6a63868a08b..e630bd85114 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -101,7 +101,9 @@ class LiteLLMDatabase: usage_rows: Final = ( db_response.to_dicts() if isinstance(db_response, pl.DataFrame) - else db_response if isinstance(db_response, list) else [] + else db_response + if isinstance(db_response, list) + else [] ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index da6458c9369..02b1e9e944b 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -103,7 +103,9 @@ class FocusLiteLLMDatabase: usage_rows: Final = ( db_response.to_dicts() if isinstance(db_response, pl.DataFrame) - else db_response if isinstance(db_response, list) else [] + else db_response + if isinstance(db_response, list) + else [] ) # v1.99 double-hashed DailyUserSpend.api_key values miss the # VerificationToken join above; recover alias/team for those rows. From d426b99f562672c930a473ec545dc8167e6345a0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 15:44:18 +0000 Subject: [PATCH 25/35] fix(spend): page reverse-hash recovery past the first 10k keys Historical dirty spend on large installs was still unlabeled when the matching token sat past the first page. Keep scanning until the digest matches or the table ends. Co-authored-by: Mateo Wang --- .../spend_tracking/key_metadata_recovery.py | 95 ++++++++++++++----- .../test_key_metadata_recovery.py | 73 ++++++++++++++ 2 files changed, 142 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index fae091d4948..895f97e1a05 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -16,7 +16,7 @@ from litellm.repositories.verification_token_repository import ( _T = TypeVar("_T") -_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 +_TOKEN_SCAN_PAGE: Final = 10_000 _SPEND_LOGS_KEY_METADATA_SQL: Final = """ SELECT DISTINCT ON (api_key) @@ -108,46 +108,86 @@ def _token_digest_metadata( ) +async def _paginate_token_digest_metadata( + load_page: Callable[[int], Awaitable[Sequence[_TokenAliasRecord] | None]], + wanted: AbstractSet[str], + *, + page_size: int, + skip: int = 0, + accumulated: Mapping[str, KeyMetadataDict] = _EMPTY_KEY_METADATA, +) -> Mapping[str, KeyMetadataDict]: + if not wanted: + return accumulated + records: Final = await load_page(skip) + if records is None: + return accumulated + page_hits: Final = _token_digest_metadata(records, wanted) + combined: Final[Mapping[str, KeyMetadataDict]] = ( + MappingProxyType({**accumulated, **page_hits}) if page_hits else accumulated + ) + still_wanted: Final = wanted - frozenset(page_hits) + if not still_wanted or len(records) < page_size: + return combined + return await _paginate_token_digest_metadata( + load_page, + still_wanted, + page_size=page_size, + skip=skip + page_size, + accumulated=combined, + ) + + async def _reverse_hash_active_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], + *, + page_size: int, ) -> Mapping[str, KeyMetadataDict]: - active_records: Final = await _db_or_empty( - lambda: VerificationTokenRepository(prisma_client).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN), - "Failed reverse-hash recovery against active keys for %d missing keys: %s", - len(wanted), - ) - if active_records is None: - return _EMPTY_KEY_METADATA - return _token_digest_metadata(active_records, wanted) + async def load_page(skip: int) -> Sequence[_TokenAliasRecord] | None: + return await _db_or_empty( + lambda: VerificationTokenRepository(prisma_client).table.find_many( + take=page_size, + skip=skip, + order={"token": "asc"}, # mutable-ok: Prisma find_many order= is a dict + ), + "Failed reverse-hash recovery against active keys for %d missing keys: %s", + len(wanted), + ) + + return await _paginate_token_digest_metadata(load_page, wanted, page_size=page_size) async def _reverse_hash_deleted_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], + *, + page_size: int, ) -> Mapping[str, KeyMetadataDict]: - deleted_records: Final = await _db_or_empty( - lambda: DeletedVerificationTokenRepository(prisma_client).table.find_many( - take=_MAX_DOUBLE_HASH_TOKEN_SCAN, - order={"deleted_at": "desc"}, # mutable-ok: Prisma find_many order= is a dict - ), - "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", - len(wanted), - ) - if deleted_records is None: - return _EMPTY_KEY_METADATA - return _token_digest_metadata(deleted_records, wanted) + async def load_page(skip: int) -> Sequence[_TokenAliasRecord] | None: + return await _db_or_empty( + lambda: DeletedVerificationTokenRepository(prisma_client).table.find_many( + take=page_size, + skip=skip, + order=[{"deleted_at": "desc"}, {"id": "asc"}], # mutable-ok: Prisma find_many order= is a dict + ), + "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + len(wanted), + ) + + return await _paginate_token_digest_metadata(load_page, wanted, page_size=page_size) async def _reverse_hash_key_metadata( prisma_client: PrismaClient, wanted: AbstractSet[str], + *, + page_size: int, ) -> Mapping[str, KeyMetadataDict]: - from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) + from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted, page_size=page_size) still_wanted: Final = wanted - frozenset(from_active) if not still_wanted: return from_active - from_deleted: Final = await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted) + from_deleted: Final = await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted, page_size=page_size) return MappingProxyType({**from_active, **from_deleted}) @@ -228,21 +268,24 @@ async def attach_user_emails( async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], + *, + token_scan_page_size: int = _TOKEN_SCAN_PAGE, ) -> Mapping[str, KeyMetadataDict]: """ Recover key_alias/team_id/user_email for DailyUserSpend.api_key values that were double-hashed by the v1.99 spend-log provenance gate. Those rows store hash(VerificationToken.token) instead of the token, so the - exact join misses. Prefer a bounded reverse-hash against active/deleted - tokens; fall back to SpendLogs metadata. Emails come from SpendLogs when - present, otherwise from UserTable via the recovered key's user_id. + exact join misses. Page through active then deleted tokens until every + wanted digest is found or the table ends; fall back to SpendLogs metadata. + Emails come from SpendLogs when present, otherwise from UserTable via the + recovered key's user_id. """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: return _EMPTY_KEY_METADATA - from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) + from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing, page_size=token_scan_page_size) still_missing: Final = sha_missing - frozenset(from_tokens) recovered: Final = ( from_tokens diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 6f0e912c5b1..704de049fc8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -111,3 +111,76 @@ async def test_recover_falls_back_to_spend_logs_when_token_scan_raises_prisma_er assert result[double_hashed]["key_alias"] == "from-spend-logs" assert result[double_hashed]["team_id"] == "team-sl" assert result[double_hashed]["user_email"] == "carol@example.com" + + +@pytest.mark.asyncio +async def test_recover_double_hashed_key_metadata_scans_past_first_page(): + token = "z" * 64 + double_hashed = hash_token(token) + decoys = ( + SimpleNamespace(token="1" * 64, key_alias="decoy-1", team_id=None, user_id=None), + SimpleNamespace(token="2" * 64, key_alias="decoy-2", team_id=None, user_id=None), + ) + match = SimpleNamespace(token=token, key_alias="late-key", team_id="team-late", user_id="dana") + + async def find_many(*, take: int | None = None, skip: int | None = None, order: object = None): + if skip == 0: + return list(decoys) + if skip == 2: + return [match] + return [] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=find_many) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="dana", user_email="dana@example.com")] + ) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}, token_scan_page_size=2) + + assert result[double_hashed]["key_alias"] == "late-key" + assert result[double_hashed]["team_id"] == "team-late" + assert result[double_hashed]["user_email"] == "dana@example.com" + assert [call.kwargs["skip"] for call in mock_prisma.db.litellm_verificationtoken.find_many.call_args_list] == [0, 2] + mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called() + mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_double_hashed_key_metadata_pages_deleted_tokens(): + token = "y" * 64 + double_hashed = hash_token(token) + decoys = ( + SimpleNamespace(token="3" * 64, key_alias="deleted-decoy-1", team_id=None, user_id=None), + SimpleNamespace(token="4" * 64, key_alias="deleted-decoy-2", team_id=None, user_id=None), + ) + match = SimpleNamespace(token=token, key_alias="deleted-late-key", team_id="team-del", user_id="erin") + + async def find_deleted(*, take: int | None = None, skip: int | None = None, order: object = None): + if skip == 0: + return list(decoys) + if skip == 2: + return [match] + return [] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(side_effect=find_deleted) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="erin", user_email="erin@example.com")] + ) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}, token_scan_page_size=2) + + assert result[double_hashed]["key_alias"] == "deleted-late-key" + assert result[double_hashed]["user_email"] == "erin@example.com" + assert [ + call.kwargs["skip"] for call in mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list + ] == [ + 0, + 2, + ] + mock_prisma.db.query_raw.assert_not_called() From ce95afe2bdfb47f4325c59fea89b8c8f8fb0a5a6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:58:17 -0700 Subject: [PATCH 26/35] fix(spend-tracking): reverse-hash dirty spend keys in Postgres instead of paging token tables --- litellm/integrations/cloudzero/database.py | 2 - litellm/integrations/focus/database.py | 2 - litellm/proxy/_types.py | 1 - .../spend_tracking/key_metadata_recovery.py | 232 ++++------------- .../spend_tracking/spend_tracking_utils.py | 1 - .../integrations/cloudzero/test_cloudzero.py | 2 +- .../test_common_daily_activity.py | 58 ++--- .../test_key_metadata_recovery.py | 234 ++++++++---------- .../test_spend_management_endpoints.py | 1 - .../test_spend_tracking_utils.py | 2 - 10 files changed, 189 insertions(+), 346 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index e630bd85114..4adf725fd0f 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -105,8 +105,6 @@ class LiteLLMDatabase: if isinstance(db_response, list) else [] ) - # v1.99 double-hashed DailyUserSpend.api_key values miss the - # VerificationToken join above; recover alias/team for those rows. recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as e: diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 02b1e9e944b..f214aa02b5b 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -107,8 +107,6 @@ class FocusLiteLLMDatabase: if isinstance(db_response, list) else [] ) - # v1.99 double-hashed DailyUserSpend.api_key values miss the - # VerificationToken join above; recover alias/team for those rows. recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) except Exception as exc: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b4a55d3510..5d5a25e7cd6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3659,7 +3659,6 @@ class SpendLogsMetadata(TypedDict): user_api_key_project_alias: str | None user_api_key_org_id: str | None user_api_key_user_id: str | None - user_api_key_user_email: ReadOnly[str | None] user_api_key_team_alias: str | None spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call requester_ip_address: str | None diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 895f97e1a05..d524b158c9c 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,38 +1,30 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from types import MappingProxyType -from typing import Final, Protocol, TypeVar +from typing import Final, TypeVar +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash -from litellm.proxy.utils import PrismaClient, hash_token -from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.proxy.utils import PrismaClient from litellm.repositories.user_repository import UserRepository -from litellm.repositories.verification_token_repository import ( - VerificationTokenRepository, -) _T = TypeVar("_T") -_TOKEN_SCAN_PAGE: Final = 10_000 +_ACTIVE_TOKEN_DIGEST_SQL: Final = """ +SELECT encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id +FROM "LiteLLM_VerificationToken" +WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) +""" -_SPEND_LOGS_KEY_METADATA_SQL: Final = """ -SELECT DISTINCT ON (api_key) - api_key, - metadata->>'user_api_key_alias' AS key_alias, - metadata->>'user_api_key_team_id' AS team_id, - metadata->>'user_api_key_user_id' AS user_id, - metadata->>'user_api_key_user_email' AS user_email -FROM "LiteLLM_SpendLogs" -WHERE api_key = ANY($1::text[]) - AND ( - NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL - OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL - OR NULLIF(metadata->>'user_api_key_user_email', '') IS NOT NULL - ) -ORDER BY api_key, "startTime" DESC NULLS LAST +_DELETED_TOKEN_DIGEST_SQL: Final = """ +SELECT DISTINCT ON (token) + encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id +FROM "LiteLLM_DeletedVerificationToken" +WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) +ORDER BY token, deleted_at DESC """ @@ -43,24 +35,18 @@ class KeyMetadataDict(TypedDict, total=False): user_email: ReadOnly[str | None] +class _TokenDigestRow(BaseModel): + digest: str + key_alias: str | None = None + team_id: str | None = None + user_id: str | None = None + + +_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) _EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) -class _TokenAliasRecord(Protocol): - @property - def token(self) -> str: ... - - @property - def key_alias(self) -> str | None: ... - - @property - def team_id(self) -> str | None: ... - - @property - def user_id(self) -> str | None: ... - - async def _db_or_empty( load: Callable[[], Awaitable[_T]], warning: str, @@ -75,142 +61,25 @@ async def _db_or_empty( return None -def _record_metadata(record: _TokenAliasRecord) -> KeyMetadataDict: - meta: Final[KeyMetadataDict] = { - "key_alias": record.key_alias, - "team_id": record.team_id, - "user_id": getattr(record, "user_id", None), - } - return meta - - -def _spend_log_row_metadata(row: Mapping[str, object]) -> KeyMetadataDict: - meta: Final[KeyMetadataDict] = { - "key_alias": row.get("key_alias") if isinstance(row.get("key_alias"), str) else None, - "team_id": row.get("team_id") if isinstance(row.get("team_id"), str) else None, - "user_id": row.get("user_id") if isinstance(row.get("user_id"), str) else None, - "user_email": row.get("user_email") if isinstance(row.get("user_email"), str) else None, - } - return meta - - -def _token_digest_metadata( - records: Sequence[_TokenAliasRecord], - wanted: AbstractSet[str], -) -> Mapping[str, KeyMetadataDict]: - return MappingProxyType( - { - digested: _record_metadata(record) - for record in records - for digested in (hash_token(record.token),) - if digested in wanted - } - ) - - -async def _paginate_token_digest_metadata( - load_page: Callable[[int], Awaitable[Sequence[_TokenAliasRecord] | None]], - wanted: AbstractSet[str], - *, - page_size: int, - skip: int = 0, - accumulated: Mapping[str, KeyMetadataDict] = _EMPTY_KEY_METADATA, -) -> Mapping[str, KeyMetadataDict]: - if not wanted: - return accumulated - records: Final = await load_page(skip) - if records is None: - return accumulated - page_hits: Final = _token_digest_metadata(records, wanted) - combined: Final[Mapping[str, KeyMetadataDict]] = ( - MappingProxyType({**accumulated, **page_hits}) if page_hits else accumulated - ) - still_wanted: Final = wanted - frozenset(page_hits) - if not still_wanted or len(records) < page_size: - return combined - return await _paginate_token_digest_metadata( - load_page, - still_wanted, - page_size=page_size, - skip=skip + page_size, - accumulated=combined, - ) - - -async def _reverse_hash_active_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], - *, - page_size: int, -) -> Mapping[str, KeyMetadataDict]: - async def load_page(skip: int) -> Sequence[_TokenAliasRecord] | None: - return await _db_or_empty( - lambda: VerificationTokenRepository(prisma_client).table.find_many( - take=page_size, - skip=skip, - order={"token": "asc"}, # mutable-ok: Prisma find_many order= is a dict - ), - "Failed reverse-hash recovery against active keys for %d missing keys: %s", - len(wanted), - ) - - return await _paginate_token_digest_metadata(load_page, wanted, page_size=page_size) - - -async def _reverse_hash_deleted_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], - *, - page_size: int, -) -> Mapping[str, KeyMetadataDict]: - async def load_page(skip: int) -> Sequence[_TokenAliasRecord] | None: - return await _db_or_empty( - lambda: DeletedVerificationTokenRepository(prisma_client).table.find_many( - take=page_size, - skip=skip, - order=[{"deleted_at": "desc"}, {"id": "asc"}], # mutable-ok: Prisma find_many order= is a dict - ), - "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", - len(wanted), - ) - - return await _paginate_token_digest_metadata(load_page, wanted, page_size=page_size) - - async def _reverse_hash_key_metadata( prisma_client: PrismaClient, + sql: str, wanted: AbstractSet[str], *, - page_size: int, + warning: str, ) -> Mapping[str, KeyMetadataDict]: - from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted, page_size=page_size) - still_wanted: Final = wanted - frozenset(from_active) - if not still_wanted: - return from_active - from_deleted: Final = await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted, page_size=page_size) - return MappingProxyType({**from_active, **from_deleted}) - - -async def _spend_logs_key_metadata( - prisma_client: PrismaClient, - wanted: AbstractSet[str], -) -> Mapping[str, KeyMetadataDict]: - spend_log_rows: Final = await _db_or_empty( - lambda: prisma_client.db.query_raw( - _SPEND_LOGS_KEY_METADATA_SQL, - tuple(wanted), - ), - "Failed SpendLogs metadata recovery for %d missing keys: %s", + rows: Final = await _db_or_empty( + lambda: prisma_client.db.query_raw(sql, sorted(wanted)), + warning, len(wanted), ) - if not isinstance(spend_log_rows, list): + if rows is None: return _EMPTY_KEY_METADATA - return MappingProxyType( { - row["api_key"]: _spend_log_row_metadata(row) - for row in spend_log_rows - if isinstance(row, dict) and isinstance(row.get("api_key"), str) and row["api_key"] in wanted + row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id) + for row in _TOKEN_DIGEST_ROWS.validate_python(rows) + if row.digest in wanted } ) @@ -223,7 +92,7 @@ async def _emails_for_user_ids( return _EMPTY_EMAILS users: Final = await _db_or_empty( lambda: UserRepository(prisma_client).table.find_many( - where={"user_id": {"in": tuple(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict + where={"user_id": {"in": list(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict ), "Failed user_email recovery for %d user ids: %s", len(user_ids), @@ -268,31 +137,35 @@ async def attach_user_emails( async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], - *, - token_scan_page_size: int = _TOKEN_SCAN_PAGE, ) -> Mapping[str, KeyMetadataDict]: """ - Recover key_alias/team_id/user_email for DailyUserSpend.api_key values that + Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that were double-hashed by the v1.99 spend-log provenance gate. Those rows store hash(VerificationToken.token) instead of the token, so the - exact join misses. Page through active then deleted tokens until every - wanted digest is found or the table ends; fall back to SpendLogs metadata. - Emails come from SpendLogs when present, otherwise from UserTable via the - recovered key's user_id. + exact join misses. Postgres hashes the token column itself, one pass over + active keys and one over deleted keys, so no key row crosses the wire. """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: return _EMPTY_KEY_METADATA - from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing, page_size=token_scan_page_size) - still_missing: Final = sha_missing - frozenset(from_tokens) - recovered: Final = ( - from_tokens - if not still_missing - else MappingProxyType({**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}) + from_active: Final = await _reverse_hash_key_metadata( + prisma_client, + _ACTIVE_TOKEN_DIGEST_SQL, + sha_missing, + warning="Failed reverse-hash recovery against active keys for %d missing keys: %s", ) - return await attach_user_emails(prisma_client, recovered) + still_missing: Final = sha_missing - frozenset(from_active) + if not still_missing: + return from_active + from_deleted: Final = await _reverse_hash_key_metadata( + prisma_client, + _DELETED_TOKEN_DIGEST_SQL, + still_missing, + warning="Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + ) + return MappingProxyType({**from_active, **from_deleted}) def _row_with_recovered_fields( @@ -345,7 +218,10 @@ async def fill_missing_api_key_aliases( if not missing_keys: return tuple(rows) - recovered: Final = await recover_double_hashed_key_metadata(prisma_client, missing_keys) + recovered: Final = await attach_user_emails( + prisma_client, + await recover_double_hashed_key_metadata(prisma_client, missing_keys), + ) if not recovered: return tuple(rows) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 4c7333143c5..a37c3ba4405 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -139,7 +139,6 @@ def _get_spend_logs_metadata( user_api_key_project_alias=None, user_api_key_org_id=None, user_api_key_user_id=None, - user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index c543156eedd..1b6e8ca513c 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -74,7 +74,7 @@ class TestCloudZeroHourlyExport: fake_db = MagicMock() async def query_raw_mock(query: str, *params): - if "LiteLLM_SpendLogs" in query: + if "sha256(" in query: return [] start_time_utc = params[0] if len(params) > 0 else None end_time_utc = params[1] if len(params) > 1 else None diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index b9c0d953086..37a54c4901a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -459,33 +459,23 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( """ v1.99 spend logging re-hashed already-hashed api_key values when provenance was missing. Usage joins DailyUserSpend.api_key to VerificationToken.token, so those - rows looked like key-hash-... with a null alias. Reverse-hash recovery must map - hash(token) back to the key's alias for historical dirty spend. + rows looked like key-hash-... with a null alias. Recovery asks Postgres for the + key whose hashed token matches the dirty value and maps it back to its alias. """ from litellm.proxy.utils import hash_token - token = "a" * 64 - double_hashed = hash_token(token) + double_hashed = hash_token("a" * 64) mock_prisma = MagicMock() - - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - side_effect=[ - [], # exact join miss - [ - SimpleNamespace( - token=token, - key_alias="batch-worker", - team_id="team-1", - user_id="alice", - ) - ], - ] - ) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_usertable.find_many = AsyncMock( return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] ) - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + {"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"} + ] + ) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -495,37 +485,37 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( assert result[double_hashed]["key_alias"] == "batch-worker" assert result[double_hashed]["team_id"] == "team-1" assert result[double_hashed]["user_email"] == "alice@example.com" - mock_prisma.db.query_raw.assert_not_called() + ((digest_sql, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + assert '"LiteLLM_VerificationToken"' in digest_sql + assert digests == [double_hashed] @pytest.mark.asyncio -async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): - """When the token tables cannot reverse-hash the dirty key, use SpendLogs metadata.""" +async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs(): + """A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan.""" from litellm.proxy.utils import hash_token double_hashed = hash_token("b" * 64) mock_prisma = MagicMock() mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.query_raw = AsyncMock( - return_value=[ - { - "api_key": double_hashed, - "key_alias": "from-spend-log", - "team_id": "team-spend", - } - ] - ) mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) result = await get_api_key_metadata( prisma_client=mock_prisma, api_keys={double_hashed}, ) - assert result[double_hashed]["key_alias"] == "from-spend-log" - assert result[double_hashed]["team_id"] == "team-spend" - mock_prisma.db.query_raw.assert_called_once() + assert double_hashed not in result + issued_sql = [call.args[0] for call in mock_prisma.db.query_raw.call_args_list] + assert len(issued_sql) == 2 + assert not any("LiteLLM_SpendLogs" in sql for sql in issued_sql) + token_lookups = ( + mock_prisma.db.litellm_verificationtoken.find_many.call_args_list + + mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list + ) + assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups) def test_key_metadata_includes_recovered_user_email(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 704de049fc8..43f7d20cf13 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -11,55 +12,126 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import ( from litellm.proxy.utils import hash_token +def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]: + return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id} + + +def _query_raw_by_table( + active_rows: Sequence[dict[str, str | None]], + deleted_rows: Sequence[dict[str, str | None]], +) -> AsyncMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: + if '"LiteLLM_VerificationToken"' in sql: + return list(active_rows) + if '"LiteLLM_DeletedVerificationToken"' in sql: + return list(deleted_rows) + raise AssertionError(f"unexpected query: {sql}") + + return AsyncMock(side_effect=query_raw) + + @pytest.mark.asyncio -async def test_recover_double_hashed_key_metadata_via_reverse_hash(): - token = "a" * 64 - double_hashed = hash_token(token) +async def test_recover_double_hashed_key_metadata_via_active_token_digest(): + double_hashed = hash_token("a" * 64) mock_prisma = MagicMock() - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[ - SimpleNamespace( - token=token, - key_alias="batch-worker", - team_id="team-1", - user_id="alice", - ) - ] + mock_prisma.db.query_raw = _query_raw_by_table( + active_rows=[_digest_row(double_hashed, "batch-worker", "team-1", "alice")], + deleted_rows=[], ) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] - ) - mock_prisma.db.query_raw = AsyncMock(return_value=[]) result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) assert result[double_hashed]["key_alias"] == "batch-worker" assert result[double_hashed]["team_id"] == "team-1" - assert result[double_hashed]["user_email"] == "alice@example.com" + assert result[double_hashed]["user_id"] == "alice" + ((_, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + assert digests == [double_hashed] + + +@pytest.mark.asyncio +async def test_recover_double_hashed_key_metadata_falls_back_to_deleted_tokens(): + double_hashed = hash_token("y" * 64) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_by_table( + active_rows=[], + deleted_rows=[_digest_row(double_hashed, "deleted-key", "team-del", "erin")], + ) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) + + assert result[double_hashed]["key_alias"] == "deleted-key" + assert result[double_hashed]["team_id"] == "team-del" + assert result[double_hashed]["user_id"] == "erin" + assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [[double_hashed], [double_hashed]] + + +@pytest.mark.asyncio +async def test_recover_only_asks_deleted_tokens_for_digests_active_keys_missed(): + found_active = hash_token("1" * 64) + found_deleted = hash_token("2" * 64) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_by_table( + active_rows=[_digest_row(found_active, "active-key", None, None)], + deleted_rows=[_digest_row(found_deleted, "deleted-key", None, None)], + ) + + result = await recover_double_hashed_key_metadata(mock_prisma, {found_active, found_deleted}) + + assert result[found_active]["key_alias"] == "active-key" + assert result[found_deleted]["key_alias"] == "deleted-key" + assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [ + sorted((found_active, found_deleted)), + [found_deleted], + ] + + +@pytest.mark.asyncio +async def test_recover_permanent_miss_costs_two_digest_lookups_and_no_table_walk(): + double_hashed = hash_token("b" * 64) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_by_table(active_rows=[], deleted_rows=[]) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) + + assert result == {} + assert len(mock_prisma.db.query_raw.call_args_list) == 2 + mock_prisma.db.litellm_verificationtoken.find_many.assert_not_called() + mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_skips_keys_that_are_not_sha256_digests(): + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_double_hashed_key_metadata(mock_prisma, {"sk-plain-key", "key-hash-short"}) + + assert result == {} mock_prisma.db.query_raw.assert_not_called() @pytest.mark.asyncio -async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows(): - token = "c" * 64 - double_hashed = hash_token(token) +async def test_recover_returns_empty_when_digest_lookup_raises_prisma_error(): + double_hashed = hash_token("c" * 64) mock_prisma = MagicMock() - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[ - SimpleNamespace( - token=token, - key_alias="recovered-alias", - team_id="team-9", - user_id="bob", - ) - ] + mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) + + result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) + + assert result == {} + + +@pytest.mark.asyncio +async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows(): + double_hashed = hash_token("d" * 64) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_by_table( + active_rows=[_digest_row(double_hashed, "recovered-alias", "team-9", "bob")], + deleted_rows=[], ) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_usertable.find_many = AsyncMock( return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")] ) - mock_prisma.db.query_raw = AsyncMock(return_value=[]) rows = ( { @@ -83,104 +155,18 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows(): assert filled[0]["api_key_alias"] == "recovered-alias" assert filled[0]["team_id"] == "team-9" assert filled[0]["user_email"] == "bob@example.com" + assert filled[0]["spend"] == 12.5 assert filled[1]["api_key_alias"] == "named-key" + assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["bob"]}} @pytest.mark.asyncio -async def test_recover_falls_back_to_spend_logs_when_token_scan_raises_prisma_error(): - token = "b" * 64 - double_hashed = hash_token(token) +async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_is_missing(): mock_prisma = MagicMock() - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=PrismaError("db down")) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(side_effect=PrismaError("db down")) - mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) - mock_prisma.db.query_raw = AsyncMock( - return_value=[ - { - "api_key": double_hashed, - "key_alias": "from-spend-logs", - "team_id": "team-sl", - "user_id": "carol", - "user_email": "carol@example.com", - } - ] - ) - - result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}) - - assert result[double_hashed]["key_alias"] == "from-spend-logs" - assert result[double_hashed]["team_id"] == "team-sl" - assert result[double_hashed]["user_email"] == "carol@example.com" - - -@pytest.mark.asyncio -async def test_recover_double_hashed_key_metadata_scans_past_first_page(): - token = "z" * 64 - double_hashed = hash_token(token) - decoys = ( - SimpleNamespace(token="1" * 64, key_alias="decoy-1", team_id=None, user_id=None), - SimpleNamespace(token="2" * 64, key_alias="decoy-2", team_id=None, user_id=None), - ) - match = SimpleNamespace(token=token, key_alias="late-key", team_id="team-late", user_id="dana") - - async def find_many(*, take: int | None = None, skip: int | None = None, order: object = None): - if skip == 0: - return list(decoys) - if skip == 2: - return [match] - return [] - - mock_prisma = MagicMock() - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=find_many) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=[SimpleNamespace(user_id="dana", user_email="dana@example.com")] - ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) + rows = ({"api_key": hash_token("e" * 64), "api_key_alias": "named", "user_email": "x@example.com"},) - result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}, token_scan_page_size=2) + filled = await fill_missing_api_key_aliases(mock_prisma, rows) - assert result[double_hashed]["key_alias"] == "late-key" - assert result[double_hashed]["team_id"] == "team-late" - assert result[double_hashed]["user_email"] == "dana@example.com" - assert [call.kwargs["skip"] for call in mock_prisma.db.litellm_verificationtoken.find_many.call_args_list] == [0, 2] - mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called() - mock_prisma.db.query_raw.assert_not_called() - - -@pytest.mark.asyncio -async def test_recover_double_hashed_key_metadata_pages_deleted_tokens(): - token = "y" * 64 - double_hashed = hash_token(token) - decoys = ( - SimpleNamespace(token="3" * 64, key_alias="deleted-decoy-1", team_id=None, user_id=None), - SimpleNamespace(token="4" * 64, key_alias="deleted-decoy-2", team_id=None, user_id=None), - ) - match = SimpleNamespace(token=token, key_alias="deleted-late-key", team_id="team-del", user_id="erin") - - async def find_deleted(*, take: int | None = None, skip: int | None = None, order: object = None): - if skip == 0: - return list(decoys) - if skip == 2: - return [match] - return [] - - mock_prisma = MagicMock() - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(side_effect=find_deleted) - mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=[SimpleNamespace(user_id="erin", user_email="erin@example.com")] - ) - mock_prisma.db.query_raw = AsyncMock(return_value=[]) - - result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed}, token_scan_page_size=2) - - assert result[double_hashed]["key_alias"] == "deleted-late-key" - assert result[double_hashed]["user_email"] == "erin@example.com" - assert [ - call.kwargs["skip"] for call in mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list - ] == [ - 0, - 2, - ] + assert filled == rows mock_prisma.db.query_raw.assert_not_called() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index dfdad545003..73a29afd9b9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -511,7 +511,6 @@ ignored_keys = [ "metadata.user_api_key_project_alias", "metadata.user_api_key_org_id", "metadata.user_api_key_user_id", - "metadata.user_api_key_user_email", "metadata.user_api_key_team_alias", "metadata.spend_logs_metadata", "metadata.requester_ip_address", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 3cfc16eedec..fb0cdc175b1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2729,7 +2729,6 @@ def test_get_logging_payload_batch_attribution_keeps_verification_token_hash(): "user_api_key_hash": token_hash, "user_api_key_alias": "batch-creator", "user_api_key_user_id": "alice", - "user_api_key_user_email": "alice@example.com", "user_api_key_team_id": "team-1", } }, @@ -2746,7 +2745,6 @@ def test_get_logging_payload_batch_attribution_keeps_verification_token_hash(): parsed_meta = json.loads(payload["metadata"]) assert parsed_meta["user_api_key"] == token_hash assert parsed_meta["user_api_key_alias"] == "batch-creator" - assert parsed_meta["user_api_key_user_email"] == "alice@example.com" def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): From 33815682bc20ec5b32d9b363995608fc29ee357a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 4 Sep 2026 01:38:26 +0000 Subject: [PATCH 27/35] fix(spend): keep CloudZero and Focus spend-user email when filling alias Export rows already join user_email from DailyUserSpend.user_id. Recovered key-owner email must not replace that when only the alias join missed. Co-authored-by: Mateo Wang --- .../spend_tracking/key_metadata_recovery.py | 2 +- .../test_key_metadata_recovery.py | 29 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index d524b158c9c..2940c9a8337 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -186,7 +186,7 @@ def _row_with_recovered_fields( **row, alias_field: meta.get("key_alias") or row.get(alias_field), team_id_field: meta.get("team_id") or row.get(team_id_field), - user_email_field: meta.get("user_email") or row.get(user_email_field), + user_email_field: row.get(user_email_field) or meta.get("user_email"), } ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 43f7d20cf13..d76d355ec0e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -170,3 +170,32 @@ async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_i assert filled == rows mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_fill_missing_api_key_aliases_keeps_spend_user_email_when_alias_is_missing(): + double_hashed = hash_token("f" * 64) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_by_table( + active_rows=[_digest_row(double_hashed, "team-key", "team-9", "key-owner")], + deleted_rows=[], + ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="key-owner", user_email="owner@example.com")] + ) + + rows = ( + { + "api_key": double_hashed, + "api_key_alias": None, + "team_id": None, + "user_email": "spender@example.com", + "spend": 4.0, + }, + ) + + filled = await fill_missing_api_key_aliases(mock_prisma, rows) + + assert filled[0]["api_key_alias"] == "team-key" + assert filled[0]["team_id"] == "team-9" + assert filled[0]["user_email"] == "spender@example.com" From 048499cdf55959753c04aae7d4be995cc3606a2c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 4 Sep 2026 01:44:18 +0000 Subject: [PATCH 28/35] fix(spend): only reverse-hash export rows whose key alias join missed Team and service keys often have no user_email after a successful token join. Treating empty email as a miss hashed every verification token on routine CloudZero and Focus exports. Co-authored-by: Mateo Wang --- .../spend_tracking/key_metadata_recovery.py | 4 +--- .../test_key_metadata_recovery.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 2940c9a8337..7de18521edd 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -211,9 +211,7 @@ async def fill_missing_api_key_aliases( key for row in rows for key in (row.get(api_key_field),) - if isinstance(key, str) - and key - and (row.get(alias_field) in (None, "") or row.get(user_email_field) in (None, "")) + if isinstance(key, str) and key and row.get(alias_field) in (None, "") ) if not missing_keys: return tuple(rows) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index d76d355ec0e..7a80319239d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -199,3 +199,22 @@ async def test_fill_missing_api_key_aliases_keeps_spend_user_email_when_alias_is assert filled[0]["api_key_alias"] == "team-key" assert filled[0]["team_id"] == "team-9" assert filled[0]["user_email"] == "spender@example.com" + + +@pytest.mark.asyncio +async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email(): + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + rows = ( + { + "api_key": hash_token("g" * 64), + "api_key_alias": "service-key", + "team_id": "team-svc", + "user_email": None, + }, + ) + + filled = await fill_missing_api_key_aliases(mock_prisma, rows) + + assert filled == rows + mock_prisma.db.query_raw.assert_not_called() From 2b7e14872f7d4bb776ef60c8168bd7e78778a86b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 18:46:01 -0700 Subject: [PATCH 29/35] fix(spend-tracking): hand plain dict rows to polars in the CloudZero and Focus exports --- litellm/integrations/cloudzero/database.py | 2 +- litellm/integrations/focus/database.py | 2 +- .../integrations/cloudzero/test_cloudzero.py | 24 +++++++++++++++++++ .../integrations/focus/test_focus_database.py | 22 +++++++++++++++++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 4adf725fd0f..2fb10ad8a96 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -106,6 +106,6 @@ class LiteLLMDatabase: else [] ) recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) - return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) + return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index f214aa02b5b..657c7e0d264 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -108,7 +108,7 @@ class FocusLiteLLMDatabase: else [] ) recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows) - return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None) + return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 1b6e8ca513c..6ddb8cbaa7c 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -1,3 +1,4 @@ +import hashlib from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -165,3 +166,26 @@ class TestCloudZeroHourlyExport: logger = CloudZeroLogger(api_key="test", connection_id="test") await logger._hourly_usage_data_export() + + +class TestLiteLLMDatabaseUsageData: + @pytest.mark.asyncio + async def test_builds_frame_from_rows_recovered_for_double_hashed_keys(self, monkeypatch: pytest.MonkeyPatch): + double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest() + joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1} + dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5} + + async def query_raw(query: str, *params): + if "sha256(" in query: + return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}] + return [joined_row, dirty_row] + + fake_client = MagicMock() + fake_client.db.query_raw = AsyncMock(side_effect=query_raw) + db = LiteLLMDatabase() + monkeypatch.setattr(db, "_ensure_prisma_client", lambda: fake_client) + + result = await db.get_usage_data() + + assert result["api_key_alias"].to_list() == ["joined", "batch-worker"] + assert result["team_id"].to_list() == ["team-0", "team-1"] diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index 5c13665f1f1..06240eac387 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -1,5 +1,6 @@ """Tests for FocusLiteLLMDatabase query construction.""" +import hashlib from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import AsyncMock @@ -87,3 +88,24 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch): ) assert "ot.organization_alias as organization_alias" in query_text assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text + + +@pytest.mark.asyncio +async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch): + double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest() + joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1} + dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5} + + async def query_raw(query: str, *params): + if "sha256(" in query: + return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}] + return [joined_row, dirty_row] + + mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw))) + db = FocusLiteLLMDatabase() + monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client) + + result = await db.get_usage_data() + + assert result["api_key_alias"].to_list() == ["joined", "batch-worker"] + assert result["team_id"].to_list() == ["team-0", "team-1"] From 7276caecd42c13562ad833bbeba4cb9460c71466 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 4 Sep 2026 08:17:07 -0700 Subject: [PATCH 30/35] refactor(rust): extract config crate (#39706) * refactor(rust): extract config crate * refactor(config): split crate modules * refactor(gateway): remove gil health counter --- litellm-rust/AGENTS.md | 5 +- litellm-rust/CLAUDE.md | 1 + litellm-rust/Cargo.lock | 12 ++- litellm-rust/Cargo.toml | 2 + litellm-rust/README.md | 4 +- litellm-rust/crates/ai-gateway/AGENTS.md | 9 +-- .../crates/ai-gateway/ARCHITECTURE.md | 2 + litellm-rust/crates/ai-gateway/Cargo.toml | 4 +- litellm-rust/crates/ai-gateway/README.md | 22 ++++-- litellm-rust/crates/ai-gateway/config.yaml | 6 +- litellm-rust/crates/ai-gateway/src/gil.rs | 58 -------------- litellm-rust/crates/ai-gateway/src/lib.rs | 10 +-- litellm-rust/crates/ai-gateway/src/main.rs | 10 +-- .../crates/ai-gateway/src/python/AGENTS.md | 27 ------- .../crates/ai-gateway/src/python/config.rs | 37 --------- .../crates/ai-gateway/src/python/mod.rs | 4 - .../crates/ai-gateway/src/routes/AGENTS.md | 2 +- .../crates/ai-gateway/src/routes/gil.rs | 30 -------- .../crates/ai-gateway/src/routes/mod.rs | 4 +- litellm-rust/crates/config/Cargo.toml | 16 ++++ litellm-rust/crates/config/src/error.rs | 11 +++ litellm-rust/crates/config/src/lib.rs | 7 ++ litellm-rust/crates/config/src/python.rs | 76 +++++++++++++++++++ .../core/tests/workspace_crate_allowlist.rs | 14 +++- 24 files changed, 173 insertions(+), 200 deletions(-) delete mode 100644 litellm-rust/crates/ai-gateway/src/gil.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/python/AGENTS.md delete mode 100644 litellm-rust/crates/ai-gateway/src/python/config.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/python/mod.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/routes/gil.rs create mode 100644 litellm-rust/crates/config/Cargo.toml create mode 100644 litellm-rust/crates/config/src/error.rs create mode 100644 litellm-rust/crates/config/src/lib.rs create mode 100644 litellm-rust/crates/config/src/python.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index b8b6291283d..b8d1f2db4d7 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,17 +1,18 @@ # AGENTS.md -litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. +litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | +| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. +Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate. ## Where a route lives diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index d9c944529df..dfacf37b6cd 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -24,6 +24,7 @@ the base when behavior is genuinely different, and say so explicitly in the PR. ## Crates (see AGENTS.md) `litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. +`litellm-config` is the config-loading boundary and returns resolved core types. `litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and `litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 803df633c27..62e943d0f42 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1412,8 +1412,8 @@ dependencies = [ "base64", "futures-channel", "futures-util", + "litellm-config", "litellm-core", - "pyo3", "reqwest", "serde", "serde_json", @@ -1425,6 +1425,16 @@ dependencies = [ "tracing", ] +[[package]] +name = "litellm-config" +version = "0.1.0" +dependencies = [ + "litellm-core", + "pyo3", + "serde_json", + "thiserror 2.0.19", +] + [[package]] name = "litellm-core" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 62f62872dd7..720c4545181 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/core", + "crates/config", "crates/ai-gateway", "crates/python-interop", "crates/python-bridge", @@ -17,6 +18,7 @@ repository = "https://github.com/BerriAI/litellm" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } +litellm-config = { path = "crates/config" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } axum = "0.7" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index e43dc7ea6ad..650d38753e7 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -25,11 +25,12 @@ coverage and production evidence. | Crate | Role | |-------|------| | litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | +| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | | litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. +Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. ## Layout @@ -38,6 +39,7 @@ crates/ core/ The SDK: route modules + provider transforms. src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client src/providers/anthropic/messages/transformation.rs + config/ Config loading and resolved deployments. ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. python-interop/ Domain-neutral PyO3 conversion and GIL primitives. python-bridge/ PyO3 API adapter for Python LiteLLM. diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md index 92567091cd3..b2fd583316b 100644 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -9,19 +9,15 @@ such as `litellm_core::messages::messages`. No provider handler lives here. src/ main.rs # entrypoint: build AppState (router + master key), bind, serve state.rs # AppState — shared Arc + master_key - gil.rs # GIL-activity tracker (records Python acquisitions) auth/ # authentication as an axum extractor — added to handler args mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY) routes/ # one module per route, all matching the same template AGENTS.md # ← the route template (read this before adding a route) mod.rs # app(): merges every module's router() health.rs # simple route (one file): router() + liveness/readiness - gil.rs # simple route (one file): router() + GET /health/gil realtime/ # route with logic → axum surface + a no-axum service: mod.rs # router() + handler + WS<->events adapter (the axum surface) service.rs # business logic (select deployment, call provider) — no axum, testable - python/ # Python interop (feature: python-config) — load-time only - mod.rs, config.rs, AGENTS.md ``` ## Rules @@ -53,5 +49,6 @@ proxy in a later phase. Health routes don't add the extractor (unauthenticated). ## Python interop -Anything that calls into Python lives in `python/` and is **load-time only** — see -`python/AGENTS.md`. The realtime data path never takes the GIL. +Python-backed loading lives in `litellm-config` and is **load-time only**. The +gateway's `python-config` feature forwards to that crate. The realtime data path +never takes the GIL. diff --git a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md index 733953bbdb3..6d090cf4c8e 100644 --- a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md +++ b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md @@ -9,4 +9,6 @@ flowchart LR C[client] <--> G[Rust ai-gateway
LLM inference] G <--> O[OpenAI realtime] G -. spend tracking callback .-> P[litellm proxy] + F[litellm-config
load-time only] --> G + F -. Python backend .-> P ``` diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index eef2bf55a07..10369fa3bfd 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -16,6 +16,7 @@ required-features = ["server"] [dependencies] tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-config.workspace = true # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true @@ -31,7 +32,6 @@ subtle = { workspace = true, optional = true } # sha2 hashes the master key into user_api_key_hash (matches the proxy's # SHA-256 hash_token) so the plaintext credential never enters a log payload. sha2 = { workspace = true, optional = true } -pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } tower = { version = "0.5.3", features = ["util"], optional = true } [features] @@ -39,7 +39,7 @@ default = [] server = ["dep:axum", "dep:subtle", "dep:sha2"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). -python-config = ["dep:pyo3"] +python-config = ["litellm-config/python"] trace-parity = ["server", "dep:tower", "litellm-core/observability"] [dev-dependencies] diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 1675e6f1b16..9fef59a277d 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,25 +6,30 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route: +`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | +| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. +Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) -- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil` +- **Health:** `GET /health/readiness`, `GET /health/liveness` - **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging)) > **Realtime serving is pure Rust.** Python is used at **load time only** — to > read the config once at boot. The realtime hot path never touches Python. +The former `/health/gil` route and its acquisition counter were removed. They +only observed the single startup config load and did not prove that every GIL +acquisition was instrumented + ## Configuration (config.yaml) The gateway loads its `model_list` from a **config.yaml**, the same as the @@ -43,9 +48,10 @@ model_list: LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway ``` -At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the -**real proxy config reader** (`ProxyConfig.get_config`). That means everything -the proxy supports in config.yaml works here too: +At boot `litellm-config` calls into `litellm.proxy.read_model_list` and returns +resolved deployments to the gateway, which constructs the router. The Python +backend still reuses the **real proxy config reader** (`ProxyConfig.get_config`), +so everything the proxy supports in config.yaml works here too: - `include:` to merge in other config files, - `os.environ/VAR` secret references (resolved via the secret manager, never @@ -82,8 +88,8 @@ stand-in built from the environment: |---|---|---| | `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). | -This mode links no libpython and needs no config file, but it only supports one -hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the +The default workspace build links no libpython and needs no config file. This +fallback mode only supports one hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the stand-in only for the leanest possible build. ## Request logging diff --git a/litellm-rust/crates/ai-gateway/config.yaml b/litellm-rust/crates/ai-gateway/config.yaml index ac598c220dd..321801f6862 100644 --- a/litellm-rust/crates/ai-gateway/config.yaml +++ b/litellm-rust/crates/ai-gateway/config.yaml @@ -1,8 +1,8 @@ # Sample realtime config for the LiteLLM Rust AI Gateway. # -# The gateway loads this model_list at boot via the embedded python config -# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader — -# so include:, os.environ/ secrets, and DB-stored models all work here too. +# litellm-config resolves this model_list at boot through the Python config +# reader (litellm.proxy.read_model_list), then the gateway builds its router. +# Includes, environment secrets, and database-stored models still work. # # Secrets are referenced (never inlined) via os.environ/. A real deploy can # override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH). diff --git a/litellm-rust/crates/ai-gateway/src/gil.rs b/litellm-rust/crates/ai-gateway/src/gil.rs deleted file mode 100644 index c749f722c73..00000000000 --- a/litellm-rust/crates/ai-gateway/src/gil.rs +++ /dev/null @@ -1,58 +0,0 @@ -//! GIL-activity tracking. -//! -//! Every acquisition of the Python GIL is recorded here so the `/health/gil` -//! endpoint can report whether Python was touched recently. The design goal is -//! that the GIL is acquired **only at load time** (config read) and never on the -//! realtime hot path — polling this endpoint during traffic should show the -//! count holding steady and `acquired_last_30s` falling to `false`. - -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{SystemTime, UNIX_EPOCH}; - -/// Window (seconds) for the "recently acquired" signal. -pub const RECENT_WINDOW_SECS: u64 = 30; - -static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0); -/// Unix seconds of the last acquisition; `0` means "never". -static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0); - -fn now_unix_secs() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) -} - -/// Record that the GIL was just acquired. Call immediately before taking the GIL. -/// -/// Only invoked under the `python-config` feature; without it the gateway never -/// touches Python, so the recorder is unused (and the endpoint reports zero). -#[cfg_attr(not(feature = "python-config"), allow(dead_code))] -pub fn record_acquisition() { - GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed); - LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed); -} - -/// Point-in-time view of GIL activity. -pub struct GilSnapshot { - pub total_acquisitions: u64, - pub seconds_since_last: Option, - pub acquired_last_30s: bool, -} - -/// Read the current GIL-activity snapshot. -pub fn snapshot() -> GilSnapshot { - let total = GIL_ACQUISITIONS.load(Ordering::Relaxed); - let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed); - let seconds_since_last = if last == 0 { - None - } else { - Some(now_unix_secs().saturating_sub(last)) - }; - let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS); - GilSnapshot { - total_acquisitions: total, - seconds_since_last, - acquired_last_30s, - } -} diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index a2950748afc..08fbde564ed 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -10,18 +10,13 @@ //! - [`io`]: compatibility exports and realtime WebSocket splice helpers. //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` -//! binary turns on. The `python-config` feature additionally pulls in [`python`] -//! for the load-time config reader. +//! binary turns on. pub mod audio_transcription; mod client; pub mod io; pub mod ocr; -/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and -/// the `python-config` reader, so it is available without either feature. -pub mod gil; - #[cfg(feature = "server")] pub mod auth; #[cfg(feature = "server")] @@ -35,6 +30,3 @@ mod constants; pub mod integrations; #[cfg(feature = "server")] mod realtime; - -#[cfg(feature = "python-config")] -pub mod python; diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index da3a486d4ee..88d7b1dbcf8 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -14,12 +14,12 @@ use std::sync::Arc; use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key}; use litellm_ai_gateway::routes; use litellm_ai_gateway::state::AppState; +#[cfg(feature = "python-config")] +use litellm_config::load_model_list; use litellm_core::router::{Deployment, LiteLLMParams, Router}; use litellm_ai_gateway::integrations::custom_logger::CustomLogger; use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger; -#[cfg(feature = "python-config")] -use litellm_ai_gateway::python; /// Bind to localhost by default so the gateway is not a public, unauthenticated /// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`). @@ -124,10 +124,10 @@ fn resolve_port() -> u16 { fn build_router() -> Router { #[cfg(feature = "python-config")] if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") { - match python::config::load_router_from_config(&config_path) { - Ok(router) => { + match load_model_list(std::path::Path::new(&config_path)) { + Ok(deployments) => { eprintln!("loaded model_list from {config_path} via python config reader"); - return router; + return Router::new(deployments); } Err(err) => { eprintln!("config load failed ({err}); falling back to env deployment"); diff --git a/litellm-rust/crates/ai-gateway/src/python/AGENTS.md b/litellm-rust/crates/ai-gateway/src/python/AGENTS.md deleted file mode 100644 index 47aa117e0b9..00000000000 --- a/litellm-rust/crates/ai-gateway/src/python/AGENTS.md +++ /dev/null @@ -1,27 +0,0 @@ -# ai-gateway/src/python — Python interop (load-time only) - -Functions here embed the Python interpreter (pyo3) and take the GIL to call into -`litellm` (e.g. read the proxy `model_list`). Compiled only under the -`python-config` feature. - -## Hard rule: non-hot-path functions only - -Everything in this folder MUST run **at most once per process lifetime — at -startup / load time** (config read, warm-up). NEVER call into Python on the -request path: - -- No GIL acquisition per request, per connection, or per realtime event. -- No Python call inside a route handler, the router's hot path, or any loop that - scales with traffic. - -**Why:** the GIL serializes execution and would cap throughput; the realtime data -path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll -`GET /health/gil`, and `total_acquisitions` MUST stay flat under load. - -## How to add one - -Resolve whatever Python-derived data you need **once at boot** and hand the rest -of the gateway an owned, plain-Rust value (e.g. build a `Router` from the -resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()` -immediately before taking the GIL. If a function would need to run per request, -it does not belong here — move the work to Rust, or pre-resolve it at startup. diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs deleted file mode 100644 index d5a4dd69c8d..00000000000 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Build the router by calling the Python proxy config reader (load time only). -//! -//! Embeds the interpreter via pyo3 and calls -//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's -//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot** -//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python. -//! -//! Compiled only under the `python-config` feature. -use litellm_core::error::Error; -use litellm_core::router::{Deployment, Router}; -use pyo3::prelude::*; - -use crate::gil; - -/// Load the router's `model_list` from `config_path` via the Python reader. -pub fn load_router_from_config(config_path: &str) -> Result { - gil::record_acquisition(); - Python::attach(|py| { - let model_list = py - .import("litellm.proxy.read_model_list") - .and_then(|module| module.getattr("read_model_list")) - .and_then(|reader| reader.call1((config_path,))) - .map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?; - - let model_list_json: String = py - .import("json") - .and_then(|json| json.getattr("dumps")) - .and_then(|dumps| dumps.call1((model_list,))) - .and_then(|encoded| encoded.extract()) - .map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?; - - let deployments: Vec = serde_json::from_str(&model_list_json) - .map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?; - - Ok(Router::new(deployments)) - }) -} diff --git a/litellm-rust/crates/ai-gateway/src/python/mod.rs b/litellm-rust/crates/ai-gateway/src/python/mod.rs deleted file mode 100644 index a677bade676..00000000000 --- a/litellm-rust/crates/ai-gateway/src/python/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path -//! only.** Compiled only under the `python-config` feature. - -pub mod config; diff --git a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md index 3eee43e7a2f..c675916f71a 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md @@ -13,7 +13,7 @@ private). This is the norm — don't split until it hurts. pub fn router() -> Router { Router::new().route(PATH, get(handle)) } async fn handle(...) -> impl IntoResponse { ... } ``` -`health.rs` and `gil.rs` are examples. +`health.rs` is the example. ## Split out `service` when there's real logic When a route has business logic worth testing without axum, put it in a sibling diff --git a/litellm-rust/crates/ai-gateway/src/routes/gil.rs b/litellm-rust/crates/ai-gateway/src/routes/gil.rs deleted file mode 100644 index 0db0c6f0b14..00000000000 --- a/litellm-rust/crates/ai-gateway/src/routes/gil.rs +++ /dev/null @@ -1,30 +0,0 @@ -//! `GET /health/gil` — poll to confirm Python is only touched at load time. -//! Simple-route template: a `router()` plus its handler, in one file. - -use axum::routing::get; -use axum::{Json, Router}; -use serde::Serialize; - -use crate::gil; -use crate::state::AppState; - -/// This route's contribution to the app router. -pub fn router() -> Router { - Router::new().route("/health/gil", get(status)) -} - -#[derive(Debug, Serialize)] -struct GilStatusResponse { - gil_acquired_last_30s: bool, - total_acquisitions: u64, - seconds_since_last: Option, -} - -async fn status() -> Json { - let snapshot = gil::snapshot(); - Json(GilStatusResponse { - gil_acquired_last_30s: snapshot.acquired_last_30s, - total_acquisitions: snapshot.total_acquisitions, - seconds_since_last: snapshot.seconds_since_last, - }) -} diff --git a/litellm-rust/crates/ai-gateway/src/routes/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/mod.rs index c26be8ffee3..71b05c7d64b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/mod.rs @@ -2,10 +2,9 @@ //! //! **Template:** every route module exposes `pub fn router() -> Router` //! that mounts its own paths; [`app`] merges them. A trivial route is a single -//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with +//! file (`health.rs`); a non-trivial one is a folder (`realtime/`) with //! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md. -pub mod gil; pub mod health; pub mod messages; pub mod realtime; @@ -19,7 +18,6 @@ use crate::state::AppState; pub fn app(state: AppState) -> Router { Router::new() .merge(health::router()) - .merge(gil::router()) .merge(messages::router()) .merge(realtime::router()) .merge(responses::router()) diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml new file mode 100644 index 00000000000..ae9710266a3 --- /dev/null +++ b/litellm-rust/crates/config/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-config" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core.workspace = true +pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } +serde_json.workspace = true +thiserror.workspace = true + +[features] +default = [] +python = ["dep:pyo3"] diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs new file mode 100644 index 00000000000..cec7bc5c110 --- /dev/null +++ b/litellm-rust/crates/config/src/error.rs @@ -0,0 +1,11 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("read_model_list failed: {0}")] + PythonLoading(String), + #[error("serializing model_list failed: {0}")] + Serialization(String), + #[error("parsing model_list failed: {0}")] + ModelListParsing(#[source] serde_json::Error), +} diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs new file mode 100644 index 00000000000..655affbb0b7 --- /dev/null +++ b/litellm-rust/crates/config/src/lib.rs @@ -0,0 +1,7 @@ +mod error; +#[cfg(feature = "python")] +mod python; + +pub use error::Error; +#[cfg(feature = "python")] +pub use python::load_model_list; diff --git a/litellm-rust/crates/config/src/python.rs b/litellm-rust/crates/config/src/python.rs new file mode 100644 index 00000000000..fdad5027baa --- /dev/null +++ b/litellm-rust/crates/config/src/python.rs @@ -0,0 +1,76 @@ +use std::path::Path; + +use litellm_core::router::Deployment; +use pyo3::prelude::*; + +use crate::Error; + +pub fn load_model_list(config_path: &Path) -> Result, Error> { + Python::attach(|python| { + let model_list = python + .import("litellm.proxy.read_model_list") + .and_then(|module| module.getattr("read_model_list")) + .and_then(|reader| reader.call1((config_path.to_string_lossy().as_ref(),))) + .map_err(|error| Error::PythonLoading(error.to_string()))?; + + let model_list_json = python + .import("json") + .and_then(|json| json.getattr("dumps")) + .and_then(|dumps| dumps.call1((model_list,))) + .and_then(|encoded| encoded.extract::()) + .map_err(|error| Error::Serialization(error.to_string()))?; + + parse_model_list(&model_list_json) + }) +} + +fn parse_model_list(model_list_json: &str) -> Result, Error> { + serde_json::from_str(model_list_json).map_err(Error::ModelListParsing) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_resolved_model_list() { + let deployments = parse_model_list( + r#"[ + { + "model_name": "realtime", + "litellm_params": { + "model": "openai/gpt-realtime", + "api_key": "resolved-secret", + "api_base": "https://api.example.test/v1" + } + }, + { + "model_name": "without-optional-values", + "litellm_params": {"model": "openai/gpt-4.1"} + } + ]"#, + ) + .expect("resolved model list should parse"); + + assert_eq!(deployments.len(), 2); + assert_eq!(deployments[0].model_name, "realtime"); + assert_eq!( + deployments[0].litellm_params.api_key.as_deref(), + Some("resolved-secret") + ); + assert_eq!( + deployments[0].litellm_params.api_base.as_deref(), + Some("https://api.example.test/v1") + ); + assert_eq!(deployments[1].litellm_params.api_key, None); + assert_eq!(deployments[1].litellm_params.api_base, None); + } + + #[test] + fn malformed_model_list_returns_parsing_error() { + let error = parse_model_list(r#"[{"model_name":"missing-params"}]"#) + .expect_err("missing litellm_params should fail"); + + assert!(matches!(error, Error::ModelListParsing(_))); + } +} diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index 8a8a5ea263a..fc0ab2b62a3 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,6 +1,7 @@ -//! Enforcement: the litellm-rust workspace has exactly four crates. +//! Enforcement: the litellm-rust workspace has exactly five crates. //! -//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host), +//! `core` (the Rust SDK), `config` (the config-loading boundary), +//! `ai-gateway` (the HTTP/WebSocket host), //! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the //! PyO3 cdylib). Adding or removing a crate must be a //! deliberate act: this test fails until the allowlist here is updated, forcing @@ -19,13 +20,20 @@ use std::path::{Path, PathBuf}; /// workspace legitimately gains or loses a crate. const EXPECTED_MEMBERS: &[&str] = &[ "crates/core", + "crates/config", "crates/ai-gateway", "crates/python-interop", "crates/python-bridge", ]; /// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"]; +const EXPECTED_CRATE_DIRS: &[&str] = &[ + "core", + "config", + "ai-gateway", + "python-interop", + "python-bridge", +]; const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; From b75ac5cf525489c59fa0e5ffe73ab6726f6f28c9 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 4 Sep 2026 08:40:44 -0700 Subject: [PATCH 31/35] feat(python): rename Rust rollout API (#39704) --- .../PROVIDER_CODING_STANDARDS.md | 2 +- litellm/__init__.py | 2 +- litellm/rust_bridge/__init__.py | 4 +- litellm/rust_bridge/configuration.py | 87 +++--------------- litellm/rust_bridge/ocr.py | 2 +- .../e2e_parity/sdk/ocr/fixtures/migrate.py | 5 +- .../e2e_parity/sdk/ocr/fixtures/record.py | 5 +- .../unit_tests_mapping/cases/ocr.py | 2 +- .../test_rust_bridge_messages.py | 45 ++++++---- tests/test_litellm/ocr/test_rust_bridge.py | 89 +++++++++++-------- .../responses/test_rust_bridge_websocket.py | 4 +- .../rust_bridge/test_chat_completions.py | 4 +- .../rust_bridge/test_configuration.py | 52 ++++------- 13 files changed, 122 insertions(+), 181 deletions(-) diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index c0a29ab14bc..952bbc38b43 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -45,7 +45,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` 22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust. 23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms///` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method. -24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. +24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. ## Checks before push diff --git a/litellm/__init__.py b/litellm/__init__.py index 41a3789ab0d..42c0ea881fd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1421,7 +1421,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .rust_bridge import use_litellm_rust +from .rust_bridge import rust from .rag.main import * from .sandbox.main import * from .search.main import * diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index 9e8558bbf7d..8f6f4390b8a 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -1,10 +1,10 @@ """LiteLLM Rust bridge package.""" -from litellm.rust_bridge.configuration import use_litellm_rust +from litellm.rust_bridge.configuration import rust from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, reset_native_bridge_cache, ) -__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "use_litellm_rust"] +__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"] diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index d54b15f060c..515ab6edef1 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -2,13 +2,7 @@ from __future__ import annotations import os import warnings -from typing import TYPE_CHECKING, Final - -if TYPE_CHECKING: - from litellm.rust_bridge.messages import RustAmessages, RustMessages - from litellm.rust_bridge.ocr import RustAocr, RustOcr - from litellm.rust_bridge.responses_websocket import RustResponsesWebSocketConnection - from litellm.rust_bridge.transcription import RustAtranscription, RustTranscription +from typing import Final DEFAULT_RUST_ENABLED: Final = False _TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) @@ -16,13 +10,6 @@ _GLOBAL_ENV_NAME: Final = "LITELLM_RUST" _LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" -class _Unset: - pass - - -_UNSET: Final = _Unset() - - class _RustConfiguration: def __init__(self) -> None: self.override: bool | None = None @@ -42,7 +29,7 @@ def resolve_rust_enabled( request_override: bool | None, process_override: bool | None, environment_override: bool | None, - legacy_ocr_override: bool | None = None, + legacy_environment_override: bool | None = None, release_default: bool = DEFAULT_RUST_ENABLED, ) -> bool: if request_override is not None: @@ -51,25 +38,12 @@ def resolve_rust_enabled( return process_override if environment_override is not None: return environment_override - if legacy_ocr_override is not None: - return legacy_ocr_override + if legacy_environment_override is not None: + return legacy_environment_override return release_default def rust_enabled(*, request_override: bool | None = None) -> bool: - if request_override is not None: - return request_override - process_override: Final = _CONFIGURATION.override - if process_override is not None: - return process_override - return resolve_rust_enabled( - request_override=None, - process_override=None, - environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), - ) - - -def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: if request_override is not None: return request_override process_override: Final = _CONFIGURATION.override @@ -87,62 +61,21 @@ def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: request_override=None, process_override=None, environment_override=global_override, - legacy_ocr_override=legacy_override, + legacy_environment_override=legacy_override, ) +def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: + return rust_enabled(request_override=request_override) + + def reset_rust_configuration() -> None: _CONFIGURATION.override = None -def use_litellm_rust( - enabled: bool = True, - *, - ocr: RustOcr | None | _Unset = _UNSET, - aocr: RustAocr | None | _Unset = _UNSET, - messages: RustMessages | None | _Unset = _UNSET, - amessages: RustAmessages | None | _Unset = _UNSET, - responses_websocket: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, - transcription: RustTranscription | None | _Unset = _UNSET, - atranscription: RustAtranscription | None | _Unset = _UNSET, -) -> None: +def rust(enabled: bool) -> None: """Set the process override for optional Rust paths. Rust-only paths, including Bedrock transcription, are not controlled by this switch. """ _CONFIGURATION.override = enabled - bindings: Final = (ocr, aocr, messages, amessages, responses_websocket, transcription, atranscription) - if all(isinstance(binding, _Unset) for binding in bindings): - return - warnings.warn( - "Injecting Rust bridge implementations through use_litellm_rust() is deprecated; " - "use the internal bridge setters in tests", - DeprecationWarning, - stacklevel=2, - ) - - if not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset): - from litellm.rust_bridge.ocr import set_rust_ocr - - if not isinstance(ocr, _Unset): - set_rust_ocr(ocr=ocr) - if not isinstance(aocr, _Unset): - set_rust_ocr(aocr=aocr) - if not isinstance(messages, _Unset) or not isinstance(amessages, _Unset): - from litellm.rust_bridge.messages import set_rust_messages - - if not isinstance(messages, _Unset): - set_rust_messages(messages=messages) - if not isinstance(amessages, _Unset): - set_rust_messages(amessages=amessages) - if not isinstance(responses_websocket, _Unset): - from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket - - set_rust_responses_websocket(connection=responses_websocket) - if not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset): - from litellm.rust_bridge.transcription import configure_rust_transcription - - if not isinstance(transcription, _Unset): - configure_rust_transcription(transcription=transcription) - if not isinstance(atranscription, _Unset): - configure_rust_transcription(atranscription=atranscription) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index b5b0a35a498..86038438f57 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -11,7 +11,7 @@ from litellm.rust_bridge import configuration as _configuration from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds rust_ocr_enabled = _configuration.rust_ocr_enabled -use_litellm_rust = _configuration.use_litellm_rust +rust = _configuration.rust class RustOcr(Protocol): diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/migrate.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/migrate.py index 08f3cc66a42..c0e32123872 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/migrate.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/migrate.py @@ -5,7 +5,7 @@ from pathlib import Path from typing import Final, cast import litellm -from litellm.rust_bridge.ocr import use_litellm_rust +from litellm.rust_bridge.ocr import rust, set_rust_ocr from ......shared.parity.fixtures.recording import ( RecordedInteraction, UpstreamEndpoint, @@ -53,7 +53,8 @@ def main() -> None: parser.add_argument("--fixture-dir", type=Path, default=configured_fixture_directory()) args: Final = parser.parse_args() directory: Final = cast(Path, args.fixture_dir) - use_litellm_rust(False, ocr=None, aocr=None) + rust(False) + set_rust_ocr(ocr=None, aocr=None) paths: Final = tuple(sorted(directory.rglob("*.json"))) for path in paths: print(f"Migrated {path.name} to {migrate_fixture(path).name}") diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/record.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/record.py index ba1ea63aa81..19022324aa0 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/record.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/record.py @@ -8,7 +8,7 @@ from typing import Final, cast from dotenv import load_dotenv import litellm -from litellm.rust_bridge.ocr import use_litellm_rust +from litellm.rust_bridge.ocr import rust, set_rust_ocr from ......shared.parity.fixtures.cli import parse_recording_args from ......shared.parity.fixtures.media import structured_image_data_uri from ......shared.parity.fixtures.pipeline import record_fixtures @@ -67,7 +67,8 @@ def main() -> int: os.environ.get(FIXTURE_DIR_ENV), DEFAULT_FIXTURE_DIRECTORY, ) - use_litellm_rust(False, ocr=None, aocr=None) + rust(False) + set_rust_ocr(ocr=None, aocr=None) summary: Final = record_fixtures(targets, root, args.examples, args.concurrency, OcrParityCase) return summary.exit_code diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index dc3167017d9..3e6c4060134 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -406,7 +406,7 @@ OCR_CONTRACT: Final = UnitTestContract( ), exclusions=( UnitParityExclusionSpec( - nodeid="tests/test_litellm/ocr/test_rust_bridge.py::test_use_litellm_rust_toggles_flag", + nodeid="tests/test_litellm/ocr/test_rust_bridge.py::test_rust_toggles_flag", reason="This test asserts the process-level backend flag selected by the parity runner.", ), ), diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 293f75b7592..b2cf253d164 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -121,23 +121,25 @@ def _reset_rust_flag(): def test_load_rust_messages_returns_injected_impl(): bridge = RecordingMessages() - litellm.use_litellm_rust(True, messages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(messages=bridge) assert rust_messages.load_rust_messages() is bridge -def test_bare_use_litellm_rust_still_toggles_ocr(): +def test_bare_rust_still_toggles_ocr(): from litellm.rust_bridge.ocr import rust_ocr_enabled - litellm.use_litellm_rust(True) + litellm.rust(True) assert rust_ocr_enabled() is True - litellm.use_litellm_rust(False) + litellm.rust(False) assert rust_ocr_enabled() is False def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) assert rust_messages.load_rust_amessages() is bridge @@ -147,7 +149,7 @@ def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): "get_native_bridge", lambda: None, ) - litellm.use_litellm_rust(True) + litellm.rust(True) assert rust_messages.load_rust_messages() is None result = rust_messages.messages( model="claude", @@ -163,7 +165,8 @@ def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): def test_messages_wrapper_forwards_args_and_converts_timeout(): bridge = RecordingMessages() - litellm.use_litellm_rust(True, messages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(messages=bridge) response = rust_messages.messages( model="claude-sonnet-4-5", @@ -190,7 +193,8 @@ def test_messages_wrapper_forwards_args_and_converts_timeout(): @pytest.mark.asyncio async def test_amessages_wrapper_forwards_args(): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await rust_messages.amessages( model="claude-sonnet-4-5", @@ -226,7 +230,8 @@ def _gate(**overrides): @pytest.mark.asyncio async def test_gate_invokes_rust_and_marks_response_header(): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate() @@ -245,7 +250,8 @@ async def test_gate_invokes_rust_and_marks_response_header(): @pytest.mark.asyncio async def test_gate_falls_back_to_python_when_bridge_raises(): bridge = RaisingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate() @@ -268,7 +274,7 @@ async def test_gate_skips_rust_when_flag_absent(): async def test_gate_uses_process_enable_without_request_override(): bridge = RecordingAsyncMessages() rust_messages.set_rust_messages(amessages=bridge) - litellm.use_litellm_rust(True) + litellm.rust(True) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) @@ -279,7 +285,8 @@ async def test_gate_uses_process_enable_without_request_override(): @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_false(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) @@ -290,7 +297,8 @@ async def test_gate_skips_rust_when_flag_false(): @pytest.mark.asyncio async def test_gate_invokes_rust_for_native_anthropic_provider(): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate( custom_llm_provider="anthropic", @@ -339,7 +347,8 @@ async def test_gate_env_var_falsey_does_not_enable(monkeypatch): @pytest.mark.asyncio async def test_gate_skips_rust_for_unsupported_provider(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate(custom_llm_provider="openai") @@ -350,7 +359,8 @@ async def test_gate_skips_rust_for_unsupported_provider(): @pytest.mark.asyncio async def test_gate_skips_rust_for_agentic_hook(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate(has_agentic_hook=True) @@ -361,7 +371,8 @@ async def test_gate_skips_rust_for_agentic_hook(): @pytest.mark.asyncio async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) streaming_body = {**REQUEST_BODY, "stream": True} response = await _gate( @@ -398,7 +409,7 @@ async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): "get_native_bridge", lambda: None, ) - litellm.use_litellm_rust(True) + litellm.rust(True) response = await _gate() diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 4afb8303d03..1c2e07e0d24 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -228,7 +228,8 @@ def _reset_rust_flag(): def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) return bridge @@ -236,15 +237,16 @@ def fake_bridge(): def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() - litellm.use_litellm_rust(True, aocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(aocr=bridge) return bridge -def test_use_litellm_rust_toggles_flag(): +def test_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False - litellm.use_litellm_rust() + litellm.rust(True) assert rust_bridge.rust_ocr_enabled() is True - litellm.use_litellm_rust(False) + litellm.rust(False) assert rust_bridge.rust_ocr_enabled() is False @@ -255,14 +257,15 @@ def test_env_var_enables_rust_ocr(monkeypatch): def test_explicit_false_overrides_process_enable(): - litellm.use_litellm_rust(True) + litellm.rust(True) assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -325,25 +328,22 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() - litellm.use_litellm_rust(True, aocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(aocr=bridge) assert rust_bridge.load_rust_aocr() is bridge def test_toggle_without_ocr_arg_preserves_injected_impl(): - """Regression: routine enable/disable calls must not clobber a prior injection. - - Earlier, ``use_litellm_rust()`` unconditionally assigned the keyword default - of ``None`` to ``_rust_ocr_impl``, silently dropping a custom bridge whenever - a caller toggled the flag without re-passing ``ocr=``. - """ + """The public flag must not clobber an internal test binding.""" bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() - litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) - litellm.use_litellm_rust(False) + litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge assert rust_bridge.load_rust_aocr() is async_bridge - litellm.use_litellm_rust(True) + litellm.rust(True) assert rust_bridge.load_rust_ocr() is bridge assert rust_bridge.load_rust_aocr() is async_bridge @@ -356,9 +356,10 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch): ) bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() - litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) - litellm.use_litellm_rust(True, ocr=None, aocr=None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -371,7 +372,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch): "get_native_bridge", lambda: None, ) - litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI + litellm.rust(True) # no impl injected; extension isn't built in CI assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -389,7 +390,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): lambda: fake_module, ) - litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension + litellm.rust(True) # enabled, no impl injected -> import the extension assert rust_bridge.load_rust_ocr() is fake_module.ocr assert rust_bridge.load_rust_aocr() is fake_module.aocr @@ -403,7 +404,9 @@ def test_timeout_to_seconds_handles_float_timeout_and_none(): def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + + rust_bridge.set_rust_ocr(ocr=bridge) response = rust_bridge.ocr( model="mistral-ocr-latest", document=DOCUMENT, @@ -436,7 +439,9 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): bridge = RecordingAsyncBridge() - litellm.use_litellm_rust(True, aocr=bridge) + litellm.rust(True) + + rust_bridge.set_rust_ocr(aocr=bridge) response = await rust_bridge.aocr( model="mistral-ocr-maas", document=DOCUMENT, @@ -464,7 +469,8 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) response = ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -496,7 +502,8 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), @@ -508,7 +515,8 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") @@ -527,7 +535,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name): resolver_calls.append(name) @@ -549,7 +558,8 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -575,7 +585,8 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) def _resolver(name: str) -> str | None: return { @@ -598,7 +609,8 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -615,7 +627,8 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -635,7 +648,8 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -722,7 +736,8 @@ def test_ocr_exception_type_uses_resolved_provider_context( return CapturedException("wrapped") monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.use_litellm_rust(True, ocr=RaisingBridge()) + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -767,7 +782,8 @@ async def test_aocr_exception_type_uses_resolved_provider_context( return CapturedException("wrapped") monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge()) + litellm.rust(True) + rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -795,7 +811,8 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): def test_ocr_does_not_route_to_rust_when_disabled(): """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() - litellm.use_litellm_rust(False, ocr=bridge) + litellm.rust(False) + rust_bridge.set_rust_ocr(ocr=bridge) assert rust_bridge.rust_ocr_enabled() is False # The impl stays available for injection, but the disabled flag gates usage, @@ -807,7 +824,7 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI + litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 1233ddf1785..4b446368dbe 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -55,13 +55,13 @@ def test_rust_websocket_bridge_is_disabled_without_flag() -> None: def test_explicit_false_overrides_process_enable() -> None: - configuration.use_litellm_rust(True) + configuration.rust(True) assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) def test_process_enable_applies_without_request_override() -> None: - configuration.use_litellm_rust(True) + configuration.rust(True) assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 03921133c77..0489f4ff017 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -139,13 +139,13 @@ class TestGate: def test_explicit_false_overrides_process_enable(self): bridge.set_rust_chat_completions(decline=_RecordingDecline()) - configuration.use_litellm_rust(True) + configuration.rust(True) assert _accepts(litellm_params={"rust": False}) is False def test_process_enable_applies_without_request_override(self): bridge.set_rust_chat_completions(decline=_RecordingDecline()) - configuration.use_litellm_rust(True) + configuration.rust(True) assert _accepts(litellm_params={}) is True diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index 1c81c1fb624..15f69f95335 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -13,21 +13,6 @@ from litellm.rust_bridge import configuration from litellm.rust_bridge import ocr as rust_ocr -class _OcrBridge: - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - return {} - - @pytest.fixture(autouse=True) def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest discovers fixtures dynamically monkeypatch: pytest.MonkeyPatch, @@ -42,7 +27,7 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest @pytest.mark.parametrize( - ("request_override", "process", "environment", "legacy_ocr", "release_default", "expected"), + ("request_override", "process", "environment", "legacy_environment", "release_default", "expected"), ( (False, True, True, True, True, False), (True, False, False, False, False, True), @@ -60,7 +45,7 @@ def test_resolution_precedence( request_override: bool | None, process: bool | None, environment: bool | None, - legacy_ocr: bool | None, + legacy_environment: bool | None, release_default: bool, expected: bool, ) -> None: @@ -69,7 +54,7 @@ def test_resolution_precedence( request_override=request_override, process_override=process, environment_override=environment, - legacy_ocr_override=legacy_ocr, + legacy_environment_override=legacy_environment, release_default=release_default, ) is expected @@ -83,7 +68,7 @@ def test_release_default_remains_disabled() -> None: def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_RUST", "0") - configuration.use_litellm_rust(True) + configuration.rust(True) assert configuration.rust_enabled() is True assert configuration.rust_enabled(request_override=False) is False @@ -105,11 +90,11 @@ def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch @pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) -def test_invalid_legacy_environment_value_disables_ocr(monkeypatch: pytest.MonkeyPatch, value: str) -> None: +def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_ocr_enabled() is False + assert configuration.rust_enabled() is False def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: @@ -117,7 +102,7 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes with ThreadPoolExecutor(max_workers=1) as executor: assert executor.submit(configuration.rust_enabled).result() is True - configuration.use_litellm_rust(False) + configuration.rust(False) assert executor.submit(configuration.rust_enabled).result() is False assert executor.submit(configuration.rust_ocr_enabled).result() is False configuration.reset_rust_configuration() @@ -129,37 +114,30 @@ def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.Monk monkeypatch.setenv("LITELLM_RUST", "sometimes") assert configuration.rust_enabled(request_override=False) is False - configuration.use_litellm_rust(True) + configuration.rust(True) assert configuration.rust_enabled() is True -def test_legacy_ocr_environment_is_deprecated_and_ocr_only(monkeypatch: pytest.MonkeyPatch) -> None: +def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert configuration.rust_enabled() is True with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): assert configuration.rust_ocr_enabled() is True - assert configuration.rust_enabled() is False def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - assert configuration.rust_ocr_enabled() is False - - -def test_deprecated_public_injection_delegates_to_internal_binding() -> None: - bridge: Final = _OcrBridge() - - with pytest.warns(DeprecationWarning, match="Injecting Rust bridge implementations"): - configuration.use_litellm_rust(True, ocr=bridge) - - assert rust_ocr.load_rust_ocr() is bridge + assert configuration.rust_enabled() is False +@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR")) @pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) -def test_environment_controls_startup(value: str, expected: str) -> None: - environment: Final = {**os.environ, "LITELLM_RUST": value} +def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None: + environment: Final = {**os.environ, environment_name: value} result: Final = subprocess.run( ( sys.executable, From 6675e1fd9c5a5d21f9621b3e6ed206d3bc7caf11 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 4 Sep 2026 09:13:05 -0700 Subject: [PATCH 32/35] fix: preserve Python 3.10 harness compatibility --- basedpyright-code-budget.json | 10 +++++----- ruff-strict-budget.json | 4 ++-- test-quality-budget.json | 2 +- tests/rust-python-harness/shared/reporting/models.py | 4 +++- .../rust-python-harness/shared/reporting/rendering.py | 4 +++- tests/rust-python-harness/shared/tracing/profiler.py | 11 +++++++++++ .../shared/tracing/pytest_usage.py | 8 ++++++-- .../strategies/trace_parity/runner.py | 5 +++-- type-discipline-budget.json | 2 +- 9 files changed, 35 insertions(+), 15 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index dc51a44b5fc..1b0650c70d8 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -18,7 +18,7 @@ "limit": 40 }, "reportDeprecated": { - "limit": 210 + "limit": 209 }, "reportDuplicateImport": { "limit": 19 @@ -45,7 +45,7 @@ "limit": 24 }, "reportInvalidTypeForm": { - "limit": 32 + "limit": 30 }, "reportInvalidTypeVarUse": { "limit": 2 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38311 + "limit": 38309 }, "reportUnknownParameterType": { - "limit": 19624 + "limit": 19622 }, "reportUnknownVariableType": { - "limit": 29847 + "limit": 29846 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 8d5fc365d1d..4aac1756af4 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 809 }, "ANN201": { - "limit": 2000 + "limit": 1999 }, "ANN202": { "limit": 835 @@ -87,7 +87,7 @@ "limit": 2 }, "DTZ003": { - "limit": 25 + "limit": 24 }, "DTZ005": { "limit": 233 diff --git a/test-quality-budget.json b/test-quality-budget.json index 281c17748ee..7ca563d25af 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11069 + "limit": 11003 } } diff --git a/tests/rust-python-harness/shared/reporting/models.py b/tests/rust-python-harness/shared/reporting/models.py index f78141fa552..1ebba6c9793 100644 --- a/tests/rust-python-harness/shared/reporting/models.py +++ b/tests/rust-python-harness/shared/reporting/models.py @@ -5,7 +5,9 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path from time import monotonic -from typing import TYPE_CHECKING, Final, Literal, TypeAlias, assert_never +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from typing_extensions import assert_never if TYPE_CHECKING: from .strategy import CaseSpec, StrategyDefinition diff --git a/tests/rust-python-harness/shared/reporting/rendering.py b/tests/rust-python-harness/shared/reporting/rendering.py index 217caa48526..109f1e6cc1d 100644 --- a/tests/rust-python-harness/shared/reporting/rendering.py +++ b/tests/rust-python-harness/shared/reporting/rendering.py @@ -2,7 +2,9 @@ from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass -from typing import Final, Protocol, assert_never +from typing import Final, Protocol + +from typing_extensions import assert_never from .models import CaseDisposition, CaseResult diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py index ead856ee58f..abfb6a2425d 100644 --- a/tests/rust-python-harness/shared/tracing/profiler.py +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -80,6 +80,17 @@ def _qualified_name(frame: FrameType) -> str: native: Final = getattr(code, "co_qualname", None) if isinstance(native, str): return native + enclosing: Final = next( + ( + name + for ancestor in _frame_ancestors(frame) + for declared_code, name in _declared_functions(ancestor.f_locals, frozenset()) + if declared_code is code + ), + None, + ) + if enclosing is not None: + return enclosing module_name: Final = frame.f_globals.get("__name__") if not isinstance(module_name, str): return code.co_name diff --git a/tests/rust-python-harness/shared/tracing/pytest_usage.py b/tests/rust-python-harness/shared/tracing/pytest_usage.py index 58af174df38..285514a5239 100644 --- a/tests/rust-python-harness/shared/tracing/pytest_usage.py +++ b/tests/rust-python-harness/shared/tracing/pytest_usage.py @@ -11,6 +11,7 @@ import tempfile import warnings from collections.abc import Generator, Sequence from pathlib import Path +from types import CodeType from typing import TYPE_CHECKING, Final from pluggy import HookimplMarker @@ -62,9 +63,12 @@ class PythonFunctionReference(BaseModel): value: object = importlib.import_module(self.module) for component in self.qualname.split("."): value = getattr(value, component) + if not callable(value): + raise ValueError(f"Python function is not callable: {self.module}:{self.qualname}") function: Final = inspect.unwrap(value) code: Final = getattr(function, "__code__", None) - if code is None: + qualname: Final = getattr(function, "__qualname__", None) + if not isinstance(code, CodeType) or not isinstance(qualname, str): raise ValueError(f"Python function has no code object: {self.module}:{self.qualname}") source: Final = Path(code.co_filename).resolve() try: @@ -74,7 +78,7 @@ class PythonFunctionReference(BaseModel): return PythonFunctionIdentity( file=relative.as_posix(), line=code.co_firstlineno, - qualname=code.co_qualname, + qualname=qualname, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/runner.py b/tests/rust-python-harness/strategies/trace_parity/runner.py index ef373a6f2f6..b78a3c7da3f 100644 --- a/tests/rust-python-harness/strategies/trace_parity/runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/runner.py @@ -165,9 +165,10 @@ def run_trace_cases( ) -> tuple[int, HarnessRun]: selected_scenarios: Final = frozenset(runner_args) run: Final = HarnessRun.from_cases(cases) - bridge_error: Final = ensure_trace_bridge(repo_root) + runnable_cases: Final = tuple(case for case in cases if isinstance(case.spec, ModuleCaseSpec)) + bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases else None if bridge_error is not None: - for harness_case in cases: + for harness_case in runnable_cases: _record_setup_failure(run, harness_case, bridge_error, "bridge") run.finished_at = monotonic() on_update(run) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ba871b5f69b..094b9749d98 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -33,6 +33,6 @@ "limit": 5514 }, "LIT012": { - "limit": 4488 + "limit": 4487 } } From dbf8fe0f4e19f9425a7dd2753c3d96410efe6beb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 4 Sep 2026 10:10:47 -0700 Subject: [PATCH 33/35] test: repair four chronically failing CI tests test_no_linear_scans_in_router: #39468 added config_deployments() and heuristic_v2_router_limit_violation(), which both scan the whole model_list from admin-only paths (model add/upsert), so add them to the allowlist. The allowlist becomes a mapping so each exemption carries its reason as data. test_missing_model_parameter_curl: a request with no model is rejected by the proxy when nothing can serve it and by the router when a wildcard or default deployment exists, and by the upstream provider when a wildcard forwards it, so the message text is not a stable contract. Assert the contract that holds in every case: HTTP 400 with a non-empty error message. test_model_group_info_e2e: /model_group/info resolves wildcards, so it can never return "anthropic/*" verbatim. cc3f9cd65b7 rewrote the assertion to expect the raw pattern after claude-3-5-haiku-20241022 left the price map, which made it unsatisfiable. Assert the expansion instead. test_should_derive_ocr_mapping_status_from_live_tests: the audit needs a native bridge built with the trace-parity feature, which CI never builds, so skip with the harness's own diagnostic instead of erroring. Extract that check out of ensure_trace_bridge as trace_bridge_error so a pytest run reports the state without kicking off a maturin rebuild. --- .../test_router_index_management.py | 10 ++-- .../shared/native_build.py | 17 ++++--- .../test_openai_error_handling.py | 13 ++++-- tests/test_models.py | 18 ++++---- tests/test_rust_python_harness.py | 46 +++++++++++++++++++ 5 files changed, 82 insertions(+), 22 deletions(-) diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index 87ddaadaf3d..35d295d581a 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -237,10 +237,12 @@ class TestRouterIndexManagement: - model_name_to_deployment_indices for O(1) + O(k) model_name lookups """ # Methods that are allowed to iterate through self.model_list - ALLOWED_METHODS = [ - "_get_deployment_by_litellm_model", # Edge case: lookup by litellm_params.model (not indexed) - "_finalize_adaptive_router_if_configured", # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match) - ] + ALLOWED_METHODS = { + "_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed", + "_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match', + "config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)", + "heuristic_v2_router_limit_violation": "counts heuristic_v2 routers across the whole list; admin path only (auto-router init/upsert)", + } # Get path to router.py router_file = os.path.join( diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index 2ca7131c2c1..7df0f999847 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -73,6 +73,16 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]: return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:]) +def trace_bridge_error() -> str | None: + """Why the installed bridge cannot serve trace parity, or None when it can. Never rebuilds.""" + bridge: Final = get_native_bridge() + if bridge is None: + return "native Rust bridge is not importable" + if getattr(bridge, "_trace", None) is None: + return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature" + return None + + def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None @@ -84,9 +94,4 @@ def ensure_trace_bridge(repo_root: Path) -> str | None: if not succeeded: return f"native Rust bridge rebuild failed:\n{output}" _drop_imported_bridge() - bridge: Final = get_native_bridge() - if bridge is None: - return "native Rust bridge is not importable" - if getattr(bridge, "_trace", None) is None: - return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature" - return None + return trace_bridge_error() diff --git a/tests/store_model_in_db_tests/test_openai_error_handling.py b/tests/store_model_in_db_tests/test_openai_error_handling.py index 554ddf49cce..9a18d7f3420 100644 --- a/tests/store_model_in_db_tests/test_openai_error_handling.py +++ b/tests/store_model_in_db_tests/test_openai_error_handling.py @@ -106,15 +106,22 @@ def test_missing_model_parameter_curl(curl_command): # Run the curl command and capture the output key = generate_key_sync() curl_command = curl_command.replace("sk-1234", key) - result = subprocess.run(curl_command, shell=True, capture_output=True, text=True) + result = subprocess.run( + f'{curl_command} -s -w "\\n%{{http_code}}"', + shell=True, + capture_output=True, + text=True, + ) + body, _, status_code = result.stdout.rpartition("\n") # Parse the JSON response - response = json.loads(result.stdout) + response = json.loads(body) # Check that we got an error response assert "error" in response print("error in response", json.dumps(response, indent=4)) - assert "litellm.BadRequestError" in response["error"]["message"] + assert status_code == "400", f"expected HTTP 400, got {status_code}: {response}" + assert isinstance(response["error"]["message"], str) and response["error"]["message"] @pytest.mark.asyncio diff --git a/tests/test_models.py b/tests/test_models.py index 151fb70b665..186752af2bc 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -487,6 +487,9 @@ async def test_get_personal_models_for_user(): async def test_model_group_info_e2e(): """ Test /model/group/info endpoint + + The proxy config declares a wildcard "anthropic/*" deployment, and the endpoint resolves + wildcards into the concrete models they cover, so the raw pattern is never returned. """ async with aiohttp.ClientSession() as session: models = await get_models(session=session, key="sk-1234") @@ -495,16 +498,13 @@ async def test_model_group_info_e2e(): model_group_info = await get_model_group_info(session=session, key="sk-1234") print(model_group_info) - # Check that the endpoint returns data and contains the wildcard - # anthropic model group from the proxy config - has_anthropic_wildcard = False - for model in model_group_info["data"]: - if model["model_group"] == "anthropic/*": - has_anthropic_wildcard = True + model_groups = [m["model_group"] for m in model_group_info["data"]] - assert has_anthropic_wildcard, ( - f"Expected 'anthropic/*' in model groups, got: " - f"{[m['model_group'] for m in model_group_info['data']]}" + assert "anthropic/*" not in model_groups, ( + f"Expected 'anthropic/*' to be expanded, but it was returned verbatim: {model_groups}" + ) + assert any(m.startswith("anthropic/") for m in model_groups), ( + f"Expected concrete anthropic models from the 'anthropic/*' config entry, got: {model_groups}" ) diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index b27d1c83597..179660dd4e9 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -2,6 +2,7 @@ from __future__ import annotations import importlib from pathlib import Path +from types import SimpleNamespace from typing import Final import pytest @@ -13,6 +14,7 @@ mapping_validator = importlib.import_module("tests.rust-python-harness.strategie mappings = importlib.import_module("tests.rust-python-harness.strategies.unit_tests_mapping.mappings") ocr_mapping = importlib.import_module("tests.rust-python-harness.strategies.unit_tests_mapping.cases.ocr") cli = importlib.import_module("tests.rust-python-harness.cli") +native_build = importlib.import_module("tests.rust-python-harness.shared.native_build") audit_mapping = mapping_validator.audit_mapping UNIT_TEST_CONTRACTS = mappings.UNIT_TEST_CONTRACTS @@ -119,7 +121,51 @@ def test_should_leave_functions_without_mapping_contracts_unimplemented() -> Non assert "messages" not in UNIT_TEST_CONTRACTS +def test_should_report_a_bridge_that_cannot_be_imported() -> None: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_build, "get_native_bridge", lambda: None) + message = native_build.trace_bridge_error() + + assert message is not None + assert "not importable" in message + + +def test_should_report_a_bridge_built_without_the_trace_feature() -> None: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) + message = native_build.trace_bridge_error() + + assert message is not None + assert native_build.BRIDGE_FEATURE in message + + +def test_should_accept_a_bridge_built_with_the_trace_feature() -> None: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=object())) + + assert native_build.trace_bridge_error() is None + + +def test_should_not_rebuild_the_bridge_while_reporting_its_state() -> None: + rebuilds: list[object] = [] + + def fake_rebuild(repo_root: object) -> tuple[bool, str]: + rebuilds.append(repo_root) + return True, "" + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_build, "_rebuild", fake_rebuild) + patch.setattr(native_build, "get_native_bridge", lambda: None) + native_build.trace_bridge_error() + + assert rebuilds == [] + + def test_should_derive_ocr_mapping_status_from_live_tests() -> None: + bridge_error: Final = native_build.trace_bridge_error() + if bridge_error is not None: + pytest.skip(bridge_error) + report = audit_mapping(OCR_CONTRACT, repo_root=REPO_ROOT) assert report.is_valid, ( From a5a78670d047011f38dd8c11e7d2df818d2fba32 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 4 Sep 2026 10:18:15 -0700 Subject: [PATCH 34/35] test: address review notes on the chronic-test repairs Drop the two new docstrings, annotate the new locals Final, and replace the mutable call recorder with a rebuild stub that fails the test if it is ever reached. --- tests/rust-python-harness/shared/native_build.py | 1 - tests/test_models.py | 6 ++---- tests/test_rust_python_harness.py | 16 ++++++---------- 3 files changed, 8 insertions(+), 15 deletions(-) diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index 7df0f999847..8693cf3bac2 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -74,7 +74,6 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]: def trace_bridge_error() -> str | None: - """Why the installed bridge cannot serve trace parity, or None when it can. Never rebuilds.""" bridge: Final = get_native_bridge() if bridge is None: return "native Rust bridge is not importable" diff --git a/tests/test_models.py b/tests/test_models.py index 186752af2bc..64c7dcd83da 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -6,6 +6,7 @@ import asyncio import aiohttp import os import dotenv +from typing import Final from dotenv import load_dotenv load_dotenv() @@ -487,9 +488,6 @@ async def test_get_personal_models_for_user(): async def test_model_group_info_e2e(): """ Test /model/group/info endpoint - - The proxy config declares a wildcard "anthropic/*" deployment, and the endpoint resolves - wildcards into the concrete models they cover, so the raw pattern is never returned. """ async with aiohttp.ClientSession() as session: models = await get_models(session=session, key="sk-1234") @@ -498,7 +496,7 @@ async def test_model_group_info_e2e(): model_group_info = await get_model_group_info(session=session, key="sk-1234") print(model_group_info) - model_groups = [m["model_group"] for m in model_group_info["data"]] + model_groups: Final = [m["model_group"] for m in model_group_info["data"]] assert "anthropic/*" not in model_groups, ( f"Expected 'anthropic/*' to be expanded, but it was returned verbatim: {model_groups}" diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index 179660dd4e9..85b45c07bc2 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -124,7 +124,7 @@ def test_should_leave_functions_without_mapping_contracts_unimplemented() -> Non def test_should_report_a_bridge_that_cannot_be_imported() -> None: with pytest.MonkeyPatch.context() as patch: patch.setattr(native_build, "get_native_bridge", lambda: None) - message = native_build.trace_bridge_error() + message: Final = native_build.trace_bridge_error() assert message is not None assert "not importable" in message @@ -133,7 +133,7 @@ def test_should_report_a_bridge_that_cannot_be_imported() -> None: def test_should_report_a_bridge_built_without_the_trace_feature() -> None: with pytest.MonkeyPatch.context() as patch: patch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) - message = native_build.trace_bridge_error() + message: Final = native_build.trace_bridge_error() assert message is not None assert native_build.BRIDGE_FEATURE in message @@ -147,18 +147,14 @@ def test_should_accept_a_bridge_built_with_the_trace_feature() -> None: def test_should_not_rebuild_the_bridge_while_reporting_its_state() -> None: - rebuilds: list[object] = [] - - def fake_rebuild(repo_root: object) -> tuple[bool, str]: - rebuilds.append(repo_root) - return True, "" + def forbidden_rebuild(repo_root: object) -> tuple[bool, str]: + raise AssertionError("trace_bridge_error must not rebuild the native bridge") with pytest.MonkeyPatch.context() as patch: - patch.setattr(native_build, "_rebuild", fake_rebuild) + patch.setattr(native_build, "_rebuild", forbidden_rebuild) patch.setattr(native_build, "get_native_bridge", lambda: None) - native_build.trace_bridge_error() - assert rebuilds == [] + assert native_build.trace_bridge_error() is not None def test_should_derive_ocr_mapping_status_from_live_tests() -> None: From dd01abc4390f494eedc5fe448b70c7b3de06a1c7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:00:47 -0700 Subject: [PATCH 35/35] feat(team): report per-user spend within a team for JWT traffic (#39771) * feat(team): report per-user spend within a team for JWT traffic Add GET /team/spend/by_user, which groups raw spend logs by (team_id, user) so JWT/SSO requests with no virtual key are attributed to the user inside each selected team. Team admins see every member, plain members see only their own row. The Team Usage page gets a Spend Per User Within Team card with CSV export backed by the same endpoint. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(team): cover /team/spend/by_user in behavior suite, tf audit allowlist and EntityUsage unit test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(team): drop explanatory docstrings from /team/spend/by_user and regen schema.d.ts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 + .../management_endpoints/team_endpoints.py | 123 +++++++++++++ .../management_endpoints/team_endpoints.py | 21 +++ .../endpointaudit/coverage_allowlist.txt | 1 + .../management/test_team_spend_by_user.py | 58 ++++++ .../proxy/auth/test_route_checks.py | 35 ++++ .../test_team_endpoints.py | 172 ++++++++++++++++++ .../EntityUsage/EntityUsage.test.tsx | 25 +++ .../components/EntityUsage/EntityUsage.tsx | 19 ++ .../EntityUsage/TeamUserSpendCard.tsx | 109 +++++++++++ .../EntityUsage/teamUserSpend.test.ts | 93 ++++++++++ .../components/EntityUsage/teamUserSpend.ts | 55 ++++++ .../src/components/networking.tsx | 18 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 114 ++++++++++++ 14 files changed, 845 insertions(+) create mode 100644 tests/proxy_behavior/management/test_team_spend_by_user.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 832d941f5b5..b33e2fe7ff6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -670,6 +670,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only "/gateway/daily/activity", # model @@ -832,6 +833,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/spend/by_user", "/team/{team_id}/members/me", "/model/new", "/model/update", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 90d7539b38d..a504c1c5e43 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -170,6 +170,8 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamMemberAddResult, TeamMemberInfoResponse, TeamMetadataSchemaResponse, + TeamUserSpendResponse, + TeamUserSpendRow, UpdateTeamMemberPermissionsRequest, ) @@ -6231,3 +6233,124 @@ async def get_team_daily_activity_aggregated( timezone_offset_minutes=timezone, include_entity_breakdown=True, ) + + +def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: + team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count)) + user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else "" + return f""" + SELECT + sl.team_id, + sl."user" AS user_id, + u.user_email, + u.user_alias, + SUM(sl.spend)::float AS spend, + SUM(sl.prompt_tokens)::bigint AS prompt_tokens, + SUM(sl.completion_tokens)::bigint AS completion_tokens, + SUM(sl.total_tokens)::bigint AS total_tokens, + COUNT(*)::bigint AS api_requests, + COUNT(*) FILTER (WHERE sl.status IS DISTINCT FROM 'failure')::bigint AS successful_requests, + COUNT(*) FILTER (WHERE sl.status = 'failure')::bigint AS failed_requests + FROM "LiteLLM_SpendLogs" sl + LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = sl."user" + WHERE sl."startTime" >= $1::timestamp + AND sl."startTime" < $2::timestamp + INTERVAL '1 day' + AND sl.team_id IN ({team_placeholders}){user_clause} + GROUP BY sl.team_id, sl."user", u.user_email, u.user_alias + ORDER BY spend DESC, sl.team_id, sl."user" + """ + + +class _TeamUserSpendDbRow(TypedDict): + team_id: ReadOnly[str] + user_id: ReadOnly[str | None] + user_email: ReadOnly[str | None] + user_alias: ReadOnly[str | None] + spend: ReadOnly[float] + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + api_requests: ReadOnly[int] + successful_requests: ReadOnly[int] + failed_requests: ReadOnly[int] + + +@router.get( + "/team/spend/by_user", + response_model=TeamUserSpendResponse, + tags=["team management"], # mutable-ok: fastapi route tags must be a list +) +async def get_team_spend_by_user( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + team_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, +) -> TeamUserSpendResponse: + """ + Spend per user within the given teams, attributed per request from spend logs. + + Proxy admins may query any team. Team admins and members holding the + `/team/daily/activity` permission see every user of the requested teams; + other members only see their own row. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) + + range_error: Final = _aggregated_date_range_error(start_date, end_date) + if range_error is not None or start_date is None or end_date is None: + raise _daily_activity_error(status_code=400, message=range_error or "Please provide start_date and end_date") + + if not team_ids: + raise _daily_activity_error(status_code=400, message="Please provide team_ids") + + scope: Final = await _resolve_team_daily_activity_scope( + team_ids=team_ids, + exclude_team_ids=None, + api_key=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + scoped_team_ids: Final = tuple(scope.team_ids or ()) + if not scoped_team_ids: + return TeamUserSpendResponse(start_date=start_date, end_date=end_date, results=()) + + own_user_only: Final = scope.api_key_filter is not None + user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else () + rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw( + _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), + start_date, + end_date, + *scoped_team_ids, + *user_param, + ) + results: Final = tuple( + TeamUserSpendRow( + team_id=row["team_id"], + team_alias=_team_alias_or_none(scope.team_alias_metadata.get(row["team_id"])), + user_id=row["user_id"] or "", + user_email=row["user_email"], + user_alias=row["user_alias"], + spend=row["spend"], + prompt_tokens=row["prompt_tokens"], + completion_tokens=row["completion_tokens"], + total_tokens=row["total_tokens"], + api_requests=row["api_requests"], + successful_requests=row["successful_requests"], + failed_requests=row["failed_requests"], + ) + for row in rows + ) + return TeamUserSpendResponse(start_date=start_date, end_date=end_date, results=results) + + +def _team_alias_or_none(metadata: Mapping[str, object] | None) -> str | None: + alias: Final = metadata.get("team_alias") if metadata is not None else None + return alias if isinstance(alias, str) else None diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 2417868fb29..a282430bb11 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -143,3 +143,24 @@ class TeamMetadataSchemaResponse(BaseModel): """Response for GET /team/metadata_schema; ``fields`` is empty when no schema is configured.""" fields: tuple[TeamMetadataFieldSchema, ...] + + +class TeamUserSpendRow(BaseModel): + team_id: str + team_alias: str | None = None + user_id: str + user_email: str | None = None + user_alias: str | None = None + spend: float = 0.0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + api_requests: int = 0 + successful_requests: int = 0 + failed_requests: int = 0 + + +class TeamUserSpendResponse(BaseModel): + start_date: str + end_date: str + results: tuple[TeamUserSpendRow, ...] diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 052962e078e..6bc8947e89f 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated +GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity GET /user/daily/activity/aggregated diff --git a/tests/proxy_behavior/management/test_team_spend_by_user.py b/tests/proxy_behavior/management/test_team_spend_by_user.py new file mode 100644 index 00000000000..1d6aab04003 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_spend_by_user.py @@ -0,0 +1,58 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/spend/by_user shares the team-scope resolver with +# /team/daily/activity, so the membership matrix must hold here too. team_ids +# is mandatory on this route (a per-user rollup with no team is meaningless), +# so the bare query is 400 for everyone instead of defaulting to own teams. +_MEMBERS = { + "alpha": { + Actor.TEAM_ADMIN, + Actor.INTERNAL_USER, + Actor.OWNER, + Actor.UNRELATED_SAME_ORG, + Actor.SERVICE_ACCOUNT, + }, + "beta": {Actor.CROSS_ORG_USER}, +} + + +def _expected(actor: Actor, team: str) -> int: + if team == "none": + return 400 + if actor == Actor.PROXY_ADMIN: + return 200 + return 200 if actor in _MEMBERS.get(team, set()) else 404 + + +_CASES = [ + (f"{team}/{actor.value}", actor, team, _expected(actor, team)) + for team in ("none", "alpha", "beta") + for actor in Actor +] + +_DATES = "start_date=2024-01-01&end_date=2024-12-31" + + +@pytest.mark.parametrize( + "actor,team,expected_status", + [(a, t, s) for (_id, a, t, s) in _CASES], + ids=[c[0] for c in _CASES], +) +async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, proxy_client, world): + team_id = {"alpha": world.team_alpha_id, "beta": world.team_beta_id}.get(team) + query = _DATES if team_id is None else f"{_DATES}&team_ids={team_id}" + + resp = await proxy_client.get( + f"/team/spend/by_user?{query}", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}" + if expected_status == 200: + body = resp.json() + assert (body["start_date"], body["end_date"]) == ("2024-01-01", "2024-12-31") + assert all(row["team_id"] == team_id for row in body["results"]) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 71ccef620e5..48926cb7bc2 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3312,6 +3312,41 @@ def test_user_daily_activity_routes_reachable_by_non_admin(route, user_role): ) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_team_spend_by_user_reachable_by_non_admin(user_role): + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + def outcome(route: str) -> str: + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as exc: + return f"denied: {exc}" + return "allowed" + + assert outcome("/team/spend/by_user") == "allowed" + assert outcome("/team/spend/by_key").startswith("denied: Only proxy admin") + + def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): """check_route_access is exact-match plus explicit wildcards, so listing the parent /user/daily/activity does not implicitly cover the /aggregated diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 019ebc9807c..ab4cd74e092 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13598,3 +13598,175 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent( assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership" assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5 + + +def _team_spend_by_user_team(team_id: str, team_alias: str, member: Member, permissions: list[str]) -> MagicMock: + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = team_id + team.team_alias = team_alias + team.members_with_roles = [member] + team.team_member_permissions = permissions + team.model_dump.return_value = { + "team_id": team_id, + "team_alias": team_alias, + "members_with_roles": [{"user_id": member.user_id, "role": member.role}], + "team_member_permissions": permissions, + } + return team + + +def _team_spend_by_user_caller(user_id: str, teams: list[str]) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id=user_id, user_email=f"{user_id}@example.com", teams=teams, user_role="internal_user" + ) + + +def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, requests: int) -> dict: + return { + "team_id": team_id, + "user_id": user_id, + "user_email": f"{user_id}@example.com", + "user_alias": None, + "spend": spend, + "prompt_tokens": 10 * requests, + "completion_tokens": 5 * requests, + "total_tokens": 15 * requests, + "api_requests": requests, + "successful_requests": requests - 1, + "failed_requests": 1, + } + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) + beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), []) + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta]) + mock_db_client.db.query_raw = AsyncMock( + return_value=[ + _team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3), + _team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2), + _team_spend_by_user_db_row("team-beta", "alice", 0.1, 1), + ] + ) + + response = await get_team_spend_by_user( + user_api_key_dict=admin, + team_ids="team-alpha,team-beta", + start_date="2026-09-01", + end_date="2026-09-04", + ) + + sql, *params = mock_db_client.db.query_raw.call_args.args + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"] + assert 'FROM "LiteLLM_SpendLogs" sl' in sql + assert 'sl."startTime" >= $1::timestamp' in sql + assert "sl.\"startTime\" < $2::timestamp + INTERVAL '1 day'" in sql + assert "sl.team_id IN ($3, $4)" in sql + assert 'GROUP BY sl.team_id, sl."user"' in sql + assert 'sl."user" = $' not in sql + + assert response.start_date == "2026-09-01" + assert response.end_date == "2026-09-04" + assert [(r.team_id, r.team_alias, r.user_id, r.user_email, r.spend, r.api_requests) for r in response.results] == [ + ("team-alpha", "Team Alpha", "alice", "alice@example.com", 0.5, 3), + ("team-alpha", "Team Alpha", "bob", "bob@example.com", 0.25, 2), + ("team-beta", "Team Beta", "alice", "alice@example.com", 0.1, 1), + ] + assert (response.results[0].successful_requests, response.results[0].failed_requests) == (2, 1) + assert (response.results[0].prompt_tokens, response.results[0].completion_tokens) == (30, 15) + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_client): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) + mock_db_client.db.query_raw = AsyncMock(return_value=[]) + mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) + ) + + await get_team_spend_by_user( + user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" + ) + + sql, *params = mock_db_client.db.query_raw.call_args.args + assert params == ["2026-09-01", "2026-09-04", "team-alpha"] + assert 'sl."user" = $' not in sql + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_client): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) + alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"]) + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) + mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) + ) + + response = await get_team_spend_by_user( + user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" + ) + + sql, *params = mock_db_client.db.query_raw.call_args.args + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"] + assert "sl.team_id IN ($3)" in sql + assert 'AND sl."user" = $4' in sql + assert [(r.user_id, r.spend) for r in response.results] == [("bob", 0.25)] + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_client): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) + mock_db_client.db.query_raw = AsyncMock(return_value=[]) + mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) + ) + + with pytest.raises(HTTPException) as exc_info: + await get_team_spend_by_user( + user_api_key_dict=caller, team_ids="team-beta", start_date="2026-09-01", end_date="2026-09-04" + ) + + assert exc_info.value.status_code == 404 + mock_db_client.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_ids,start_date,end_date,expected_error", + [ + (None, "2026-09-01", "2026-09-04", "team_ids"), + ("", "2026-09-01", "2026-09-04", "team_ids"), + ("team-alpha", None, "2026-09-04", "start_date and end_date"), + ("team-alpha", "2026-09-04", "2026-09-01", "on or after"), + ("team-alpha", "2020-01-01", "2026-12-31", "at most 400 days"), + ("team-alpha", "nope", "2026-09-04", "valid YYYY-MM-DD"), + ], +) +async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + mock_db_client.db.query_raw = AsyncMock(return_value=[]) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(HTTPException) as exc_info: + await get_team_spend_by_user( + user_api_key_dict=admin, team_ids=team_ids, start_date=start_date, end_date=end_date + ) + + assert exc_info.value.status_code == 400 + assert expected_error in str(exc_info.value.detail) + mock_db_client.db.query_raw.assert_not_called() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 5bb48a78437..2a6c2ede478 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import type { ReactNode } from "react"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; +import useTeams from "@/app/(dashboard)/hooks/useTeams"; import * as networking from "@/components/networking"; import EntityUsage from "./EntityUsage"; @@ -60,6 +61,10 @@ vi.mock("./TopModelView", () => ({ ), })); +vi.mock("./TeamUserSpendCard", () => ({ + default: ({ teamIds }: { teamIds: string[] }) =>
{`team-user-spend:${teamIds.join("|")}`}
, +})); + vi.mock("@/components/EntityUsageExport/EntityUsageExportModal", () => ({ default: () =>
Entity Usage Export Modal
, })); @@ -460,6 +465,26 @@ describe("EntityUsage", () => { }); }); + it("feeds the per-user spend card every visible team except the dashboard team, only for teams", async () => { + const mockUseTeams = vi.mocked(useTeams); + const teamsResult = (teams: { team_id: string }[]) => + ({ teams, setTeams: vi.fn() }) as unknown as ReturnType; + mockUseTeams.mockReturnValue( + teamsResult([{ team_id: "team-alpha" }, { team_id: "litellm-dashboard" }, { team_id: "team-beta" }]), + ); + + render(); + expect(await screen.findByText("team-user-spend:team-alpha|team-beta")).toBeInTheDocument(); + + cleanup(); + mockUseTeams.mockReturnValue(teamsResult([])); + render(); + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + expect(screen.queryByText(/^team-user-spend:/)).not.toBeInTheDocument(); + }); + it("should render with organization entity type and call organization API", async () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index ef3943e5b71..273e478528e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -43,6 +43,7 @@ import EndpointUsage from "../EndpointUsage/EndpointUsage"; import ModelViewToggle, { ModelViewType } from "../ModelViewToggle"; import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; import TopModelView from "./TopModelView"; +import TeamUserSpendCard from "./TeamUserSpendCard"; interface EntityMetrics { metrics: { @@ -275,6 +276,13 @@ const EntityUsage: React.FC = ({ const capitalizedEntityLabel = entityType.charAt(0).toUpperCase() + entityType.slice(1); const showFlatCost = entityType === "team" && hasFlatCost(spendData.metadata); + const userSpendTeamIds = useMemo( + () => + selectedTags.length > 0 + ? selectedTags + : (teams ?? []).map((team) => team.team_id).filter((id) => id !== "litellm-dashboard"), + [selectedTags, teams], + ); const providerSpend = useMemo(() => getProviderSpend(spendData.results), [spendData.results]); const entityBreakdownColumns = useMemo[]>( () => [ @@ -530,6 +538,17 @@ const EntityUsage: React.FC = ({ + {entityType === "team" && ( +
+ +
+ )} + {/* Top API Keys */}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx new file mode 100644 index 00000000000..ed90e144efb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx @@ -0,0 +1,109 @@ +import { useQuery } from "@tanstack/react-query"; +import type { ColumnDef } from "@tanstack/react-table"; +import { Download } from "lucide-react"; +import React, { useMemo } from "react"; + +import { teamSpendByUserCall } from "@/components/networking"; +import { DataTable } from "@/components/shared/DataTable"; +import { MoneyCell } from "@/components/shared/table_cells"; +import { Button } from "@/components/ui/button"; +import { Card as ShadcnCard, CardContent } from "@/components/ui/card"; + +import { + buildTeamUserSpendCsv, + downloadCsv, + sortBySpendDesc, + teamLabel, + teamUserSpendCsvFileName, + teamUserSpendRowId, + userLabel, + type TeamUserSpendRow, +} from "./teamUserSpend"; + +interface TeamUserSpendCardProps { + accessToken: string | null; + startTime: Date | null; + endTime: Date | null; + teamIds: string[]; +} + +const columns: ColumnDef[] = [ + { header: "Team", accessorFn: teamLabel, id: "team", cell: ({ row }) => teamLabel(row.original) }, + { header: "User", accessorFn: userLabel, id: "user", cell: ({ row }) => userLabel(row.original) }, + { + header: "Spend", + accessorKey: "spend", + meta: { numeric: true }, + cell: ({ row }) => , + }, + { + header: "Requests", + accessorKey: "api_requests", + meta: { numeric: true }, + cell: ({ row }) => row.original.api_requests.toLocaleString(), + }, + { + header: "Successful", + accessorKey: "successful_requests", + meta: { numeric: true, className: "text-success" }, + cell: ({ row }) => row.original.successful_requests.toLocaleString(), + }, + { + header: "Failed", + accessorKey: "failed_requests", + meta: { numeric: true, className: "text-destructive" }, + cell: ({ row }) => row.original.failed_requests.toLocaleString(), + }, + { + header: "Tokens", + accessorKey: "total_tokens", + meta: { numeric: true }, + cell: ({ row }) => row.original.total_tokens.toLocaleString(), + }, +]; + +const TeamUserSpendCard: React.FC = ({ accessToken, startTime, endTime, teamIds }) => { + const hasTeams = teamIds.length > 0; + const { data, isLoading } = useQuery({ + queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds], + queryFn: () => + accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null, + enabled: Boolean(accessToken && startTime && endTime) && hasTeams, + }); + const rows = useMemo(() => sortBySpendDesc(data?.results ?? []), [data]); + + return ( + + +
+
+

Spend Per User Within Team

+

+ Attributed per request from spend logs, so it includes JWT/SSO traffic that does not use a virtual key +

+
+ +
+ + + + ); +}; + +export default TeamUserSpendCard; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.test.ts new file mode 100644 index 00000000000..36d442c617f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, it } from "vitest"; + +import type { TeamUserSpendResponse } from "@/components/networking"; + +import { + buildTeamUserSpendCsv, + sortBySpendDesc, + teamUserSpendCsvFileName, + teamUserSpendRowId, + userLabel, + type TeamUserSpendRow, +} from "./teamUserSpend"; + +const row = (overrides: Partial): TeamUserSpendRow => ({ + team_id: "team-alpha", + team_alias: "Team Alpha", + user_id: "alice@example.com", + user_email: "alice@example.com", + user_alias: null, + spend: 0.5, + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + api_requests: 3, + successful_requests: 2, + failed_requests: 1, + ...overrides, +}); + +const aliceInBeta: Partial = { + team_id: "team-beta", + team_alias: "Team Beta", + spend: 0.1, + api_requests: 1, +}; +const bobInAlpha: Partial = { + user_id: "bob", + user_email: null, + user_alias: "Bob", + spend: 0.25, + api_requests: 2, +}; + +const response: TeamUserSpendResponse = { + start_date: "2026-09-01", + end_date: "2026-09-04", + results: [row(aliceInBeta), row({}), row(bobInAlpha)], +}; + +describe("teamUserSpend", () => { + it("keeps the same user as separate rows per team", () => { + const ids = response.results.map(teamUserSpendRowId); + expect(new Set(ids).size).toBe(3); + expect(ids[0]).not.toBe(ids[1]); + }); + + it("labels a user by email, then alias, then id, then a placeholder", () => { + expect(userLabel(row({}))).toBe("alice@example.com"); + expect(userLabel(row({ user_email: null, user_alias: "Bob", user_id: "u1" }))).toBe("Bob"); + expect(userLabel(row({ user_email: null, user_alias: null, user_id: "u1" }))).toBe("u1"); + expect(userLabel(row({ user_email: null, user_alias: null, user_id: "" }))).toBe("(no user)"); + }); + + it("sorts by spend descending without mutating the input", () => { + const before = [...response.results]; + expect(sortBySpendDesc(response.results).map((r) => r.spend)).toEqual([0.5, 0.25, 0.1]); + expect(response.results).toEqual(before); + }); + + it("writes one CSV line per (team, user) with the team kept on every line", () => { + const lines = buildTeamUserSpendCsv(response).split(/\r?\n/); + expect(lines[0]).toBe( + "Start Date,End Date,Team,Team ID,User,User ID,User Email,Spend (USD),Requests,Successful,Failed,Prompt Tokens,Completion Tokens,Total Tokens", + ); + expect(lines.slice(1)).toEqual([ + "2026-09-01,2026-09-04,Team Alpha,team-alpha,alice@example.com,alice@example.com,alice@example.com,0.5,3,2,1,10,5,15", + "2026-09-01,2026-09-04,Team Alpha,team-alpha,Bob,bob,,0.25,2,2,1,10,5,15", + "2026-09-01,2026-09-04,Team Beta,team-beta,alice@example.com,alice@example.com,alice@example.com,0.1,1,2,1,10,5,15", + ]); + }); + + it("neutralises spreadsheet formulas in user-controlled cells", () => { + const csv = buildTeamUserSpendCsv({ + ...response, + results: [row({ user_alias: null, user_email: "=HYPERLINK(1)" })], + }); + expect(csv).toContain("'=HYPERLINK(1)"); + }); + + it("names the file after the exported range", () => { + expect(teamUserSpendCsvFileName(response)).toBe("team_user_spend_2026-09-01_to_2026-09-04.csv"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.ts new file mode 100644 index 00000000000..d0b47a4e5c0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/teamUserSpend.ts @@ -0,0 +1,55 @@ +import Papa from "papaparse"; + +import type { TeamUserSpendResponse } from "@/components/networking"; + +export type TeamUserSpendRow = TeamUserSpendResponse["results"][number]; + +export const NO_USER_LABEL = "(no user)"; + +export const userLabel = (row: TeamUserSpendRow): string => { + const identity = row.user_email || row.user_alias; + return identity || row.user_id || NO_USER_LABEL; +}; + +export const teamLabel = (row: TeamUserSpendRow): string => row.team_alias || row.team_id; + +export const teamUserSpendRowId = (row: TeamUserSpendRow): string => `${row.team_id}\u0000${row.user_id}`; + +export const sortBySpendDesc = (rows: readonly TeamUserSpendRow[]): TeamUserSpendRow[] => + [...rows].sort((a, b) => b.spend - a.spend || teamLabel(a).localeCompare(teamLabel(b))); + +export const buildTeamUserSpendCsv = (response: TeamUserSpendResponse): string => + Papa.unparse( + sortBySpendDesc(response.results).map((row) => ({ + "Start Date": response.start_date, + "End Date": response.end_date, + Team: teamLabel(row), + "Team ID": row.team_id, + User: userLabel(row), + "User ID": row.user_id, + "User Email": row.user_email ?? "", + "Spend (USD)": row.spend, + Requests: row.api_requests, + Successful: row.successful_requests, + Failed: row.failed_requests, + "Prompt Tokens": row.prompt_tokens, + "Completion Tokens": row.completion_tokens, + "Total Tokens": row.total_tokens, + })), + { escapeFormulae: true }, + ); + +export const teamUserSpendCsvFileName = (response: TeamUserSpendResponse): string => + `team_user_spend_${response.start_date}_to_${response.end_date}.csv`; + +export const downloadCsv = (csv: string, fileName: string): void => { + const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" }); + const url = window.URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = fileName; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + window.URL.revokeObjectURL(url); +}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 1384679a88a..697216c5254 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -81,6 +81,7 @@ import { EmailEventSettingsResponse, EmailEventSettingsUpdateRequest } from "./e import type { SkillRegisterRequest } from "./claude_code_plugins/types"; import type { ModelBudgetUsage, ModelMaxBudget } from "./key_team_helpers/ModelMaxBudgetEditor"; import type { ObjectPermission } from "./object_permission_types"; +import type { components } from "@/lib/http/schema"; import { jsonFields } from "./common_components/check_openapi_schema"; import type { MCPUserEnvVarsStatus } from "./mcp_tools/types"; import type { @@ -1535,6 +1536,23 @@ export const teamDailyActivityAggregatedCall = async ( } }; +export type TeamUserSpendResponse = components["schemas"]["TeamUserSpendResponse"]; + +export const teamSpendByUserCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + teamIds: string[], +): Promise => + apiClient.get(`/team/spend/by_user`, { + accessToken, + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + team_ids: teamIds.join(","), + }, + }); + export const organizationDailyActivityCall = async ( accessToken: string, startTime: Date, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8c72872bdd2..f4cb88bbae1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15531,6 +15531,30 @@ export interface paths { patch?: never; trace?: never; }; + "/team/spend/by_user": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Team Spend By User + * @description Spend per user within the given teams, attributed per request from spend logs. + * + * Proxy admins may query any team. Team admins and members holding the + * `/team/daily/activity` permission see every user of the requested teams; + * other members only see their own row. + */ + get: operations["get_team_spend_by_user_team_spend_by_user_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/spend/report": { parameters: { query?: never; @@ -36769,6 +36793,63 @@ export interface components { /** Team Id */ team_id: string; }; + /** TeamUserSpendResponse */ + TeamUserSpendResponse: { + /** End Date */ + end_date: string; + /** Results */ + results: components["schemas"]["TeamUserSpendRow"][]; + /** Start Date */ + start_date: string; + }; + /** TeamUserSpendRow */ + TeamUserSpendRow: { + /** + * Api Requests + * @default 0 + */ + api_requests: number; + /** + * Completion Tokens + * @default 0 + */ + completion_tokens: number; + /** + * Failed Requests + * @default 0 + */ + failed_requests: number; + /** + * Prompt Tokens + * @default 0 + */ + prompt_tokens: number; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Successful Requests + * @default 0 + */ + successful_requests: number; + /** Team Alias */ + team_alias?: string | null; + /** Team Id */ + team_id: string; + /** + * Total Tokens + * @default 0 + */ + total_tokens: number; + /** User Alias */ + user_alias?: string | null; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id: string; + }; /** * TestCustomCodeGuardrailRequest * @description Request model for testing custom code guardrails. @@ -58550,6 +58631,39 @@ export interface operations { }; }; }; + get_team_spend_by_user_team_spend_by_user_get: { + parameters: { + query?: { + team_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamUserSpendResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_team_spend_report_team_spend_report_get: { parameters: { query?: {