mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
commit
49ab201e75
2030 changed files with 80626 additions and 25440 deletions
|
|
@ -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
|
||||
|
|
|
|||
40
.github/actions/cache-prisma-binaries/action.yml
vendored
Normal file
40
.github/actions/cache-prisma-binaries/action.yml
vendored
Normal 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
152
.github/ci-coverage-allowlist.yml
vendored
Normal 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
|
||||
33
.github/pull_request_template.md
vendored
33
.github/pull_request_template.md
vendored
|
|
@ -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
262
.github/scripts/assert_ci_coverage.py
vendored
Normal 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())
|
||||
47
.github/workflows/_test-unit-base.yml
vendored
47
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
4
.github/workflows/check-schema-sync.yml
vendored
4
.github/workflows/check-schema-sync.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
50
.github/workflows/check-ui-api-types.yml
vendored
50
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -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
42
.github/workflows/ci-coverage.yml
vendored
Normal 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
|
||||
4
.github/workflows/conventional-commits.yml
vendored
4
.github/workflows/conventional-commits.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
35
.github/workflows/helm_unit_test.yml
vendored
35
.github/workflows/helm_unit_test.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
98
.github/workflows/image-scan.yml
vendored
98
.github/workflows/image-scan.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
5
.github/workflows/mutation-test.yml
vendored
5
.github/workflows/mutation-test.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
63
.github/workflows/publish-basedpyright-base-counts.yml
vendored
Normal file
63
.github/workflows/publish-basedpyright-base-counts.yml
vendored
Normal 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
|
||||
6
.github/workflows/test-code-quality.yml
vendored
6
.github/workflows/test-code-quality.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
61
.github/workflows/test-linting.yml
vendored
61
.github/workflows/test-linting.yml
vendored
|
|
@ -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"
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-build.yml
vendored
4
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
11
.github/workflows/test-litellm-ui-lint.yml
vendored
11
.github/workflows/test-litellm-ui-lint.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
24
.github/workflows/test-litellm-ui-unit.yml
vendored
24
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -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"
|
||||
|
|
|
|||
4
.github/workflows/test-mcp.yml
vendored
4
.github/workflows/test-mcp.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
4
.github/workflows/test-model-map.yaml
vendored
4
.github/workflows/test-model-map.yaml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
4
.github/workflows/test-unit-misc.yml
vendored
4
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
4
.github/workflows/test-unit-proxy-db.yml
vendored
4
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-infra.yml
vendored
2
.github/workflows/test-unit-proxy-infra.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
6
.github/workflows/test-unit-proxy-legacy.yml
vendored
6
.github/workflows/test-unit-proxy-legacy.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
5
.github/workflows/weekly_load_anomaly.yml
vendored
5
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -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
1
.gitignore
vendored
|
|
@ -1,5 +1,6 @@
|
|||
.python-version
|
||||
.venv
|
||||
.venv-typecheck
|
||||
.venv_policy_test
|
||||
.env
|
||||
.claude
|
||||
|
|
|
|||
22
CLAUDE.md
22
CLAUDE.md
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
33
Makefile
33
Makefile
|
|
@ -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/
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/integrations/__init__.py
Normal file
0
enterprise/litellm_enterprise/integrations/__init__.py
Normal 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(),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/proxy/hooks/__init__.py
Normal file
0
enterprise/litellm_enterprise/proxy/hooks/__init__.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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(
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/py.typed
Normal file
0
enterprise/litellm_enterprise/py.typed
Normal file
0
enterprise/litellm_enterprise/types/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/proxy/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/proxy/__init__.py
Normal 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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
|
@ -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 '[]';
|
||||
|
|
@ -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");
|
||||
|
|
@ -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");
|
||||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';
|
||||
181
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
Normal file
181
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
Normal 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)
|
||||
0
litellm-proxy-extras/litellm_proxy_extras/py.typed
Normal file
0
litellm-proxy-extras/litellm_proxy_extras/py.typed
Normal 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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
//
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
276
litellm/caching/evicted_client_closer.py
Normal file
276
litellm/caching/evicted_client_closer.py
Normal 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()
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue