Merge remote-tracking branch 'berri/litellm_internal_staging' into litellm_mcp_e2e_coverage_registry

# Conflicts:
#	tests/e2e/mcp/mcp_client.py
#	tests/e2e/models.py
This commit is contained in:
mubashir1osmani 2026-08-10 22:49:01 -07:00
commit 49ab201e75
2030 changed files with 80626 additions and 25440 deletions

View file

@ -17,3 +17,24 @@
# style: unify ruff format width on 120 (#31518)
48b5a5a0cc5a694a11219416ee0b6eb6e620e74e
# refactor(imports): move collections.abc names out of typing (#35495)
397e8e4918777e4e60a7f5e88699e0a9a7dabb3d
# refactor(lint): apply every safe ruff autofix and zero 28 strict-rule budgets (#35495)
b604e2b20c6db2099085a2f0e59b7e99e87eed6f
# refactor(logging): drop redundant !s conversion flags from f-strings (#35546)
7b2d3440cba3160277470f7a0180098ae9b87864
# perf: build log messages lazily so filtered-out log records cost nothing (#35703)
c9887a1f94bc1e7e4bdfe64d640f0509a0bc19dd
# feat(lint): enforce Final on locals and freeze function parameters (#35807)
2708620d6a599cc73c1950a942d26ac26a7ed3d4
# chore(lint): remove litellm/types from the ruff lint exclusion (#35926)
4e32a8bf6a1e1af1e04b67c759841ccef44b2235
# chore(lint): strip inert type: ignore comments and zero LIT009/LIT010/LIT011 headroom (#35928)
338e411103ad5d7003e97f34f04fa36bca542dbe

View file

@ -0,0 +1,40 @@
name: "Cache Prisma binaries"
description: >-
Cache the Prisma CLI and engine binaries that `prisma generate` downloads, so
only the first job on a given prisma-client-py version pays for the download.
prisma-client-py shells out to `npm install prisma@<version>` whenever its
binary cache directory has no CLI entrypoint, which pulls ~85 MB of query and
schema engines over the network. That normally takes a few seconds, but it is
unbounded: one shard of a proxy-db run took 5m18s on that single step versus
3.8s on its eleven siblings, which pushed the job past its timeout and got a
fully passing test run cancelled.
Callers must not set PRISMA_BINARY_CACHE_DIR. The prisma-client-py default
(~/.cache/prisma-python/binaries/<prisma-version>/<engine-version>) is already
keyed by both versions, so a cache entry can never be served to a run that
expects different binaries.
runs:
using: composite
steps:
- name: Resolve prisma-client-py version
id: version
shell: bash
run: |
version="$(grep -A1 '^name = "prisma"$' uv.lock | sed -n 's/^version = "\(.*\)"$/\1/p' | head -1)"
if [ -z "${version}" ]; then
echo "could not resolve the prisma package version from uv.lock" >&2
exit 1
fi
echo "version=${version}" >> "$GITHUB_OUTPUT"
- name: Restore Prisma binaries
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
# ~/.cache/prisma-python holds the npm install tree prisma-client-py
# drives; ~/.cache/prisma is where @prisma/engines stages its downloads.
path: |
~/.cache/prisma-python
~/.cache/prisma
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}

152
.github/ci-coverage-allowlist.yml vendored Normal file
View file

@ -0,0 +1,152 @@
description: >-
Paths deliberately outside CI coverage, each with the reason it is exempt.
assert_ci_coverage.py fails when a test file or Dockerfile is neither invoked
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds
paths:
- tests/e2e
- reason: >-
The documentation and code-quality workflows execute four files in this directory by name as
scripts and pytest never collects the directory, so these six run nowhere; listed individually
so a seventh cannot inherit the exemption
paths:
- tests/documentation_tests/test_exception_types.py
- tests/documentation_tests/test_general_setting_keys.py
- tests/documentation_tests/test_optional_params.py
- tests/documentation_tests/test_readme_providers.py
- tests/documentation_tests/test_requests_lib_usage.py
- tests/documentation_tests/test_standard_logging_payload.py
- reason: >-
Sibling files here are executed by name from the code-quality workflow; this one is referenced
by no job
paths:
- tests/code_coverage_tests/test_aio_http_image_conversion.py
- reason: >-
A second mirror of the package tree living beside tests/test_litellm, which is the mirror the
repo convention names; only test_no_hardcoded_secrets.py is invoked, from the linting
workflow, and whether this directory should exist at all is unresolved
paths:
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
- tests/litellm/integrations/helicone/test_helicone_gemini.py
- tests/litellm/litellm_core_utils/test_json_schema_validation.py
- tests/litellm/llms/anthropic/test_anthropic_reasoning_effort.py
- tests/litellm/llms/anthropic/test_anthropic_schema_filter.py
- tests/litellm/llms/azure/test_azure_embedding.py
- tests/litellm/llms/bedrock/embed/test_embedding.py
- tests/litellm/llms/bedrock/test_nova_imported_models.py
- tests/litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py
- tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
- tests/litellm/llms/oci/chat/test_oci_chat_transformation.py
- tests/litellm/llms/openai_like/test_abliteration_provider.py
- tests/litellm/llms/openai_like/test_assemblyai_provider.py
- tests/litellm/llms/openai_like/test_empiriolabs_provider.py
- tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
- tests/litellm/llms/vertex_ai/gemini/test_transformation.py
- tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py
- tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
- tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
- tests/litellm/proxy/common_utils/test_rbac_utils.py
- tests/litellm/proxy/management_endpoints/test_common_utils.py
- tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py
- tests/litellm/proxy/test_claude_code_marketplace.py
- tests/litellm/proxy/test_init_litellm_callbacks.py
- tests/litellm/proxy/test_prisma_engine_watchdog.py
- tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py
- tests/litellm/test_bedrock_extended_beta_models.py
- tests/litellm/test_bedrock_nemotron_super.py
- tests/litellm/test_proxy_auth.py
- tests/litellm/test_router_retry_backoff_headers.py
- tests/litellm/test_sambanova_model_metadata.py
- tests/litellm/test_stream_chunk_builder_images.py
- reason: >-
Legacy proxy suite superseded by the proxy shards; no job invokes it and whether it still
describes supported behaviour is unresolved
paths:
- tests/old_proxy_tests/tests/test_anthropic_context_caching.py
- tests/old_proxy_tests/tests/test_anthropic_sdk.py
- tests/old_proxy_tests/tests/test_async.py
- tests/old_proxy_tests/tests/test_gemini_context_caching.py
- tests/old_proxy_tests/tests/test_langchain_embedding.py
- tests/old_proxy_tests/tests/test_langchain_request.py
- tests/old_proxy_tests/tests/test_llamaindex.py
- tests/old_proxy_tests/tests/test_mistral_sdk.py
- tests/old_proxy_tests/tests/test_openai_embedding.py
- tests/old_proxy_tests/tests/test_openai_exception_request.py
- tests/old_proxy_tests/tests/test_openai_request.py
- tests/old_proxy_tests/tests/test_openai_request_with_traceparent.py
- tests/old_proxy_tests/tests/test_openai_simple_embedding.py
- tests/old_proxy_tests/tests/test_openai_tts_request.py
- tests/old_proxy_tests/tests/test_pass_through_langfuse.py
- tests/old_proxy_tests/tests/test_q.py
- tests/old_proxy_tests/tests/test_simple_traceparent_openai.py
- tests/old_proxy_tests/tests/test_vertex_sdk_forward_headers.py
- tests/old_proxy_tests/tests/test_vtx_embedding.py
- tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py
- reason: >-
No job invokes this suite and its files mix pure transformation tests with ones driving live
vendor vector stores, so assigning them needs a per-file decision
paths:
- tests/vector_store_tests/rag/test_rag_bedrock.py
- tests/vector_store_tests/rag/test_rag_openai.py
- tests/vector_store_tests/rag/test_rag_s3_vectors.py
- tests/vector_store_tests/rag/test_rag_vertex_ai.py
- tests/vector_store_tests/test_azure_ai_vector_store.py
- tests/vector_store_tests/test_azure_vector_store.py
- tests/vector_store_tests/test_bedrock_vector_store.py
- tests/vector_store_tests/test_gemini_vector_store.py
- tests/vector_store_tests/test_milvus_vector_store.py
- tests/vector_store_tests/test_openai_vector_store.py
- tests/vector_store_tests/test_ragflow_vector_store.py
- tests/vector_store_tests/test_s3_vectors_vector_store.py
- tests/vector_store_tests/test_vertex_ai_search_api_vector_store.py
- tests/vector_store_tests/test_vertex_ai_vector_store.py
- reason: >-
Throughput and memory-growth measurements whose runtime and variance make them unsuitable for
a per-pull-request job
paths:
- tests/load_tests/test_datadog_load_test.py
- tests/load_tests/test_langsmith_load_test.py
- tests/load_tests/test_linear_memory_growth.py
- tests/load_tests/test_memory_usage.py
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
Third-party integration tests that skip themselves without OCI configuration or sandbox
credentials, neither of which a pull request job holds
paths:
- tests/integration/sandbox/test_e2b_sandbox.py
- tests/integration/test_oci_integration.py
- tests/integration/test_oci_proxy_integration.py
- reason: >-
Two prompt-factory tests sitting at the top level of tests/ instead of under the
tests/test_litellm mirror the shards enumerate; they need moving rather than a shard entry
paths:
- tests/litellm_core_utils/test_anthropic_dedup_factory.py
- tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py
- reason: >-
A unit test for the proxy-extras package that no job invokes, while the package's other tests
live under tests/proxy_migration_tests
paths:
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
dockerfiles:
- reason: >-
The dashboard container is a static Next.js export served by nginx, and the dashboard build
and lint workflows already exercise that output, so building the image adds no signal about it
paths:
- ui/Dockerfile
- reason: >-
The Rust gateway ships as its own chart and package with a separate release pipeline, so its
image is not part of this repo's Python image set
paths:
- litellm-rust/crates/ai-gateway/Dockerfile
- reason: >-
An example image under cookbook/ that is documentation rather than a shipped artifact
paths:
- cookbook/litellm-ollama-docker-image/Dockerfile

View file

@ -13,6 +13,33 @@ How it solves it:
- <blah>
- ...
## User Flow
<!-- Two ordered lists, Before and After, walking the same end user through the same task, written strictly from that user's seat
Read the linked issue, ticket, or customer thread first so the flow reflects the real application and the routes its users actually hit; don't invent a generic scenario
Lead each list with one plain sentence saying where the flow fails (Before) or succeeds (After), then number the steps
Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
Keep the two lists step-for-step identical until they diverge, so the changed step is obvious
If the bug had a security or authorization consequence, end each list with what another user could or could no longer do
Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision
Example:
Before: a developer whose app streams chat completions gets no token counts back, so their cost dashboard reads zero
1. They send POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
2. The last SSE chunk arrives with `"usage": null`, so their app records 0 prompt and 0 completion tokens
3. They open https://litellm-domain/ui/?page=logs and see the request logged at $0 spend
After: the same request comes back with real token counts, so the dashboard shows real spend
1. The proxy admin sets `always_include_stream_usage: true` and restarts the proxy
2. The developer sends the same POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
3. The last SSE chunk now carries a `usage` object with real prompt and completion token counts
4. https://litellm-domain/ui/?page=logs shows that request at non-zero spend
-->
## Relevant issues
<!-- e.g., "Fixes #000" -->
@ -56,7 +83,11 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
🚄 Infrastructure
✅ Test
## Changes
## Caveats (if any)
<!-- Short bullet points, just like the TLDR: one line per bullet, roughly 10 words max
Call out known limitations, follow-up work, or anything a reviewer should watch out for
Leave this section empty if there are none -->
## QA runbook

262
.github/scripts/assert_ci_coverage.py vendored Normal file
View file

@ -0,0 +1,262 @@
from __future__ import annotations
import pathlib
import re
import sys
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
import yaml
REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
WORKFLOW_DIR = REPO_ROOT / ".github" / "workflows"
CIRCLECI_CONFIG = REPO_ROOT / ".circleci" / "config.yml"
ALLOWLIST_FILE = REPO_ROOT / ".github" / "ci-coverage-allowlist.yml"
TESTS_ROOT = REPO_ROOT / "tests"
ALLOWLIST_KEYS = frozenset({"description", "test_paths", "dockerfiles"})
PATH_FILTER_KEYS = frozenset({"paths", "paths-ignore"})
TEST_PATH_KEYS = frozenset({"test-path", "test-paths"})
DOCKERFILE_INPUT_KEYS = frozenset({"file", "dockerfile"})
TEST_RUNNER_RE = re.compile(r"\bpytest\b|\bcircleci tests\b|\bhelm unittest\b|\bplaywright test\b|\bpython[0-9.]*\s")
IMAGE_BUILD_RE = re.compile(r"\bdocker\s+(?:buildx\s+)?build\b")
TEST_TOKEN_RE = re.compile(r"tests/[A-Za-z0-9_./*?-]+")
DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
GLOB_CHARS = frozenset("*?")
@dataclass(frozen=True, slots=True)
class AllowEntry:
paths: tuple[str, ...]
reason: str
@dataclass(frozen=True, slots=True)
class Allowlist:
test_paths: tuple[AllowEntry, ...]
dockerfiles: tuple[AllowEntry, ...]
def covers_test(self, relative_path: str) -> bool:
return any(_token_covers(path, relative_path) for entry in self.test_paths for path in entry.paths)
def covers_dockerfile(self, relative_path: str) -> bool:
return any(relative_path == path for entry in self.dockerfiles for path in entry.paths)
@dataclass(frozen=True, slots=True)
class Scalar:
key: str
value: str
@dataclass(frozen=True, slots=True)
class Finding:
subject: str
detail: str
def _scalars(node: object, key: str) -> tuple[Scalar, ...]:
if isinstance(node, str):
return (Scalar(key=key, value=node),)
if isinstance(node, Mapping):
return tuple(
scalar
for child_key, value in node.items()
if child_key not in PATH_FILTER_KEYS
for scalar in _scalars(value, str(child_key))
)
if isinstance(node, Sequence):
return tuple(scalar for item in node for scalar in _scalars(item, key))
return ()
def _config_files() -> tuple[pathlib.Path, ...]:
workflows = tuple(sorted(path for path in WORKFLOW_DIR.iterdir() if path.suffix in (".yml", ".yaml")))
circleci = (CIRCLECI_CONFIG,) if CIRCLECI_CONFIG.is_file() else ()
return workflows + circleci
def _all_scalars() -> tuple[Scalar, ...]:
return tuple(
scalar
for path in _config_files()
for scalar in _scalars(yaml.safe_load(path.read_text(encoding="utf-8")), path.name)
)
def _uncommented(value: str) -> str:
return COMMENT_RE.sub("", value)
def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
return frozenset(
match.group(0).rstrip("/")
for scalar in scalars
if scalar.key in TEST_PATH_KEYS or TEST_RUNNER_RE.search(scalar.value)
for match in TEST_TOKEN_RE.finditer(_uncommented(scalar.value))
)
def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
return frozenset(
match.group(0)
for scalar in scalars
if scalar.key in DOCKERFILE_INPUT_KEYS or IMAGE_BUILD_RE.search(scalar.value)
for match in DOCKERFILE_TOKEN_RE.finditer(_uncommented(scalar.value))
)
def _glob_to_regex(token: str) -> re.Pattern[str]:
parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
translated = "".join(
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
)
return re.compile(rf"{translated}(?:/.*)?$")
def _token_covers(token: str, relative_path: str) -> bool:
if GLOB_CHARS & set(token):
return _glob_to_regex(token).match(relative_path) is not None
return relative_path == token or relative_path.startswith(f"{token}/")
def _test_files() -> tuple[str, ...]:
return tuple(
sorted(
path.relative_to(REPO_ROOT).as_posix()
for path in TESTS_ROOT.rglob("test_*.py")
if path.is_file() and "node_modules" not in path.parts
)
)
def _dockerfiles() -> tuple[str, ...]:
return tuple(
sorted(
path.relative_to(REPO_ROOT).as_posix()
for path in REPO_ROOT.rglob("Dockerfile*")
if path.is_file()
and ".git" not in path.parts
and "node_modules" not in path.parts
and not path.name.endswith(".dockerignore")
)
)
def _uncovered_tests(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
uncovered = tuple(
relative_path
for relative_path in _test_files()
if not any(_token_covers(token, relative_path) for token in tokens) and not allowlist.covers_test(relative_path)
)
directories = tuple(dict.fromkeys(path.rsplit("/", 1)[0] for path in uncovered))
return tuple(
Finding(
subject=directory,
detail=_describe(tuple(p for p in uncovered if p.rsplit("/", 1)[0] == directory)),
)
for directory in directories
)
def _describe(paths: tuple[str, ...]) -> str:
names = ", ".join(path.rsplit("/", 1)[1] for path in paths[:3])
suffix = f", +{len(paths) - 3} more" if len(paths) > 3 else ""
return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
return tuple(
Finding(subject=relative_path, detail="built by no job")
for relative_path in _dockerfiles()
if relative_path not in tokens and not allowlist.covers_dockerfile(relative_path)
)
def _parse_entry(item: object, section: str) -> AllowEntry:
if not isinstance(item, dict):
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' entries must be mappings")
paths = item.get("paths")
reason = item.get("reason")
if (
not isinstance(paths, list)
or not paths
or not all(isinstance(path, str) for path in paths)
or not isinstance(reason, str)
or not reason.strip()
):
raise SystemExit(
f"{ALLOWLIST_FILE.name}: every '{section}' entry needs a non-empty 'paths' "
"list of strings and a non-empty 'reason'"
)
return AllowEntry(paths=tuple(paths), reason=reason)
def _parse_entries(raw: object, section: str) -> tuple[AllowEntry, ...]:
if not isinstance(raw, list):
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' must be a list")
return tuple(_parse_entry(item, section) for item in raw)
def _load_allowlist() -> Allowlist:
if not ALLOWLIST_FILE.is_file():
return Allowlist(test_paths=(), dockerfiles=())
raw = yaml.safe_load(ALLOWLIST_FILE.read_text(encoding="utf-8")) or {}
if not isinstance(raw, dict):
raise SystemExit(f"{ALLOWLIST_FILE.name}: top level must be a mapping")
unknown = sorted(str(key) for key in raw if key not in ALLOWLIST_KEYS)
if unknown:
raise SystemExit(
f"{ALLOWLIST_FILE.name}: unknown top-level key(s) {unknown}; expected only {sorted(ALLOWLIST_KEYS)}"
)
return Allowlist(
test_paths=_parse_entries(raw.get("test_paths", []), "test_paths"),
dockerfiles=_parse_entries(raw.get("dockerfiles", []), "dockerfiles"),
)
def _write(message: str) -> None:
sys.stdout.write(f"{message}\n")
def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
_write(f"ERROR: {title}")
for finding in findings:
_write(f" - {finding.subject}: {finding.detail}")
_write("")
_write(remedy)
_write("")
def main() -> int:
allowlist = _load_allowlist()
scalars = _all_scalars()
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
if test_findings:
_report(
"test files that no CI job invokes",
test_findings,
"Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if dockerfile_findings:
_report(
"Dockerfiles that no CI job builds",
dockerfile_findings,
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if test_findings or dockerfile_findings:
return 1
_write(
f"OK: {len(_test_files())} test files and {len(_dockerfiles())} Dockerfiles are each "
"invoked by at least one job or carry an explicit allowlist entry."
)
return 0
if __name__ == "__main__":
sys.exit(main())

View file

@ -18,10 +18,25 @@ on:
type: number
default: 2
timeout-minutes:
description: "Job timeout in minutes"
description: >-
Timeout for the test step alone. Setup (checkout, dependency install,
Prisma client generation) gets its own allowance on top, so a slow
runner or a cold binary download can never cancel passing tests.
required: false
type: number
default: 20
job-timeout-minutes:
description: >-
Backstop for the whole job. Keep it >= `timeout-minutes` plus 35: 30 for
the per-step ceilings on the setup steps below, and 5 for the runner
overhead the job clock charges but no step owns (job init, step
transitions, post-job cleanup). That headroom is what makes the test
budget a floor rather than a hope, since setup cannot overrun into it
without failing its own step first. GitHub expressions have no
arithmetic, so the sum is passed in rather than computed.
required: false
type: number
default: 55
max-failures:
description: "Stop after this many failures"
required: false
@ -44,30 +59,35 @@ jobs:
run:
name: Run tests
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.timeout-minutes }}
timeout-minutes: ${{ inputs.job-timeout-minutes }}
outputs:
decision: ${{ steps.changes.outputs.decision }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: 3
with:
persist-credentials: false
- name: Detect backend-relevant changes
id: changes
timeout-minutes: 2
uses: ./.github/actions/detect-backend-changes
- name: Set up Python
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
timeout-minutes: 3
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache uv dependencies
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -79,18 +99,24 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
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
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 3
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
timeout-minutes: 3
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Run tests
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: ${{ inputs.timeout-minutes }}
env:
TEST_PATH: ${{ inputs.test-path }}
MAX_FAILURES: ${{ inputs.max-failures }}
@ -154,6 +180,19 @@ jobs:
merge-multiple: true
- name: Upload to Codecov
id: codecov-upload
continue-on-error: true
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
with:
use_oidc: true
directory: coverage-reports
root_dir: ${{ github.workspace }}
flags: ${{ inputs.artifact-name }}
fail_ci_if_error: false
- name: Upload to Codecov (retry)
if: steps.codecov-upload.outcome == 'failure'
continue-on-error: true
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
with:
use_oidc: true

View file

@ -10,6 +10,10 @@ on:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
check-sync:
name: Verify schema.prisma copies match root

View file

@ -2,18 +2,19 @@ name: Check UI API Types Sync
on:
pull_request:
paths:
- "litellm/proxy/**"
- "litellm/types/**"
- "ui/litellm-dashboard/src/lib/http/schema.d.ts"
- "ui/litellm-dashboard/scripts/gen-api-types.mjs"
- "ui/litellm-dashboard/package.json"
- "ui/litellm-dashboard/package-lock.json"
- ".github/workflows/check-ui-api-types.yml"
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
check-sync:
name: Verify schema.d.ts matches the proxy OpenAPI spec
@ -24,18 +25,39 @@ jobs:
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
fetch-depth: 2
- name: Detect changes that can affect the generated types
id: changes
run: |
set -euo pipefail
if ! base="$(git rev-parse --verify --quiet HEAD^2 >/dev/null && git rev-parse HEAD^1)"; then
echo "Not a pull request merge commit, running the full check."
echo "relevant=true" >> "$GITHUB_OUTPUT"
exit 0
fi
files="$(git diff --name-only "$base" HEAD)"
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)' <<< "$files"; then
echo "relevant=true" >> "$GITHUB_OUTPUT"
else
echo "No proxy, types or generator changes in this pull request, nothing to verify."
echo "relevant=false" >> "$GITHUB_OUTPUT"
fi
- name: Set up Python
if: steps.changes.outputs.relevant == 'true'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
if: steps.changes.outputs.relevant == 'true'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache uv dependencies
if: steps.changes.outputs.relevant == 'true'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -46,14 +68,19 @@ jobs:
${{ runner.os }}-uv-
- name: Install backend dependencies
if: steps.changes.outputs.relevant == 'true'
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Cache Prisma binaries
if: steps.changes.outputs.relevant == 'true'
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Set up Node.js
if: steps.changes.outputs.relevant == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
@ -61,16 +88,19 @@ jobs:
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dashboard dependencies
if: steps.changes.outputs.relevant == 'true'
working-directory: ui/litellm-dashboard
run: npm ci
- name: Regenerate types from the live spec
if: steps.changes.outputs.relevant == 'true'
working-directory: ui/litellm-dashboard
env:
LITELLM_PYTHON: "uv run --no-sync python"
run: npm run gen:api
- name: Fail if types are stale
if: steps.changes.outputs.relevant == 'true'
run: |
if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."

42
.github/workflows/ci-coverage.yml vendored Normal file
View file

@ -0,0 +1,42 @@
name: "CI Coverage"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
assert-ci-coverage:
name: assert-ci-coverage
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Assert every test file and Dockerfile is invoked by a job
run: |
python -m pip install "pyyaml==6.0.3"
python .github/scripts/assert_ci_coverage.py

View file

@ -14,6 +14,10 @@ on:
permissions:
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
lint-pr-title:
name: Validate PR title

View file

@ -13,35 +13,16 @@ jobs:
contents: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create daily oss-agent-shin branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
# Configure Git user
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
# Generate branch name with MM_DD_YYYY format
BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
# Fetch all branches
git fetch --all
# Check if the branch already exists
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
else
echo "Creating new branch: $BRANCH_NAME"
# Create the new branch from main
git checkout -b $BRANCH_NAME origin/main
# Push the new branch
git push origin $BRANCH_NAME
echo "Successfully created and pushed branch: $BRANCH_NAME"
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"

View file

@ -13,38 +13,19 @@ jobs:
contents: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create daily staging branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
# Configure Git user
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
# Generate branch name with MM_DD_YYYY format
BRANCH_NAME="litellm_oss_staging_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
# Fetch all branches
git fetch --all
# Check if the branch already exists
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
else
echo "Creating new branch: $BRANCH_NAME"
# Create the new branch from main
git checkout -b $BRANCH_NAME origin/main
# Push the new branch
git push origin $BRANCH_NAME
echo "Successfully created and pushed branch: $BRANCH_NAME"
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"
create-internal-dev-branch:
if: github.repository == 'BerriAI/litellm'
@ -53,35 +34,16 @@ jobs:
contents: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create internal dev branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
# Configure Git user
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
# Generate branch name with MM_DD_YYYY format
BRANCH_NAME="litellm_internal_dev_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
# Fetch all branches
git fetch --all
# Check if the branch already exists
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
else
echo "Creating new branch: $BRANCH_NAME"
# Create the new branch from main
git checkout -b $BRANCH_NAME origin/main
# Push the new branch
git push origin $BRANCH_NAME
echo "Successfully created and pushed branch: $BRANCH_NAME"
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"

View file

@ -15,6 +15,10 @@ on:
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
guard:
name: Block fork dependency changes

View file

@ -9,6 +9,10 @@ on:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
unit-test:
runs-on: ubuntu-latest
@ -23,21 +27,28 @@ jobs:
with:
version: "3.11.1"
- name: Download and verify Helm Unit Test Plugin
run: |
curl -fsSLo "$RUNNER_TEMP/helm-unittest.tgz" https://github.com/helm-unittest/helm-unittest/releases/download/v0.8.2/helm-unittest-linux-amd64-0.8.2.tgz
echo "56ab3091e6fa52a7c92ee951def9bed957f295d9ce98483aed404e748d7b3a94 $RUNNER_TEMP/helm-unittest.tgz" | sha256sum -c -
- name: Install Helm Unit Test Plugin
run: |
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4
- name: Verify Helm Unit Test Plugin integrity
run: |
EXPECTED_SHA="e251ba198448629678ff2168e1a469249d998155"
PLUGIN_DIR="$(helm env HELM_PLUGINS)/helm-unittest"
ACTUAL_SHA="$(git -C "$PLUGIN_DIR" rev-parse HEAD)"
if [ "$ACTUAL_SHA" != "$EXPECTED_SHA" ]; then
echo "::error::Helm unittest plugin checksum mismatch! Expected $EXPECTED_SHA but got $ACTUAL_SHA"
exit 1
fi
echo "Helm unittest plugin integrity verified: $ACTUAL_SHA"
mkdir -p "$PLUGIN_DIR"
tar -xzf "$RUNNER_TEMP/helm-unittest.tgz" -C "$PLUGIN_DIR"
helm plugin list
- name: Run unit tests
run: |
helm unittest -f 'tests/*.yaml' helm/litellm-helm
helm unittest -f 'tests/*.yaml' helm/litellm
for chart in helm/litellm-helm helm/litellm; do
declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')"
output="$(mktemp)"
helm unittest -f 'tests/*.yaml' "$chart" | tee "$output"
executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")"
if [ "$declared" != "$executed" ]; then
echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute."
exit 1
fi
echo "$chart: all $declared declared test suites ran"
done

View file

@ -8,10 +8,17 @@ on:
- litellm_oss_branch
- "litellm_**"
paths:
- Dockerfile
- docker/Dockerfile.non_root
- migrations/Dockerfile
- migrations/run.py
- tests/proxy_migration_tests/test_offline_image_migration.py
- gateway/Dockerfile
- gateway/main.py
- backend/Dockerfile
- backend/main.py
- docker/component_entrypoint.sh
- litellm-proxy-extras/**
- tests/proxy_migration_tests/**
- uv.lock
- ui/litellm-dashboard/package-lock.json
- .github/workflows/image-scan.yml
@ -86,6 +93,35 @@ jobs:
--fail-on high \
--output table
runtime-image:
name: runtime-image
runs-on: ubuntu-latest
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build runtime image
run: docker build -f Dockerfile -t litellm-runtime-scan:${{ github.sha }} .
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify offline migration as a non-root uid
env:
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
migrations-image:
name: migrations-image
runs-on: ubuntu-latest
@ -116,3 +152,63 @@ jobs:
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
gateway-image:
name: gateway-image
runs-on: ubuntu-latest
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build gateway image
run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} .
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify the gateway serves offline as a non-root uid
env:
LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }}
LITELLM_COMPONENT_PORT: "4000"
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
backend-image:
name: backend-image
runs-on: ubuntu-latest
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build backend image
run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} .
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify the backend serves offline as a non-root uid
env:
LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }}
LITELLM_COMPONENT_PORT: "4001"
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v

View file

@ -57,9 +57,10 @@ jobs:
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma

View file

@ -0,0 +1,63 @@
name: Publish basedpyright base counts
# Every commit on litellm_internal_staging is some branch's future merge-base.
# Publishing its per-rule basedpyright counts as an artifact lets
# scripts/type_check_gate.py download them in seconds instead of paying a
# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
# No concurrency group on purpose: runs must never cancel each other, because
# every sha's artifact matters (any of them can become a merge-base).
on:
push:
branches:
- litellm_internal_staging
workflow_dispatch:
inputs:
ref:
description: "Ref to compute and publish base counts for"
required: false
default: litellm_internal_staging
permissions:
contents: read
jobs:
publish:
runs-on: ubuntu-latest
timeout-minutes: 20
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ inputs.ref || github.sha }}
clean: true
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
# The gate provisions its own measurement env (.venv-typecheck: a frozen
# uv sync of its canonical dependency groups plus a generated Prisma
# client), so no install step here can drift from what local runs measure.
- name: Emit basedpyright counts for HEAD
run: |
python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json)
echo "COUNTS_ARTIFACT_NAME=$(basename "$counts_file" .json)" >> "$GITHUB_ENV"
- name: Upload counts artifact
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: ${{ env.COUNTS_ARTIFACT_NAME }}
path: ${{ runner.temp }}/basedpyright-counts/
if-no-files-found: error

View file

@ -65,6 +65,12 @@ jobs:
- name: check_provider_folders_documented
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
- name: check_prisma_binary_cache
run: uv run --no-sync python ./tests/code_coverage_tests/check_prisma_binary_cache.py
- name: check_workflow_startup_safety
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

View file

@ -11,10 +11,20 @@ on:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
lint:
runs-on: ubuntu-latest
timeout-minutes: 15
# actions: read lets scripts/type_check_gate.py download the base-counts
# artifact published by publish-basedpyright-base-counts.yml instead of
# re-running basedpyright over the merge-base tree.
permissions:
contents: read
actions: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@ -23,10 +33,21 @@ jobs:
# Any-discipline) would otherwise blame on this branch.
with:
ref: ${{ github.event.pull_request.head.sha }}
fetch-depth: 0
fetch-depth: 1
clean: true
persist-credentials: false
- name: Fetch gate base (merge-base with target branch)
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
MERGE_BASE=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
test -n "$MERGE_BASE"
git fetch --no-tags --depth=1 origin "$MERGE_BASE"
echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV"
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
@ -50,20 +71,19 @@ jobs:
run: |
uv sync --frozen --group proxy-dev --group e2e-dev
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
# basedpyright resolves Prisma's generated client (litellm/proxy/schema.prisma)
# only after `prisma generate` writes prisma/client.py et al. Without this the
# DB wrappers typed against the generated client would degrade to Unknown.
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Check ruff format
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
git diff --name-only --diff-filter=ACMR "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
echo "No changed litellm Python files to check with ruff format."
exit 0
@ -86,16 +106,12 @@ jobs:
cd ..
- name: Check strict-rule budget (delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
uv run --no-sync python scripts/ruff_strict_gate.py --base "$BASE_SHA"
uv run --no-sync python scripts/ruff_strict_gate.py --base "$GATE_BASE_SHA"
- name: Check type-discipline budget (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
uv run --no-sync python scripts/type_discipline_gate.py --base "$BASE_SHA"
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
- name: Print OpenAI version
run: |
@ -103,15 +119,13 @@ jobs:
- name: Check basedpyright budget (delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
GH_TOKEN: ${{ github.token }}
run: |
uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"
- name: Check tests/e2e basedpyright (zero errors)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
if git diff --name-only --diff-filter=ACMRD "$BASE_SHA"...HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
uv run --no-sync basedpyright tests/e2e
else
echo "No changed tests/e2e Python files; skipping."
@ -140,9 +154,15 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
fetch-depth: 1
persist-credentials: false
- name: Fetch ratchet base
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
git fetch --no-tags --depth=1 origin "$BASE_SHA"
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
@ -163,7 +183,7 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
fetch-depth: 1
persist-credentials: false
- name: Set up Python
@ -178,13 +198,14 @@ jobs:
- name: Run secret scan test
run: |
uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
uv run --no-project --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
- name: Run ggshield secret scan
env:
GITGUARDIAN_API_KEY: ${{ secrets.GITGUARDIAN_API_KEY }}
run: |
if [ -n "$GITGUARDIAN_API_KEY" ]; then
git fetch --no-tags --unshallow origin
uv tool run --from 'ggshield==1.48.0' ggshield secret scan repo .
else
echo "GITGUARDIAN_API_KEY not set, skipping ggshield scan"

View file

@ -10,6 +10,10 @@ on:
- litellm_oss_staging
- "litellm_**"
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
build-ui:
runs-on: ubuntu-latest

View file

@ -10,6 +10,10 @@ on:
- litellm_oss_staging
- "litellm_**"
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
frontend-lint:
runs-on: ubuntu-latest
@ -22,12 +26,13 @@ jobs:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
fetch-depth: 1
persist-credentials: false
- name: Collect changed files
id: changed
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
@ -37,7 +42,9 @@ jobs:
# landed since, so a PR that touches no UI file still gets linted
# against hundreds of other people's files. Diff the PR head against its
# own merge base instead, which is exactly what this PR changed.
merge_base=$(git merge-base "$BASE_SHA" "$HEAD_SHA")
merge_base=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
test -n "$merge_base"
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
: > "$RUNNER_TEMP/prettier_files.txt"
: > "$RUNNER_TEMP/eslint_files.txt"
while IFS= read -r f; do

View file

@ -29,7 +29,7 @@ jobs:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
fetch-depth: 1
persist-credentials: false
- name: Setup Node.js
@ -42,14 +42,32 @@ jobs:
- name: Install dependencies
run: npm ci
- name: Run UI type tests (Vitest)
env:
CI: "true"
run: npm run test:types
- name: Run UI unit tests (Vitest)
env:
CI: "true"
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
if [ -n "$BASE_SHA" ]; then
echo "Pull request: running only tests related to changes since $BASE_SHA"
npm run test -- --run --changed "$BASE_SHA" --passWithNoTests \
merge_base=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
test -n "$merge_base"
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
changed_files=()
while IFS= read -r f; do
changed_files+=("$f")
done < <(git diff --name-only --relative "$merge_base" "$HEAD_SHA" -- .)
if [ ${#changed_files[@]} -eq 0 ]; then
echo "No UI files changed in this PR; skipping unit tests."
exit 0
fi
echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
--pool forks --poolOptions.forks.maxForks=14
else
echo "Push to $GITHUB_REF_NAME: running the full suite"

View file

@ -11,6 +11,10 @@ on:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
test:
runs-on: ubuntu-latest

View file

@ -11,6 +11,10 @@ on:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
validate-model-prices-json:
runs-on: ubuntu-latest

View file

@ -92,9 +92,10 @@ jobs:
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma

View file

@ -65,10 +65,12 @@ jobs:
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma

View file

@ -40,7 +40,11 @@ jobs:
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/realtime_api
tests/test_litellm/rerank_api
tests/test_litellm/sandbox
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py

View file

@ -28,6 +28,10 @@ concurrency:
# Most of a shard's time is pytest plugin load + xdist worker imports +
# pytest-cov instrumentation, not the tests themselves. Keeping per-shard
# work low and matching worker count to runner cores is what controls it.
# * `timeout` bounds the pytest step only. Checkout, dependency install, and
# Prisma client generation draw on a separate allowance in the base
# workflow, so slow setup shows up as a slow job rather than as a
# cancelled shard whose tests were passing.
# * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores
# oversubscribes 2x and workers fight for CPU during their cold-start
# imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective).

View file

@ -29,19 +29,24 @@ jobs:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
@ -71,4 +76,5 @@ jobs:
workers: 4
reruns: 2
timeout-minutes: 60
job-timeout-minutes: 95
artifact-name: proxy-server

View file

@ -33,6 +33,8 @@ jobs:
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
workers: 2

View file

@ -82,10 +82,12 @@ jobs:
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma

View file

@ -51,9 +51,10 @@ jobs:
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma

1
.gitignore vendored
View file

@ -1,5 +1,6 @@
.python-version
.venv
.venv-typecheck
.venv_policy_test
.env
.claude

View file

@ -1,4 +1,12 @@
Do not write any comments (existing comments can stay) unless explicitly asked to in a user (not system) prompt
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
@ -9,7 +17,7 @@ Don't assume that the existing code is correct or the right way of doing things
- easy to maintain/change
- modern
In that order of importance
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
@ -29,7 +37,7 @@ Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We pref
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
@ -41,9 +49,11 @@ Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
@ -57,7 +67,7 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages
When working on a PR, keep the PR description in sync with new commits being made
Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
@ -70,7 +80,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc.
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>` explaining why
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed

View file

@ -134,7 +134,8 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
EXPOSE 4000/tcp

View file

@ -8,7 +8,7 @@
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety pre-commit \
install-helm-unittest check-circular-imports check-import-safety check pre-commit \
lint-install lint-fetch-base bootstrap
# Default target
@ -22,7 +22,8 @@ help:
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)"
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
@echo " make pre-commit - Legacy alias for make check"
@echo " make format - Apply ruff format code formatting"
@echo " make format-check - Check ruff format code formatting (matches CI)"
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
@ -75,7 +76,7 @@ install-dev:
bootstrap:
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
cd ui/litellm-dashboard && npm install --no-audit --no-fund
cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
@ -99,7 +100,10 @@ install-test-deps: install-proxy-dev
$(UV_RUN) prisma generate --schema litellm/proxy/schema.prisma
install-helm-unittest:
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists"
@helm plugin list | grep -qE '^unittest[[:space:]]+0\.8\.2([[:space:]]|$$)' || { \
helm plugin uninstall unittest >/dev/null 2>&1 || true; \
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
}
# Install git hooks that enforce Conventional Commits and Conventional Branches.
# Opt-in: not chained into install-dev.
@ -121,10 +125,10 @@ lint-fetch-base:
git fetch origin litellm_internal_staging
# Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
# Prisma client, so basedpyright resolves the same modules CI does (without the generated
# client the DB wrappers typed against it degrade to Unknown, drifting the budget from
# CI's). --inexact tops up the venv instead of pruning the proxy extras gen:api and the
# running proxy need.
# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
# budget gate itself no longer measures here (scripts/type_check_gate.py provisions its
# own .venv-typecheck). --inexact tops up the venv instead of pruning the proxy extras
# gen:api and the running proxy need.
lint-install:
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
@ -233,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety
# Run the gating CI checks against your staged files right before committing. Mirrors
# Run the gating CI checks against your changes. Scopes to staged files when anything
# is staged (warning about changed files left unstaged); with nothing staged it falls
# back to the working tree's diff against the merge base with the base branch, so a
# fresh merge commit or an unstaged working tree still gets checked. Mirrors
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope.
# Not auto-installed as a git hook so it never slows an unrelated human commit.
pre-commit: bootstrap
check: bootstrap
./scripts/pre_commit_lint.sh
pre-commit:
@echo "make pre-commit is a legacy alias; use make check" >&2
@$(MAKE) check
# Testing targets
test: install-test-deps
$(UV_RUN) pytest tests/

View file

@ -59,9 +59,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra semantic-router \
--python python3
RUN mkdir -p /home/nonroot && \
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
chown -R nonroot:nonroot /home/nonroot/.cache
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
@ -83,13 +83,16 @@ ENV HOME=/home/nonroot \
PATH="/app/.venv/bin:${PATH}" \
PYTHONPATH="/app" \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /home/nonroot/.cache /home/nonroot/.cache
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
USER nonroot

View file

@ -82,6 +82,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/user_agent",
"/usage/",
"/daily/",
# Deployment-wide gateway request counts. Scoped to the analytics read rather
# than all of /gateway/, which stays free for data-plane routes.
"/gateway/daily/",
# CloudZero cost-export admin (init / settings / export / dry-run / delete)
"/cloudzero/",
# Caching admin

View file

@ -1,45 +1,45 @@
{
"reportAny": {
"limit": 29809
"limit": 26391
},
"reportArgumentType": {
"limit": 2645
"limit": 2614
},
"reportAssignmentType": {
"limit": 329
"limit": 327
},
"reportAttributeAccessIssue": {
"limit": 516
"limit": 514
},
"reportCallIssue": {
"limit": 123
"limit": 114
},
"reportConstantRedefinition": {
"limit": 40
},
"reportDeprecated": {
"limit": 325
"limit": 215
},
"reportDuplicateImport": {
"limit": 24
"limit": 19
},
"reportExplicitAny": {
"limit": 9473
"limit": 8319
},
"reportFunctionMemberAccess": {
"limit": 11
"limit": 7
},
"reportGeneralTypeIssues": {
"limit": 157
},
"reportIncompatibleMethodOverride": {
"limit": 77
"limit": 56
},
"reportIncompatibleVariableOverride": {
"limit": 12
"limit": 8
},
"reportInconsistentOverload": {
"limit": 18
"limit": 12
},
"reportIndexIssue": {
"limit": 35
@ -48,16 +48,16 @@
"limit": 35
},
"reportInvalidTypeVarUse": {
"limit": 5
"limit": 2
},
"reportMatchNotExhaustive": {
"limit": 0
},
"reportMissingParameterType": {
"limit": 5855
"limit": 5825
},
"reportMissingTypeArgument": {
"limit": 15849
"limit": 15695
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1079
"limit": 1077
},
"reportOptionalOperand": {
"limit": 0
@ -84,49 +84,49 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 2436
"limit": 1824
},
"reportRedeclaration": {
"limit": 8
},
"reportReturnType": {
"limit": 219
"limit": 213
},
"reportTypedDictNotRequiredAccess": {
"limit": 27
"limit": 26
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45262
"limit": 44996
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40452
"limit": 39643
},
"reportUnknownParameterType": {
"limit": 20309
"limit": 20132
},
"reportUnknownVariableType": {
"limit": 31978
"limit": 31153
},
"reportUnnecessaryCast": {
"limit": 124
"limit": 118
},
"reportUnnecessaryComparison": {
"limit": 703
"limit": 701
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 866
"limit": 857
},
"reportUntypedBaseClass": {
"limit": 165
"limit": 0
},
"reportUntypedFunctionDecorator": {
"limit": 33
@ -138,9 +138,9 @@
"limit": 139
},
"reportUnusedImport": {
"limit": 588
"limit": 555
},
"reportUnusedVariable": {
"limit": 147
"limit": 146
}
}

View file

@ -133,7 +133,8 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
EXPOSE 4000/tcp

View file

@ -185,7 +185,8 @@ RUN mkdir -p /nonexistent /app/.cache /var/lib/litellm/assets /var/lib/litellm/u
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1
ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1 && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
USER 65534

View file

@ -47,7 +47,13 @@ RUN uv venv --python python && \
"prisma==0.11.0" \
"openai==2.24.0"
RUN prisma generate --schema=./schema.prisma
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma && \
chmod -R a+rX /opt/prisma && \
python -c "import sys; from prisma.client import BINARY_PATHS; bad = sorted(p for group in BINARY_PATHS.model_dump().values() for p in group.values() if not p.startswith('/opt/prisma/')); sys.exit('prisma engines baked outside /opt/prisma: %r' % bad) if bad else None"
ENV PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
EXPOSE 4000/tcp

View file

@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, List, Optional, Tuple
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -43,11 +43,15 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
async def _get_user_info(self, batch_id, user_id) -> dict:
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
"""
Look up user email and key alias by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
Returns an empty dict when user_id is None: batches created by a team or service
account key carry no user id, and find_unique(where={"user_id": None}) raises.
"""
if not user_id:
return {}
try:
user_row = await self.prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
@ -62,6 +66,66 @@ class CheckBatchCost:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
return {}
async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None:
"""Resolve the creating virtual key's alias from its hashed token."""
if not api_key:
return None
try:
key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
return getattr(key_row, "key_alias", None) if key_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
return None
async def _get_team_alias(self, team_id: str | None) -> str | None:
"""Resolve a team's alias from its id."""
if not team_id:
return None
try:
team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
return getattr(team_row, "team_alias", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> Dict[str, Any]:
"""
Rebuild the spend-tracking metadata for the key, team, and tags that created the
batch so the batch-cost spend log is attributed the same way a non-batch request
is. Rows created before api_key and request_tags were persisted carry only
created_by and team_id, and fall back to those. A named creating key owns
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
request_tags = getattr(job, "request_tags", None)
metadata: Dict[str, Any] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}
key_alias = await self._get_key_alias(batch_id, api_key)
if key_alias is not None:
metadata["user_api_key_alias"] = key_alias
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
return metadata
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@ -296,17 +360,13 @@ class CheckBatchCost:
underlying provider model (e.g. ``gpt-5.5``), which no key is allowed to call.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
get_models_from_unified_file_id,
resolve_managed_output_file_model_name,
)
input_file_id = cls._get_input_file_id(job)
target_model_names = (
get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(input_file_id)) if input_file_id else []
return resolve_managed_output_file_model_name(
unified_input_file_id=cls._get_input_file_id(job),
fallback_model_name=deployment_info.model_name or None,
)
if target_model_names:
return ",".join(target_model_names)
return deployment_info.model_name or None
@staticmethod
def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]:
@ -489,9 +549,6 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
creator_user_id = job.created_by
user_info = await self._get_user_info(batch_id, job.created_by)
logging_obj.update_environment_variables(
litellm_params={
# set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@ -500,10 +557,7 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
"metadata": {
"user_api_key_user_id": creator_user_id,
**user_info,
},
"metadata": await self._build_creator_attribution_metadata(job, batch_id),
},
optional_params={},
)
@ -660,6 +714,20 @@ class CheckBatchCost:
elif response.status in ("failed", "expired", "cancelled"):
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
ensure_batch_response_managed_file_ids,
)
response.id = job.unified_object_id
await ensure_batch_response_managed_file_ids(
response=response,
managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"),
prisma_client=self.prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data = {
"status": response.status,
"file_object": response.model_dump_json(),

View file

@ -1,10 +1,10 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by litellm.aget_responses().
Cost tracking is handled automatically by the get-responses call.
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Dict, Optional, cast
import litellm
from litellm._logging import verbose_proxy_logger
@ -13,11 +13,15 @@ from litellm.constants import (
MAX_OBJECTS_PER_POLL_CYCLE,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
class CheckResponsesCost:
def __init__(
@ -33,6 +37,28 @@ class CheckResponsesCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
async def _get_response(
self,
response_id: str,
litellm_metadata: Dict[str, str],
) -> ResponsesAPIResponse:
"""Fetch the upstream response, using deployment credentials when available.
LiteLLM-encoded response IDs carry the ``model_id`` of the deployment that
served the original request, so routing through ``llm_router`` applies that
deployment's ``api_base`` / ``api_key`` / ``api_version``, exactly like
``GET /v1/responses/{id}`` does. ``litellm.aget_responses`` on its own only
sees provider env vars, so it fails for every deployment whose credentials
live in the config; the row then never leaves ``queued``.
"""
model_id: Optional[str] = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
if model_id is None or self.llm_router.get_deployment(model_id=model_id) is None:
return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
router_response = await self.llm_router.aget_responses(
response_id=response_id, litellm_metadata=litellm_metadata
)
return cast(ResponsesAPIResponse, router_response)
async def _expire_stale_rows(
self, cutoff: datetime, batch_size: int
) -> int:
@ -87,8 +113,8 @@ class CheckResponsesCost:
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is automatically tracked by litellm.aget_responses()
- Mark completed/failed/cancelled responses as complete in the database
- Cost is automatically tracked by the get-responses call
- Mark responses in a terminal state as complete in the database
"""
try:
await self._cleanup_stale_managed_objects()
@ -134,7 +160,7 @@ class CheckResponsesCost:
litellm_metadata["model"] = model_name
litellm_metadata["model_group"] = model_name # Use same value for model_group
response = await litellm.aget_responses(
response = await self._get_response(
response_id=responses_id_security,
litellm_metadata=litellm_metadata,
)
@ -144,21 +170,14 @@ class CheckResponsesCost:
)
except Exception as e:
verbose_proxy_logger.info(
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
continue
# Check if response is in a terminal state
if response.status == "completed":
if response.status in TERMINAL_RESPONSE_STATUSES:
verbose_proxy_logger.info(
f"Response {unified_object_id} is complete. Cost automatically tracked by aget_responses."
)
completed_jobs.append(job)
elif response.status in ["failed", "cancelled"]:
verbose_proxy_logger.info(
f"Response {unified_object_id} has status {response.status}, marking as complete"
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
)
completed_jobs.append(job)

File diff suppressed because it is too large Load diff

View file

@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import delete_cached_project_object
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_helpers.utils import (
@ -514,6 +515,7 @@ async def update_project(
litellm_proxy_admin_name,
premium_user,
prisma_client,
user_api_key_cache,
)
try:
@ -672,6 +674,11 @@ async def update_project(
include={"litellm_budget_table": True, "object_permission": True},
)
await delete_cached_project_object(
project_id=data.project_id,
user_api_key_cache=user_api_key_cache,
)
return updated_project
except Exception as e:
verbose_proxy_logger.exception(
@ -710,7 +717,7 @@ async def delete_project(
}'
```
"""
from litellm.proxy.proxy_server import premium_user, prisma_client
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
try:
if not premium_user:
@ -773,6 +780,11 @@ async def delete_project(
prisma_models.LiteLLM_ProjectTable | None
) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
await delete_cached_project_object(
project_id=project_id,
user_api_key_cache=user_api_key_cache,
)
deleted_projects.append(deleted_project)
return deleted_projects
@ -831,7 +843,7 @@ async def project_info(
)
# Check if user has access to this project (admin or team member)
is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
is_admin = user_api_key_has_admin_view(user_api_key_dict)
is_team_member = False
if project.team_id and user_api_key_dict.user_id:
@ -886,7 +898,7 @@ async def list_projects(
)
# If proxy admin, get all projects
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
if user_api_key_has_admin_view(user_api_key_dict):
projects: Sequence[
prisma_models.LiteLLM_ProjectTable
] = await prisma_client.db.litellm_projecttable.find_many(

View file

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.53"
version = "0.1.54"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.53"
version = "0.1.54"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -61,9 +61,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3
RUN mkdir -p /home/nonroot && \
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
chown -R nonroot:nonroot /home/nonroot/.cache
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
@ -85,13 +85,16 @@ ENV HOME=/home/nonroot \
PATH="/app/.venv/bin:${PATH}" \
PYTHONPATH="/app" \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /home/nonroot/.cache /home/nonroot/.cache
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
USER nonroot

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -0,0 +1,5 @@
-- Add api_key and request_tags columns to LiteLLM_ManagedObjectTable
-- Captured at batch-create time so CheckBatchCost can attribute batch-cost spend
-- back to the creating virtual key (and its tags) even when created_by is null.
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT;
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB DEFAULT '[]';

View file

@ -0,0 +1,15 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGatewayRequests" (
"date" TEXT NOT NULL,
"category" TEXT NOT NULL,
"route" TEXT NOT NULL,
"successful_requests" BIGINT NOT NULL DEFAULT 0,
"failed_requests" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyGatewayRequests_pkey" PRIMARY KEY ("date","category","route")
);
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGatewayRequests_date_idx" ON "LiteLLM_DailyGatewayRequests"("date");

View file

@ -0,0 +1,31 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterSession" (
"api_key" TEXT NOT NULL,
"session_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"router_type" TEXT NOT NULL,
"first_turn_at" TIMESTAMP(3) NOT NULL,
"last_turn_at" TIMESTAMP(3) NOT NULL,
"last_model" TEXT NOT NULL,
"models" JSONB NOT NULL DEFAULT '{}',
"turns" INTEGER NOT NULL DEFAULT 0,
"unordered_turns" INTEGER NOT NULL DEFAULT 0,
"covered_turns" INTEGER NOT NULL DEFAULT 0,
"cache_hits" INTEGER NOT NULL DEFAULT 0,
"same_model_turns" INTEGER NOT NULL DEFAULT 0,
"same_model_hits" INTEGER NOT NULL DEFAULT 0,
"first_visit_turns" INTEGER NOT NULL DEFAULT 0,
"first_visit_hits" INTEGER NOT NULL DEFAULT 0,
"return_turns" INTEGER NOT NULL DEFAULT 0,
"return_hits" INTEGER NOT NULL DEFAULT 0,
"return_expired_misses" INTEGER NOT NULL DEFAULT 0,
"return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0,
"ttl_5m_turns" INTEGER NOT NULL DEFAULT 0,
"ttl_1h_turns" INTEGER NOT NULL DEFAULT 0,
"total_tokens" BIGINT NOT NULL DEFAULT 0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_AutoRouterSession_pkey" PRIMARY KEY ("api_key", "session_id", "router_name")
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_session_last_turn" ON "LiteLLM_AutoRouterSession"("last_turn_at");

View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';

View file

@ -0,0 +1,181 @@
"""Prepare the Node toolchain the Prisma CLI needs, separately from migrations.
The Prisma CLI is a Node program. The first invocation inside a fresh
container installs a private Node runtime and npm-installs the CLI itself,
which can take minutes on a cold or slow machine. Sharing one timeout between
that one-time bootstrap and the migration commands makes a slow bootstrap
indistinguishable from a slow migration, so the bootstrap gets killed long
before it can finish.
A killed bootstrap does not correct itself. The installer leaves its cache
directory behind, and Prisma decides whether to install by testing that
directory for existence alone, so every later attempt skips the install and
then fails on a Node binary that was never written. Deleting a cache directory
that exists without a Node binary is what turns a killed bootstrap back into a
recoverable one.
Both budgets are overridable so an operator can widen them without a release:
``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command.
"""
import math
import os
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from litellm_proxy_extras._logging import logger
try:
from prisma import config as prisma_config
except ImportError:
prisma_config = None
PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
BOOTSTRAP_ARG = "--version"
@dataclass(frozen=True)
class ToolchainBootstrap:
"""Outcome of preparing the Prisma toolchain."""
healed_incomplete_cache: bool
ready: bool
def _timeout_from_env(env_var: str, default: float) -> float:
raw = os.getenv(env_var)
if raw is None:
return default
try:
seconds = float(raw)
except ValueError:
logger.warning(
"%s=%r is not a number, falling back to %ss", env_var, raw, default
)
return default
if not math.isfinite(seconds) or seconds <= 0:
logger.warning(
"%s=%r is not a finite positive number, falling back to %ss",
env_var,
raw,
default,
)
return default
return seconds
def prisma_command_timeout() -> float:
"""Seconds any single Prisma command may run for."""
return _timeout_from_env(
PRISMA_COMMAND_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_COMMAND_TIMEOUT
)
def prisma_bootstrap_timeout() -> float:
"""Seconds the one-time Node toolchain install may run for."""
return _timeout_from_env(
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT
)
def nodeenv_cache_dir() -> Optional[Path]:
"""Where Prisma installs its private Node runtime, or None if unknowable."""
override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
if override:
return Path(override).absolute()
if prisma_config is not None:
try:
return Path(prisma_config.nodeenv_cache_dir).absolute()
except (OSError, ValueError) as e:
logger.warning("Could not read the Prisma nodeenv cache dir: %s", e)
try:
return Path.home() / ".cache" / "prisma-python" / "nodeenv"
except RuntimeError:
logger.warning(
"No resolvable home directory, cannot locate the Prisma nodeenv cache"
)
return None
def node_binary_path(cache_dir: Path) -> Path:
"""Path the Node binary occupies once the toolchain is fully installed."""
if os.name == "nt":
return cache_dir / "Scripts" / "node.exe"
return cache_dir / "bin" / "node"
def heal_incomplete_nodeenv_cache() -> bool:
"""Delete a nodeenv cache directory left without a Node binary.
Returns True when a half-installed toolchain was removed, so the next
Prisma invocation reinstalls it instead of failing on a missing binary.
"""
cache_dir = nodeenv_cache_dir()
if cache_dir is None:
return False
try:
if not cache_dir.is_dir() or node_binary_path(cache_dir).exists():
return False
except OSError as e:
logger.warning("Could not inspect the Node toolchain at %s: %s", cache_dir, e)
return False
logger.warning(
"Node toolchain at %s has no %s, so a previous install was interrupted. "
"Removing it so it can be reinstalled.",
cache_dir,
node_binary_path(cache_dir).name,
)
try:
shutil.rmtree(cache_dir)
except OSError as e:
logger.warning("Could not remove %s: %s", cache_dir, e)
return False
return True
def ensure_prisma_toolchain(
prisma_command: str, prisma_env: dict[str, str]
) -> ToolchainBootstrap:
"""Install whatever the Prisma CLI needs to run, under its own timeout.
Never raises. A toolchain that cannot be prepared is reported so the
caller can go on and let the real Prisma command produce the real error.
"""
healed = heal_incomplete_nodeenv_cache()
timeout = prisma_bootstrap_timeout()
logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout)
try:
subprocess.run(
[prisma_command, BOOTSTRAP_ARG],
timeout=timeout,
check=True,
capture_output=True,
text=True,
env=prisma_env,
)
except subprocess.TimeoutExpired:
logger.warning(
"Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "
"if this machine needs longer to install it.",
timeout,
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR,
)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
except subprocess.CalledProcessError as e:
logger.warning("Preparing the Prisma CLI toolchain failed: %s", e.stderr)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
except OSError as e:
logger.warning("Could not run the Prisma CLI: %s", e)
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
logger.info("Prisma CLI toolchain ready")
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=True)

View file

@ -16,6 +16,7 @@ import tempfile
from pathlib import Path
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
@ -75,7 +76,7 @@ def apply_replica_identity_full(
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,

View file

@ -30,7 +30,7 @@ model LiteLLM_BudgetTable {
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
tags LiteLLM_TagTable[] // multiple tags can have the same budget
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
// Models on proxy
@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
ptu_flat_cost Float @default(0.0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -985,6 +986,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt
updated_by String?
@ -1118,6 +1121,26 @@ model LiteLLM_DailyToolSpend {
@@id([date, tool_name])
}
// Gateway request counts recorded at the ASGI edge by
// BillableRequestMetricsMiddleware. This is the source of truth for SGR
// (successful gateway requests): it counts what the proxy actually answered,
// independent of whether the request reached litellm's logging callbacks.
// The key carries no deployment or caller dimension. Every part of it is
// chosen by the proxy and drawn from a closed set, so the table is bounded by
// (days x categories x routes) rather than by anything a caller can vary.
model LiteLLM_DailyGatewayRequests {
date String
category String
route String
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@id([date, category, route])
@@index([date])
}
// Prompt table for storing prompt configurations
model LiteLLM_PromptTable {
id String @id @default(uuid())
@ -1393,6 +1416,38 @@ model LiteLLM_AdaptiveRouterSession {
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
model LiteLLM_AutoRouterSession {
api_key String
session_id String
router_name String
router_type String
first_turn_at DateTime
last_turn_at DateTime
last_model String
models Json @default("{}")
turns Int @default(0)
unordered_turns Int @default(0)
covered_turns Int @default(0)
cache_hits Int @default(0)
same_model_turns Int @default(0)
same_model_hits Int @default(0)
first_visit_turns Int @default(0)
first_visit_hits Int @default(0)
return_turns Int @default(0)
return_hits Int @default(0)
return_expired_misses Int @default(0)
return_within_ttl_misses Int @default(0)
ttl_5m_turns Int @default(0)
ttl_1h_turns Int @default(0)
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//

View file

@ -14,6 +14,10 @@ from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
ensure_prisma_toolchain,
prisma_command_timeout,
)
def str_to_bool(value: Optional[str]) -> bool:
@ -142,7 +146,7 @@ class ProxyExtrasDBManager:
],
stdout=open(migration_file, "w"),
check=True,
timeout=30,
timeout=prisma_command_timeout(),
env=prisma_env,
)
@ -157,7 +161,7 @@ class ProxyExtrasDBManager:
"0_init",
],
check=True,
timeout=30,
timeout=prisma_command_timeout(),
env=prisma_env,
)
@ -193,7 +197,7 @@ class ProxyExtrasDBManager:
"--rolled-back",
migration_name,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
@ -205,7 +209,7 @@ class ProxyExtrasDBManager:
prisma_env = _get_prisma_env()
subprocess.run(
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
@ -303,7 +307,7 @@ class ProxyExtrasDBManager:
"--script",
],
check=True,
timeout=60,
timeout=prisma_command_timeout(),
stdout=f,
env=_get_prisma_env(),
)
@ -335,7 +339,7 @@ class ProxyExtrasDBManager:
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -364,7 +368,7 @@ class ProxyExtrasDBManager:
"--schema",
schema_path,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -393,7 +397,7 @@ class ProxyExtrasDBManager:
"--applied",
migration_name,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -530,7 +534,7 @@ class ProxyExtrasDBManager:
try:
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
env=_get_prisma_env(),
)
@ -555,7 +559,7 @@ class ProxyExtrasDBManager:
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -731,6 +735,9 @@ class ProxyExtrasDBManager:
Returns:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(
prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()
)
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
@ -757,7 +764,7 @@ class ProxyExtrasDBManager:
# Set migrations directory for Prisma
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -840,7 +847,7 @@ class ProxyExtrasDBManager:
"--rolled-back",
failed_migration,
],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
@ -968,7 +975,7 @@ class ProxyExtrasDBManager:
# Use prisma db push with increased timeout
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=60,
timeout=prisma_command_timeout(),
check=True,
)
return True

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.83"
version = "0.4.84"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.83"
version = "0.4.84"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -197,6 +197,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
)
log_raw_request_response: bool = False
request_correlation_in_logs: bool = False
redact_messages_in_exceptions: Optional[bool] = False
redact_user_api_key_info: Optional[bool] = False
# When True (default — preserves historical behavior), the Router appends
@ -244,6 +245,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
# Or via `litellm_settings.strip_anthropic_total_tokens: true` in
# config.yaml.
strip_anthropic_total_tokens: bool = False
anthropic_sse_ping_interval_seconds: float = 15.0
route_all_chat_openai_to_responses: bool = (
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge
@ -704,9 +706,8 @@ def is_openai_finetune_model(key: str) -> bool:
return key.startswith("ft:") and not key.count(":") > 1
def add_known_models(model_cost_map: Optional[Dict] = None):
_map: Final = model_cost_map if model_cost_map is not None else model_cost
for key, value in _map.items():
def _populate_provider_model_sets(model_cost_map: Dict) -> None:
for key, value in model_cost_map.items():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":
@ -949,7 +950,16 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
bedrock_mantle_models.add(key)
add_known_models()
def add_known_models(model_cost_map: Optional[Dict] = None):
"""Fold `model_cost_map` (defaults to `litellm.model_cost`) into the per-provider model sets,
then refresh `models_by_provider` from those sets so the additions reach wildcard expansion.
The refresh updates the dict in place, so references captured before a reload stay live.
"""
_populate_provider_model_sets(model_cost_map if model_cost_map is not None else model_cost)
models_by_provider.update(_build_models_by_provider())
_populate_provider_model_sets(model_cost)
# known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary
# this is maintained for Exception Mapping
@ -1071,112 +1081,116 @@ model_list_set = set(model_list)
# provider_list is lazy-loaded via __getattr__ to avoid importing LlmProviders at import time
models_by_provider: dict = {
"openai": open_ai_chat_completion_models | open_ai_text_completion_models,
"text-completion-openai": open_ai_text_completion_models,
"cohere": cohere_models | cohere_chat_models,
"cohere_chat": cohere_chat_models,
"anthropic": anthropic_models,
"replicate": replicate_models,
"huggingface": huggingface_models,
"together_ai": together_ai_models,
"baseten": baseten_models,
"openrouter": openrouter_models,
"vercel_ai_gateway": vercel_ai_gateway_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
| vertex_text_models
| vertex_anthropic_models
| vertex_vision_models
| vertex_language_models
| vertex_deepseek_models
| vertex_minimax_models
| vertex_moonshot_models
| vertex_zai_models,
"ai21": ai21_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
"ollama": ollama_models,
"ollama_chat": ollama_models,
"deepinfra": deepinfra_models,
"perplexity": perplexity_models,
"maritalk": maritalk_models,
"watsonx": watsonx_models,
"gemini": gemini_models,
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
"aleph_alpha": aleph_alpha_models,
"text-completion-codestral": text_completion_codestral_models,
"text-completion-inception": text_completion_inception_models,
"xai": xai_models,
"zai": zai_models,
"fal_ai": fal_ai_models,
"deepseek": deepseek_models,
"tencent": tencent_models,
"runwayml": runwayml_models,
"mistral": mistral_chat_models,
"azure_ai": azure_ai_models,
"voyage": voyage_models,
"infinity": infinity_models,
"databricks": databricks_models,
"cloudflare": cloudflare_models,
"codestral": codestral_models,
"nlp_cloud": nlp_cloud_models,
"friendliai": friendliai_models,
"palm": palm_models,
"groq": groq_models,
"azure": azure_models | azure_text_models,
"azure_anthropic": azure_anthropic_models,
"azure_text": azure_text_models,
"anyscale": anyscale_models,
"cerebras": cerebras_models,
"galadriel": galadriel_models,
"nvidia_nim": nvidia_nim_models,
"nvidia_riva": nvidia_riva_models,
"soniox": soniox_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
"aiml": aiml_models,
"assemblyai": assemblyai_models,
"jina_ai": jina_ai_models,
"snowflake": snowflake_models,
"gradient_ai": gradient_ai_models,
"meta_llama": llama_models,
"nscale": nscale_models,
"featherless_ai": featherless_ai_models,
"deepgram": deepgram_models,
"elevenlabs": elevenlabs_models,
"heroku": heroku_models,
"dashscope": dashscope_models,
"modelscope": modelscope_models,
"moonshot": moonshot_models,
"publicai": publicai_models,
"darkbloom": darkbloom_models,
"v0": v0_models,
"morph": morph_models,
"lambda_ai": lambda_ai_models,
"inception": inception_models,
"hyperbolic": hyperbolic_models,
"black_forest_labs": black_forest_labs_models,
"recraft": recraft_models,
"cometapi": cometapi_models,
"oci": oci_models,
"volcengine": volcengine_models,
"wandb": wandb_models,
"ovhcloud": ovhcloud_models | ovhcloud_embedding_models,
"lemonade": lemonade_models,
"clarifai": clarifai_models,
"amazon_nova": amazon_nova_models,
"stability": stability_models,
"github_copilot": github_copilot_models,
"chatgpt": chatgpt_models,
"minimax": minimax_models,
"aws_polly": aws_polly_models,
"gigachat": gigachat_models,
"llamagate": llamagate_models,
"reducto": reducto_models,
"bedrock_mantle": bedrock_mantle_models,
}
def _build_models_by_provider() -> dict:
return {
"openai": open_ai_chat_completion_models | open_ai_text_completion_models,
"text-completion-openai": open_ai_text_completion_models,
"cohere": cohere_models | cohere_chat_models,
"cohere_chat": cohere_chat_models,
"anthropic": anthropic_models,
"replicate": replicate_models,
"huggingface": huggingface_models,
"together_ai": together_ai_models,
"baseten": baseten_models,
"openrouter": openrouter_models,
"vercel_ai_gateway": vercel_ai_gateway_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
| vertex_text_models
| vertex_anthropic_models
| vertex_vision_models
| vertex_language_models
| vertex_deepseek_models
| vertex_minimax_models
| vertex_moonshot_models
| vertex_zai_models,
"ai21": ai21_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
"ollama": ollama_models,
"ollama_chat": ollama_models,
"deepinfra": deepinfra_models,
"perplexity": perplexity_models,
"maritalk": maritalk_models,
"watsonx": watsonx_models,
"gemini": gemini_models,
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
"aleph_alpha": aleph_alpha_models,
"text-completion-codestral": text_completion_codestral_models,
"text-completion-inception": text_completion_inception_models,
"xai": xai_models,
"zai": zai_models,
"fal_ai": fal_ai_models,
"deepseek": deepseek_models,
"tencent": tencent_models,
"runwayml": runwayml_models,
"mistral": mistral_chat_models,
"azure_ai": azure_ai_models,
"voyage": voyage_models,
"infinity": infinity_models,
"databricks": databricks_models,
"cloudflare": cloudflare_models,
"codestral": codestral_models,
"nlp_cloud": nlp_cloud_models,
"friendliai": friendliai_models,
"palm": palm_models,
"groq": groq_models,
"azure": azure_models | azure_text_models,
"azure_anthropic": azure_anthropic_models,
"azure_text": azure_text_models,
"anyscale": anyscale_models,
"cerebras": cerebras_models,
"galadriel": galadriel_models,
"nvidia_nim": nvidia_nim_models,
"nvidia_riva": nvidia_riva_models,
"soniox": soniox_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
"aiml": aiml_models,
"assemblyai": assemblyai_models,
"jina_ai": jina_ai_models,
"snowflake": snowflake_models,
"gradient_ai": gradient_ai_models,
"meta_llama": llama_models,
"nscale": nscale_models,
"featherless_ai": featherless_ai_models,
"deepgram": deepgram_models,
"elevenlabs": elevenlabs_models,
"heroku": heroku_models,
"dashscope": dashscope_models,
"modelscope": modelscope_models,
"moonshot": moonshot_models,
"publicai": publicai_models,
"darkbloom": darkbloom_models,
"v0": v0_models,
"morph": morph_models,
"lambda_ai": lambda_ai_models,
"inception": inception_models,
"hyperbolic": hyperbolic_models,
"black_forest_labs": black_forest_labs_models,
"recraft": recraft_models,
"cometapi": cometapi_models,
"oci": oci_models,
"volcengine": volcengine_models,
"wandb": wandb_models,
"ovhcloud": ovhcloud_models | ovhcloud_embedding_models,
"lemonade": lemonade_models,
"clarifai": clarifai_models,
"amazon_nova": amazon_nova_models,
"stability": stability_models,
"github_copilot": github_copilot_models,
"chatgpt": chatgpt_models,
"minimax": minimax_models,
"aws_polly": aws_polly_models,
"gigachat": gigachat_models,
"llamagate": llamagate_models,
"reducto": reducto_models,
"bedrock_mantle": bedrock_mantle_models,
}
models_by_provider: dict = _build_models_by_provider()
# mapping for those models which have larger equivalents
longer_context_model_fallback_dict: dict = {
@ -1269,8 +1283,8 @@ from .llms.xai.common_utils import XAIModelInfo
from litellm.types.utils import LlmProviders
## Lazy loading this is not straightforward, will leave it here for now.
from .main import * # type: ignore
from .compression import compress # type: ignore[no-redef]
from .main import *
from .compression import compress
# Skills API
from .skills.main import (
@ -1341,7 +1355,7 @@ from .assistants.main import *
from .batches.main import *
from .images.main import *
from .videos.main import *
from .batch_completion.main import * # type: ignore
from .batch_completion.main import *
from .rerank_api.main import *
from .llms.anthropic.experimental_pass_through.messages.handler import *
from .responses.main import *
@ -2054,7 +2068,7 @@ if TYPE_CHECKING:
supports_reasoning: Callable[..., bool]
acreate: Callable[..., Any]
get_max_tokens: Callable[..., int]
get_model_info: Callable[..., _ModelInfoType] # type: ignore[no-redef]
get_model_info: Callable[..., _ModelInfoType]
register_prompt_template: Callable[..., None]
validate_environment: Callable[..., dict]
check_valid_key: Callable[..., bool]
@ -2150,9 +2164,9 @@ def __getattr__(name: str) -> Any:
# Lazy load encoding from main.py to avoid heavy tiktoken import
if name == "encoding":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "encoding" not in _globals:
from .main import encoding as _encoding
@ -2162,9 +2176,9 @@ def __getattr__(name: str) -> Any:
# Lazy load bedrock_tool_name_mappings instance
if name == "bedrock_tool_name_mappings":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "bedrock_tool_name_mappings" not in _globals:
from .llms.bedrock.chat.invoke_handler import (
@ -2176,9 +2190,9 @@ def __getattr__(name: str) -> Any:
# Lazy load AzureOpenAIError exception class
if name == "AzureOpenAIError":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "AzureOpenAIError" not in _globals:
from .llms.azure.common_utils import AzureOpenAIError as _AzureOpenAIError
@ -2188,9 +2202,9 @@ def __getattr__(name: str) -> Any:
# Lazy load openaiOSeriesConfig instance
if name == "openaiOSeriesConfig":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
if "openaiOSeriesConfig" not in _globals:
# Import the config class and instantiate it
config_class = __getattr__("OpenAIOSeriesConfig")
@ -2206,9 +2220,9 @@ def __getattr__(name: str) -> Any:
"nvidiaNimEmbeddingConfig": "NvidiaNimEmbeddingConfig",
}
if name in _config_instances:
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
if name not in _globals:
# Import the config class and instantiate it
config_class = __getattr__(_config_instances[name])
@ -2221,9 +2235,9 @@ def __getattr__(name: str) -> Any:
# Lazy load provider_list
if name == "provider_list":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "provider_list" not in _globals:
# LlmProviders is eagerly imported above, so we can import it directly
@ -2234,9 +2248,9 @@ def __getattr__(name: str) -> Any:
# Lazy load priority_reservation_settings instance
if name == "priority_reservation_settings":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "priority_reservation_settings" not in _globals:
# Import the class and instantiate it
@ -2246,9 +2260,9 @@ def __getattr__(name: str) -> Any:
# Lazy load logging_callback_manager instance
if name == "logging_callback_manager":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "logging_callback_manager" not in _globals:
# Import the class and instantiate it
@ -2258,9 +2272,9 @@ def __getattr__(name: str) -> Any:
# Lazy load _service_logger module
if name == "_service_logger":
from ._lazy_imports import _get_litellm_globals
from ._lazy_imports import get_litellm_globals
_globals = _get_litellm_globals()
_globals = get_litellm_globals()
# Check if already cached
if "_service_logger" not in _globals:
# Import the module lazily

View file

@ -54,7 +54,7 @@ from ._lazy_imports_registry import (
)
def _get_litellm_globals() -> dict:
def get_litellm_globals() -> dict:
"""
Get the globals dictionary of the litellm module.
@ -233,7 +233,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
raise AttributeError(f"{category} lazy import: unknown attribute {name!r}")
# Step 2: Get the cache (where we store imported things)
_globals: Final = _get_litellm_globals()
_globals: Final = get_litellm_globals()
# Step 3: If we've already imported it, just return the cached version
if name in _globals:
@ -332,7 +332,7 @@ def _lazy_import_utils_module(name: str) -> Any:
Handler for utils module lazy imports.
This uses a custom implementation because utils module needs to use
_get_utils_globals() instead of _get_litellm_globals() for caching.
_get_utils_globals() instead of get_litellm_globals() for caching.
"""
# Check if this attribute exists in our map
if name not in _UTILS_MODULE_IMPORT_MAP:
@ -379,7 +379,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
- "in_memory_llm_clients_cache" is a singleton instance of that class
So we need custom logic to handle both cases.
"""
_globals: Final = _get_litellm_globals()
_globals: Final = get_litellm_globals()
# If already cached, return it
if name in _globals:
@ -412,7 +412,7 @@ def _lazy_import_http_handlers(name: str) -> Any:
- They need configuration (timeout, etc.) from the module globals
- They use factory functions instead of direct instantiation
"""
_globals: Final = _get_litellm_globals()
_globals: Final = get_litellm_globals()
if name == "module_level_aclient":
# Create an async HTTP client using the factory function

View file

@ -1461,32 +1461,30 @@ _UTILS_MODULE_IMPORT_MAP: Final = {
# Export all name tuples and import maps for use in _lazy_imports.py
__all__ = [
# Name tuples
"COST_CALCULATOR_NAMES",
"LITELLM_LOGGING_NAMES",
"UTILS_NAMES",
"TOKEN_COUNTER_NAMES",
"LLM_CLIENT_CACHE_NAMES",
"BEDROCK_TYPES_NAMES",
"TYPES_UTILS_NAMES",
"CACHING_NAMES",
"HTTP_HANDLER_NAMES",
"COST_CALCULATOR_NAMES",
"DOTPROMPT_NAMES",
"HTTP_HANDLER_NAMES",
"LITELLM_LOGGING_NAMES",
"LLM_CLIENT_CACHE_NAMES",
"LLM_CONFIG_NAMES",
"TYPES_NAMES",
"LLM_PROVIDER_LOGIC_NAMES",
"TOKEN_COUNTER_NAMES",
"TYPES_NAMES",
"TYPES_UTILS_NAMES",
"UTILS_MODULE_NAMES",
# Import maps
"_UTILS_IMPORT_MAP",
"_COST_CALCULATOR_IMPORT_MAP",
"_TYPES_UTILS_IMPORT_MAP",
"_TOKEN_COUNTER_IMPORT_MAP",
"UTILS_NAMES",
"_BEDROCK_TYPES_IMPORT_MAP",
"_CACHING_IMPORT_MAP",
"_LITELLM_LOGGING_IMPORT_MAP",
"_COST_CALCULATOR_IMPORT_MAP",
"_DOTPROMPT_IMPORT_MAP",
"_TYPES_IMPORT_MAP",
"_LITELLM_LOGGING_IMPORT_MAP",
"_LLM_CONFIGS_IMPORT_MAP",
"_LLM_PROVIDER_LOGIC_IMPORT_MAP",
"_TOKEN_COUNTER_IMPORT_MAP",
"_TYPES_IMPORT_MAP",
"_TYPES_UTILS_IMPORT_MAP",
"_UTILS_IMPORT_MAP",
"_UTILS_MODULE_IMPORT_MAP",
]

View file

@ -1,4 +1,5 @@
import ast
import contextvars
import logging
import os
import sys
@ -6,12 +7,44 @@ from datetime import datetime
from logging import Formatter
from typing import Any, Final
import litellm
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import redact_string
set_verbose = False
session_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("session_id", default="")
trace_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("trace_id", default="")
_MAX_CORRELATION_ID_LENGTH: Final = 256
def _sanitize_correlation_id(value: str) -> str:
"""Strip control characters, bound length, and redact credential-shaped
content before a caller-controlled trace_id/session_id (e.g.
litellm_session_id, x-litellm-trace-id) is stamped into log lines.
Without the first two, a caller could embed \\r/\\n or terminal escape
sequences to forge fake log entries, or submit an oversized value repeated
across every log line for the request. Without the redaction, a caller
could smuggle a real credential (e.g. an sk-... key) through this field:
CorrelationContextFilter stamps trace_id/session_id onto the record after
SecretRedactionFilter has already run, so those two fields never otherwise
pass through credential redaction.
"""
stripped: Final = "".join(ch for ch in value if ch.isprintable())
return _redact_string(stripped[:_MAX_CORRELATION_ID_LENGTH])
def set_session_id(session_id: str) -> "contextvars.Token[str]":
return session_id_var.set(_sanitize_correlation_id(session_id))
def set_trace_id(trace_id: str) -> "contextvars.Token[str]":
return trace_id_var.set(_sanitize_correlation_id(trace_id))
if set_verbose is True:
logging.warning(
"`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
@ -77,6 +110,28 @@ class SecretRedactionFilter(logging.Filter):
_secret_filter: Final = SecretRedactionFilter()
class CorrelationContextFilter(logging.Filter):
"""Stamps each log record with the current request's trace_id and session_id from contextvars.
Works in tandem with JsonFormatter: the formatter's record.__dict__ loop picks up these
attributes as first-class JSON fields without any formatter-level code.
"""
def filter(self, record: logging.LogRecord) -> bool:
if not litellm.request_correlation_in_logs:
return True
trace_id: Final = trace_id_var.get()
if trace_id:
record.trace_id = trace_id # rebind-ok: stamping the LogRecord is the Filter interface's contract
session_id: Final = session_id_var.get()
if session_id:
record.session_id = session_id # rebind-ok: stamping the LogRecord is the Filter interface's contract
return True
_correlation_filter: Final = CorrelationContextFilter()
json_logs = bool(os.getenv("JSON_LOGS", False))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
@ -84,6 +139,7 @@ numeric_level: Final[str] = getattr(logging, log_level.upper())
handler: Final = logging.StreamHandler()
handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)
handler.addFilter(_correlation_filter)
def _try_parse_json_message(message: str) -> dict[str, Any] | None:
@ -146,6 +202,11 @@ def _get_standard_record_attrs() -> frozenset:
_STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs()
# CorrelationContextFilter is the only legitimate source for these two JSON fields;
# see JsonFormatter.format() for why they're excluded from the generic message-content
# and extra-attribute promotion paths.
_RESERVED_CORRELATION_FIELDS: Final = frozenset(("trace_id", "session_id"))
class JsonFormatter(Formatter):
def __init__(self):
@ -164,13 +225,18 @@ class JsonFormatter(Formatter):
"timestamp": self.formatTime(record),
}
# Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties
# Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties.
# trace_id/session_id are excluded here unconditionally (not just "if not already
# set") - CorrelationContextFilter is the only legitimate source for these two
# fields, and a message that merely happens to parse as JSON/dict (e.g. a proxy
# log line dumping raw request headers) must never be able to claim them, even on
# a record the filter hasn't stamped yet (no correlation context active for it).
parsed = _try_parse_json_message(message_str)
if parsed is None:
parsed = _try_parse_embedded_python_dict(message_str)
if parsed is not None:
for key, value in parsed.items():
if key not in json_record:
if key not in json_record and key not in _RESERVED_CORRELATION_FIELDS:
json_record[key] = value
# Include extra attributes passed via logger.debug("msg", extra={...})
@ -178,6 +244,18 @@ class JsonFormatter(Formatter):
if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
json_record[key] = value
# trace_id/session_id are reserved: CorrelationContextFilter is the only
# legitimate source for these two fields. Without this, a message string
# that happens to parse as JSON/dict (e.g. a proxy log line dumping raw
# request headers) with a "trace_id"/"session_id" key would have already
# claimed the key at the parsed-message step above, and the extra-attributes
# loop's "key not in json_record" guard would then skip the real value -
# letting a caller-supplied header spoof another request's correlation ids.
for reserved_key in _RESERVED_CORRELATION_FIELDS:
value = getattr(record, reserved_key, None)
if value:
json_record[reserved_key] = value
# Set component/logger only if not already supplied via extra={...}
if "component" not in json_record:
json_record["component"] = record.name
@ -190,12 +268,34 @@ class JsonFormatter(Formatter):
return safe_dumps(json_record)
class CorrelationPlainFormatter(logging.Formatter):
"""Appends trace_id/session_id to plain-text log lines stamped by CorrelationContextFilter.
Mirrors JsonFormatter's handling of these two fields so request_correlation_in_logs
behaves the same whether or not json_logs is enabled.
"""
def format(self, record: logging.LogRecord) -> str:
formatted: Final = super().format(record)
trace_id: Final = getattr(record, "trace_id", None)
session_id: Final = getattr(record, "session_id", None)
if not trace_id and not session_id:
return formatted
parts: Final = tuple(
p
for p in (f"trace_id={trace_id}" if trace_id else None, f"session_id={session_id}" if session_id else None)
if p
)
return f"{formatted} [{' '.join(parts)}]"
# Function to set up exception handlers for JSON logging
def _setup_json_exception_handlers(formatter):
# Create a handler with JSON formatting for exceptions
error_handler: Final = logging.StreamHandler()
error_handler.setFormatter(formatter)
error_handler.addFilter(_secret_filter)
error_handler.addFilter(_correlation_filter)
# Setup excepthook for uncaught exceptions
def json_excepthook(exc_type, exc_value, exc_traceback):
@ -243,7 +343,7 @@ if json_logs:
handler.setFormatter(JsonFormatter())
_setup_json_exception_handlers(JsonFormatter())
else:
formatter: Final = logging.Formatter(
formatter: Final = CorrelationPlainFormatter(
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
datefmt="%H:%M:%S",
)
@ -273,8 +373,42 @@ def _suppress_loggers():
apscheduler_scheduler_logger.setLevel(logging.WARNING)
_REDACTED_THIRD_PARTY_LOGGERS: Final[tuple[str, ...]] = (
"apscheduler.executors.default",
"apscheduler.scheduler",
"asyncio",
"backoff",
"httpx",
"uvicorn.error",
)
def _redact_third_party_loggers() -> None:
"""Extend secret redaction to records litellm does not emit directly.
litellm's own loggers are covered by the filter on their shared handler, but a
litellm value can also reach a log record through a dependency that logs on its
own logger. Those records never pass through a litellm handler.
The filter is attached to each emitting logger rather than to the root logger or
to root's handlers. `Logger.handle` applies the emitting logger's filters before
any handler runs, so redaction happens once, at the earliest point in the
record's life, and covers every downstream handler regardless of who owns it.
The alternatives do not hold: `callHandlers` consults ancestors for handlers but
never for filters, so a filter on the root logger never sees these records at
all, and a filter on a root handler only covers that one handler, leaving
handlers registered earlier or on the emitting logger itself untouched.
Each name is the exact logger a dependency emits on; a parent name would not
cover its children, for the same reason the root logger does not.
"""
for name in _REDACTED_THIRD_PARTY_LOGGERS:
logging.getLogger(name).addFilter(_secret_filter)
# Call the suppression function
_suppress_loggers()
_redact_third_party_loggers()
ALL_LOGGERS: Final = [
logging.getLogger(),
@ -312,6 +446,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
- Prevents bubbling to parent/root (critical to prevent duplicate JSON logs)
"""
handler.addFilter(_secret_filter)
handler.addFilter(_correlation_filter)
for lg in _get_loggers_to_initialize():
lg.handlers.clear() # remove any existing handlers
lg.addHandler(handler) # add JSON formatter handler

View file

@ -15,8 +15,8 @@ import os
from collections.abc import Callable
from typing import Final
import redis # type: ignore
import redis.asyncio as async_redis # type: ignore
import redis
import redis.asyncio as async_redis
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
@ -153,7 +153,7 @@ def _redis_kwargs_from_environment():
return_dict: Final = {}
for k, v in mapping.items():
value = get_secret(k, default_value=None) # type: ignore
value = get_secret(k, default_value=None)
if value is not None:
return_dict[v] = value
return return_dict
@ -317,7 +317,7 @@ def create_azure_ad_redis_connect_func(
# AzureADCredentialProvider for refresh-aware token retrieval. The raw
# client_id/tenant_id/secret are intentionally NOT exposed here — the
# credential closure already holds them.
ad_connect._azure_credential = credential # type: ignore[attr-defined]
ad_connect._azure_credential = credential
return ad_connect
@ -351,7 +351,7 @@ def _get_redis_client_logic(**env_overrides):
for k, v in env_overrides.items():
if isinstance(v, str) and v.startswith("os.environ/"):
v = v.replace("os.environ/", "")
value = get_secret(v) # type: ignore
value = get_secret(v)
env_overrides[k] = value
environment_kwargs: Final = _redis_kwargs_from_environment()
@ -370,7 +370,7 @@ def _get_redis_client_logic(**env_overrides):
**env_overrides,
}
_startup_nodes: Final[str | list | None] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
_startup_nodes: Final[str | list | None] = redis_kwargs.get("startup_nodes", None) or get_secret(
"REDIS_CLUSTER_NODES"
)
@ -381,7 +381,7 @@ def _get_redis_client_logic(**env_overrides):
elif _startup_nodes is None:
redis_kwargs.pop("startup_nodes", None)
_sentinel_nodes: Final[str | list | None] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
_sentinel_nodes: Final[str | list | None] = redis_kwargs.get("sentinel_nodes", None) or get_secret(
"REDIS_SENTINEL_NODES"
)
@ -395,9 +395,7 @@ def _get_redis_client_logic(**env_overrides):
if _sentinel_password is not None:
redis_kwargs["sentinel_password"] = _sentinel_password
_service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
"REDIS_SERVICE_NAME"
)
_service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret("REDIS_SERVICE_NAME")
if _service_name is not None:
redis_kwargs["service_name"] = _service_name
@ -412,7 +410,7 @@ def _get_redis_client_logic(**env_overrides):
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
)
# Store GCP service account in redis_connect_func for async cluster access
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account # type: ignore[attr-defined]
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("gcp_service_account", None)
@ -449,7 +447,7 @@ def _get_redis_client_logic(**env_overrides):
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
# are intentionally NOT exposed on the function to avoid leaking
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True # type: ignore[attr-defined]
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("azure_redis_ad_token", None)
@ -481,7 +479,7 @@ def _get_redis_client_logic(**env_overrides):
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
_redis_cluster_nodes_in_env: Final[str | None] = get_secret("REDIS_CLUSTER_NODES") # type: ignore
_redis_cluster_nodes_in_env: Final[str | None] = get_secret("REDIS_CLUSTER_NODES")
if _redis_cluster_nodes_in_env is not None:
try:
redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env)
@ -505,7 +503,7 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
new_startup_nodes.append(ClusterNode(**item))
cluster_kwargs.pop("startup_nodes", None)
return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs) # type: ignore
return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs)
def _get_redis_sentinel_connection_kwargs(redis_kwargs: dict) -> dict:
@ -638,7 +636,7 @@ def get_redis_async_client(
# Create async RedisCluster with IAM token as password if available
cluster_client: Final = async_redis.RedisCluster(
startup_nodes=new_startup_nodes,
**cluster_kwargs, # type: ignore
**cluster_kwargs,
)
return cluster_client

View file

@ -3,7 +3,7 @@ import threading
import time
from typing import Any, Final
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
from redis.credentials import CredentialProvider
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"

View file

@ -1,6 +1,6 @@
import asyncio
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm._logging import verbose_logger
@ -16,7 +16,7 @@ if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
Span = Union[_Span, Any]
Span = _Span | Any
OTELClass = OpenTelemetry
else:
Span = Any

View file

@ -4,7 +4,7 @@ Internal unified UUID helper.
Always uses fastuuid for performance.
"""
import fastuuid as _uuid # type: ignore
import fastuuid as _uuid
# Expose a module-like alias so callers can use: uuid.uuid4()
uuid = _uuid

View file

@ -55,19 +55,15 @@ from litellm.a2a_protocol.main import (
from litellm.types.agents import LiteLLMSendMessageResponse
__all__ = [
# Client
"A2AClient",
# Functions
"asend_message",
"send_message",
"asend_message_streaming",
"aget_agent_card",
"create_a2a_client",
# Response types
"LiteLLMSendMessageResponse",
# Exceptions
"A2AError",
"A2AConnectionError",
"A2AAgentCardError",
"A2AClient",
"A2AConnectionError",
"A2AError",
"A2ALocalhostURLError",
"LiteLLMSendMessageResponse",
"aget_agent_card",
"asend_message",
"asend_message_streaming",
"create_a2a_client",
"send_message",
]

View file

@ -18,8 +18,8 @@ AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json"
PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json"
try:
from a2a.client import A2ACardResolver as _A2ACardResolver # type: ignore[no-redef]
from a2a.utils.constants import ( # type: ignore[no-redef]
from a2a.client import A2ACardResolver as _A2ACardResolver
from a2a.utils.constants import (
AGENT_CARD_WELL_KNOWN_PATH,
PREV_AGENT_CARD_WELL_KNOWN_PATH,
)
@ -102,7 +102,7 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
return agent_card
class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
class LiteLLMA2ACardResolver(_A2ACardResolver):
"""
Custom A2A card resolver that supports multiple well-known paths.

View file

@ -29,9 +29,9 @@ try:
A2A_SDK_AVAILABLE = True
except ImportError:
A2A_SDK_AVAILABLE = False
Client = None # type: ignore[misc, assignment]
ClientConfig = None # type: ignore[misc, assignment]
create_client = None # type: ignore[misc, assignment]
Client = None
ClientConfig = None
create_client = None
class A2AExceptionCheckers:
@ -202,9 +202,9 @@ async def handle_a2a_localhost_retry(
# Fix the agent card URL
set_agent_card_url(agent_card, error.base_url)
# Reuse the httpx client LiteLLM attached at creation. It carries this agent's
# trace-id and auth headers, so a fresh client would drop them. Only clients built
# by ``create_a2a_client`` have it; an externally-supplied client cannot be retried.
# Reuse the httpx client and call context LiteLLM attached at creation, since the
# context carries this agent's trace-id/auth headers. Only clients built by
# ``create_a2a_client`` have them; an externally-supplied client cannot be retried.
httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None)
if httpx_client is None:
raise RuntimeError(
@ -219,6 +219,9 @@ async def handle_a2a_localhost_retry(
streaming=is_streaming,
),
)
new_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined]
new_client._litellm_agent_card = agent_card # type: ignore[attr-defined]
new_client._litellm_httpx_client = httpx_client
new_client._litellm_call_context = getattr( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash
a2a_client, "_litellm_call_context", None
)
new_client._litellm_agent_card = agent_card
return new_client

View file

@ -10,7 +10,7 @@ A2A Streaming Events (in order):
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
import litellm
@ -21,6 +21,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
)
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
from litellm.interactions.agents.utils import merge_agent_headers
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponse
# litellm_params key carrying the authenticated principal (hashed virtual key) so
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
@ -45,47 +47,14 @@ class A2ACompletionBridgeHandler:
"""
@staticmethod
async def handle_non_streaming(
request_id: str,
def _build_completion_params(
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
litellm_params: Mapping[str, Any],
api_base: str | None,
agent_extra_headers: Mapping[str, str] | None,
*,
_skip_a2a_provider_routing: bool = False,
) -> dict[str, Any]:
"""
Handle non-streaming A2A request via litellm.acompletion.
Args:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Returns:
A2A SendMessageResponse dict
"""
custom_llm_provider = litellm_params.get("custom_llm_provider")
if not _skip_a2a_provider_routing:
a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
if a2a_provider_config is not None:
verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider)
return await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
)
stream: bool,
) -> Mapping[str, Any]:
# Extract message from params
message: Final = params.get("message", {})
@ -93,7 +62,7 @@ class A2ACompletionBridgeHandler:
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
model: Final = litellm_params.get("model", "agent")
# Build full model string if provider specified
@ -103,14 +72,17 @@ class A2ACompletionBridgeHandler:
else:
full_model = model
verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base)
if stream:
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
else:
verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: Final[dict[str, Any]] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": False,
"stream": stream,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add: Final = {
@ -134,8 +106,64 @@ class A2ACompletionBridgeHandler:
static_headers=completion_params.get("extra_headers"),
)
return completion_params
@staticmethod
async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper:
return await litellm.acompletion(**completion_params)
@staticmethod
async def handle_non_streaming(
request_id: str,
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> dict[str, object]:
"""
Handle non-streaming A2A request via litellm.acompletion.
Args:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Returns:
A2A SendMessageResponse dict
"""
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
if not _skip_a2a_provider_routing:
a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
if a2a_provider_config is not None:
verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider)
return await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
)
completion_params: Final = A2ACompletionBridgeHandler._build_completion_params(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
stream=False,
)
# Call litellm.acompletion
response: Final = await litellm.acompletion(**completion_params)
response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params)
# Transform response to A2A format
a2a_response: Final = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
@ -156,7 +184,7 @@ class A2ACompletionBridgeHandler:
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
@ -177,7 +205,7 @@ class A2ACompletionBridgeHandler:
Yields:
A2A streaming response events
"""
custom_llm_provider = litellm_params.get("custom_llm_provider")
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
if not _skip_a2a_provider_routing:
a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
@ -198,60 +226,20 @@ class A2ACompletionBridgeHandler:
return
# Extract message from params
message: Final = params.get("message", {})
# Create streaming context
ctx: Final = A2AStreamingContext(
request_id=request_id,
input_message=message,
input_message=params.get("message", {}),
)
# Transform A2A message to OpenAI format
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model: Final = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: Final[dict[str, Any]] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": True,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add: Final = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
completion_params: Final = A2ACompletionBridgeHandler._build_completion_params(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
stream=True,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event: Final = A2ACompletionBridgeTransformation.create_task_event(ctx)
yield task_event
@ -266,12 +254,12 @@ class A2ACompletionBridgeHandler:
yield working_event
# Call litellm.acompletion with streaming
response: Final = await litellm.acompletion(**completion_params)
response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""
chunk_count = 0
async for chunk in response: # type: ignore[union-attr]
async for chunk in response:
chunk_count += 1
# Extract delta content
@ -312,7 +300,7 @@ async def handle_a2a_completion(
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
@ -329,7 +317,7 @@ async def handle_a2a_completion_streaming(
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,

View file

@ -30,6 +30,7 @@ from litellm.utils import client
if TYPE_CHECKING:
from a2a.client import Client as A2AClientType
from a2a.client import ClientCallContext as A2ACallContextType
from a2a.compat.v0_3.types import (
AgentCard,
Message,
@ -45,7 +46,7 @@ A2A_SDK_AVAILABLE = False
_a2a_conversions: Any = None
try:
from a2a.client import Client, ClientConfig, create_client
from a2a.client import Client, ClientCallContext, ClientConfig, create_client
from a2a.compat.v0_3 import conversions as _a2a_conversions
from a2a.compat.v0_3.types import (
Message,
@ -59,9 +60,10 @@ try:
A2A_SDK_AVAILABLE = True
except ImportError:
Client = None # type: ignore[misc, assignment]
ClientConfig = None # type: ignore[misc, assignment]
create_client = None # type: ignore[misc, assignment]
Client = None
ClientCallContext = None
ClientConfig = None
create_client = None
# Import our custom card resolver that supports multiple well-known paths
from litellm.a2a_protocol.card_resolver import (
@ -218,6 +220,10 @@ async def _send_message_via_completion_bridge(
return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id))
def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]:
return getattr(a2a_client, "_litellm_call_context", None)
async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse":
"""Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response."""
if _a2a_conversions is None:
@ -227,7 +233,7 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
last_event = None
async for event in a2a_client.send_message(pb_request):
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
last_event = event
if last_event is None:
raise RuntimeError("A2A send_message failed: no response received from agent.")
@ -301,7 +307,7 @@ async def _stream_messages(
)
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
async for event in a2a_client.send_message(pb_request):
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
compat_chunk = _a2a_conversions.to_compat_stream_response(
event,
request_id=request.id,
@ -756,26 +762,12 @@ async def create_a2a_client(
verbose_logger.info("Creating A2A client for %s", base_url)
# Use get_async_httpx_client with per-agent params so that different agents
# (with different extra_headers) get separate cached clients. The params
# dict is hashed into the cache key, keeping agent auth isolated while
# still reusing connections within the same agent.
#
# Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout).
# Use "disable_aiohttp_transport" key for cache-key-only data (it's
# filtered out before reaching the constructor).
_client_params: Final[dict] = {"timeout": timeout}
if extra_headers:
# Encode headers into a cache-key-only param so each unique header
# set produces a distinct cache key.
_client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
_async_handler: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2AProvider,
params=_client_params,
params={"timeout": timeout},
)
httpx_client: Final = _async_handler.client
if extra_headers:
httpx_client.headers.update(extra_headers)
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
@ -784,14 +776,20 @@ async def create_a2a_client(
httpx_client=httpx_client,
streaming=streaming,
),
resolver_http_kwargs={"headers": extra_headers} if extra_headers else None,
)
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
# the configured httpx client (with this agent's trace-id/auth headers) without
# excavating a2a-sdk private internals.
a2a_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined]
# the configured httpx client and this agent's headers without excavating
# a2a-sdk private internals.
a2a_client._litellm_httpx_client = httpx_client
a2a_client._litellm_call_context = ( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash
ClientCallContext(service_parameters=extra_headers) # pyright: ignore[reportOptionalCall] # SDK checked above
if extra_headers
else None
)
agent_card: Final = getattr(a2a_client, "_card", None)
if agent_card is not None:
a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined]
a2a_client._litellm_agent_card = agent_card
verbose_logger.info("A2A client created for %s", base_url)

View file

@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s
"""
import asyncio
from collections.abc import AsyncIterator
from typing import Any, Final, cast
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any, Final, Protocol, cast, runtime_checkable
from uuid import uuid4
from pydantic import TypeAdapter
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_LIST_ADAPTER: Final = TypeAdapter(list[object])
_TEXT_ADAPTER: Final = TypeAdapter(str)
@runtime_checkable
class _SupportsModelDump(Protocol):
def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ...
@runtime_checkable
class _SupportsPydanticDict(Protocol):
def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ...
class PydanticAITransformation:
"""
@ -28,7 +45,7 @@ class PydanticAITransformation:
"""
@staticmethod
def _remove_none_values(obj: Any) -> Any:
def _remove_none_values(obj: object) -> object:
"""
Recursively remove None values from a dict/list structure.
@ -42,14 +59,18 @@ class PydanticAITransformation:
Cleaned object with None values removed
"""
if isinstance(obj, dict):
return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None}
typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj)
return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None}
elif isinstance(obj, list):
return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None]
typed_list: Final = _LIST_ADAPTER.validate_python(obj)
return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None]
else:
return obj
@staticmethod
def _params_to_dict(params: Any) -> dict[str, Any]:
def _params_to_dict(
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
) -> Mapping[str, object]:
"""
Convert params to a dict, handling Pydantic models.
@ -59,10 +80,10 @@ class PydanticAITransformation:
Returns:
Dict representation of params
"""
if hasattr(params, "model_dump"):
if isinstance(params, _SupportsModelDump):
# Pydantic v2 model
return params.model_dump(mode="python", exclude_none=True)
elif hasattr(params, "dict"):
elif isinstance(params, _SupportsPydanticDict):
# Pydantic v1 model
return params.dict(exclude_none=True)
elif isinstance(params, dict):
@ -75,12 +96,12 @@ class PydanticAITransformation:
async def _poll_for_completion(
client: AsyncHTTPHandler,
endpoint: str,
task_id: str,
task_id: object,
request_id: str,
max_attempts: int = 30,
poll_interval: float = 0.5,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Poll for task completion using tasks/get method.
@ -112,10 +133,10 @@ class PydanticAITransformation:
},
)
response.raise_for_status()
poll_data = response.json()
poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
result = poll_data.get("result", {})
status = result.get("status", {})
result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {}))
status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
state = status.get("state", "")
verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state)
@ -133,10 +154,10 @@ class PydanticAITransformation:
async def _send_and_poll_raw(
api_base: str,
request_id: str,
params: Any,
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
timeout: float = 60.0,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -153,14 +174,16 @@ class PydanticAITransformation:
Raw Pydantic AI task response (with history/artifacts)
"""
# Convert params to dict if it's a Pydantic model
params_dict = PydanticAITransformation._params_to_dict(params)
# Remove None values - FastA2A doesn't accept null for optional fields
params_dict = PydanticAITransformation._remove_none_values(params_dict)
params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(
PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params))
)
# Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
if "message" in params_dict:
params_dict["message"]["kind"] = "message"
message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"])
message_value["kind"] = "message"
params_dict["message"] = message_value
# Build A2A JSON-RPC request using message/send method for FastA2A compatibility
a2a_request: Final = {
@ -189,11 +212,11 @@ class PydanticAITransformation:
},
)
response.raise_for_status()
response_data = response.json()
response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
# Check if task is already completed
result: Final = response_data.get("result", {})
status: Final = result.get("status", {})
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
state: Final = status.get("state", "")
if state != "completed":
@ -217,10 +240,10 @@ class PydanticAITransformation:
async def send_non_streaming_request(
api_base: str,
request_id: str,
params: Any,
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
timeout: float = 60.0,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
@ -253,10 +276,10 @@ class PydanticAITransformation:
async def send_and_get_raw_response(
api_base: str,
request_id: str,
params: Any,
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
timeout: float = 60.0,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -282,9 +305,9 @@ class PydanticAITransformation:
@staticmethod
def _transform_to_a2a_response(
response_data: dict[str, Any],
response_data: Mapping[str, object],
request_id: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform Pydantic AI task response to standard A2A non-streaming format.
@ -328,7 +351,7 @@ class PydanticAITransformation:
}
@staticmethod
def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]:
"""
Extract response text from completed task response.
@ -342,52 +365,53 @@ class PydanticAITransformation:
Returns:
Tuple of (full_text, message_id, parts)
"""
result: Final = response_data.get("result", {})
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
# Try to extract from artifacts first (preferred for results)
artifacts: Final = result.get("artifacts", [])
if artifacts:
for artifact in artifacts:
parts = artifact.get("parts", [])
for artifact in _LIST_ADAPTER.validate_python(artifacts):
parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", []))
for part in parts:
if part.get("kind") == "text":
text = part.get("text", "")
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
text = part_dict.get("text", "")
if text:
return text, str(uuid4()), parts
# Fall back to history - get the last agent message
history: Final = result.get("history", [])
history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
for msg in reversed(history):
if msg.get("role") == "agent":
parts = msg.get("parts", [])
message_id = msg.get("messageId", str(uuid4()))
if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent":
parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", []))
message_id = msg_dict.get("messageId", str(uuid4()))
full_text = ""
for part in parts:
if part.get("kind") == "text":
full_text += part.get("text", "")
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
if full_text:
return full_text, message_id, parts
# Fall back to message field (original format)
message: Final = result.get("message", {})
if message:
parts = message.get("parts", [])
message_id = message.get("messageId", str(uuid4()))
message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", []))
message_id = message_dict.get("messageId", str(uuid4()))
full_text = ""
for part in parts:
if part.get("kind") == "text":
full_text += part.get("text", "")
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
return full_text, message_id, parts
return "", str(uuid4()), []
@staticmethod
async def fake_streaming_from_response(
response_data: dict[str, Any],
response_data: Mapping[str, object],
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""
Convert a non-streaming A2A response into fake streaming chunks.
@ -410,12 +434,12 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
# Extract input message from raw response for history
result: Final = response_data.get("result", {})
history: Final = result.get("history", [])
input_message = {}
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
input_message = _STR_KEY_DICT_ADAPTER.validate_python({})
for msg in history:
if msg.get("role") == "user":
input_message = msg
if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user":
input_message = msg_dict
break
# Generate IDs for streaming events
@ -426,45 +450,49 @@ class PydanticAITransformation:
# 1. Emit initial task event (kind: "task", status: "submitted")
# Format matches A2ACompletionBridgeTransformation.create_task_event
task_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"history": [
{
"contextId": context_id,
"kind": "message",
"messageId": input_message_id,
"parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
"role": "user",
"taskId": task_id,
}
],
"id": task_id,
"kind": "task",
"status": {
"state": "submitted",
task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
{
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"history": [
{
"contextId": context_id,
"kind": "message",
"messageId": input_message_id,
"parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
"role": "user",
"taskId": task_id,
}
],
"id": task_id,
"kind": "task",
"status": {
"state": "submitted",
},
},
},
}
}
)
yield task_event
# 2. Emit status update (kind: "status-update", status: "working")
# Format matches A2ACompletionBridgeTransformation.create_status_update_event
working_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": False,
"kind": "status-update",
"status": {
"state": "working",
working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
{
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": False,
"kind": "status-update",
"status": {
"state": "working",
},
"taskId": task_id,
},
"taskId": task_id,
},
}
}
)
yield working_event
# Small delay to simulate processing
@ -473,29 +501,32 @@ class PydanticAITransformation:
# 3. Emit artifact update chunks (kind: "artifact-update")
# Format matches A2ACompletionBridgeTransformation.create_artifact_update_event
if full_text:
full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text)
# Split text into chunks
for i in range(0, len(full_text), chunk_size):
chunk_text = full_text[i : i + chunk_size]
is_last_chunk = (i + chunk_size) >= len(full_text)
for i in range(0, len(full_text_str), chunk_size):
chunk_text = full_text_str[i : i + chunk_size]
is_last_chunk = (i + chunk_size) >= len(full_text_str)
artifact_event = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"kind": "artifact-update",
"taskId": task_id,
"artifact": {
"artifactId": artifact_id,
"parts": [
{
"kind": "text",
"text": chunk_text,
}
],
artifact_event = _STR_KEY_DICT_ADAPTER.validate_python(
{
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"kind": "artifact-update",
"taskId": task_id,
"artifact": {
"artifactId": artifact_id,
"parts": [
{
"kind": "text",
"text": chunk_text,
}
],
},
},
},
}
}
)
yield artifact_event
# Add delay between chunks (except for last chunk)
@ -503,19 +534,21 @@ class PydanticAITransformation:
await asyncio.sleep(delay_ms / 1000.0)
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
completed_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": True,
"kind": "status-update",
"status": {
"state": "completed",
completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
{
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": True,
"kind": "status-update",
"status": {
"state": "completed",
},
"taskId": task_id,
},
"taskId": task_id,
},
}
}
)
yield completed_event
verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id)

View file

@ -153,7 +153,7 @@ class AnthropicExceptionMapping:
# Optionally add request_id if provided and not present
if request_id and "request_id" not in parsed:
parsed["request_id"] = request_id
return parsed # type: ignore
return parsed
# Extract message - use parsed dict if available, otherwise raw string
if parsed is not None:

View file

@ -51,9 +51,7 @@ async def aget_assistants(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -61,7 +59,7 @@ async def aget_assistants(
response = await init_response
else:
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -98,7 +96,7 @@ def get_assistants(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -132,12 +130,12 @@ def get_assistants(
max_retries=optional_params.max_retries,
organization=organization,
client=client,
aget_assistants=aget_assistants, # type: ignore
) # type: ignore
aget_assistants=aget_assistants,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -145,14 +143,14 @@ def get_assistants(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
response = azure_assistants_api.get_assistants(
api_base=api_base,
@ -162,7 +160,7 @@ def get_assistants(
timeout=timeout,
max_retries=optional_params.max_retries,
client=client,
aget_assistants=aget_assistants, # type: ignore
aget_assistants=aget_assistants,
litellm_params=litellm_params_dict,
)
else:
@ -173,7 +171,7 @@ def get_assistants(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
@ -185,7 +183,7 @@ def get_assistants(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
@ -210,9 +208,7 @@ async def acreate_assistants(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model=model, custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -220,7 +216,7 @@ async def acreate_assistants(
response = await init_response
else:
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model=model,
@ -267,7 +263,7 @@ def create_assistants(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -318,12 +314,12 @@ def create_assistants(
organization=organization,
create_assistant_data=create_assistant_data,
client=client,
async_create_assistants=async_create_assistants, # type: ignore
) # type: ignore
async_create_assistants=async_create_assistants,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -331,14 +327,14 @@ def create_assistants(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
if isinstance(client, OpenAI):
client = None # only pass client if it's AzureOpenAI
@ -363,7 +359,7 @@ def create_assistants(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
if response is None:
@ -392,9 +388,7 @@ async def adelete_assistant(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -402,7 +396,7 @@ async def adelete_assistant(
response = await init_response
else:
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -442,7 +436,7 @@ def delete_assistant(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -472,9 +466,9 @@ def delete_assistant(
async_delete_assistants=async_delete_assistants,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -482,14 +476,14 @@ def delete_assistant(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
if isinstance(client, OpenAI):
client = None # only pass client if it's AzureOpenAI
@ -541,9 +535,7 @@ async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwar
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -551,7 +543,7 @@ async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwar
response = await init_response
else:
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -608,7 +600,7 @@ def create_thread(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -649,7 +641,7 @@ def create_thread(
acreate_thread=acreate_thread,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_key = (
optional_params.api_key
@ -657,16 +649,16 @@ def create_thread(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
if isinstance(client, OpenAI):
client = None # only pass client if it's AzureOpenAI
@ -692,10 +684,10 @@ def create_thread(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response # type: ignore
return response
async def aget_thread(
@ -715,9 +707,7 @@ async def aget_thread(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -725,7 +715,7 @@ async def aget_thread(
response = await init_response
else:
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -758,7 +748,7 @@ def get_thread(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
api_base: str | None = None
@ -797,9 +787,9 @@ def get_thread(
aget_thread=aget_thread,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -807,14 +797,14 @@ def get_thread(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
if isinstance(client, OpenAI):
client = None # only pass client if it's AzureOpenAI
@ -839,10 +829,10 @@ def get_thread(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response # type: ignore
return response
### MESSAGES ###
@ -879,9 +869,7 @@ async def a_add_message(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -890,7 +878,7 @@ async def a_add_message(
else:
# Call the synchronous function using run_in_executor
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -937,7 +925,7 @@ def add_message(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
api_key: str | None = None
@ -976,9 +964,9 @@ def add_message(
a_add_message=a_add_message,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -986,14 +974,14 @@ def add_message(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
response = azure_assistants_api.add_message(
thread_id=thread_id,
@ -1016,11 +1004,11 @@ def add_message(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response # type: ignore
return response
async def aget_messages(
@ -1046,9 +1034,7 @@ async def aget_messages(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -1057,7 +1043,7 @@ async def aget_messages(
else:
# Call the synchronous function using run_in_executor
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -1090,7 +1076,7 @@ def get_messages(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -1129,9 +1115,9 @@ def get_messages(
aget_messages=aget_messages,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -1139,14 +1125,14 @@ def get_messages(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
response = azure_assistants_api.get_messages(
thread_id=thread_id,
@ -1168,11 +1154,11 @@ def get_messages(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response # type: ignore
return response
### RUNS ###
@ -1182,7 +1168,7 @@ def arun_thread_stream(
**kwargs,
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
kwargs["arun_thread"] = True
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
return run_thread(stream=True, event_handler=event_handler, **kwargs)
async def arun_thread(
@ -1222,9 +1208,7 @@ async def arun_thread(
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
# Await normally
init_response: Final = await loop.run_in_executor(None, func_with_context)
@ -1233,7 +1217,7 @@ async def arun_thread(
else:
# Call the synchronous function using run_in_executor
response = init_response
return response # type: ignore
return response
except Exception as e:
raise exception_type(
model="",
@ -1249,7 +1233,7 @@ def run_thread_stream(
event_handler: AssistantEventHandler | None = None,
**kwargs,
) -> AssistantStreamManager[AssistantEventHandler]:
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
return run_thread(stream=True, event_handler=event_handler, **kwargs)
def run_thread(
@ -1283,7 +1267,7 @@ def run_thread(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -1329,9 +1313,9 @@ def run_thread(
event_handler=event_handler,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -1339,14 +1323,14 @@ def run_thread(
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
)
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
azure_ad_token = get_secret("AZURE_AD_TOKEN")
response = azure_assistants_api.run_thread(
thread_id=thread_id,
@ -1366,7 +1350,7 @@ def run_thread(
client=client,
arun_thread=arun_thread,
litellm_params=litellm_params_dict,
) # type: ignore
)
else:
raise litellm.exceptions.BadRequestError(
message=f"LiteLLM doesn't support {custom_llm_provider} for 'run_thread'. Only 'openai' is supported.",
@ -1375,7 +1359,7 @@ def run_thread(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response # type: ignore
return response

View file

@ -8,8 +8,8 @@ from ..types.llms.openai import *
def get_optional_params_add_message(
role: str | None,
content: str | List[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
attachments: List[Attachment] | None,
content: str | list[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
attachments: list[Attachment] | None,
metadata: dict | None,
custom_llm_provider: str,
**kwargs,
@ -57,7 +57,7 @@ def get_optional_params_add_message(
optional_params = litellm.AzureOpenAIAssistantsAPIConfig().map_openai_params_create_message_params(
non_default_params=non_default_params, optional_params=optional_params
)
for k in passed_params.keys():
for k in passed_params:
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params
@ -128,7 +128,7 @@ def get_optional_params_image_gen(
if n is not None:
optional_params["sampleCount"] = int(n)
for k in passed_params.keys():
for k in passed_params:
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params

View file

@ -274,7 +274,7 @@ async def _fetch_batch_output_file_content(
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content: Final = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
_file_content: Final = await afile_content(**file_content_kwargs)
return _file_content.content
@ -432,7 +432,11 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
usage_object=response_body.get("usage", None) or {},
reasoning_content=None,
)
from litellm.responses.utils import ResponseAPILoggingUtils
_usage_dict: Final = response_body.get("usage", None) or {}
if ResponseAPILoggingUtils._is_response_api_usage(_usage_dict):
return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_usage_dict)
usage: Final[Usage] = Usage(**_usage_dict)
return usage

View file

@ -104,7 +104,7 @@ def _resolve_timeout(
@client
async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: dict[str, str] | None = None,
@ -154,7 +154,7 @@ async def acreate_batch(
@client
def create_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: dict[str, str] | None = None,
@ -287,7 +287,7 @@ def create_batch(
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
get_secret_str("AZURE_AD_TOKEN")
response = azure_batches_instance.create_batch(
_is_async=_is_async,
@ -327,7 +327,7 @@ def create_batch(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"),
),
)
return response
@ -370,7 +370,7 @@ async def aretrieve_batch(
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response # type: ignore
response = init_response
return response
except Exception as e:
@ -436,7 +436,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
get_secret_str("AZURE_AD_TOKEN")
response = azure_batches_instance.retrieve_batch(
_is_async=_is_async,
@ -498,7 +498,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"),
),
)
return response
@ -545,7 +545,7 @@ def retrieve_batch(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -677,7 +677,7 @@ async def alist_batches(
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response # type: ignore
response = init_response
return response
except Exception as e:
@ -723,7 +723,7 @@ def list_batches(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -755,7 +755,7 @@ def list_batches(
max_retries=optional_params.max_retries,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
api_key = (
@ -770,7 +770,7 @@ def list_batches(
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
get_secret_str("AZURE_AD_TOKEN")
response = azure_batches_instance.list_batches(
_is_async=_is_async,
@ -813,7 +813,7 @@ def list_batches(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
),
)
return response
@ -909,7 +909,7 @@ def cancel_batch(
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
timeout = float(timeout)
elif timeout is None:
timeout = 600.0
@ -959,7 +959,7 @@ def cancel_batch(
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
get_secret_str("AZURE_AD_TOKEN")
response = azure_batches_instance.cancel_batch(
_is_async=_is_async,
@ -999,7 +999,7 @@ def cancel_batch(
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"),
),
)
return response

View file

@ -9,12 +9,12 @@ Has 4 methods:
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
Span = Union[_Span, Any]
Span = _Span | Any
else:
Span = Any

View file

@ -534,11 +534,9 @@ class Cache:
if isinstance(cached_response, dict):
pass
else:
cached_response = json.loads(
cached_response # type: ignore
) # Convert string to dictionary
cached_response = json.loads(cached_response) # Convert string to dictionary
except Exception:
cached_response = ast.literal_eval(cached_response) # type: ignore
cached_response = ast.literal_eval(cached_response)
return cached_response
return cached_result

View file

@ -242,7 +242,7 @@ class LLMCachingHandler:
or litellm.cache.get_cache_key(**self.request_kwargs)
)
if hasattr(cached_result, "_hidden_params"):
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
cached_result._hidden_params["cache_key"] = cache_key
return CachingHandlerResponse(cached_result=cached_result)
elif (
call_type == CallTypes.aembedding.value
@ -356,7 +356,7 @@ class LLMCachingHandler:
or litellm.cache.get_cache_key(**self.request_kwargs)
)
if hasattr(cached_result, "_hidden_params"):
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
cached_result._hidden_params["cache_key"] = cache_key
return CachingHandlerResponse(cached_result=cached_result)
return CachingHandlerResponse(cached_result=cached_result)

View file

@ -1,12 +1,12 @@
import json
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final
from .base_cache import BaseCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
Span = Union[_Span, Any]
Span = _Span | Any
else:
Span = Any
@ -44,7 +44,7 @@ class DiskCache(BaseCache):
original_cached_response: Final = self.disk_cache.get(key)
if original_cached_response:
try:
cached_response = json.loads(original_cached_response) # type: ignore
cached_response = json.loads(original_cached_response)
except Exception:
cached_response = original_cached_response
return cached_response

View file

@ -13,7 +13,7 @@ import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from threading import Lock
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@ -29,7 +29,7 @@ from .redis_cache import RedisCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
Span = Union[_Span, Any]
Span = _Span | Any
else:
Span = Any

View file

@ -0,0 +1,276 @@
"""
Deferred close of HTTP/SDK clients that the LLM client cache has evicted.
Eviction only drops the cache's reference to a client. Every OpenAI/Azure SDK
client is a reference cycle (each resource namespace holds the client back), so
an evicted client and its pooled TCP connections survive until a generational
collection runs, which under load is thousands of requests later.
Closing at eviction time is not an option: a request that was handed the client
just before it was evicted is still using it, and closing it underneath that
request raises ``RuntimeError: Cannot send a request, as the client has been
closed.``
So an evicted client is closed once two conditions hold. A grace window must
have passed since its eviction, which covers a request that holds the client
but is momentarily not on the wire, and the client must report no connection in
flight. The second condition is what keeps the first honest: a request may run
for ``litellm.request_timeout`` seconds, 6000 by default, and a streaming
response is bounded only by how long the upstream keeps sending, so no deadline
on its own can promise that a request has finished.
Only clients litellm itself created are closed; a client the caller supplied is
left alone because litellm does not own its lifecycle.
A client that closes synchronously is closed from wherever the cache is next
used. One whose close is a coroutine needs the event loop it was evicted on, so
it waits for a call from that loop rather than having work scheduled onto a loop
it does not belong to. Queued clients are therefore bucketed by what it takes to
close them, and each bucket is ordered by deadline, so a reap walks the entries
that are due rather than the whole queue.
The queue holds its clients weakly, so waiting out a grace window never keeps
alive anything the collector would have reclaimed first.
"""
import asyncio
import contextlib
import inspect
import threading
import time
import weakref
from collections import deque
from collections.abc import Awaitable, Callable, Iterator
from dataclasses import dataclass, replace
from typing import Final
from litellm.constants import (
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
)
_CLOSABLE_ANYWHERE: Final = "closable-anywhere"
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
_BucketKey = str | int
@dataclass(frozen=True, slots=True)
class _PendingClose:
"""A queued close.
The client is held weakly, so queueing one never keeps alive anything the
collector would otherwise have reclaimed first.
``needs_loop`` is set for a client whose close is a coroutine; those can only
be closed from the event loop they were evicted on, recorded in ``loop_id``.
A client that closes synchronously carries neither constraint.
"""
client_ref: "weakref.ref[object]"
loop_id: int | None
needs_loop: bool
close_after: float
def _bucket_key(pending: _PendingClose) -> _BucketKey:
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
if not pending.needs_loop:
return _CLOSABLE_ANYWHERE
if pending.loop_id is None:
return _CLOSABLE_ON_ANY_LOOP
return pending.loop_id
def _running_loop_id() -> int | None:
try:
return id(asyncio.get_running_loop())
except RuntimeError:
return None
def _close_function(client: object) -> Callable[[], object] | None:
close_fn: Final[Callable[[], object] | None] = getattr(client, "aclose", None) or getattr(client, "close", None)
return close_fn
def _transport_of(client: object) -> object:
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
transport: object = getattr(holder, "_transport", None)
if transport is not None:
return transport
return None
def _connection_is_idle(connection: object) -> bool:
"""A pooled connection is idle unless it is servicing a request."""
is_idle: Final[object] = getattr(connection, "is_idle", None)
return bool(is_idle()) if callable(is_idle) else True
def _pool_has_busy_connection(transport: object) -> bool | None:
"""Whether the httpcore pool behind the transport is servicing a request.
``None`` when there is no such pool, so the caller can ask the other backend.
"""
pooled: Final[object] = getattr(getattr(transport, "_pool", None), "connections", None)
if not isinstance(pooled, (list, tuple)):
return None
return any(
not _connection_is_idle(connection) # pyright: ignore[reportUnknownArgumentType] # untyped pool list
for connection in pooled # pyright: ignore[reportUnknownVariableType] # untyped pool list
)
def _has_connection_in_flight(client: object) -> bool:
"""Whether the client is servicing a request right now.
Both connection backends litellm uses already account for the connections
they have handed out, so this reads the client's own lease accounting rather
than inferring it from elapsed time: httpcore reports a non-idle connection
for the whole of a response including a stream, and aiohttp holds the
connection in ``_acquired`` over the same span.
A client that cannot answer is reported as idle, which leaves the grace
window as the only guard, exactly as it was before this check existed.
"""
try:
transport: Final = _transport_of(client)
pooled_busy: Final = _pool_has_busy_connection(transport)
if pooled_busy is not None:
return pooled_busy
session: Final[object] = getattr(transport, "client", None)
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
return False
async def _close_quietly(closing: Awaitable[object]) -> None:
with contextlib.suppress(Exception):
await closing
class EvictedClientCloser:
"""Closes evicted, litellm-owned clients once they are idle and out of grace."""
def __init__(
self,
grace_seconds: float = EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
max_pending: int = EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
clock: Callable[[], float] = time.monotonic,
) -> None:
self._grace_seconds = grace_seconds
self._max_pending = max_pending
self._clock = clock
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
self._pending_count = 0
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
def mark_owned(self, client: object) -> None:
"""Record that litellm created this client, so it may be closed on eviction."""
try:
self._owned.add(client)
except TypeError:
pass # values that cannot be weak-referenced are never litellm clients
def _is_owned(self, client: object) -> bool:
try:
return client in self._owned
except TypeError:
return False # unhashable values are never litellm clients
def schedule(self, client: object) -> None:
"""Queue an evicted client for closing once it is idle and out of grace.
Past ``max_pending`` the client is left to the collector instead, so a
workload that churns the cache cannot grow this queue without bound.
Every queued entry comes due within one grace window, so the capacity it
occupies is returned within that window rather than held.
"""
if client is None or not self._is_owned(client):
return
close_fn: Final = _close_function(client)
if close_fn is None:
return
if self._pending_count >= self._max_pending:
return
self._enqueue(
_PendingClose(
client_ref=weakref.ref(client),
loop_id=_running_loop_id(),
needs_loop=inspect.iscoroutinefunction(close_fn),
close_after=self._clock() + self._grace_seconds,
)
)
def reap(self) -> None:
"""Close every queued client that is due, idle, and closable from here.
Called from the cache's read path, so the empty-queue exit comes first and
the work done past it is proportional to what is due, not to the queue.
"""
if not self._pending_count:
return
now: Final = self._clock()
for pending in self._take_due(_running_loop_id(), now):
client = pending.client_ref()
if client is None:
continue
if _has_connection_in_flight(client):
self._enqueue(replace(pending, close_after=now + self._grace_seconds))
continue
self._close(client)
@property
def pending_count(self) -> int:
return self._pending_count
def _enqueue(self, pending: _PendingClose) -> None:
"""Append to the entry's bucket, dropping any dead entries it queues behind.
Deadlines only ever move forward, so appending keeps each bucket ordered
by deadline, and entries whose client the collector already took sit at
the front rather than having to be searched for.
"""
with self._queue_lock:
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design
while bucket and bucket[0].client_ref() is None:
bucket.popleft()
self._pending_count -= 1
bucket.append(pending)
self._pending_count += 1
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
with self._queue_lock:
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
bucket: Final = self._buckets.get(key)
if bucket is None:
return
while bucket and bucket[0].close_after <= now:
self._pending_count -= 1
yield bucket.popleft()
if not bucket:
del self._buckets[key]
def _close(self, client: object) -> None:
close_fn: Final = _close_function(client)
if close_fn is None:
return
try:
closing: Final = close_fn()
except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers
return
if not inspect.isawaitable(closing):
return
task: Final = asyncio.get_running_loop().create_task(_close_quietly(closing))
self._close_tasks.add(task)
task.add_done_callback(self._close_tasks.discard)
default_evicted_client_closer: Final = EvictedClientCloser()

View file

@ -5,21 +5,44 @@ Add the event loop to the cache key, to prevent event loop closed errors.
import asyncio
from typing import Final
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
from .in_memory_cache import InMemoryCache
class LLMClientCache(InMemoryCache):
"""Cache for LLM HTTP clients (OpenAI, Azure, httpx, etc.).
IMPORTANT: This cache intentionally does NOT close clients on eviction.
Evicted clients may still be in use by in-flight requests. Closing them
eagerly causes ``RuntimeError: Cannot send a request, as the client has
been closed.`` errors in production after the TTL (1 hour) expires.
An evicted client is never closed on the spot: a request handed the client
just before eviction is still using it, and closing it there raises
``RuntimeError: Cannot send a request, as the client has been closed.``
Clients that are no longer referenced will be garbage-collected normally.
For explicit shutdown cleanup, use ``close_litellm_async_clients()``.
Nor can eviction be left to rely on garbage collection. The SDK clients are
reference cycles, so an evicted client and its open TCP connections survive
until a generational collection runs. Instead a client litellm created is
handed to ``EvictedClientCloser``, which closes it once a grace window has
passed. Clients the caller supplied are left untouched.
"""
def __init__(
self,
max_size_in_memory: int | None = 200,
default_ttl: int | None = 600,
max_size_per_item: int | None = 1024,
evicted_client_closer: EvictedClientCloser | None = None,
) -> None:
super().__init__(
max_size_in_memory=max_size_in_memory,
default_ttl=default_ttl,
max_size_per_item=max_size_per_item,
)
self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer
def _remove_key(self, key: str) -> None:
evicted: Final[object] = self.cache_dict.get(key)
super()._remove_key(key)
self.evicted_client_closer.schedule(evicted)
self.evicted_client_closer.reap()
def update_cache_key_with_event_loop(self, key):
"""
Add the event loop to the cache key, to prevent event loop closed errors.
@ -32,16 +55,22 @@ class LLMClientCache(InMemoryCache):
except RuntimeError: # handle no current running event loop
return key
def set_cache(self, key, value, **kwargs):
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return super().set_cache(key, value, **kwargs)
async def async_set_cache(self, key, value, **kwargs):
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return await super().async_set_cache(key, value, **kwargs)
def get_cache(self, key, **kwargs):
key = self.update_cache_key_with_event_loop(key)
self.evicted_client_closer.reap()
return super().get_cache(key, **kwargs)

View file

@ -18,7 +18,7 @@ import time
from collections.abc import Awaitable, Callable, Sequence
from contextvars import ContextVar
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, TypeVar, Union, cast
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -49,7 +49,7 @@ if TYPE_CHECKING:
cluster_pipeline = ClusterPipeline
async_redis_client = Redis
async_redis_cluster_client = RedisCluster
Span = Union[_Span, Any]
Span = _Span | Any
else:
pipeline = Any
cluster_pipeline = Any
@ -242,7 +242,7 @@ async def _run_under_circuit_breaker(
return result
def _redis_circuit_breaker_guard(method): # type: ignore
def _redis_circuit_breaker_guard(method):
"""
Decorator for RedisCache async methods.
Checks the circuit breaker before each call; records success/failure after.
@ -256,7 +256,7 @@ def _redis_circuit_breaker_guard(method): # type: ignore
"""
@functools.wraps(method)
async def wrapper(self, *args, **kwargs): # type: ignore
async def wrapper(self, *args, **kwargs):
return await _run_under_circuit_breaker(
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
)
@ -319,7 +319,7 @@ class RedisCache(BaseCache):
self.redis_version = "Unknown"
try:
if not coroutine_checker.is_async_callable(self.redis_client):
self.redis_version = self.redis_client.info()["redis_version"] # type: ignore
self.redis_version = self.redis_client.info()["redis_version"]
except Exception:
pass
@ -355,7 +355,7 @@ class RedisCache(BaseCache):
# SYNC HEALTH PING
try:
if hasattr(self.redis_client, "ping"):
self.redis_client.ping() # type: ignore
self.redis_client.ping()
except Exception as e:
verbose_logger.error("Error connecting to Sync Redis client", extra={"error": str(e)})
self._handle_sync_ping_error(e)
@ -423,7 +423,7 @@ class RedisCache(BaseCache):
redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs)
in_memory_llm_clients_cache.set_cache(key=cache_key, value=redis_async_client)
self.redis_async_client = redis_async_client # type: ignore
self.redis_async_client = redis_async_client
return redis_async_client
def check_and_fix_namespace(self, key: str) -> str:
@ -431,7 +431,7 @@ class RedisCache(BaseCache):
Make sure each key starts with the given namespace
"""
if key is None:
return key # type: ignore[return-value]
return key
if self.namespace is not None and not key.startswith(self.namespace):
key = self.namespace + ":" + key
@ -493,7 +493,7 @@ class RedisCache(BaseCache):
key = self.check_and_fix_namespace(key=key)
try:
start_time = time.time()
result: Final[int] = _redis_client.incr(name=key, amount=value) # type: ignore
result: Final[int] = _redis_client.incr(name=key, amount=value)
end_time = time.time()
_duration = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -520,7 +520,7 @@ class RedisCache(BaseCache):
if current_ttl == -1:
# Key has no expiration
start_time = time.time()
_redis_client.expire(key, set_ttl) # type: ignore
_redis_client.expire(key, set_ttl)
end_time = time.time()
_duration = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -555,7 +555,7 @@ class RedisCache(BaseCache):
return []
pattern = self.check_and_fix_namespace(key=pattern)
async for key in _redis_client.scan_iter(match=pattern + "*", count=count): # type: ignore
async for key in _redis_client.scan_iter(match=pattern + "*", count=count):
keys.append(key)
if len(keys) >= count:
break
@ -680,7 +680,7 @@ class RedisCache(BaseCache):
start_time: Final = time.time()
try:
_redis_client: Final[Redis] = self.init_async_client() # type: ignore
_redis_client: Final[Redis] = self.init_async_client()
except Exception as e:
end_time = time.time()
_duration = end_time - start_time
@ -773,7 +773,7 @@ class RedisCache(BaseCache):
_td: timedelta | None = None
if ttl is not None:
_td = timedelta(seconds=ttl)
pipe.set( # type: ignore
pipe.set(
name=cache_key,
value=json_cache_value,
ex=_td,
@ -849,7 +849,7 @@ class RedisCache(BaseCache):
"""Helper function for async_set_cache_sadd. Separated for testing."""
ttl = self.get_ttl(ttl=ttl)
try:
await redis_client.sadd(key, *value) # type: ignore
await redis_client.sadd(key, *value)
if ttl is not None:
_td: Final = timedelta(seconds=ttl)
await redis_client.expire(key, _td)
@ -862,7 +862,7 @@ class RedisCache(BaseCache):
start_time: Final = time.time()
try:
_redis_client: Final[Redis] = self.init_async_client() # type: ignore
_redis_client: Final[Redis] = self.init_async_client()
except Exception as e:
end_time = time.time()
_duration = end_time - start_time
@ -945,7 +945,7 @@ class RedisCache(BaseCache):
) -> float:
from redis.asyncio import Redis
_redis_client: Final[Redis] = self.init_async_client() # type: ignore
_redis_client: Final[Redis] = self.init_async_client()
start_time: Final = time.time()
_used_ttl: Final = self.get_ttl(ttl=ttl)
key = self.check_and_fix_namespace(key=key)
@ -1080,7 +1080,7 @@ class RedisCache(BaseCache):
We use a wrapper so RedisCluster can override this method
"""
return self.redis_client.mget(keys=keys) # type: ignore
return self.redis_client.mget(keys=keys)
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
@ -1089,7 +1089,7 @@ class RedisCache(BaseCache):
We use a wrapper so RedisCluster can override this method
"""
async_redis_client: Final = self.init_async_client()
return await async_redis_client.mget(keys=keys) # type: ignore
return await async_redis_client.mget(keys=keys)
def batch_get_cache(
self,
@ -1147,7 +1147,7 @@ class RedisCache(BaseCache):
async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
from redis.asyncio import Redis
_redis_client: Final[Redis] = self.init_async_client() # type: ignore
_redis_client: Final[Redis] = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
@ -1269,7 +1269,7 @@ class RedisCache(BaseCache):
print_verbose("Pinging Sync Redis Cache")
start_time: Final = time.time()
try:
response: Final[bool] = self.redis_client.ping() # type: ignore
response: Final[bool] = self.redis_client.ping()
print_verbose(f"Redis Cache PING: {response}")
## LOGGING ##
end_time = time.time()
@ -1339,7 +1339,7 @@ class RedisCache(BaseCache):
await _redis_client.delete(*keys)
def client_list(self) -> list:
client_list: Final[list] = self.redis_client.client_list() # type: ignore
client_list: Final[list] = self.redis_client.client_list()
return client_list
def info(self):
@ -1376,10 +1376,10 @@ class RedisCache(BaseCache):
redis_client: Final = redis_async.Redis(**self.redis_kwargs)
# Test the connection
ping_result: Final = await redis_client.ping() # type: ignore[misc]
ping_result: Final = await redis_client.ping()
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
await redis_client.aclose()
if ping_result:
return {
@ -1448,7 +1448,7 @@ class RedisCache(BaseCache):
from redis.asyncio import Redis
_redis_client: Final[Redis] = self.init_async_client() # type: ignore
_redis_client: Final[Redis] = self.init_async_client()
start_time: Final = time.time()
print_verbose(f"Increment Async Redis Cache Pipeline: increment list: {increment_list}")
@ -1769,7 +1769,7 @@ class RedisCache(BaseCache):
or None
)
except Exception:
decoded_results.append(r) # type: ignore
decoded_results.append(r)
else:
decoded_results.append(None)
return decoded_results

View file

@ -5,7 +5,7 @@ Key differences:
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
"""
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final
from litellm.caching.redis_cache import RedisCache
@ -16,7 +16,7 @@ if TYPE_CHECKING:
pipeline = Pipeline
async_redis_client = Redis
Span = Union[_Span, Any]
Span = _Span | Any
else:
pipeline = Any
async_redis_client = Any
@ -47,14 +47,14 @@ class RedisClusterCache(RedisCache):
"""
Overrides `_run_redis_mget_operation` in redis_cache.py
"""
return self.redis_client.mget_nonatomic(keys=keys) # type: ignore
return self.redis_client.mget_nonatomic(keys=keys)
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Overrides `_async_run_redis_mget_operation` in redis_cache.py
"""
async_redis_cluster_client: Final = self.init_async_client()
return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore
return await async_redis_cluster_client.mget_nonatomic(keys=keys)
async def test_connection(self) -> dict:
"""
@ -78,14 +78,14 @@ class RedisClusterCache(RedisCache):
# Create a fresh Redis Cluster client with current settings
redis_client: Final = redis_async.RedisCluster(
startup_nodes=new_startup_nodes,
**cluster_kwargs, # type: ignore
**cluster_kwargs,
)
# Test the connection
ping_result: Final = await redis_client.ping() # type: ignore[attr-defined, misc]
ping_result: Final = await redis_client.ping()
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
await redis_client.aclose()
if ping_result:
return {

View file

@ -13,6 +13,7 @@ import ast
import asyncio
import json
import os
from collections.abc import Callable, Mapping
from typing import Any, Final, cast
import litellm
@ -47,7 +48,7 @@ class RedisSemanticCache(BaseCache):
similarity_threshold: float | None = None,
embedding_model: str = "text-embedding-ada-002",
index_name: str | None = None,
**kwargs,
**kwargs: object,
):
"""
Initialize the Redis Semantic Cache.
@ -126,8 +127,8 @@ class RedisSemanticCache(BaseCache):
# CustomTextVectorizer probes its embedding dimension at construction by
# embedding "dimension test", so the first cache request issues one extra
# billable embedding on top of the request's own.
from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped]
from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped]
from redisvl.extensions.llmcache import SemanticCache
from redisvl.utils.vectorize import CustomTextVectorizer
try:
cache_vectorizer: Final = CustomTextVectorizer(self._get_embedding)
@ -150,11 +151,11 @@ class RedisSemanticCache(BaseCache):
def _init_semantic_cache(
self,
semantic_cache_cls: Any,
semantic_cache_cls: Callable[..., object],
index_name: str,
redis_url: str,
cache_vectorizer: Any,
) -> Any:
cache_vectorizer: object,
) -> object:
def _is_schema_mismatch(exc: ValueError) -> bool:
error_message: Final = str(exc).lower()
return any(phrase in error_message for phrase in ("schema does not match", "index schema"))
@ -206,12 +207,12 @@ class RedisSemanticCache(BaseCache):
def _get_cache_filters(self, key: str) -> dict[str, str]:
return {self.CACHE_KEY_FIELD_NAME: str(key)}
def _get_cache_key_filter_expression(self, key: str) -> Any:
from redisvl.query.filter import Tag # type: ignore[import-not-found, import-untyped]
def _get_cache_key_filter_expression(self, key: str) -> object:
from redisvl.query.filter import Tag
return Tag(self.CACHE_KEY_FIELD_NAME) == str(key)
def _cache_hit_matches_key(self, cache_hit: dict[str, Any], key: str) -> bool:
def _cache_hit_matches_key(self, cache_hit: Mapping[str, object], key: str) -> bool:
# Pre-isolation entries with no ``litellm_cache_key`` field cannot be
# safely reassigned to a caller's scope and are treated as misses.
cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME)
@ -297,7 +298,7 @@ class RedisSemanticCache(BaseCache):
return
@staticmethod
def _coerce_response_input_value(value: Any) -> Any:
def _coerce_response_input_value(value: object) -> object:
model_dump: Final = getattr(value, "model_dump", None)
if callable(model_dump):
return model_dump()
@ -340,7 +341,7 @@ class RedisSemanticCache(BaseCache):
)
return embedding_response["data"][0]["embedding"]
def _get_cache_logic(self, cached_response: Any) -> Any:
def _get_cache_logic(self, cached_response: Any) -> object:
"""
Process the cached response to prepare it for use.
@ -369,7 +370,7 @@ class RedisSemanticCache(BaseCache):
return cached_response
def set_cache(self, key: str, value: Any, **kwargs) -> None:
def set_cache(self, key: str, value: object, **kwargs) -> None:
"""
Store a value in the semantic cache.
@ -405,7 +406,7 @@ class RedisSemanticCache(BaseCache):
except Exception as e:
print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e}")
def get_cache(self, key: str, **kwargs) -> Any:
def get_cache(self, key: str, **kwargs) -> object:
"""
Retrieve a semantically similar cached response.
@ -428,7 +429,7 @@ class RedisSemanticCache(BaseCache):
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata"))
check_kwargs: Final[dict[str, Any]] = {
check_kwargs: Final[Mapping[str, object]] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
@ -508,7 +509,7 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Error generating async embedding: {e}")
raise ValueError(f"Failed to generate embedding: {e}") from e
async def async_set_cache(self, key: str, value: Any, **kwargs) -> None:
async def async_set_cache(self, key: str, value: object, **kwargs) -> None:
"""
Asynchronously store a value in the semantic cache.
@ -548,7 +549,7 @@ class RedisSemanticCache(BaseCache):
except Exception as e:
print_verbose(f"Error in async_set_cache: {e}")
async def async_get_cache(self, key: str, **kwargs) -> Any:
async def async_get_cache(self, key: str, **kwargs) -> object:
"""
Asynchronously retrieve a semantically similar cached response.
@ -573,7 +574,7 @@ class RedisSemanticCache(BaseCache):
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
check_kwargs: Final[dict[str, Any]] = {
check_kwargs: Final[Mapping[str, object]] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
@ -615,7 +616,7 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Error in async_get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _index_info(self) -> dict[str, Any]:
async def _index_info(self) -> Mapping[str, object]:
"""
Get information about the Redis index.
@ -625,7 +626,7 @@ class RedisSemanticCache(BaseCache):
aindex: Final = await self.llmcache._get_async_index()
return await aindex.info()
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: object) -> None:
"""
Asynchronously store multiple values in the semantic cache.

View file

@ -146,7 +146,7 @@ class S3Cache(BaseCache):
)
return cached_response
except botocore.exceptions.ClientError as e: # type: ignore
except botocore.exceptions.ClientError as e:
if e.response["Error"]["Code"] == "NoSuchKey":
verbose_logger.debug("S3 Cache: The specified key '%s' does not exist in the S3 bucket.", key)
return None

Some files were not shown because too many files have changed in this diff Show more