mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #37721 from BerriAI/litellm_internal_staging
Some checks failed
CI Coverage / assert-ci-coverage (push) Has been cancelled
CodeQL / Analyze (actions) (push) Has been cancelled
CodeQL / Analyze (javascript-typescript) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Unit Tests / core-utils (push) Has been cancelled
Unit Tests / enterprise-routing (push) Has been cancelled
Unit Tests / integrations (push) Has been cancelled
Unit Tests / All Other Providers (push) Has been cancelled
Unit Tests / Vertex AI (push) Has been cancelled
Unit Tests / misc (push) Has been cancelled
Unit Tests / proxy-auth (push) Has been cancelled
Unit Tests / proxy-endpoints (push) Has been cancelled
Unit Tests / proxy-infra (push) Has been cancelled
Unit Tests / proxy-server (push) Has been cancelled
Unit Tests / responses-caching-types (push) Has been cancelled
GitHub Actions Security Analysis / zizmor (push) Has been cancelled
Code Quality Checks / code-quality (push) Has been cancelled
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
Unit Tests: Documentation Validation / documentation (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
CodSpeed Benchmarks / benchmarks (push) Has been cancelled
Helm unit test / unit-test (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Some checks failed
CI Coverage / assert-ci-coverage (push) Has been cancelled
CodeQL / Analyze (actions) (push) Has been cancelled
CodeQL / Analyze (javascript-typescript) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Unit Tests / core-utils (push) Has been cancelled
Unit Tests / enterprise-routing (push) Has been cancelled
Unit Tests / integrations (push) Has been cancelled
Unit Tests / All Other Providers (push) Has been cancelled
Unit Tests / Vertex AI (push) Has been cancelled
Unit Tests / misc (push) Has been cancelled
Unit Tests / proxy-auth (push) Has been cancelled
Unit Tests / proxy-endpoints (push) Has been cancelled
Unit Tests / proxy-infra (push) Has been cancelled
Unit Tests / proxy-server (push) Has been cancelled
Unit Tests / responses-caching-types (push) Has been cancelled
GitHub Actions Security Analysis / zizmor (push) Has been cancelled
Code Quality Checks / code-quality (push) Has been cancelled
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
Unit Tests: Documentation Validation / documentation (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
CodSpeed Benchmarks / benchmarks (push) Has been cancelled
Helm unit test / unit-test (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
chore(ci): promote internal staging to main
This commit is contained in:
commit
418c7c6012
1363 changed files with 88066 additions and 28345 deletions
|
|
@ -1025,7 +1025,7 @@ jobs:
|
|||
name: Run tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
TEST_FILES=$(circleci tests glob "tests/agent_tests/**/test_*.py" | grep -v "^tests/agent_tests/local_only_agent_tests/")
|
||||
TEST_FILES=$(circleci tests glob "tests/agent_tests/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
|
|
|
|||
|
|
@ -1,15 +1,17 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
ui/* | tests/e2e/ui/*) has_client=true ;;
|
||||
docs/* | *.md | *.mdx) : ;;
|
||||
.github/* | .circleci/*) has_ci=true; has_backend=true ;;
|
||||
*) has_backend=true ;;
|
||||
esac
|
||||
done
|
||||
|
|
@ -21,6 +23,9 @@ case "$category" in
|
|||
client)
|
||||
{ [ "$has_client" = true ] || [ "$has_backend" = true ]; } && echo run || echo skip
|
||||
;;
|
||||
ui)
|
||||
{ [ "$has_client" = true ] || [ "$has_ci" = true ]; } && echo run || echo skip
|
||||
;;
|
||||
*)
|
||||
echo run
|
||||
;;
|
||||
|
|
|
|||
2
.github/CODEOWNERS
vendored
2
.github/CODEOWNERS
vendored
|
|
@ -1,3 +1,5 @@
|
|||
/ui/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/ui/litellm-dashboard/src/lib/http/schema.d.ts
|
||||
/model_prices_and_context_window.json @mateo-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
name: "Detect backend-relevant changes"
|
||||
description: >-
|
||||
Classify the pull request's changed files with .circleci/scripts/classify_changes.sh
|
||||
and expose decision=run|skip. decision=skip means only ui/**, **.md or **.mdx files
|
||||
changed, so callers can short-circuit expensive steps while the job still completes
|
||||
successfully and satisfies its required status check. The decision defaults to run for
|
||||
any non pull_request event or whenever the changed set cannot be resolved, so tests are
|
||||
never skipped when the classification is uncertain.
|
||||
|
||||
outputs:
|
||||
decision:
|
||||
description: "run when backend-relevant files changed, otherwise skip"
|
||||
value: ${{ steps.classify.outputs.decision }}
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- id: classify
|
||||
shell: bash
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
set -uo pipefail
|
||||
if [ -z "${BASE_SHA:-}" ]; then
|
||||
echo "detect-backend-changes: not a pull_request event; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
if ! git fetch --no-tags --depth=1 origin "${BASE_SHA}" >/dev/null 2>&1; then
|
||||
echo "detect-backend-changes: could not fetch base ${BASE_SHA}; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
changed="$(git diff --name-only "${BASE_SHA}" HEAD 2>/dev/null)" || {
|
||||
echo "detect-backend-changes: git diff failed; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
}
|
||||
if [ -z "${changed}" ]; then
|
||||
echo "detect-backend-changes: no changed files vs ${BASE_SHA}; skipping job"
|
||||
echo "decision=skip" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
echo "detect-backend-changes: changed files vs ${BASE_SHA}:"
|
||||
printf '%s\n' "${changed}" | sed 's/^/ /'
|
||||
decision="$(printf '%s\n' "${changed}" | bash .circleci/scripts/classify_changes.sh backend)" || decision="run"
|
||||
echo "detect-backend-changes: decision=${decision}"
|
||||
echo "decision=${decision}" >> "${GITHUB_OUTPUT}"
|
||||
41
.github/actions/detect-changes/action.yml
vendored
Normal file
41
.github/actions/detect-changes/action.yml
vendored
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
name: "Detect relevant changes"
|
||||
description: >-
|
||||
Classify the pull request's changed files with .circleci/scripts/classify_changes.sh
|
||||
and expose decision=run|skip for one category. backend means anything outside ui/,
|
||||
docs/ and markdown; ui means the dashboard sources alone. decision=skip lets callers
|
||||
short-circuit expensive steps while the job still completes successfully and satisfies
|
||||
its required status check, which a paths: filter cannot do because a workflow that
|
||||
never starts never reports. The file list comes from the pull request itself rather
|
||||
than from a git diff, because the checked-out merge ref is recomputed as the base
|
||||
branch advances and would otherwise attribute the base branch's own commits to the
|
||||
pull request. The decision defaults to run for any non pull_request event or whenever
|
||||
the changed set cannot be resolved, so jobs are never skipped when the classification
|
||||
is uncertain.
|
||||
|
||||
inputs:
|
||||
category:
|
||||
description: "Which classification to apply: backend, client or ui"
|
||||
required: false
|
||||
default: backend
|
||||
github-token:
|
||||
description: "Token used to list the pull request's files; needs pull-requests: read"
|
||||
required: false
|
||||
default: ${{ github.token }}
|
||||
|
||||
outputs:
|
||||
decision:
|
||||
description: "run when category-relevant files changed, otherwise skip"
|
||||
value: ${{ steps.classify.outputs.decision }}
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- id: classify
|
||||
shell: bash
|
||||
env:
|
||||
GH_TOKEN: ${{ inputs.github-token }}
|
||||
CATEGORY: ${{ inputs.category }}
|
||||
REPO: ${{ github.repository }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
CHANGED_FILE_COUNT: ${{ github.event.pull_request.changed_files }}
|
||||
run: bash "${GITHUB_ACTION_PATH}/../../scripts/detect_changes.sh"
|
||||
117
.github/ci-coverage-allowlist.yml
vendored
117
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -4,6 +4,25 @@ description: >-
|
|||
by a job nor listed here, so every entry below is a decision on the record.
|
||||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
The caching suite in tests/local_testing, which runs nowhere. Every job that globs that
|
||||
directory either deselects it (local_testing_part1 and part2 carry `-k "... and not caching
|
||||
and not cache"`) or keeps only another keyword (langfuse, router, assistants), and no job
|
||||
names these files the way redis_caching_unit_tests names test_dual_cache.py. Measured
|
||||
2026-08-20 by collecting the directory under each job's own selector: 118 tests across
|
||||
these eight files are selected by none of them. Listed so the gap is a decision rather
|
||||
than an accident, and so the --slices guard has a baseline to ratchet down from. Revisit
|
||||
when tests/local_testing is ported off CircleCI, where the keyless part of this suite
|
||||
belongs in a real job
|
||||
paths:
|
||||
- tests/local_testing/test_cache_preset_key.py
|
||||
- tests/local_testing/test_caching.py
|
||||
- tests/local_testing/test_caching_handler.py
|
||||
- tests/local_testing/test_disk_cache_unit_tests.py
|
||||
- tests/local_testing/test_gcs_cache_unit_tests.py
|
||||
- tests/local_testing/test_prompt_caching.py
|
||||
- tests/local_testing/test_responses_stream_cache_keys.py
|
||||
- tests/local_testing/test_unit_test_caching.py
|
||||
- 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
|
||||
|
|
@ -21,72 +40,24 @@ test_paths:
|
|||
- 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
|
||||
Named like a test but shaped like a benchmark: it fetches live image URLs, times aiohttp
|
||||
against httpx, prints the ratio, and asserts nothing, so pytest cannot collect it (its
|
||||
functions take arguments, not fixtures) and running it beside its siblings in the
|
||||
code-quality workflow would add a network dependency for a number nothing reads. Exempt
|
||||
as a script rather than as an unresolved gap; revisit by deleting it once the aiohttp
|
||||
choice it informed is settled
|
||||
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
|
||||
The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its
|
||||
other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging
|
||||
their bodies into the live file of the same name. This one cannot follow either route yet:
|
||||
its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no
|
||||
counterpart while 25 assertions fail against today's code, so what survives that rewrite
|
||||
is a judgement about the endpoints, not a merge. Revisit by deciding which of the five
|
||||
behaviours still hold
|
||||
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
|
||||
|
|
@ -116,6 +87,14 @@ test_paths:
|
|||
- 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: >-
|
||||
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
|
||||
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
|
||||
pull request job. Until 2026-08-20 the CircleCI agent job hid them behind a grep -v
|
||||
that this census could not see; the glob now excludes them structurally and this entry
|
||||
is the decision on the record. Revisit when the A2A bridge gets a recorded-wire fixture
|
||||
paths:
|
||||
- tests/agent_tests/local_only_agent_tests
|
||||
- reason: >-
|
||||
Third-party integration tests that skip themselves without OCI configuration or sandbox
|
||||
credentials, neither of which a pull request job holds
|
||||
|
|
@ -124,14 +103,14 @@ test_paths:
|
|||
- 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
|
||||
A unit test for the proxy-extras package that no job invokes, while the package's other
|
||||
tests live under tests/proxy_migration_tests. Measured 2026-08-20: 24 of its 28 tests pass
|
||||
and the 4 in TestMigrationSQLIdempotency fail, because 13 migrations from 2026-03 onward use
|
||||
bare CREATE TABLE, ADD COLUMN, CREATE INDEX and ADD CONSTRAINT rather than the guarded forms
|
||||
this file requires. It also matches those keywords inside SQL comments, so two further
|
||||
migrations are reported that are in fact fine. Wiring it up means deciding what to do about
|
||||
the 13 first, and they cannot simply be edited: Prisma checksums an applied migration, so a
|
||||
changed one breaks migrate deploy for existing installs
|
||||
paths:
|
||||
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
|
||||
|
||||
|
|
|
|||
3
.github/pull_request_template.md
vendored
3
.github/pull_request_template.md
vendored
|
|
@ -53,7 +53,8 @@ After: the same request comes back with real token counts, so the dashboard show
|
|||
**Please complete all items before asking a LiteLLM maintainer to review your PR**
|
||||
|
||||
- [ ] I have added meaningful tests
|
||||
- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests)
|
||||
- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/<your_test_file>.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
|
||||
- [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
|
||||
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
|
||||
- [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes)
|
||||
|
||||
|
|
|
|||
279
.github/scripts/assert_ci_coverage.py
vendored
279
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -1,10 +1,14 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import operator
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import yaml
|
||||
|
||||
|
|
@ -25,6 +29,15 @@ 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("*?")
|
||||
|
||||
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
|
||||
# tests has to be named by some shard or it runs nowhere. A child listed here is
|
||||
# itself decomposed one level deeper and is checked through its own entry.
|
||||
SHARDED_ROOTS: tuple[str, ...] = (
|
||||
"tests/proxy_unit_tests",
|
||||
"tests/test_litellm",
|
||||
"tests/test_litellm/proxy",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AllowEntry:
|
||||
|
|
@ -44,6 +57,14 @@ class Allowlist:
|
|||
return any(relative_path == path for entry in self.dockerfiles for path in entry.paths)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Section:
|
||||
name: str
|
||||
entries: tuple[AllowEntry, ...]
|
||||
candidates: tuple[str, ...]
|
||||
matches: Callable[[str, str], bool]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Scalar:
|
||||
key: str
|
||||
|
|
@ -107,20 +128,34 @@ def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
|
|||
)
|
||||
|
||||
|
||||
def _glob_to_regex(token: str) -> re.Pattern[str]:
|
||||
parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
|
||||
def _glob_to_regex(token: str, *, subtree: bool) -> re.Pattern[str]:
|
||||
parts = re.split(r"(\*\*/|\*\*|\*|\?|\[[^\]]*\])", token)
|
||||
translated = "".join(
|
||||
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
|
||||
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part)
|
||||
or (part if part.startswith("[") and part.endswith("]") else re.escape(part))
|
||||
for part in parts
|
||||
)
|
||||
return re.compile(rf"{translated}(?:/.*)?$")
|
||||
return re.compile(rf"{translated}(?:/.*)?$" if subtree else 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 _glob_to_regex(token, subtree=True).match(relative_path) is not None
|
||||
return relative_path == token or relative_path.startswith(f"{token}/")
|
||||
|
||||
|
||||
def _token_names(token: str, relative_path: str) -> bool:
|
||||
"""Whether the token names this path itself, rather than merely containing it.
|
||||
|
||||
A sharded tree has no catch-all bucket, so the ancestor token the census is happy
|
||||
with (`tests/x` standing in for everything below it) is exactly what would let a
|
||||
newly added child ride along without a shard.
|
||||
"""
|
||||
if GLOB_CHARS & set(token):
|
||||
return _glob_to_regex(token, subtree=False).match(relative_path) is not None
|
||||
return token == relative_path
|
||||
|
||||
|
||||
def _test_files() -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
|
|
@ -166,6 +201,174 @@ def _describe(paths: tuple[str, ...]) -> str:
|
|||
return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
|
||||
|
||||
|
||||
GLOB_CALL_RE = re.compile(r'circleci tests glob "([^"]+)"')
|
||||
KEYWORD_RE = re.compile(r"-k\s+\\?[\"']([^\"'\\]+)")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Slice:
|
||||
"""One job's selection: the files it globs, narrowed by its `-k` expression."""
|
||||
|
||||
job: str
|
||||
globs: tuple[str, ...]
|
||||
named: frozenset[str]
|
||||
required: tuple[str, ...]
|
||||
excluded: tuple[str, ...]
|
||||
understood: bool
|
||||
|
||||
def claims(self, relative_path: str, inner_names: frozenset[str]) -> bool:
|
||||
"""Whether this job runs any test in the file.
|
||||
|
||||
The question is deliberately per-file, not per-test. An excluded term is only
|
||||
honoured when it appears in the path, because that is the case where it takes
|
||||
the whole module with it; a term matching one function inside drops that test
|
||||
and leaves the file claimed. Losing a whole file is the failure worth a gate,
|
||||
and answering per-test would mean a baseline of test ids that churns on every
|
||||
rename.
|
||||
"""
|
||||
if relative_path in self.named:
|
||||
return True
|
||||
if not any(_token_covers(glob, relative_path) for glob in self.globs):
|
||||
return False
|
||||
if not self.understood:
|
||||
return True # a `-k` this parser cannot model is assumed to claim everything
|
||||
if any(term.lower() in relative_path.lower() for term in self.excluded):
|
||||
return False
|
||||
return not self.required or any(
|
||||
term.lower() in name.lower() for term in self.required for name in inner_names
|
||||
)
|
||||
|
||||
|
||||
def _strings(node: object) -> Iterable[str]:
|
||||
if isinstance(node, str):
|
||||
yield node
|
||||
elif isinstance(node, dict):
|
||||
for value in node.values():
|
||||
yield from _strings(value)
|
||||
elif isinstance(node, list):
|
||||
for value in node:
|
||||
yield from _strings(value)
|
||||
|
||||
|
||||
def _keyword_terms(
|
||||
expressions: Sequence[str], *, attributable: bool = True
|
||||
) -> tuple[tuple[str, ...], tuple[str, ...], bool]:
|
||||
"""A `-k` expression as (required, excluded, understood).
|
||||
|
||||
Only flat `and` chains of bare terms are modelled. Anything with `or`, parentheses
|
||||
or negation of a group is left unmodelled, and its job is then treated as claiming
|
||||
every file it globs, so an unparsed selector can never raise a false alarm.
|
||||
|
||||
`attributable` is False when a job runs several pytest commands, since a selector
|
||||
read out of the job's text cannot then be tied to the glob it belongs to, and
|
||||
pairing one command's exclusion with another's glob would invent a gap.
|
||||
"""
|
||||
terms: Final = tuple(part.strip() for expression in expressions for part in expression.split(" and "))
|
||||
if not attributable and terms:
|
||||
return (), (), False
|
||||
if any(("or " in term) or ("(" in term) or (term.startswith("not ") and " " in term[4:]) for term in terms):
|
||||
return (), (), False
|
||||
return (
|
||||
tuple(term for term in terms if term and not term.startswith("not ")),
|
||||
tuple(term[4:].strip() for term in terms if term.startswith("not ")),
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
def _slices() -> tuple[Slice, ...]:
|
||||
if not CIRCLECI_CONFIG.exists():
|
||||
return ()
|
||||
jobs: Final = yaml.safe_load(CIRCLECI_CONFIG.read_text()).get("jobs", {})
|
||||
return tuple(
|
||||
Slice(job=job, globs=globs, named=named, required=required, excluded=excluded, understood=understood)
|
||||
for job, body in jobs.items()
|
||||
for text in ("\n".join(_strings(body)),)
|
||||
if "pytest" in text
|
||||
for globs in (tuple(GLOB_CALL_RE.findall(text)),)
|
||||
for named in (frozenset(TEST_TOKEN_RE.findall(text)) & frozenset(_test_files()),)
|
||||
for required, excluded, understood in (
|
||||
_keyword_terms(tuple(KEYWORD_RE.findall(text)), attributable=len(globs) < 2),
|
||||
)
|
||||
if globs or named
|
||||
)
|
||||
|
||||
|
||||
def _matchable_names(relative_path: str) -> frozenset[str]:
|
||||
"""Every name a `-k` term can match for this file: its path, plus the names inside it.
|
||||
|
||||
pytest matches a keyword against an item's own name and each of its parents', so a
|
||||
positive term hits a file when it appears in the path or in a class or function name.
|
||||
"""
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore") # test files carry stray escapes; their names still parse
|
||||
tree: Final = ast.parse((REPO_ROOT / relative_path).read_text())
|
||||
except (OSError, SyntaxError):
|
||||
return frozenset({relative_path})
|
||||
return frozenset({relative_path}) | frozenset(
|
||||
node.name
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
|
||||
)
|
||||
|
||||
|
||||
def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
|
||||
slices: Final = _slices()
|
||||
globbed: Final = tuple(
|
||||
path
|
||||
for path in _test_files()
|
||||
if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
|
||||
)
|
||||
return tuple(
|
||||
Finding(
|
||||
subject=path,
|
||||
detail="globbed by a job, then deselected by every one of their -k expressions",
|
||||
)
|
||||
for path in globbed
|
||||
if not allowlist.covers_test(path)
|
||||
and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices)
|
||||
)
|
||||
|
||||
|
||||
def _holds_tests(directory: pathlib.Path) -> bool:
|
||||
return any(directory.rglob("test_*.py"))
|
||||
|
||||
|
||||
def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str, ...]:
|
||||
"""Children of a sharded root that carry tests, so each one needs its own shard.
|
||||
|
||||
A directory earns an entry by containing a test file rather than by being named
|
||||
`test_*`, which is what keeps fixture directories (`test_configs`, `expected_*`)
|
||||
out without a hand-maintained list of exceptions.
|
||||
"""
|
||||
return tuple(
|
||||
sorted(
|
||||
child.relative_to(repo_root).as_posix()
|
||||
for child in (repo_root / root).iterdir()
|
||||
if not child.name.startswith(".")
|
||||
and (
|
||||
_holds_tests(child)
|
||||
if child.is_dir()
|
||||
else child.name.startswith("test_") and child.suffix == ".py"
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _unassigned_shard_children(
|
||||
tokens: frozenset[str],
|
||||
roots: tuple[str, ...] = SHARDED_ROOTS,
|
||||
repo_root: pathlib.Path = REPO_ROOT,
|
||||
) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
Finding(subject=child, detail=f"holds tests but no shard of {root} names it")
|
||||
for root in roots
|
||||
if (repo_root / root).is_dir()
|
||||
for child in _shard_children(root, repo_root)
|
||||
if child not in roots and not any(_token_names(token, child) for token in tokens)
|
||||
)
|
||||
|
||||
|
||||
def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
Finding(subject=relative_path, detail="built by no job")
|
||||
|
|
@ -174,6 +377,25 @@ def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tupl
|
|||
)
|
||||
|
||||
|
||||
def _stale_allowlist_paths(
|
||||
allowlist: Allowlist,
|
||||
*,
|
||||
test_files: tuple[str, ...],
|
||||
dockerfiles: tuple[str, ...],
|
||||
) -> tuple[Finding, ...]:
|
||||
sections: Final[tuple[Section, ...]] = (
|
||||
Section("test_paths", allowlist.test_paths, test_files, _token_covers),
|
||||
Section("dockerfiles", allowlist.dockerfiles, dockerfiles, operator.eq),
|
||||
)
|
||||
return tuple(
|
||||
Finding(subject=path, detail=f"listed under '{section.name}' but matches no file the census looks at")
|
||||
for section in sections
|
||||
for entry in section.entries
|
||||
for path in entry.paths
|
||||
if not any(section.matches(path, candidate) for candidate in section.candidates)
|
||||
)
|
||||
|
||||
|
||||
def _parse_entry(item: object, section: str) -> AllowEntry:
|
||||
if not isinstance(item, dict):
|
||||
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' entries must be mappings")
|
||||
|
|
@ -229,13 +451,56 @@ def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
|
|||
_write("")
|
||||
|
||||
|
||||
def _check_slices() -> int:
|
||||
findings: Final = _deselected_everywhere(_load_allowlist())
|
||||
if findings:
|
||||
_report(
|
||||
"test files a -k expression removes from every job that globs them",
|
||||
findings,
|
||||
"Give each one a job whose -k keeps it, or list it in "
|
||||
".github/ci-coverage-allowlist.yml with the reason it may stay unrun.",
|
||||
)
|
||||
return 1
|
||||
|
||||
_write(f"OK: no test file is globbed by a job and then deselected by every -k across {len(_slices())} slices.")
|
||||
return 0
|
||||
|
||||
|
||||
def _check_shards() -> int:
|
||||
findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
|
||||
if findings:
|
||||
_report(
|
||||
"test directories and files that no shard claims",
|
||||
findings,
|
||||
"Add each to the shard it belongs to. A directory that is itself split across "
|
||||
"several shards belongs in SHARDED_ROOTS instead, so its own children get checked.",
|
||||
)
|
||||
return 1
|
||||
|
||||
counted = sum(len(_shard_children(root)) for root in SHARDED_ROOTS if (REPO_ROOT / root).is_dir())
|
||||
_write(f"OK: all {counted} test children across {len(SHARDED_ROOTS)} sharded trees are assigned to a shard.")
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if "--shards" in sys.argv[1:]:
|
||||
return _check_shards()
|
||||
if "--slices" in sys.argv[1:]:
|
||||
return _check_slices()
|
||||
|
||||
allowlist = _load_allowlist()
|
||||
scalars = _all_scalars()
|
||||
|
||||
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
|
||||
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
|
||||
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
|
||||
|
||||
if stale_findings:
|
||||
_report(
|
||||
"allowlist entries that exempt nothing",
|
||||
stale_findings,
|
||||
"Delete each from .github/ci-coverage-allowlist.yml; the file it named is gone or was renamed.",
|
||||
)
|
||||
if test_findings:
|
||||
_report(
|
||||
"test files that no CI job invokes",
|
||||
|
|
@ -248,7 +513,7 @@ def main() -> int:
|
|||
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:
|
||||
if stale_findings or test_findings or dockerfile_findings:
|
||||
return 1
|
||||
|
||||
_write(
|
||||
|
|
|
|||
149
.github/scripts/assert_workflow_dir_hygiene.py
vendored
Normal file
149
.github/scripts/assert_workflow_dir_hygiene.py
vendored
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Three invariants about what lives in .github/workflows/ and what its names mean.
|
||||
|
||||
`.github/workflows/` is a directory GitHub reads, not a place to keep things. Every
|
||||
file at its top level is parsed as a workflow, so a script or a data file parked there
|
||||
is either an invalid workflow or an orphan nobody can find. A subdirectory is not read
|
||||
at all, so helper files may live in one. GitHub accepts both `.yml` and `.yaml`, and
|
||||
this repo spells them `.yml`, which is a naming rule rather than a validity one and is
|
||||
reported separately. And the `_` prefix is the repo's only signal that a workflow is a
|
||||
reusable building block rather than something that runs on its own, which is worth
|
||||
nothing unless it is true both ways.
|
||||
|
||||
WF001 a top-level file in .github/workflows/ that is not a workflow at all
|
||||
WF002 a workflow whose only trigger is `workflow_call` but is not `_`-prefixed
|
||||
WF003 a `_`-prefixed workflow that no other workflow can call
|
||||
WF004 a real workflow spelled `.yaml` where this directory spells them `.yml`
|
||||
|
||||
A workflow with `workflow_call` alongside a human trigger is deliberately dual-mode
|
||||
and belongs under its plain name, so only the call-only ones are held to WF002.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python assert_workflow_dir_hygiene.py
|
||||
|
||||
Exit code 1 if any violation is found.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import yaml
|
||||
|
||||
REPO_ROOT: Final = pathlib.Path(__file__).resolve().parents[2]
|
||||
WORKFLOW_DIR: Final = REPO_ROOT / ".github" / "workflows"
|
||||
SCRIPT_HOME: Final = ".github/scripts/"
|
||||
REUSABLE_PREFIX: Final = "_"
|
||||
CALL_TRIGGER: Final = "workflow_call"
|
||||
CANONICAL_SUFFIX: Final = ".yml"
|
||||
WORKFLOW_SUFFIXES: Final = frozenset((CANONICAL_SUFFIX, ".yaml"))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Finding:
|
||||
subject: str
|
||||
code: str
|
||||
detail: str
|
||||
|
||||
def render(self) -> str:
|
||||
return f" - {self.subject}: {self.code} {self.detail}"
|
||||
|
||||
|
||||
def _triggers(document: object) -> frozenset[str]:
|
||||
if not isinstance(document, dict):
|
||||
return frozenset()
|
||||
raw: Final = document.get("on", document.get(True))
|
||||
if isinstance(raw, str):
|
||||
return frozenset({raw})
|
||||
if isinstance(raw, dict):
|
||||
return frozenset(str(key) for key in raw)
|
||||
if isinstance(raw, list):
|
||||
return frozenset(str(item) for item in raw)
|
||||
return frozenset()
|
||||
|
||||
|
||||
def _workflows(directory: pathlib.Path) -> tuple[pathlib.Path, ...]:
|
||||
return tuple(
|
||||
path
|
||||
for path in sorted(directory.iterdir())
|
||||
if path.is_file() and path.suffix in WORKFLOW_SUFFIXES
|
||||
)
|
||||
|
||||
|
||||
def _strays(directory: pathlib.Path) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
Finding(
|
||||
path.name,
|
||||
"WF001",
|
||||
f"is not a workflow, and GitHub parses every top-level file here as one; "
|
||||
f"move it to {SCRIPT_HOME} or into a subdirectory, which GitHub does not read",
|
||||
)
|
||||
for path in sorted(directory.iterdir())
|
||||
if path.is_file() and path.suffix not in WORKFLOW_SUFFIXES
|
||||
)
|
||||
|
||||
|
||||
def _misspelled(directory: pathlib.Path) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
Finding(
|
||||
path.name,
|
||||
"WF004",
|
||||
f"is a real workflow and GitHub reads it, but this directory spells them "
|
||||
f"{CANONICAL_SUFFIX}; rename it to {path.stem}{CANONICAL_SUFFIX}",
|
||||
)
|
||||
for path in _workflows(directory)
|
||||
if path.suffix != CANONICAL_SUFFIX
|
||||
)
|
||||
|
||||
|
||||
def _misnamed(directory: pathlib.Path) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
finding
|
||||
for path in _workflows(directory)
|
||||
for finding in _naming_findings(path, _triggers(yaml.safe_load(path.read_text(encoding="utf-8"))))
|
||||
)
|
||||
|
||||
|
||||
def _naming_findings(path: pathlib.Path, triggers: frozenset[str]) -> tuple[Finding, ...]:
|
||||
underscored: Final = path.name.startswith(REUSABLE_PREFIX)
|
||||
if triggers == frozenset({CALL_TRIGGER}) and not underscored:
|
||||
return (
|
||||
Finding(
|
||||
path.name,
|
||||
"WF002",
|
||||
f"is only callable by another workflow, so name it {REUSABLE_PREFIX}{path.name}",
|
||||
),
|
||||
)
|
||||
if underscored and CALL_TRIGGER not in triggers:
|
||||
return (
|
||||
Finding(
|
||||
path.name,
|
||||
"WF003",
|
||||
f"is named as a reusable workflow but has no {CALL_TRIGGER} trigger; "
|
||||
"add one or drop the prefix",
|
||||
),
|
||||
)
|
||||
return ()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
findings: Final = _strays(WORKFLOW_DIR) + _misspelled(WORKFLOW_DIR) + _misnamed(WORKFLOW_DIR)
|
||||
if not findings:
|
||||
total: Final = len(_workflows(WORKFLOW_DIR))
|
||||
sys.stdout.write(
|
||||
f"OK: {total} workflows, every file in .github/workflows/ is one, and the "
|
||||
f"{REUSABLE_PREFIX} prefix means callable in both directions.\n"
|
||||
)
|
||||
return 0
|
||||
sys.stdout.write("ERROR: .github/workflows/ holds files that break its own conventions\n")
|
||||
for finding in findings:
|
||||
sys.stdout.write(f"{finding.render()}\n")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
42
.github/scripts/detect_changes.sh
vendored
Executable file
42
.github/scripts/detect_changes.sh
vendored
Executable file
|
|
@ -0,0 +1,42 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
readonly API_FILE_CEILING=3000
|
||||
readonly CATEGORY="${CATEGORY:-backend}"
|
||||
|
||||
decide() {
|
||||
echo "detect-changes[${CATEGORY}]: decision=$1"
|
||||
[ -z "${GITHUB_OUTPUT:-}" ] || echo "decision=$1" >>"${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
}
|
||||
|
||||
run_full() {
|
||||
echo "detect-changes[${CATEGORY}]: $1; running job"
|
||||
decide run
|
||||
}
|
||||
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
classify="${here}/../../.circleci/scripts/classify_changes.sh"
|
||||
|
||||
[ -n "${PR_NUMBER:-}" ] || run_full "not a pull_request event"
|
||||
[ -n "${REPO:-}" ] || run_full "no repository in the environment"
|
||||
|
||||
case "${CHANGED_FILE_COUNT:-}" in
|
||||
'' | *[!0-9]*) run_full "the event payload carries no changed_files count" ;;
|
||||
esac
|
||||
[ "${CHANGED_FILE_COUNT}" -le "${API_FILE_CEILING}" ] ||
|
||||
run_full "PR #${PR_NUMBER} changes ${CHANGED_FILE_COUNT} files, past the ${API_FILE_CEILING}-file listing ceiling"
|
||||
|
||||
changed="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate --jq '.[].filename')" ||
|
||||
run_full "could not list the files on PR #${PR_NUMBER}"
|
||||
[ -n "${changed}" ] || run_full "the API listed no files on PR #${PR_NUMBER}"
|
||||
|
||||
echo "detect-changes[${CATEGORY}]: files changed by PR #${PR_NUMBER}:"
|
||||
printf '%s\n' "${changed}" | sed 's/^/ /'
|
||||
|
||||
decision="$(printf '%s\n' "${changed}" | bash "${classify}" "${CATEGORY}")" ||
|
||||
run_full "classify_changes.sh failed"
|
||||
case "${decision}" in
|
||||
run | skip) decide "${decision}" ;;
|
||||
*) run_full "classify_changes.sh printed an unexpected decision: ${decision}" ;;
|
||||
esac
|
||||
15
.github/scripts/select_ui_test_scope.sh
vendored
Executable file
15
.github/scripts/select_ui_test_scope.sh
vendored
Executable file
|
|
@ -0,0 +1,15 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
has_file=false
|
||||
has_file_outside_src=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
has_file=true
|
||||
case "$file" in
|
||||
src/*) ;;
|
||||
*) has_file_outside_src=true ;;
|
||||
esac
|
||||
done
|
||||
|
||||
{ [ "$has_file" = true ] && [ "$has_file_outside_src" = false ]; } && echo related || echo full
|
||||
557
.github/scripts/triage_rollout_heads_up.py
vendored
557
.github/scripts/triage_rollout_heads_up.py
vendored
|
|
@ -1,557 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""One-shot 7-day heads-up sweep for the Agent Shin rollout.
|
||||
|
||||
Posts a friendly "the OSS triage bot kicks in next Monday" comment on every
|
||||
open external PR/issue that currently *would* fail the new rubric — i.e.,
|
||||
every PR/issue Agent Shin would close once the rollout completes. The point
|
||||
is to give contributors a full week to fix their description before the bot
|
||||
ever takes a destructive action, so nobody is surprised by an auto-close.
|
||||
|
||||
The script is designed to run **exactly once** at rollout, fired by a manual
|
||||
``workflow_dispatch`` (``dry_run=false``) on the heads-up workflow. Re-runs
|
||||
are safe: every comment is stamped with the hidden ``HEADS_UP_MARKER`` and
|
||||
PRs/issues that already carry the marker are skipped.
|
||||
|
||||
Dry-run vs. real run
|
||||
--------------------
|
||||
Defaults to dry-run. Passing ``--close`` flips into real mode. Every GitHub
|
||||
mutation goes through ``_agent_shin_actions``, which has a one-line
|
||||
``if dry_run: log else: do_it`` per call, so the only difference between a
|
||||
dry-run preview and the real run is the call site that actually hits the
|
||||
GitHub API.
|
||||
|
||||
Local preview::
|
||||
|
||||
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm
|
||||
|
||||
Real run (the manual rollout dispatch uses this)::
|
||||
|
||||
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm --close
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# Make the sibling triage_with_llm + _agent_shin_actions importable when this
|
||||
# script is invoked directly (the GitHub workflow does `python3 .github/scripts/...`).
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parent
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from _agent_shin_actions import maybe_post_comment # noqa: E402
|
||||
from agent_shin_shared import ( # noqa: E402
|
||||
AGENT_SHIN_DEFAULT_BOT_LOGIN,
|
||||
ALLOWLIST_LOGINS,
|
||||
list_open_items,
|
||||
)
|
||||
from triage_with_llm import ( # noqa: E402
|
||||
DEFAULT_MODEL,
|
||||
call_llm_judge,
|
||||
fetch_issue,
|
||||
fetch_pr,
|
||||
gh,
|
||||
is_internal_contributor,
|
||||
review_gate,
|
||||
triage,
|
||||
)
|
||||
|
||||
# Hidden marker so re-runs skip PRs/issues we've already notified. Distinct from
|
||||
# the within-grace / ready / regressed markers so it can't be confused with the
|
||||
# steady-state lifecycle comments.
|
||||
HEADS_UP_MARKER = "<!-- agent-shin:rollout-heads-up -->"
|
||||
|
||||
# Placeholder until the litellm-docs PR ships. The rollout blog post explains
|
||||
# the new rubric, the 7-day grace, and how to recover after an auto-close.
|
||||
# TODO(docs): replace with the canonical URL once the litellm-docs PR merges.
|
||||
ROLLOUT_BLOG_URL = "https://docs.litellm.ai/docs/agent_shin_triage_rollout"
|
||||
|
||||
# Default cutoff is one week from "now". Computed at runtime so the wording
|
||||
# stays correct even if the rollout is merged later than planned. The user can
|
||||
# override with --close-on YYYY-MM-DD when running the script manually.
|
||||
DEFAULT_GRACE_DAYS = 7
|
||||
|
||||
# The daily auto-close sweeps (close_low_quality_prs.yml at 09:00 UTC and
|
||||
# review_gate.yml at 09:30 UTC) are what actually close a still-failing item,
|
||||
# so the deadline we promise contributors has to name that wall-clock moment.
|
||||
ACTIVATION_TIME_UTC = "09:00 UTC"
|
||||
|
||||
|
||||
def _format_cutoff(cutoff: dt.date) -> str:
|
||||
"""Human-readable, timezone-explicit cutoff, e.g. ``Monday, June 1, 2026
|
||||
(09:00 UTC)`` — the moment a still-failing PR/issue gets closed."""
|
||||
return (
|
||||
f"{cutoff.strftime('%A, %B')} {cutoff.day}, {cutoff.year} "
|
||||
f"({ACTIVATION_TIME_UTC})"
|
||||
)
|
||||
|
||||
|
||||
def _rubric_section_pr() -> str:
|
||||
return (
|
||||
"**Going forward, every external PR needs ONE of:**\n"
|
||||
"\n"
|
||||
"- A linked GitHub issue using a closing keyword: "
|
||||
"`Fixes #1234`, `Closes #1234`, or `Resolves #1234`, OR\n"
|
||||
"- All three of: a clear **problem description**, **expected vs. "
|
||||
"actual behavior**, and **end-to-end QA proof** (at least one of a "
|
||||
"short screen recording / video, before/after screenshots, or the "
|
||||
"exact commands you ran with their real output; mocked or stubbed "
|
||||
"runs don't count).\n"
|
||||
"\n"
|
||||
"PRs also need a **Greptile confidence score of 4/5 or higher** before "
|
||||
"the bot will tag them `ready for review`. You can `@greptileai` to "
|
||||
"request a fresh review at any time, including after the PR is closed."
|
||||
)
|
||||
|
||||
|
||||
def _rubric_section_issue() -> str:
|
||||
return (
|
||||
"**Going forward, every external issue needs:**\n"
|
||||
"\n"
|
||||
"- For **bug reports**: end-to-end evidence of the bug (at least one "
|
||||
"of a screen recording / video, a screenshot, or the exact commands "
|
||||
"you ran with their real output / traceback) plus expected vs. actual "
|
||||
"behavior. Written steps with no run output don't count, and mocked "
|
||||
"or stubbed runs don't count.\n"
|
||||
"- For **feature requests**: a clear description of the proposed "
|
||||
"feature plus a use case + concrete example (config, API call, UI "
|
||||
"flow, or scenario showing what's blocked today)."
|
||||
)
|
||||
|
||||
|
||||
def _description_only_note(kind: str) -> str:
|
||||
noun = "PR" if kind == "pr" else "issue"
|
||||
return (
|
||||
f"⚠️ **The requirements must live in the {noun} *description*, not in "
|
||||
"comments.** Some PRs/issues collect 100+ comments from humans and "
|
||||
"bots; reading the entire thread on every triage run would balloon "
|
||||
"GitHub API usage (we'd start getting 429'd) and blow out the LLM "
|
||||
"judge's context. The bot only reads the description, so anything "
|
||||
"you add as a comment will be invisible to it."
|
||||
)
|
||||
|
||||
|
||||
def _missing_section(verdict: dict, greptile_score: int | None) -> str:
|
||||
"""Bullet list of what's currently missing on this PR/issue.
|
||||
|
||||
Combines the LLM judge's `missing` list (rubric items) with a Greptile
|
||||
shortfall (for PRs) so the contributor sees one list of things to fix.
|
||||
"""
|
||||
missing = list(verdict.get("missing") or [])
|
||||
if greptile_score is not None and greptile_score < 4:
|
||||
missing.insert(
|
||||
0,
|
||||
f"Greptile's most recent review scored this PR {greptile_score}/5 "
|
||||
"(below the 4/5 bar Agent Shin will require).",
|
||||
)
|
||||
if not missing:
|
||||
return (
|
||||
"_The bot couldn't articulate a specific missing piece; see the "
|
||||
"rubric link above and double-check the description includes all "
|
||||
"of it before the rollout._"
|
||||
)
|
||||
bullets = "\n".join(f"- {m}" for m in missing)
|
||||
return f"**What this one is currently missing:**\n\n{bullets}"
|
||||
|
||||
|
||||
def _recovery_section(kind: str) -> str:
|
||||
if kind == "pr":
|
||||
return (
|
||||
"**If the bot closes this PR after the rollout:** update the "
|
||||
"description with the missing pieces, then either open a fresh "
|
||||
"PR or comment `@agent-shin reconsider` on the closed PR. If "
|
||||
"Greptile re-scores you at 4/5 or higher I'll reopen and tag "
|
||||
"the PR `ready for review`. (`@greptileai` works on closed PRs "
|
||||
"too; a fresh review is one of the signals that lifts you back "
|
||||
"into the queue.) This is **not** us losing interest in your "
|
||||
"change; far from it. We just need open PRs to be a list of "
|
||||
"things a maintainer can act on, so we can get to yours faster."
|
||||
)
|
||||
return (
|
||||
"**If the bot closes this issue after the rollout:** edit the issue "
|
||||
"description to add the missing pieces, then comment `@agent-shin "
|
||||
"reconsider` on the closed issue. I'll re-evaluate and, if the rubric "
|
||||
"is met, reopen it. (GitHub doesn't let external authors reopen an "
|
||||
"issue a maintainer or bot closed, so the comment is the reliable "
|
||||
"path.) This is **not** us saying the bug isn't real or the request "
|
||||
"isn't useful; it's so the remaining open issues are a list of things "
|
||||
"a maintainer can act on."
|
||||
)
|
||||
|
||||
|
||||
def format_heads_up_comment(
|
||||
*, kind: str, verdict: dict, greptile_score: int | None, cutoff: dt.date
|
||||
) -> str:
|
||||
"""Compose the friendly 7-day heads-up comment posted on a failing PR/issue."""
|
||||
noun = "PR" if kind == "pr" else "issue"
|
||||
rubric = _rubric_section_pr() if kind == "pr" else _rubric_section_issue()
|
||||
cutoff_str = _format_cutoff(cutoff)
|
||||
explanation = (verdict.get("explanation") or "").strip()
|
||||
explanation_block = (
|
||||
f"> _(The judge's note for this one: {explanation})_\n\n" if explanation else ""
|
||||
)
|
||||
|
||||
return (
|
||||
"🚅 **Heads-up: we're turning on the OSS triage bot in "
|
||||
f"{DEFAULT_GRACE_DAYS} days, on {cutoff_str}.**\n"
|
||||
"\n"
|
||||
"We're rolling out **Agent Shin**, an LLM-as-judge triage bot for "
|
||||
f"external {noun}s. Once it's live, the bot reads each open "
|
||||
f"{noun}'s description, scores it against a small rubric, and "
|
||||
f"auto-closes any {noun} that's missing the basics, with a single "
|
||||
f"comment explaining what's missing and how to recover. Full "
|
||||
f"context: [Agent Shin rollout blog post]({ROLLOUT_BLOG_URL}).\n"
|
||||
"\n"
|
||||
f"{rubric}\n"
|
||||
"\n"
|
||||
f"{_description_only_note(kind)}\n"
|
||||
"\n"
|
||||
f"{_missing_section(verdict, greptile_score)}\n"
|
||||
"\n"
|
||||
f"{explanation_block}"
|
||||
"**Timeline (you have a week):**\n"
|
||||
"\n"
|
||||
f"- We turn the bot on in {DEFAULT_GRACE_DAYS} days, on "
|
||||
f"**{cutoff_str}**. You have until then to update this {noun}'s "
|
||||
"description with the missing pieces above.\n"
|
||||
f"- If this {noun} still fails the rubric at **{cutoff_str}**, "
|
||||
"we'll close it.\n"
|
||||
f"- From then on the bot runs daily, and every {noun} that fails "
|
||||
"the rubric gets a **2-hour lifetime**: one warning comment, then "
|
||||
"auto-close 2 hours later.\n"
|
||||
"\n"
|
||||
f"{_recovery_section(kind)}\n"
|
||||
"\n"
|
||||
f"{HEADS_UP_MARKER}"
|
||||
)
|
||||
|
||||
|
||||
def _list_open_numbers(repo: str, kind: str) -> list[int]:
|
||||
"""Return every open PR or issue number in ``repo``.
|
||||
|
||||
Delegates to ``list_open_items`` so the full backlog is fetched (no cap)
|
||||
and the `gh {pr,issue} list` invocation stays in one shared place. ``gh
|
||||
issue list`` would include PRs, but ``list_open_items`` uses the dedicated
|
||||
command per kind, so the two never mix.
|
||||
"""
|
||||
return [
|
||||
item["number"] for item in list_open_items(kind, repo=repo, fields="number")
|
||||
]
|
||||
|
||||
|
||||
def _has_heads_up_marker(item: dict) -> bool:
|
||||
"""Cheap fast-path: check the PR/issue body itself for the marker.
|
||||
|
||||
The marker is appended to the *comment* we post, not the body, so this
|
||||
will only fire if the body literally contains the marker text. We still
|
||||
do the comment-marker check separately below; this body check just lets
|
||||
us short-circuit for PRs/issues that quote the marker for any reason.
|
||||
"""
|
||||
body = item.get("body") or ""
|
||||
return HEADS_UP_MARKER in body
|
||||
|
||||
|
||||
def _comments_have_marker(repo: str, number: int) -> bool:
|
||||
"""True if the bot already posted a comment carrying the marker.
|
||||
|
||||
Used for idempotency: a re-run skips items the previous run notified.
|
||||
Filters by author (matching the sibling marker-checks in
|
||||
``triage_with_llm._has_marker`` and
|
||||
``agent_shin_shared.seconds_since_latest_marker_comment``) so a
|
||||
contributor who quotes the heads-up via GitHub's "Quote reply" — which
|
||||
preserves HTML comments in the raw markdown — can't trick the
|
||||
idempotency check into silently skipping a real heads-up.
|
||||
|
||||
Comments live on the unified issues endpoint regardless of whether the
|
||||
item is a PR or an issue, so no ``kind`` argument is required here.
|
||||
"""
|
||||
expected_login = (
|
||||
os.environ.get("AGENT_SHIN_BOT_LOGIN") or AGENT_SHIN_DEFAULT_BOT_LOGIN
|
||||
).lower()
|
||||
raw = gh(
|
||||
"api",
|
||||
"--paginate",
|
||||
f"repos/{repo}/issues/{number}/comments?per_page=100",
|
||||
)
|
||||
for line in raw.splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
payload = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
comments = payload if isinstance(payload, list) else [payload]
|
||||
for comment in comments:
|
||||
author = ((comment.get("user") or {}).get("login") or "").lower()
|
||||
if author != expected_login:
|
||||
continue
|
||||
if HEADS_UP_MARKER in (comment.get("body") or ""):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _evaluate_pr(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
|
||||
"""Run the future PR rubric (review_gate) in dry-run and return the result."""
|
||||
return review_gate(
|
||||
repo=repo,
|
||||
number=number,
|
||||
close=False, # we only want the verdict, never act here
|
||||
model=model,
|
||||
judge=judge,
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_issue(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
|
||||
"""Run the future issue rubric (triage kind='issue') in dry-run."""
|
||||
return triage(
|
||||
repo=repo,
|
||||
kind="issue",
|
||||
number=number,
|
||||
close=False,
|
||||
model=model,
|
||||
judge=judge,
|
||||
)
|
||||
|
||||
|
||||
def _would_be_closed(kind: str, result: dict) -> bool:
|
||||
"""True if the future triage would auto-close this PR/issue based on the
|
||||
rubric (regardless of grace-period gating).
|
||||
|
||||
For PRs we trust ``review_gate``'s ``passing`` field — it combines the LLM
|
||||
verdict and the Greptile score. For issues we read the LLM verdict
|
||||
directly. Both fields are ``None``/missing on skip paths
|
||||
(skip-internal-author, skip-llm-error, etc.) where the future bot would
|
||||
NOT close the item — those return False.
|
||||
"""
|
||||
if kind == "pr":
|
||||
passing = result.get("passing")
|
||||
if passing is None:
|
||||
return False # skipped — nothing for the heads-up to warn about
|
||||
return passing is False
|
||||
verdict = result.get("verdict") or {}
|
||||
return (verdict.get("verdict") or "").lower() == "fail"
|
||||
|
||||
|
||||
def _process_one(
|
||||
*,
|
||||
repo: str,
|
||||
kind: str,
|
||||
number: int,
|
||||
model: str,
|
||||
cutoff: dt.date,
|
||||
dry_run: bool,
|
||||
judge: Any = None,
|
||||
skip_marker_check: bool = False,
|
||||
allowlist: frozenset[str] = ALLOWLIST_LOGINS,
|
||||
) -> dict:
|
||||
"""Evaluate one PR/issue and post a heads-up if it would be auto-closed.
|
||||
|
||||
Returns a per-item dict for the summary table.
|
||||
"""
|
||||
base = {"kind": kind, "number": number}
|
||||
fetcher = fetch_pr if kind == "pr" else fetch_issue
|
||||
item = fetcher(repo, number)
|
||||
|
||||
if (item.get("state") or "") != "open":
|
||||
return {**base, "action": "skip-not-open"}
|
||||
if allowlist:
|
||||
login = (item.get("user") or {}).get("login") or ""
|
||||
if login.lower() not in allowlist:
|
||||
return {**base, "action": "skip-not-allowlisted"}
|
||||
elif is_internal_contributor(item):
|
||||
return {**base, "action": "skip-internal-author"}
|
||||
if not skip_marker_check and _has_heads_up_marker(item):
|
||||
return {**base, "action": "skip-already-marked-in-body"}
|
||||
if not skip_marker_check and _comments_have_marker(repo, number):
|
||||
return {**base, "action": "skip-already-notified"}
|
||||
|
||||
if kind == "pr":
|
||||
result = _evaluate_pr(repo=repo, number=number, model=model, judge=judge)
|
||||
else:
|
||||
result = _evaluate_issue(repo=repo, number=number, model=model, judge=judge)
|
||||
|
||||
if not _would_be_closed(kind, result):
|
||||
return {**base, "action": "skip-passing", "evaluator": result.get("action")}
|
||||
|
||||
verdict = result.get("verdict") or {}
|
||||
greptile_score = result.get("greptile_score") if kind == "pr" else None
|
||||
comment = format_heads_up_comment(
|
||||
kind=kind, verdict=verdict, greptile_score=greptile_score, cutoff=cutoff
|
||||
)
|
||||
maybe_post_comment(repo, number, comment, dry_run=dry_run)
|
||||
return {
|
||||
**base,
|
||||
"action": "heads-up-posted" if not dry_run else "would-post-heads-up",
|
||||
"verdict": (verdict.get("verdict") or "").lower(),
|
||||
"greptile_score": greptile_score,
|
||||
}
|
||||
|
||||
|
||||
def _print_summary(results: list[dict]) -> None:
|
||||
"""Tally per-action counts so a dry-run preview tells you at a glance how
|
||||
many comments the real run would post."""
|
||||
counts: dict[str, int] = {}
|
||||
for r in results:
|
||||
counts[r["action"]] = counts.get(r["action"], 0) + 1
|
||||
print("\n=== rollout heads-up summary ===")
|
||||
for action in sorted(counts):
|
||||
print(f" {action:35s} {counts[action]}")
|
||||
print(f" total {len(results)}")
|
||||
|
||||
|
||||
def run(
|
||||
*,
|
||||
repo: str,
|
||||
close: bool,
|
||||
cutoff: dt.date,
|
||||
model: str,
|
||||
kinds: tuple[str, ...] = ("pr", "issue"),
|
||||
judge: Any = None,
|
||||
only_numbers: dict[str, list[int]] | None = None,
|
||||
skip_marker_check: bool = False,
|
||||
) -> list[dict]:
|
||||
"""Sweep ``repo`` and post heads-up comments. Returns the per-item results."""
|
||||
dry_run = not close
|
||||
if dry_run:
|
||||
print(
|
||||
f"[DRY RUN] sweeping {repo}; --close not passed, no comments will be posted."
|
||||
)
|
||||
else:
|
||||
print(f"[REAL RUN] sweeping {repo}; comments WILL be posted.")
|
||||
print(f"Cutoff date in comment body: {cutoff.isoformat()}")
|
||||
|
||||
results: list[dict] = []
|
||||
for kind in kinds:
|
||||
if only_numbers and kind in only_numbers:
|
||||
numbers = list(only_numbers[kind])
|
||||
else:
|
||||
numbers = _list_open_numbers(repo, kind)
|
||||
print(f"\n--- {kind}s: {len(numbers)} open ---")
|
||||
for n in numbers:
|
||||
try:
|
||||
result = _process_one(
|
||||
repo=repo,
|
||||
kind=kind,
|
||||
number=n,
|
||||
model=model,
|
||||
cutoff=cutoff,
|
||||
dry_run=dry_run,
|
||||
judge=judge,
|
||||
skip_marker_check=skip_marker_check,
|
||||
)
|
||||
except (
|
||||
Exception
|
||||
) as exc: # noqa: BLE001 - per-item errors don't abort the sweep
|
||||
result = {
|
||||
"kind": kind,
|
||||
"number": n,
|
||||
"action": "error",
|
||||
"error": str(exc),
|
||||
}
|
||||
print(f"!! {kind}#{n}: {exc}", file=sys.stderr)
|
||||
print(f" {kind}#{n}: {result['action']}")
|
||||
results.append(result)
|
||||
_print_summary(results)
|
||||
return results
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--repo", required=True, help="owner/repo")
|
||||
parser.add_argument(
|
||||
"--close",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Actually post comments. Without this flag the script is in "
|
||||
"dry-run mode and only logs what it would do."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--close-on",
|
||||
type=dt.date.fromisoformat,
|
||||
default=None,
|
||||
help=(
|
||||
"Cutoff date shown in the heads-up comment as the rollout date "
|
||||
f"(default: today + {DEFAULT_GRACE_DAYS} days)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL,
|
||||
help=f"Model for the rubric LLM judge (default: {DEFAULT_MODEL}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kind",
|
||||
choices=("pr", "issue", "both"),
|
||||
default="both",
|
||||
help="Restrict the sweep to PRs or issues only (default: both).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-pr",
|
||||
type=int,
|
||||
action="append",
|
||||
default=[],
|
||||
help="Limit the PR sweep to these PR numbers (repeat for several).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-issue",
|
||||
type=int,
|
||||
action="append",
|
||||
default=[],
|
||||
help="Limit the issue sweep to these issue numbers (repeat for several).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ignore-existing-marker",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Re-post on PRs/issues that already carry the heads-up marker. "
|
||||
"Useful for testing the comment wording on a known PR."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
cutoff = args.close_on or (
|
||||
dt.datetime.now(dt.timezone.utc).date() + dt.timedelta(days=DEFAULT_GRACE_DAYS)
|
||||
)
|
||||
|
||||
kinds: tuple[str, ...]
|
||||
if args.kind == "pr":
|
||||
kinds = ("pr",)
|
||||
elif args.kind == "issue":
|
||||
kinds = ("issue",)
|
||||
else:
|
||||
kinds = ("pr", "issue")
|
||||
|
||||
only: dict[str, list[int]] = {}
|
||||
if args.only_pr:
|
||||
only["pr"] = args.only_pr
|
||||
if args.only_issue:
|
||||
only["issue"] = args.only_issue
|
||||
|
||||
# The script must NOT hit the LLM in dry-run if no key is set — we still
|
||||
# want a useful preview that says "skip-no-llm-key" for items that would
|
||||
# have been judged. Production runs require OPENAI_API_KEY.
|
||||
if args.close and not os.environ.get("OPENAI_API_KEY"):
|
||||
parser.error("OPENAI_API_KEY must be set for --close (real-run) mode.")
|
||||
|
||||
run(
|
||||
repo=args.repo,
|
||||
close=args.close,
|
||||
cutoff=cutoff,
|
||||
model=args.model,
|
||||
kinds=kinds,
|
||||
only_numbers=only or None,
|
||||
skip_marker_check=args.ignore_existing_marker,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
17
.github/workflows/_test-unit-base.yml
vendored
17
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -60,6 +60,9 @@ jobs:
|
|||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.job-timeout-minutes }}
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
outputs:
|
||||
decision: ${{ steps.changes.outputs.decision }}
|
||||
|
||||
|
|
@ -69,24 +72,27 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect backend-relevant changes
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
timeout-minutes: 2
|
||||
uses: ./.github/actions/detect-backend-changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
|
|
@ -123,6 +129,13 @@ jobs:
|
|||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
# coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has.
|
||||
# It is only the default from Python 3.14, and these shards run 3.12, so it
|
||||
# has to be asked for. Coverage refuses it when branch measurement is on
|
||||
# (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with
|
||||
# a `no-sysmon` warning, so turning on `branch = true` here means giving this
|
||||
# back until the runners move to 3.14.
|
||||
COVERAGE_CORE: sysmon
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
- name: Update JSON Data
|
||||
run: |
|
||||
uv run --frozen --with 'aiohttp==3.13.3' python ".github/workflows/auto_update_price_and_context_window_file.py"
|
||||
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
|
||||
- name: Regenerate JSON Schema
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
|
|
|
|||
9
.github/workflows/ci-coverage.yml
vendored
9
.github/workflows/ci-coverage.yml
vendored
|
|
@ -40,3 +40,12 @@ jobs:
|
|||
run: |
|
||||
python -m pip install "pyyaml==6.0.3"
|
||||
python .github/scripts/assert_ci_coverage.py
|
||||
|
||||
# The census asks whether a job names a file; this asks whether that job's -k
|
||||
# then throws it back out. A file both globbed and deselected everywhere runs
|
||||
# nowhere while counting as covered, which is how the caching suite went unrun.
|
||||
- name: Assert no -k expression deselects a file from every job that globs it
|
||||
run: python .github/scripts/assert_ci_coverage.py --slices
|
||||
|
||||
- name: Assert .github/workflows/ holds only workflows, correctly named
|
||||
run: python .github/scripts/assert_workflow_dir_hygiene.py
|
||||
|
|
|
|||
35
.github/workflows/test-linting.yml
vendored
35
.github/workflows/test-linting.yml
vendored
|
|
@ -24,6 +24,7 @@ jobs:
|
|||
# re-running basedpyright over the merge-base tree.
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
actions: read
|
||||
|
||||
steps:
|
||||
|
|
@ -37,7 +38,12 @@ jobs:
|
|||
clean: true
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Fetch gate base (merge-base with target branch)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
|
|
@ -50,39 +56,47 @@ jobs:
|
|||
echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Clean Python cache
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
find . -type d -name "__pycache__" -exec rm -rf {} + || true
|
||||
find . -name "*.pyc" -delete || true
|
||||
|
||||
- name: Check uv.lock is up to date
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv lock --check || (echo "❌ uv.lock is out of sync with pyproject.toml. Run 'uv lock' locally and commit the result." && exit 1)
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv sync --frozen --group proxy-dev --group e2e-dev
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
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
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Check ruff format
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
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
|
||||
|
|
@ -92,6 +106,7 @@ jobs:
|
|||
xargs uv run --no-sync ruff format --check --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt"
|
||||
|
||||
- name: Debug - Check file state
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
echo "Current branch:"
|
||||
git branch --show-current
|
||||
|
|
@ -101,30 +116,46 @@ jobs:
|
|||
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
|
||||
|
||||
- name: Run Ruff linting
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync ruff check .
|
||||
cd ..
|
||||
|
||||
- name: Run Ruff linting (test tree)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync ruff check --config ruff-tests.toml tests
|
||||
|
||||
- name: Check strict-rule budget (delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
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)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes, litellm global mutation, credential-gated skips, conftest snapshot inventory, delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Print OpenAI version
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
|
||||
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check tests/e2e basedpyright (zero errors)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
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
|
||||
|
|
@ -133,12 +164,14 @@ jobs:
|
|||
fi
|
||||
|
||||
- name: Check for circular imports
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync python ../tests/documentation_tests/test_circular_imports.py
|
||||
cd ..
|
||||
|
||||
- name: Check import safety
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
|
||||
|
||||
|
|
@ -200,7 +233,7 @@ jobs:
|
|||
|
||||
- name: Run secret scan test
|
||||
run: |
|
||||
uv run --no-project --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/code_coverage_tests/test_no_hardcoded_secrets.py -v
|
||||
|
||||
- name: Run ggshield secret scan
|
||||
env:
|
||||
|
|
|
|||
10
.github/workflows/test-litellm-ui-build.yml
vendored
10
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -1,6 +1,7 @@
|
|||
name: UI Build Check
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
|
|
@ -28,7 +29,14 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
with:
|
||||
category: ui
|
||||
|
||||
- name: Setup Node.js
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
|
|
@ -36,7 +44,9 @@ jobs:
|
|||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm run build
|
||||
|
|
|
|||
56
.github/workflows/test-litellm-ui-unit.yml
vendored
56
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -1,6 +1,7 @@
|
|||
name: UI Unit Tests
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
|
|
@ -32,7 +33,14 @@ jobs:
|
|||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
with:
|
||||
category: ui
|
||||
|
||||
- name: Setup Node.js
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
|
|
@ -40,36 +48,50 @@ jobs:
|
|||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm ci
|
||||
|
||||
- name: Run UI type tests (Vitest)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
CI: "true"
|
||||
run: npm run test:types
|
||||
|
||||
- name: Run UI unit tests (Vitest)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
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
|
||||
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
|
||||
full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; }
|
||||
|
||||
if [ -z "$BASE_SHA" ]; then
|
||||
echo "Push to $GITHUB_REF_NAME: running the full suite"
|
||||
npm run test -- --run --pool forks --poolOptions.forks.maxForks=14
|
||||
full_suite
|
||||
exit 0
|
||||
fi
|
||||
|
||||
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
|
||||
|
||||
scope=$(printf '%s\n' "${changed_files[@]}" | bash "$GITHUB_WORKSPACE/.github/scripts/select_ui_test_scope.sh")
|
||||
if [ "$scope" != related ]; then
|
||||
echo "Pull request: ${#changed_files[@]} changed UI files reach outside src/, so related would miss their dependents; running the full suite"
|
||||
full_suite
|
||||
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
|
||||
|
|
|
|||
9
.github/workflows/test-mcp.yml
vendored
9
.github/workflows/test-mcp.yml
vendored
|
|
@ -10,6 +10,7 @@ on:
|
|||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
|
|
@ -25,26 +26,34 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Thank You Message
|
||||
run: |
|
||||
echo "### 🙏 Thank you for contributing to LiteLLM!" >> $GITHUB_STEP_SUMMARY
|
||||
echo "Your PR is being tested now. We appreciate your help in making LiteLLM better!" >> $GITHUB_STEP_SUMMARY
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv lock --check
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
|
||||
|
||||
- name: Run MCP tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov-report=xml --durations=5
|
||||
|
|
|
|||
31
.github/workflows/test-unit-core-utils.yml
vendored
31
.github/workflows/test-unit-core-utils.yml
vendored
|
|
@ -1,31 +0,0 @@
|
|||
name: "Unit Tests: Core Utilities"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
core-utils:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/litellm_core_utils"
|
||||
workers: 2
|
||||
reruns: 1
|
||||
artifact-name: core-utils
|
||||
15
.github/workflows/test-unit-documentation.yml
vendored
15
.github/workflows/test-unit-documentation.yml
vendored
|
|
@ -23,34 +23,41 @@ jobs:
|
|||
documentation:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Checkout litellm-docs into docs/my-website (for documentation_tests)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
repository: BerriAI/litellm-docs
|
||||
path: docs/my-website
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect backend-relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-backend-changes
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
|
|||
|
|
@ -1,35 +0,0 @@
|
|||
name: "Unit Tests: Enterprise, Google GenAI & Routing"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
enterprise-routing:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: >-
|
||||
tests/test_litellm/enterprise
|
||||
tests/test_litellm/google_genai
|
||||
tests/test_litellm/router_utils
|
||||
tests/test_litellm/router_strategy
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: enterprise-routing
|
||||
31
.github/workflows/test-unit-integrations.yml
vendored
31
.github/workflows/test-unit-integrations.yml
vendored
|
|
@ -1,31 +0,0 @@
|
|||
name: "Unit Tests: Integrations (Callbacks & Logging)"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
integrations:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/integrations"
|
||||
workers: 2
|
||||
reruns: 3
|
||||
artifact-name: integrations
|
||||
47
.github/workflows/test-unit-llm-providers.yml
vendored
47
.github/workflows/test-unit-llm-providers.yml
vendored
|
|
@ -1,47 +0,0 @@
|
|||
name: "Unit Tests: LLM Provider Transformations"
|
||||
|
||||
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:
|
||||
vertex-ai:
|
||||
name: Vertex AI
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/llms/vertex_ai"
|
||||
workers: 1
|
||||
reruns: 2
|
||||
artifact-name: llm-vertex-ai
|
||||
|
||||
other-providers:
|
||||
name: All Other Providers
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: llm-other-providers
|
||||
53
.github/workflows/test-unit-misc.yml
vendored
53
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -1,53 +0,0 @@
|
|||
name: "Unit Tests: MCP, Secrets, Containers & Misc"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
misc:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: >-
|
||||
tests/test_litellm/batches
|
||||
tests/test_litellm/secret_managers
|
||||
tests/test_litellm/a2a_protocol
|
||||
tests/test_litellm/anthropic_interface
|
||||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/compression
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
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
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: misc
|
||||
31
.github/workflows/test-unit-proxy-auth.yml
vendored
31
.github/workflows/test-unit-proxy-auth.yml
vendored
|
|
@ -1,31 +0,0 @@
|
|||
name: "Unit Tests: Proxy Auth & Key Management"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-auth:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/proxy/auth tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine tests/test_litellm/proxy/client"
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: proxy-auth
|
||||
36
.github/workflows/test-unit-proxy-db.yml
vendored
36
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -42,11 +42,10 @@ concurrency:
|
|||
# pinning the whole file to one worker (the default --dist=loadscope
|
||||
# behavior for single-file targets).
|
||||
jobs:
|
||||
# Fast guard — fails the workflow if a test_*.py file under
|
||||
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
|
||||
# The semantic-shard design (no catch-all "remaining" bucket) relies on
|
||||
# every test file being explicitly assigned; this guard prevents a new
|
||||
# file from silently dropping out of CI.
|
||||
# Fast guard — fails the workflow when a test directory or file inside a sharded
|
||||
# tree is claimed by no shard. The semantic-shard design has no catch-all bucket,
|
||||
# so an unassigned child runs nowhere; assert_ci_coverage.py holds the tree list
|
||||
# and reads the same test-path keys the coverage census does.
|
||||
assert-shard-coverage:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 2
|
||||
|
|
@ -56,31 +55,8 @@ jobs:
|
|||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Assert every test_*.py is in a matrix shard
|
||||
run: |
|
||||
python3 - <<'PY'
|
||||
import pathlib, sys, yaml
|
||||
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
|
||||
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
|
||||
referenced = set()
|
||||
for entry in matrix:
|
||||
for token in entry["test-path"].split():
|
||||
if token.startswith("tests/proxy_unit_tests/"):
|
||||
referenced.add(pathlib.PurePosixPath(token).name)
|
||||
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
|
||||
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
|
||||
and p.name != "test_configs"}
|
||||
orphans = sorted(actual - referenced)
|
||||
if orphans:
|
||||
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
|
||||
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
|
||||
for o in orphans:
|
||||
print(f" - {o}")
|
||||
print()
|
||||
print("Add each to whichever semantic shard it belongs to.")
|
||||
sys.exit(1)
|
||||
print(f"OK: all {len(actual)} files assigned to a shard.")
|
||||
PY
|
||||
- name: Assert every test directory and file is claimed by a shard
|
||||
run: python3 .github/scripts/assert_ci_coverage.py --shards
|
||||
|
||||
proxy-db:
|
||||
needs: assert-shard-coverage
|
||||
|
|
|
|||
81
.github/workflows/test-unit-proxy-endpoints.yml
vendored
81
.github/workflows/test-unit-proxy-endpoints.yml
vendored
|
|
@ -1,81 +0,0 @@
|
|||
name: "Unit Tests: Proxy API Endpoints"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-endpoints:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
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/ocr_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
|
||||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
tests/test_litellm/proxy/utils
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: proxy-endpoints
|
||||
|
||||
# Behavior-pinning tests for litellm/proxy/proxy_server.py. Owns its
|
||||
# own job (not a path on the proxy-endpoints job above) so its budget
|
||||
# is independent and its coverage artifact is uploaded separately.
|
||||
# See: https://www.notion.so/36c43b8acdab81ee845fd5365128a2fc
|
||||
proxy-server:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: tests/test_litellm/proxy/proxy_server
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 60
|
||||
job-timeout-minutes: 95
|
||||
artifact-name: proxy-server
|
||||
42
.github/workflows/test-unit-proxy-infra.yml
vendored
42
.github/workflows/test-unit-proxy-infra.yml
vendored
|
|
@ -1,42 +0,0 @@
|
|||
name: "Unit Tests: Proxy Infrastructure"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-infra:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/db
|
||||
tests/test_litellm/proxy/middleware
|
||||
tests/test_litellm/proxy/spend_tracking
|
||||
tests/test_litellm/proxy/pass_through_endpoints
|
||||
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
|
||||
reruns: 2
|
||||
artifact-name: proxy-infra
|
||||
|
|
@ -1,31 +0,0 @@
|
|||
name: "Unit Tests: Responses, Caching & Types"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
responses-caching-types:
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: "tests/test_litellm/responses tests/test_litellm/caching tests/test_litellm/types"
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: responses-caching-types
|
||||
220
.github/workflows/test-unit.yml
vendored
Normal file
220
.github/workflows/test-unit.yml
vendored
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
name: "Unit Tests"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
# One caller for every tests/test_litellm shard, replacing the nine thin workflow
|
||||
# files that each wrapped a single call to _test-unit-base.yml. Adding a shard is
|
||||
# now one matrix entry rather than a new file.
|
||||
#
|
||||
# `name` is the shard id and nothing else, so each check reports as
|
||||
# "<shard> / Run tests" exactly as it did when the shard had its own file. Those
|
||||
# strings are the branch ruleset's required contexts, so they are load-bearing:
|
||||
# renaming an entry renames a required check and the ruleset stops matching it.
|
||||
#
|
||||
# Every entry states its timeouts even when they equal the base workflow's
|
||||
# defaults. An absent matrix key renders as an empty string, which is not a
|
||||
# number, so a partially-specified entry would fail the call rather than fall
|
||||
# back to the default.
|
||||
#
|
||||
# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is
|
||||
# already a matrix and carries a shard-coverage guard that reads that file by
|
||||
# name. Folding it in here is a follow-up, together with generalising that guard
|
||||
# into assert_ci_coverage.py.
|
||||
jobs:
|
||||
unit:
|
||||
name: ${{ matrix.shard }}
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- shard: core-utils
|
||||
artifact-name: core-utils
|
||||
test-path: "tests/test_litellm/litellm_core_utils"
|
||||
workers: 2
|
||||
reruns: 1
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: enterprise-routing
|
||||
artifact-name: enterprise-routing
|
||||
test-path: >-
|
||||
tests/test_litellm/enterprise
|
||||
tests/test_litellm/google_genai
|
||||
tests/test_litellm/router_utils
|
||||
tests/test_litellm/router_strategy
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: integrations
|
||||
artifact-name: integrations
|
||||
test-path: "tests/test_litellm/integrations"
|
||||
workers: 2
|
||||
reruns: 3
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: Vertex AI
|
||||
artifact-name: llm-vertex-ai
|
||||
test-path: "tests/test_litellm/llms/vertex_ai"
|
||||
workers: 1
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: All Other Providers
|
||||
artifact-name: llm-other-providers
|
||||
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: misc
|
||||
artifact-name: misc
|
||||
test-path: >-
|
||||
tests/test_litellm/batches
|
||||
tests/test_litellm/secret_managers
|
||||
tests/test_litellm/a2a_protocol
|
||||
tests/test_litellm/anthropic_interface
|
||||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/compression
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
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/rust_bridge
|
||||
tests/test_litellm/sandbox
|
||||
tests/test_litellm/test_router
|
||||
tests/test_litellm/vector_stores
|
||||
tests/test_litellm/videos
|
||||
tests/test_litellm/test_*.py
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: proxy-auth
|
||||
artifact-name: proxy-auth
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/auth
|
||||
tests/test_litellm/proxy/hooks
|
||||
tests/test_litellm/proxy/policy_engine
|
||||
tests/test_litellm/proxy/client
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: proxy-endpoints
|
||||
artifact-name: proxy-endpoints
|
||||
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/ocr_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
|
||||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
tests/test_litellm/proxy/utils
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: proxy-server
|
||||
artifact-name: proxy-server
|
||||
test-path: "tests/test_litellm/proxy/proxy_server"
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 60
|
||||
job-timeout-minutes: 95
|
||||
|
||||
- shard: proxy-infra
|
||||
artifact-name: proxy-infra
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/db
|
||||
tests/test_litellm/proxy/middleware
|
||||
tests/test_litellm/proxy/spend_tracking
|
||||
tests/test_litellm/proxy/pass_through_endpoints
|
||||
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
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
|
||||
- shard: responses-caching-types
|
||||
artifact-name: responses-caching-types
|
||||
test-path: >-
|
||||
tests/test_litellm/responses
|
||||
tests/test_litellm/caching
|
||||
tests/test_litellm/types
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: ${{ matrix.test-path }}
|
||||
workers: ${{ matrix.workers }}
|
||||
reruns: ${{ matrix.reruns }}
|
||||
timeout-minutes: ${{ matrix.timeout-minutes }}
|
||||
job-timeout-minutes: ${{ matrix.job-timeout-minutes }}
|
||||
artifact-name: ${{ matrix.artifact-name }}
|
||||
92
.github/workflows/triage_rollout_heads_up.yml
vendored
92
.github/workflows/triage_rollout_heads_up.yml
vendored
|
|
@ -1,92 +0,0 @@
|
|||
name: Agent Shin — rollout heads-up (one-shot)
|
||||
|
||||
# Fires the 7-day heads-up comment on every open external PR/issue that the
|
||||
# new triage bot would auto-close. The real sweep is a deliberate one-shot:
|
||||
# trigger it at rollout via a manual `workflow_dispatch` with `dry_run=false`.
|
||||
# The script is idempotent (skips items that already carry the
|
||||
# `<!-- agent-shin:rollout-heads-up -->` marker), so a re-run is harmless.
|
||||
#
|
||||
# The automatic push trigger runs DRY-RUN only, so merging the script to
|
||||
# `litellm_internal_staging` never posts a comment; it just confirms the
|
||||
# workflow is wired up. Posting real comments requires the manual dispatch,
|
||||
# which is also the only trigger that exposes `OPENAI_API_KEY`. The heads-up
|
||||
# is intentionally NOT gated on `AGENT_SHIN_ENABLED`: it has to warn
|
||||
# contributors while that flag is still off, ahead of the flip that turns on
|
||||
# auto-closing.
|
||||
#
|
||||
# The workflow is a thin shell over `.github/scripts/triage_rollout_heads_up.py`.
|
||||
# Dry-run vs. real run differ in EXACTLY one CLI flag (`--close`), added only
|
||||
# on a manual dispatch with `dry_run=false`.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- litellm_internal_staging
|
||||
paths:
|
||||
# The presence of this script on staging IS the rollout merge marker.
|
||||
# Editing the file later would re-fire the workflow; that's safe because
|
||||
# the script skips PRs/issues that already have the heads-up marker.
|
||||
- ".github/scripts/triage_rollout_heads_up.py"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
dry_run:
|
||||
description: "Dry run (true = preview only, false = actually post comments)."
|
||||
required: false
|
||||
default: "true"
|
||||
type: choice
|
||||
options:
|
||||
- "true"
|
||||
- "false"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
heads-up:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout triage scripts
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
sparse-checkout: .github/scripts
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Install LLM client
|
||||
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
|
||||
|
||||
- name: Run heads-up sweep
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Only the manual dispatch (the real-run trigger) needs the LLM key.
|
||||
# The automatic push trigger runs dry-run and never posts, so it gets
|
||||
# no key. Mirrors the sibling triage workflows, which expose the key
|
||||
# only on an enabled/dispatched run rather than unconditionally.
|
||||
OPENAI_API_KEY: ${{ github.event_name == 'workflow_dispatch' && secrets.OPENAI_API_KEY || '' }}
|
||||
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
|
||||
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
|
||||
# The real run is a deliberate manual dispatch with dry_run=false.
|
||||
# Use the EXACT "false" comparison so any unexpected input value
|
||||
# fail-closes to dry-run (mirrors the AGENT_SHIN_ENABLED pattern in
|
||||
# the sibling workflows). The automatic push trigger always stays
|
||||
# dry-run, so merging the script never posts.
|
||||
DRY_RUN_INPUT: ${{ github.event.inputs.dry_run }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
ARGS=(--repo "${{ github.repository }}")
|
||||
if [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${DRY_RUN_INPUT:-true}" = "false" ]; then
|
||||
ARGS+=(--close)
|
||||
echo "::notice::Manual rollout dispatch with dry_run=false -> heads-up comments WILL be posted."
|
||||
elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ]; then
|
||||
echo "::notice::Manual dispatch in dry-run mode -> previewing only, no comments will be posted."
|
||||
else
|
||||
echo "::notice::Automatic push trigger -> dry-run preview only. Fire the real rollout sweep with a manual workflow_dispatch (dry_run=false)."
|
||||
fi
|
||||
python3 .github/scripts/triage_rollout_heads_up.py "${ARGS[@]}"
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -1,9 +1,11 @@
|
|||
.python-version
|
||||
.venv
|
||||
tests/e2e/.fixtures/
|
||||
.venv-typecheck
|
||||
.venv_policy_test
|
||||
.env
|
||||
.claude
|
||||
CLAUDE.local.md
|
||||
.newenv
|
||||
newenv/*
|
||||
litellm/proxy/myenv/*
|
||||
|
|
|
|||
|
|
@ -13,8 +13,8 @@ Here are the core requirements for any PR submitted to LiteLLM:
|
|||
|
||||
- [ ] **Add testing** - Adding at least 1 test is a hard requirement - [see details](#adding-testing)
|
||||
- [ ] **Ensure your PR passes all checks**:
|
||||
- [ ] [Unit Tests](#running-unit-tests) - `make test-unit`
|
||||
- [ ] [Linting / Formatting](#running-linting-and-formatting-checks) - `make lint`
|
||||
- [ ] [The tests covering your change](#running-unit-tests) pass, e.g. `uv run pytest tests/test_litellm/<your_test_file>.py -v`. CI runs the full unit test matrix, so you don't need to run the whole suite locally
|
||||
|
||||
#### UI PRs
|
||||
|
||||
|
|
@ -71,8 +71,8 @@ make format
|
|||
# Run all linting checks (matches CI exactly)
|
||||
make lint
|
||||
|
||||
# Run unit tests to ensure nothing is broken
|
||||
make test-unit
|
||||
# Run the tests covering your change (CI runs the full suite)
|
||||
uv run pytest tests/test_litellm/<your_test_file>.py -v
|
||||
|
||||
# Commit your changes (must follow Conventional Commits — see above)
|
||||
git add .
|
||||
|
|
@ -123,12 +123,13 @@ def test_your_feature():
|
|||
|
||||
### Running Unit Tests
|
||||
|
||||
Run all unit tests (uses parallel execution for speed):
|
||||
|
||||
Run the tests covering your change:
|
||||
```bash
|
||||
make test-unit
|
||||
uv run pytest tests/test_litellm/test_your_file.py -v
|
||||
```
|
||||
|
||||
`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets) on beefier boxes, so if, for whatever reason, you must run the whole suite, it's better to rely on CI to do that.
|
||||
|
||||
If you're running broader test suites, proxy tests, or anything that touches PostgreSQL-backed fixtures/plugins, install the full local test environment first:
|
||||
|
||||
```bash
|
||||
|
|
@ -137,11 +138,6 @@ make install-test-deps
|
|||
|
||||
This syncs the locked test environment used across the repo, including `psycopg` v3 plus `psycopg-binary` (used by `pytest-postgresql`), `psycopg2-binary` (used by some proxy E2E tests), and a generated Prisma client for DB-backed proxy tests, so pytest startup matches CI without manual package installs.
|
||||
|
||||
Run specific test files:
|
||||
```bash
|
||||
uv run pytest tests/test_litellm/test_your_file.py -v
|
||||
```
|
||||
|
||||
### Running Linting and Formatting Checks
|
||||
|
||||
Run all linting checks (matches CI exactly):
|
||||
|
|
|
|||
30
Makefile
30
Makefile
|
|
@ -7,6 +7,7 @@
|
|||
info lint lint-inner lint-dev lint-checks format \
|
||||
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 \
|
||||
lint-test-quality lint-test-quality-budget-update \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
|
||||
lint-install lint-fetch-base bootstrap
|
||||
|
|
@ -35,7 +36,8 @@ help:
|
|||
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
|
||||
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
|
||||
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
|
||||
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + basedpyright)"
|
||||
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
|
||||
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
|
||||
@echo " make check-circular-imports - Check for circular imports"
|
||||
@echo " make check-import-safety - Check import safety"
|
||||
@echo " make test - Run all tests"
|
||||
|
|
@ -142,11 +144,13 @@ lint-install:
|
|||
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
|
||||
$(UV_RUN) python scripts/prisma_generate_if_needed.py
|
||||
|
||||
# Diff-scoped format check, identical to test-linting.yml's "Check ruff format" step:
|
||||
# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
|
||||
# only the litellm Python files changed vs the base are checked, so a pre-existing
|
||||
# format issue elsewhere doesn't block an unrelated commit.
|
||||
# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
|
||||
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
|
||||
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
|
||||
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
||||
@files=$$(git diff --name-only origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \
|
||||
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
|
||||
if [ -z "$$files" ]; then \
|
||||
echo "No changed litellm Python files to format-check."; \
|
||||
else \
|
||||
|
|
@ -156,6 +160,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
|||
# Linting targets
|
||||
lint-ruff: $(LINT_DEP_INSTALL)
|
||||
cd litellm && $(UV_RUN) ruff check . && cd ..
|
||||
$(UV_RUN) ruff check --config ruff-tests.toml tests
|
||||
|
||||
# faster linter for developing ...
|
||||
# inspiration from:
|
||||
|
|
@ -200,6 +205,12 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
|
|||
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
||||
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
|
||||
# litellm module-global mutation, credential-gated skips, conftest snapshot
|
||||
# inventory), counted across tests/ the same delta-vs-base way.
|
||||
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
||||
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
# --update lowers each limit by what this branch fixed since its branch point, so
|
||||
# it needs the base ref fetched to resolve the merge-base.
|
||||
lint-basedpyright-budget-update: install-dev lint-fetch-base
|
||||
|
|
@ -221,8 +232,11 @@ lint-ruff-budget-update: install-dev lint-fetch-base
|
|||
lint-type-discipline-budget-update: install-dev lint-fetch-base
|
||||
$(UV_RUN) python scripts/type_discipline_gate.py --update
|
||||
|
||||
# Ratchet all budgets in one shot (ruff strict + type-discipline + basedpyright)
|
||||
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-basedpyright-budget-update
|
||||
lint-test-quality-budget-update: install-dev lint-fetch-base
|
||||
$(UV_RUN) python scripts/test_quality_gate.py --update
|
||||
|
||||
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
|
||||
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
|
||||
|
||||
check-circular-imports: $(LINT_DEP_INSTALL)
|
||||
cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
|
||||
|
|
@ -244,7 +258,7 @@ lint:
|
|||
lint-inner: lint-install lint-fetch-base
|
||||
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
|
||||
|
||||
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
|
||||
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
|
||||
|
||||
# Faster linting for local development (only checks changed code)
|
||||
lint-dev: lint-format-changed check-circular-imports check-import-safety
|
||||
|
|
@ -314,7 +328,7 @@ test-unit-helm: install-helm-unittest
|
|||
# LLM Translation testing targets
|
||||
test-llm-translation: install-test-deps
|
||||
@echo "Running LLM translation tests..."
|
||||
@python .github/workflows/run_llm_translation_tests.py
|
||||
@python .github/scripts/run_llm_translation_tests.py
|
||||
|
||||
test-llm-translation-single: install-test-deps
|
||||
@echo "Running single LLM translation test file..."
|
||||
|
|
|
|||
|
|
@ -292,6 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |
|
||||
| [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 22343
|
||||
"limit": 19955
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2578
|
||||
"limit": 2566
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 323
|
||||
"limit": 320
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"limit": 488
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 6991
|
||||
"limit": 6049
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5681
|
||||
"limit": 5663
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15605
|
||||
"limit": 15555
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44709
|
||||
"limit": 44655
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 112
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39154
|
||||
"limit": 39011
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19944
|
||||
"limit": 19885
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30772
|
||||
"limit": 30569
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
|
|
@ -123,7 +123,7 @@
|
|||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 851
|
||||
"limit": 836
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -145,6 +145,11 @@ NUMBER_KEYS: dict[str, JsonSchema] = {
|
|||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
"regional_endpoint_uplift_multiplier": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
|
||||
},
|
||||
}
|
||||
|
||||
COST_DESCRIPTIONS: dict[str, str] = {
|
||||
|
|
|
|||
|
|
@ -60,4 +60,4 @@ if __name__ == "__main__":
|
|||
print("\n💡 Tips:")
|
||||
print("1. Run 'litellm-proxy login' to authenticate first")
|
||||
print("2. Replace 'https://your-proxy.com' with your actual proxy URL")
|
||||
print("3. The token is stored locally at ~/.litellm/token.json")
|
||||
print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none")
|
||||
|
|
|
|||
|
|
@ -255,6 +255,52 @@ class CheckBatchCost:
|
|||
"so it will no longer be polled"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool:
|
||||
"""
|
||||
Atomically flip batch_processed from false to true, returning whether this pod won
|
||||
the row. Every pod and uvicorn worker schedules its own poller against the shared
|
||||
table, so without this compare-and-swap two of them can select the same completed
|
||||
batch in one window and both emit an aretrieve_batch spend log for it. Schemas
|
||||
without the column can't be claimed, so they keep the pre-existing behavior.
|
||||
|
||||
Called immediately before the spend log is written rather than before the results
|
||||
fetch, because batch_processed is also what holds off deletion of the files that
|
||||
fetch reads and what keeps an unbilled row selectable by the next poll cycle.
|
||||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return True
|
||||
try:
|
||||
claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": job.id, "batch_processed": False},
|
||||
data={"batch_processed": True},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to claim job {job.id} for cost tracking: {db_err}"
|
||||
)
|
||||
return False
|
||||
return claimed > 0
|
||||
|
||||
async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None:
|
||||
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
|
||||
|
||||
Safe to match on batch_processed=True: while this poller is active the retrieve
|
||||
path leaves the column alone (batch_cost_poller_is_active), so a true value here
|
||||
is always this pod's own claim.
|
||||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return
|
||||
try:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": job.id, "batch_processed": True},
|
||||
data={"batch_processed": False},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to release the claim on job {job.id}, "
|
||||
f"so its cost will not be retried: {db_err}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
|
||||
"""A unified id that decodes but carries no model_id can never be routed."""
|
||||
|
|
@ -572,9 +618,10 @@ class CheckBatchCost:
|
|||
"""
|
||||
Fetch a completed batch's results, compute cost/usage, and emit the
|
||||
aretrieve_batch spend log. Returns (model_name, llm_provider) on
|
||||
success, None when the job can't be routed to a deployment. Raises on
|
||||
results-fetch or cost-computation failures so the caller can leave the
|
||||
job unprocessed and retry it on a later poll.
|
||||
success, None when the job can't be routed to a deployment or when
|
||||
another pod claimed it. Raises on results-fetch or cost-computation
|
||||
failures so the caller can leave the job unprocessed and retry it on a
|
||||
later poll.
|
||||
"""
|
||||
from litellm.batches.batch_utils import (
|
||||
_get_file_content_as_dictionary,
|
||||
|
|
@ -743,12 +790,23 @@ class CheckBatchCost:
|
|||
optional_params={},
|
||||
)
|
||||
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
)
|
||||
if not await self._claim_job_for_costing(job):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
|
||||
"in this window, so its cost is already being tracked there"
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job)
|
||||
raise
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
if prom_logger and response.completed_at and response.created_at:
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ All /vector_store management endpoints
|
|||
|
||||
import copy
|
||||
import json
|
||||
from typing import List, Optional
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final, List, Optional, Protocol
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
|
|
@ -32,9 +33,35 @@ from litellm.types.vector_stores import (
|
|||
)
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class ManagedVectorStoreRow(Protocol):
|
||||
"""A ``litellm_managedvectorstorestable`` row as returned by Prisma."""
|
||||
|
||||
def model_dump(self) -> LiteLLM_ManagedVectorStore: ...
|
||||
|
||||
|
||||
class ManagedVectorStoreTable(Protocol):
|
||||
"""The Prisma actions namespace for ``litellm_managedvectorstorestable``."""
|
||||
|
||||
async def find_unique(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
|
||||
|
||||
async def create(self, data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
|
||||
|
||||
async def delete(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
|
||||
|
||||
async def update(self, where: Mapping[str, str | None], data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
|
||||
|
||||
|
||||
def managed_vector_store_table(prisma_client: "PrismaClient") -> ManagedVectorStoreTable:
|
||||
"""The Prisma table actions for managed vector stores, behind a typed surface."""
|
||||
return prisma_client.db.litellm_managedvectorstorestable
|
||||
|
||||
|
||||
########################################################
|
||||
# Management Endpoints
|
||||
########################################################
|
||||
|
|
@ -66,7 +93,7 @@ async def new_vector_store(
|
|||
try:
|
||||
# Check if vector store already exists
|
||||
existing_vector_store = (
|
||||
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
|
||||
await managed_vector_store_table(prisma_client).find_unique(
|
||||
where={"vector_store_id": vector_store.get("vector_store_id")}
|
||||
)
|
||||
)
|
||||
|
|
@ -92,7 +119,7 @@ async def new_vector_store(
|
|||
del vector_store["litellm_params"]
|
||||
|
||||
_new_vector_store = (
|
||||
await prisma_client.db.litellm_managedvectorstorestable.create(
|
||||
await managed_vector_store_table(prisma_client).create(
|
||||
data={
|
||||
**vector_store,
|
||||
"litellm_params": litellm_params_json,
|
||||
|
|
@ -213,7 +240,7 @@ async def delete_vector_store(
|
|||
try:
|
||||
# Check if vector store exists
|
||||
existing_vector_store = (
|
||||
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
|
||||
await managed_vector_store_table(prisma_client).find_unique(
|
||||
where={"vector_store_id": data.vector_store_id}
|
||||
)
|
||||
)
|
||||
|
|
@ -224,7 +251,7 @@ async def delete_vector_store(
|
|||
)
|
||||
|
||||
# Delete vector store
|
||||
await prisma_client.db.litellm_managedvectorstorestable.delete(
|
||||
await managed_vector_store_table(prisma_client).delete(
|
||||
where={"vector_store_id": data.vector_store_id}
|
||||
)
|
||||
|
||||
|
|
@ -288,7 +315,7 @@ async def get_vector_store_info(
|
|||
return {"vector_store": vector_store_pydantic_obj}
|
||||
|
||||
vector_store = (
|
||||
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
|
||||
await managed_vector_store_table(prisma_client).find_unique(
|
||||
where={"vector_store_id": data.vector_store_id}
|
||||
)
|
||||
)
|
||||
|
|
@ -298,7 +325,7 @@ async def get_vector_store_info(
|
|||
detail=f"Vector store with ID {data.vector_store_id} not found",
|
||||
)
|
||||
|
||||
vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined]
|
||||
vector_store_dict = vector_store.model_dump()
|
||||
return {"vector_store": vector_store_dict}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}")
|
||||
|
|
@ -322,13 +349,13 @@ async def update_vector_store(
|
|||
|
||||
try:
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
vector_store_id = update_data.pop("vector_store_id")
|
||||
vector_store_id: Final[str] = update_data.pop("vector_store_id")
|
||||
if update_data.get("vector_store_metadata") is not None:
|
||||
update_data["vector_store_metadata"] = safe_dumps(
|
||||
update_data["vector_store_metadata"]
|
||||
)
|
||||
|
||||
updated = await prisma_client.db.litellm_managedvectorstorestable.update(
|
||||
updated = await managed_vector_store_table(prisma_client).update(
|
||||
where={"vector_store_id": vector_store_id},
|
||||
data=update_data,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.57"
|
||||
version = "0.1.58"
|
||||
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.57"
|
||||
version = "0.1.58"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ type: application
|
|||
# This is the chart version. This version number should be incremented each time you make changes
|
||||
# to the chart and its templates, including the app version.
|
||||
# Versions are expected to follow Semantic Versioning (https://semver.org/)
|
||||
version: 1.1.1
|
||||
version: 1.1.2
|
||||
|
||||
# This is the version number of the application being deployed. This version number should be
|
||||
# incremented each time you make changes to the application. Versions are not expected to
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
|
||||
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `docker.litellm.ai/berriai/litellm` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
|
||||
| `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` |
|
||||
| `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` |
|
||||
| `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` |
|
||||
|
|
|
|||
|
|
@ -100,6 +100,13 @@ spec:
|
|||
- name: DATABASE_URL
|
||||
value: {{ .Values.db.url | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
|
||||
- name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.readReplicaEndpointKey }}
|
||||
{{- end }}
|
||||
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
|
||||
- name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ tests:
|
|||
pattern: -litellm$
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-database:test
|
||||
value: ghcr.io/berriai/litellm:test
|
||||
- it: should work with tolerations
|
||||
template: deployment.yaml
|
||||
set:
|
||||
|
|
@ -80,6 +80,96 @@ tests:
|
|||
secretKeyRef:
|
||||
name: my-secret
|
||||
key: my-key
|
||||
- it: should inject DATABASE_READER_HOST from readReplicaEndpointKey before DATABASE_URL_READ_REPLICA
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
deployStandalone: false
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaEndpointKey: reader-host
|
||||
readReplicaUrl: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
|
||||
# $(VAR) interpolation only resolves vars defined EARLIER in the env
|
||||
# array, so the reader host must precede the composed URL
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[7].name
|
||||
value: DATABASE_READER_HOST
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[8].name
|
||||
value: DATABASE_URL_READ_REPLICA
|
||||
- it: should omit reader host when readReplicaUrl is unset
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
deployStandalone: false
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaEndpointKey: reader-host
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- it: should prefer readReplicaUrlKey over readReplicaEndpointKey composition
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaUrlKey: reader-url
|
||||
readReplicaEndpointKey: reader-host
|
||||
readReplicaUrl: postgresql://ignored
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-url
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
value: postgresql://ignored
|
||||
# the unused reader-host secret ref must be suppressed so a missing
|
||||
# key can't fail pod creation
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- it: should work with extraEnvVars
|
||||
template: deployment.yaml
|
||||
set:
|
||||
|
|
@ -337,7 +427,7 @@ tests:
|
|||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
extraInitContainers:
|
||||
- name: init-tpl
|
||||
|
|
@ -348,7 +438,7 @@ tests:
|
|||
path: spec.template.spec.initContainers
|
||||
content:
|
||||
name: init-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
command: ["echo", "hello"]
|
||||
- it: should work with extraContainers
|
||||
template: deployment.yaml
|
||||
|
|
@ -366,7 +456,7 @@ tests:
|
|||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
extraContainers:
|
||||
- name: sidecar-tpl
|
||||
|
|
@ -376,12 +466,12 @@ tests:
|
|||
path: spec.template.spec.containers
|
||||
content:
|
||||
name: sidecar-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
- it: should support tpl in podAnnotations
|
||||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
# Mirrors the real-world scenario this feature unblocks:
|
||||
# user disables the built-in ConfigMap (and its built-in checksum/config
|
||||
|
|
@ -398,7 +488,7 @@ tests:
|
|||
value: "test"
|
||||
- equal:
|
||||
path: spec.template.metadata.annotations["example.com/some-key"]
|
||||
value: "ghcr.io/berriai/litellm-database"
|
||||
value: "ghcr.io/berriai/litellm"
|
||||
- equal:
|
||||
path: spec.template.metadata.annotations["example.com/literal"]
|
||||
value: "plain-string-value"
|
||||
|
|
|
|||
|
|
@ -208,7 +208,7 @@ tests:
|
|||
template: migrations-job.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
migrationJob:
|
||||
enabled: true
|
||||
|
|
@ -221,7 +221,7 @@ tests:
|
|||
path: spec.template.spec.initContainers
|
||||
content:
|
||||
name: init-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
command: ["echo", "hello"]
|
||||
- it: should work with extraContainers
|
||||
template: migrations-job.yaml
|
||||
|
|
@ -241,7 +241,7 @@ tests:
|
|||
template: migrations-job.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
migrationJob:
|
||||
enabled: true
|
||||
|
|
@ -253,7 +253,7 @@ tests:
|
|||
path: spec.template.spec.containers
|
||||
content:
|
||||
name: sidecar-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
- it: should render the pod-level securityContext from podSecurityContext
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -6,8 +6,9 @@ replicaCount: 1
|
|||
# numWorkers: 2
|
||||
|
||||
image:
|
||||
# Use "ghcr.io/berriai/litellm-database" for optimized image with database
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
# Bundles the prisma CLI and engines, which is what lets the migrations job
|
||||
# and the proxy's own schema check run without network access.
|
||||
repository: ghcr.io/berriai/litellm
|
||||
pullPolicy: Always
|
||||
# Overrides the image tag whose default is the chart appVersion.
|
||||
# tag: "latest"
|
||||
|
|
@ -276,6 +277,14 @@ db:
|
|||
# written to db.readReplicaUrl ends up visible in the rendered pod spec
|
||||
# and the Helm release secret.
|
||||
readReplicaUrlKey: ""
|
||||
# Optional: when set, a DATABASE_READER_HOST env var is sourced from this
|
||||
# secret key, so db.readReplicaUrl can compose the reader URL from
|
||||
# individual secret components, e.g.
|
||||
# postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)
|
||||
# Use this when your secret store holds the bare reader hostname rather
|
||||
# than a full connection URL. Only takes effect when readReplicaUrl is
|
||||
# set; ignored when readReplicaUrlKey is set.
|
||||
readReplicaEndpointKey: ""
|
||||
|
||||
# Optional read-replica routing. When set, the proxy sends read-only
|
||||
# queries (find_*, count, group_by, query_raw/_first) to this URL while
|
||||
|
|
|
|||
|
|
@ -213,18 +213,21 @@ whenever the password contains a URL-reserved character (@, /, ?, %, +,
|
|||
|
||||
When `database.writer.useIAMAuth: true`, the chart injects
|
||||
IAM_TOKEN_DB_AUTH=true and omits DATABASE_PASSWORD — the entrypoint mints
|
||||
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived IAM token
|
||||
instead of a static password.
|
||||
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived AWS RDS IAM
|
||||
token instead of a static password. `database.writer.useAzureEntraAuth: true`
|
||||
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
|
||||
for Azure Database for PostgreSQL. The two are mutually exclusive.
|
||||
|
||||
The read replica is opt-in via `database.reader.host`. The chart emits
|
||||
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
|
||||
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
|
||||
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
|
||||
password auth. When `database.reader.useIAMAuth: true` it omits
|
||||
password auth. When `database.reader.useIAMAuth: true` (or
|
||||
`database.reader.useAzureEntraAuth: true`) it omits
|
||||
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
|
||||
same way. Reader IAM only takes effect when the writer also uses IAM auth
|
||||
(the proxy gates URL minting on IAM_TOKEN_DB_AUTH, which only the writer
|
||||
sets).
|
||||
same way. Reader token auth only takes effect when the writer uses the same
|
||||
token source, since the proxy gates URL minting on the single global
|
||||
IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
||||
*/}}
|
||||
{{- define "litellm.serverEnv" -}}
|
||||
{{- $root := .root -}}
|
||||
|
|
@ -254,9 +257,15 @@ sets).
|
|||
- name: DATABASE_SCHEMA
|
||||
value: {{ .schema | quote }}
|
||||
{{- end }}
|
||||
{{- if and .useIAMAuth .useAzureEntraAuth }}
|
||||
{{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
|
||||
{{- end }}
|
||||
{{- if .useIAMAuth }}
|
||||
- name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
{{- else if .useAzureEntraAuth }}
|
||||
- name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
{{- else }}
|
||||
- name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
|
|
@ -270,6 +279,9 @@ sets).
|
|||
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
|
||||
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
{{- if and .useAzureEntraAuth (not $root.Values.database.writer.useAzureEntraAuth) }}
|
||||
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
- name: DATABASE_HOST_READ_REPLICA
|
||||
value: {{ .host | quote }}
|
||||
- name: DATABASE_PORT_READ_REPLICA
|
||||
|
|
@ -280,7 +292,7 @@ sets).
|
|||
- name: DATABASE_SCHEMA_READ_REPLICA
|
||||
value: {{ .schema | quote }}
|
||||
{{- end }}
|
||||
{{- if .useIAMAuth }}
|
||||
{{- if or .useIAMAuth .useAzureEntraAuth }}
|
||||
{{- if .passwordSecret.name }}
|
||||
- name: DATABASE_USER_READ_REPLICA
|
||||
valueFrom:
|
||||
|
|
|
|||
116
helm/litellm/tests/database_auth_tests.yaml
Normal file
116
helm/litellm/tests/database_auth_tests.yaml
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
suite: test database token auth env vars
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- backend/configmap.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: writer emits DATABASE_PASSWORD and no token toggle by default
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: password
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: writer emits AZURE_POSTGRESQL_AUTH and omits DATABASE_PASSWORD under Entra auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: backend gets the same Entra toggle as the gateway
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: writer rejects both token sources at once
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useIAMAuth: true
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
|
||||
|
||||
- it: reader Entra auth without writer Entra auth is rejected
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.host: reader.example.com
|
||||
database.reader.dbname: litellm
|
||||
database.reader.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)"
|
||||
|
||||
- it: reader under Entra auth omits DATABASE_PASSWORD_READ_REPLICA
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
database.reader.host: reader.example.com
|
||||
database.reader.dbname: litellm
|
||||
database.reader.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST_READ_REPLICA
|
||||
value: reader.example.com
|
||||
any: true
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-reader-secret
|
||||
key: username
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD_READ_REPLICA
|
||||
any: true
|
||||
|
|
@ -145,6 +145,8 @@ database:
|
|||
dbname: ""
|
||||
schema: ""
|
||||
useIAMAuth: false
|
||||
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
|
||||
useAzureEntraAuth: false
|
||||
passwordSecret:
|
||||
name: litellm-writer-secret
|
||||
usernameKey: username
|
||||
|
|
@ -159,6 +161,8 @@ database:
|
|||
dbname: ""
|
||||
schema: ""
|
||||
useIAMAuth: false
|
||||
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
|
||||
useAzureEntraAuth: false
|
||||
passwordSecret:
|
||||
name: litellm-reader-secret
|
||||
usernameKey: username
|
||||
|
|
|
|||
|
|
@ -0,0 +1,9 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_ProxyWorkerHeartbeat" (
|
||||
"worker_id" TEXT NOT NULL,
|
||||
"hostname" TEXT NOT NULL,
|
||||
"started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"last_heartbeat_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_ProxyWorkerHeartbeat_pkey" PRIMARY KEY ("worker_id")
|
||||
);
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "group_id" TEXT;
|
||||
|
||||
UPDATE "LiteLLM_ShadowEvalJob" SET "group_id" = "id" WHERE "group_id" IS NULL;
|
||||
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ALTER COLUMN "group_id" SET NOT NULL;
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_group_id_idx" ON "LiteLLM_ShadowEvalJob"("group_id");
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE "LiteLLM_SpendLogs"
|
||||
ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
ADD COLUMN IF NOT EXISTS "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "stopped_by" TEXT;
|
||||
|
||||
UPDATE "LiteLLM_ShadowEvalJob" SET stopped_by = 'unknown'
|
||||
WHERE stopped_at IS NOT NULL AND ends_at > (NOW() AT TIME ZONE 'utc');
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
UPDATE "LiteLLM_SpendLogs"
|
||||
SET "created_at" = "endTime",
|
||||
"updated_at" = "endTime"
|
||||
WHERE "created_at" > "endTime" + interval '1 hour';
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "max_budget" DOUBLE PRECISION;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN "shadow_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
|
||||
|
|
@ -641,6 +641,8 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
|
|
@ -945,6 +947,17 @@ model LiteLLM_DailyTagSpend {
|
|||
}
|
||||
|
||||
|
||||
// One row per live proxy worker process. Workers upsert their row on a fixed
|
||||
// heartbeat; counting rows with a recent heartbeat tells how many workers share
|
||||
// this database, which lets the Admin UI hide its "no Redis" warning for
|
||||
// deployments that are provably a single worker.
|
||||
model LiteLLM_ProxyWorkerHeartbeat {
|
||||
worker_id String @id
|
||||
hostname String
|
||||
started_at DateTime @default(now())
|
||||
last_heartbeat_at DateTime @default(now())
|
||||
}
|
||||
|
||||
// Track the status of cron jobs running. Only allow one pod to run the job at a time
|
||||
model LiteLLM_CronJob {
|
||||
cronjob_id String @id @default(cuid()) // Unique ID for the record
|
||||
|
|
@ -1465,28 +1478,39 @@ model LiteLLM_AutoRouterSession {
|
|||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
|
||||
// direction. forward duplicates the requests the key did not route through the router
|
||||
// through it, answering whether the key should adopt it; reverse duplicates the requests
|
||||
// the router did serve against a fixed baseline model, answering whether a key already on
|
||||
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
|
||||
// compares real vs shadow responses blind. The job row is immutable config plus
|
||||
// stopped_at; every count, status, and spend figure is derived from the append-only
|
||||
// attempt rows, so nothing can disagree across pods or stop races.
|
||||
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
|
||||
// either direction. forward duplicates the requests the keys did not route through the
|
||||
// router through it, answering whether they should adopt it; reverse duplicates the
|
||||
// requests the router did serve against a fixed baseline model, answering whether a key
|
||||
// already on it still benefits. Either way a sampled slice runs in a detached task and an
|
||||
// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job:
|
||||
// immutable config plus that key's own turn budget and stop state, so one key exhausting
|
||||
// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id
|
||||
// (the id the API reports), written together by one atomic create_many with identical
|
||||
// config; single-key jobs predating group_id were backfilled group_id = id. "One active
|
||||
// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE
|
||||
// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state
|
||||
// partial indexes; it is what makes a concurrent start on another pod race-safe rather
|
||||
// than read-then-create. Every count, status, and spend figure is derived from the
|
||||
// append-only attempt rows, so nothing can disagree across pods or stop races.
|
||||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
api_key_id String // hashed virtual key whose traffic is shadowed
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
router_name String // the auto-router under evaluation, in either direction
|
||||
direction String @default("forward") // forward | reverse
|
||||
baseline_model String? // reverse only: the fixed model the router is judged against
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // sample budget: judge at most this many turns
|
||||
max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
stopped_at DateTime?
|
||||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
|
@ -1502,6 +1526,7 @@ model LiteLLM_ShadowEvalAttempt {
|
|||
shadow_model String?
|
||||
confidence Float?
|
||||
judge_cost Float @default(0)
|
||||
shadow_cost Float @default(0)
|
||||
error String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.87"
|
||||
version = "0.4.88"
|
||||
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.87"
|
||||
version = "0.4.88"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -295,6 +295,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
|
|||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -324,6 +324,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
|
|||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,12 +105,20 @@ impl IntoResponse for MessagesRouteError {
|
|||
),
|
||||
CoreError::Http { .. }
|
||||
| CoreError::Network(_)
|
||||
| CoreError::Connect(_)
|
||||
| CoreError::InvalidResponse(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider request failed".to_string(),
|
||||
),
|
||||
// The gateway has no Python implementation to decline to, so a
|
||||
// request the core cannot serve is reported to the caller. The
|
||||
// reason is a fixed internal string, never provider content.
|
||||
CoreError::Unsupported(reason) => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("messages request is not supported: {reason}"),
|
||||
),
|
||||
};
|
||||
(
|
||||
status,
|
||||
|
|
|
|||
15
litellm-rust/crates/core/src/chat_completions/client.rs
Normal file
15
litellm-rust/crates/core/src/chat_completions/client.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS};
|
||||
|
||||
pub(super) fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))
|
||||
.connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
|
||||
|
||||
use super::transformation::ChatCompletionsProviderConfig;
|
||||
|
||||
const HEADER_CONTEXT: &str = "chat completions";
|
||||
|
||||
pub(super) fn chat_completions_provider_config(
|
||||
provider: &str,
|
||||
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
|
||||
match provider {
|
||||
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
"bedrock" => Some(
|
||||
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
254
litellm-rust/crates/core/src/chat_completions/conversation.rs
Normal file
254
litellm-rust/crates/core/src/chat_completions/conversation.rs
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
//! Provider-neutral conversation shape.
|
||||
//!
|
||||
//! Both Anthropic Messages and Bedrock Converse want the same thing out of an
|
||||
//! OpenAI message list: the system prompt lifted out, consecutive same-role
|
||||
//! turns merged, and text blocks that are never empty. That normalization is
|
||||
//! shared here so a provider config only renders the result into its own wire
|
||||
//! shape.
|
||||
//!
|
||||
//! Mirrors Python's `anthropic_messages_pt` /
|
||||
//! `_bedrock_converse_messages_pt` for the text-only surface this route
|
||||
//! accepts; anything richer is declined upstream by the capability gate.
|
||||
|
||||
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
|
||||
|
||||
use super::types::{ChatMessage, ChatMessageContent};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TurnRole {
|
||||
User,
|
||||
Assistant,
|
||||
}
|
||||
|
||||
impl TurnRole {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::User => "user",
|
||||
Self::Assistant => "assistant",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Turn {
|
||||
pub role: TurnRole,
|
||||
pub texts: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct Conversation {
|
||||
pub system: Vec<String>,
|
||||
pub turns: Vec<Turn>,
|
||||
}
|
||||
|
||||
/// True when the conversation can be sent as-is.
|
||||
///
|
||||
/// Python inserts a placeholder first user turn only under
|
||||
/// `litellm.modify_params`, which the core cannot see, so a conversation that
|
||||
/// does not open on a user turn is declined rather than guessed at.
|
||||
impl Conversation {
|
||||
pub fn opens_on_user_turn(&self) -> bool {
|
||||
self.turns
|
||||
.first()
|
||||
.is_some_and(|turn| turn.role == TurnRole::User)
|
||||
}
|
||||
}
|
||||
|
||||
fn message_texts(content: &ChatMessageContent) -> Vec<String> {
|
||||
match content {
|
||||
ChatMessageContent::Text(text) => vec![text.clone()],
|
||||
ChatMessageContent::Parts(parts) => parts
|
||||
.iter()
|
||||
.filter_map(|part| part.get("text").and_then(|text| text.as_str()))
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Python rewrites empty or whitespace-only text rather than dropping it, so an
|
||||
/// entirely empty content list never reaches a provider that rejects one.
|
||||
fn sanitize(text: String) -> String {
|
||||
if text.trim().is_empty() {
|
||||
return EMPTY_TEXT_PLACEHOLDER.to_string();
|
||||
}
|
||||
text
|
||||
}
|
||||
|
||||
pub fn build_conversation(messages: &[ChatMessage]) -> Conversation {
|
||||
let system = messages
|
||||
.iter()
|
||||
.filter(|message| message.role == "system")
|
||||
.filter_map(|message| message.content.as_ref())
|
||||
.flat_map(message_texts)
|
||||
.filter(|text| !text.is_empty())
|
||||
.collect();
|
||||
|
||||
let turns = messages
|
||||
.iter()
|
||||
.filter(|message| message.role != "system")
|
||||
.fold(Vec::<Turn>::new(), |mut turns, message| {
|
||||
let role = if message.role == "assistant" {
|
||||
TurnRole::Assistant
|
||||
} else {
|
||||
TurnRole::User
|
||||
};
|
||||
let texts = message
|
||||
.content
|
||||
.as_ref()
|
||||
.map(message_texts)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(sanitize);
|
||||
match turns.last_mut() {
|
||||
Some(last) if last.role == role => last.texts.extend(texts),
|
||||
_ => turns.push(Turn {
|
||||
role,
|
||||
texts: texts.collect(),
|
||||
}),
|
||||
}
|
||||
turns
|
||||
});
|
||||
|
||||
// Anthropic and Bedrock both reject trailing whitespace on the final
|
||||
// assistant turn, so Python right-strips it there; mirror that exactly.
|
||||
let turns = match turns.split_last() {
|
||||
Some((last, rest)) if last.role == TurnRole::Assistant => rest
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([Turn {
|
||||
role: last.role,
|
||||
texts: last
|
||||
.texts
|
||||
.iter()
|
||||
.map(|text| text.trim_end().to_string())
|
||||
.collect(),
|
||||
}])
|
||||
.collect(),
|
||||
_ => turns,
|
||||
};
|
||||
|
||||
Conversation { system, turns }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: serde_json::Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifts_system_messages_out_of_the_turn_list() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
])));
|
||||
assert_eq!(conversation.system, vec!["be terse".to_string()]);
|
||||
assert_eq!(
|
||||
conversation.turns,
|
||||
vec![Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["hi".to_string()]
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_same_role_turns() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": "two"},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "three"}
|
||||
])));
|
||||
assert_eq!(
|
||||
conversation.turns,
|
||||
vec![
|
||||
Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["one".to_string(), "two".to_string()]
|
||||
},
|
||||
Turn {
|
||||
role: TurnRole::Assistant,
|
||||
texts: vec!["ack".to_string()]
|
||||
},
|
||||
Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["three".to_string()]
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flattens_text_parts_in_order() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "text", "text": "second"}
|
||||
]}
|
||||
])));
|
||||
assert_eq!(
|
||||
conversation.turns[0].texts,
|
||||
vec!["first".to_string(), "second".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrites_empty_and_whitespace_only_text_to_the_python_placeholder() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": ""},
|
||||
{"role": "assistant", "content": " "},
|
||||
{"role": "user", "content": "real"}
|
||||
])));
|
||||
assert_eq!(conversation.turns[0].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
|
||||
assert_eq!(conversation.turns[1].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn right_strips_only_the_final_assistant_turn() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "kept "},
|
||||
{"role": "user", "content": "more"},
|
||||
{"role": "assistant", "content": "stripped "}
|
||||
])));
|
||||
assert_eq!(conversation.turns[1].texts, vec!["kept ".to_string()]);
|
||||
assert_eq!(conversation.turns[3].texts, vec!["stripped".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_strip_when_the_last_turn_is_a_user_turn() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "assistant", "content": "kept "},
|
||||
{"role": "user", "content": "hi "}
|
||||
])));
|
||||
assert_eq!(conversation.turns[0].texts, vec!["kept ".to_string()]);
|
||||
assert_eq!(conversation.turns[1].texts, vec!["hi ".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_whether_the_conversation_opens_on_a_user_turn() {
|
||||
assert!(
|
||||
build_conversation(&messages(json!([{"role": "user", "content": "hi"}])))
|
||||
.opens_on_user_turn()
|
||||
);
|
||||
assert!(
|
||||
!build_conversation(&messages(json!([{"role": "assistant", "content": "hi"}])))
|
||||
.opens_on_user_turn()
|
||||
);
|
||||
assert!(!Conversation::default().opens_on_user_turn());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drops_empty_system_text_the_way_python_skips_empty_system_blocks() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "system", "content": ""},
|
||||
{"role": "system", "content": "kept"},
|
||||
{"role": "user", "content": "hi"}
|
||||
])));
|
||||
assert_eq!(conversation.system, vec!["kept".to_string()]);
|
||||
}
|
||||
}
|
||||
147
litellm-rust/crates/core/src/chat_completions/handler.rs
Normal file
147
litellm-rust/crates/core/src/chat_completions/handler.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::http_utils::truncate_error_body;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
use super::types::{
|
||||
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
|
||||
};
|
||||
|
||||
pub(super) async fn execute_chat_completions_provider_call(
|
||||
request: ProviderChatCompletionsRequest,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = serde_json::to_vec(&request.body).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"failed to serialize chat completions request: {err}"
|
||||
))
|
||||
})?;
|
||||
let headers = signed_headers(&request, &body).await?;
|
||||
|
||||
let mut request_builder = http_client().post(&request.url).body(body);
|
||||
for (key, value) in &headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request_builder.send().await.map_err(|err| {
|
||||
// Failing to establish the connection means the request never went out,
|
||||
// so the host can still serve it. Everything else here, a timeout
|
||||
// above all, may have reached the provider and been answered.
|
||||
if err.is_connect() || err.is_builder() {
|
||||
CoreError::Connect(err.to_string())
|
||||
} else {
|
||||
CoreError::Network(err.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let body: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
|
||||
})?;
|
||||
request
|
||||
.config
|
||||
.transform_response(&request.model, ProviderChatResponseData { body })
|
||||
.map_err(as_response_error)
|
||||
}
|
||||
|
||||
/// Re-tag an error raised while normalizing a response the provider already
|
||||
/// returned.
|
||||
///
|
||||
/// A config reports the same variants on either side of the call: a missing
|
||||
/// field or an unsupported block can mean "this request cannot be translated"
|
||||
/// during prepare and "this response cannot be normalized" here. Only the
|
||||
/// second kind has already been billed, and a host that keeps a reference
|
||||
/// implementation must not retry those, so collapse them to one variant that
|
||||
/// can only mean the provider was already called.
|
||||
pub(super) fn as_response_error(err: CoreError) -> CoreError {
|
||||
match err {
|
||||
already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already,
|
||||
other => CoreError::InvalidResponse(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use crate::providers::bedrock::aws_base::{
|
||||
aws_auth_config, aws_signature_headers, host_supplied_credentials,
|
||||
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
|
||||
};
|
||||
|
||||
let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else {
|
||||
return Ok(request.upstream_headers.clone());
|
||||
};
|
||||
// Reattaching a header the signer also emits would put both copies on the
|
||||
// wire, and Bedrock rejects that pair. Python instead drops the caller's
|
||||
// copy and prefers a forwarded Authorization over the signature, so leave
|
||||
// the request to Python rather than serving it a different way here.
|
||||
if request
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, _)| is_sigv4_computed_header(name))
|
||||
{
|
||||
return Err(CoreError::Unsupported(
|
||||
"request forwards a header AWS SigV4 computes",
|
||||
));
|
||||
}
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
|
||||
// A host with its own resolution chain hands the result down; only fall
|
||||
// back to deriving credentials here when it supplied none.
|
||||
let credentials = match host_supplied_credentials(&request.optional_params) {
|
||||
Some(credentials) => credentials,
|
||||
None => {
|
||||
resolve_credentials(
|
||||
aws_auth_config(&request.optional_params, &env_lookup),
|
||||
&env_lookup,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
let signature = sign_bedrock_post(
|
||||
&request.url,
|
||||
body,
|
||||
&aws_signature_headers(&unsigned),
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
// Every original header goes back on the wire alongside the computed ones,
|
||||
// as Python reattaches them. The guard above already rejected the names
|
||||
// that would collide, so no name appears twice.
|
||||
Ok(unsigned.into_iter().chain(signature).collect())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "bedrock-auth"))]
|
||||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
_body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
match &request.auth {
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported(
|
||||
"AWS SigV4 requires the bedrock-auth feature",
|
||||
)),
|
||||
_ => Ok(request.upstream_headers.clone()),
|
||||
}
|
||||
}
|
||||
59
litellm-rust/crates/core/src/chat_completions/mod.rs
Normal file
59
litellm-rust/crates/core/src/chat_completions/mod.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
//! The `/chat/completions` call, the Rust equivalent of Python's
|
||||
//! `litellm.completion()`.
|
||||
//!
|
||||
//! [`chat_completions`] is the top-level entrypoint: give it a model, the
|
||||
//! OpenAI-shaped message list, the provider-mapped optional params, and
|
||||
//! credentials, and it resolves the provider, translates the conversation,
|
||||
//! calls the provider, and returns a typed OpenAI-shaped response.
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
pub mod conversation;
|
||||
pub(crate) mod handler;
|
||||
mod prepare;
|
||||
pub mod response_utils;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use handler::execute_chat_completions_provider_call;
|
||||
use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config};
|
||||
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
|
||||
|
||||
pub async fn chat_completions(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await
|
||||
}
|
||||
|
||||
/// Whether the core would accept this request, without resolving credentials or
|
||||
/// touching the network.
|
||||
///
|
||||
/// A host that keeps the Python implementation asks this first so it can emit
|
||||
/// its pre-call logging exactly once, on whichever path is about to run.
|
||||
/// Returns the decline reason, or `None` when the request is accepted.
|
||||
pub fn chat_completions_decline_reason(
|
||||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
messages: Value,
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<&'static str> {
|
||||
let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else {
|
||||
return Some("provider is not on the rust chat completions path");
|
||||
};
|
||||
let Ok(messages) = parse_messages(messages) else {
|
||||
return Some("unreadable message list");
|
||||
};
|
||||
if messages.is_empty() {
|
||||
return Some("empty message list");
|
||||
}
|
||||
config
|
||||
.unsupported_reason(&messages, optional_params)
|
||||
.map(|reason| reason.0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
118
litellm-rust/crates/core/src/chat_completions/prepare.rs
Normal file
118
litellm-rust/crates/core/src/chat_completions/prepare.rs
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::http_utils::has_header;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::common_utils::{chat_completions_provider_config, string_headers};
|
||||
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
|
||||
use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest};
|
||||
|
||||
pub(super) fn resolve_provider_config<'a>(
|
||||
model: &'a str,
|
||||
custom_llm_provider: Option<&'a str>,
|
||||
) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> {
|
||||
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
|
||||
.or_else(|| {
|
||||
custom_llm_provider.map(|provider| CustomLlmProvider {
|
||||
model,
|
||||
custom_llm_provider: provider,
|
||||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for chat completions request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let config = chat_completions_provider_config(provider_info.custom_llm_provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
|
||||
Ok((provider_info.model.to_string(), config))
|
||||
}
|
||||
|
||||
pub(super) fn parse_messages(messages: Value) -> CoreResult<Vec<ChatMessage>> {
|
||||
serde_json::from_value(messages).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!("invalid chat completions messages: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn prepare_chat_completions_call(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ProviderChatCompletionsRequest> {
|
||||
let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let messages = parse_messages(request.messages)?;
|
||||
if messages.is_empty() {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"chat completions requires at least one message".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
|
||||
return Err(CoreError::Unsupported(reason.0));
|
||||
}
|
||||
|
||||
let mut headers = string_headers(request.extra_headers)?;
|
||||
let auth = config.auth(
|
||||
request.api_key,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
match &auth {
|
||||
ChatCompletionsAuth::Header { name, value } => {
|
||||
// The deployment's credential replaces whatever the caller forwarded
|
||||
// under the same name, mirroring Python's
|
||||
// `{**headers, **anthropic_headers}`: letting a request header win
|
||||
// would let its sender choose the principal the call bills to.
|
||||
//
|
||||
// The exception is a scheme the provider hands off to entirely, such
|
||||
// as an Anthropic OAuth bearer, where Python drops `x-api-key`
|
||||
// instead of resolving one. Re-adding it there would put the
|
||||
// credential into a header the host removed on purpose.
|
||||
if !config.defers_to_forwarded_auth(&headers) {
|
||||
headers.retain(|(header, _)| !header.eq_ignore_ascii_case(name));
|
||||
headers.push(((*name).to_string(), value.clone()));
|
||||
}
|
||||
}
|
||||
ChatCompletionsAuth::Bearer { token } => {
|
||||
// Bedrock's `get_request_headers` assigns `headers["Authorization"]`
|
||||
// unconditionally once a bearer token resolves, so the deployment's
|
||||
// identity outranks whatever the caller forwarded. Keeping the
|
||||
// caller's would bill and authorize the call as a different
|
||||
// principal than the same deployment uses on Python.
|
||||
//
|
||||
// The `Header` arm below keeps the opposite precedence on purpose:
|
||||
// Anthropic's transform honours a forwarded OAuth bearer.
|
||||
headers.retain(|(name, _)| !name.eq_ignore_ascii_case("authorization"));
|
||||
headers.push(("authorization".to_string(), format!("Bearer {token}")));
|
||||
}
|
||||
// SigV4 signs the serialized body, so the handler adds its headers.
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => {}
|
||||
}
|
||||
|
||||
for (name, value) in config.default_headers() {
|
||||
if !has_header(&headers, name) {
|
||||
headers.push(((*name).to_string(), (*value).to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let url = config.complete_url(
|
||||
request.api_base,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let transformed =
|
||||
config.transform_request(&model, messages, request.optional_params.clone())?;
|
||||
|
||||
Ok(ProviderChatCompletionsRequest {
|
||||
model,
|
||||
config,
|
||||
url,
|
||||
body: transformed.body,
|
||||
upstream_headers: headers,
|
||||
auth,
|
||||
optional_params: request.optional_params,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
101
litellm-rust/crates/core/src/chat_completions/response_utils.rs
Normal file
101
litellm-rust/crates/core/src/chat_completions/response_utils.rs
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
//! Response normalization shared by every chat completions provider config.
|
||||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::types::{ChatCompletionsUsage, PromptTokensDetails};
|
||||
|
||||
/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the
|
||||
/// reasons the providers on this route can emit. Python warns and falls back to
|
||||
/// `stop` for anything unmapped, so do the same.
|
||||
const FINISH_REASONS: &[(&str, &str)] = &[
|
||||
("end_turn", "stop"),
|
||||
("stop_sequence", "stop"),
|
||||
("max_tokens", "length"),
|
||||
("refusal", "content_filter"),
|
||||
("compaction", "length"),
|
||||
("guardrail_intervened", "content_filter"),
|
||||
("content_filtered", "content_filter"),
|
||||
("content_filter", "content_filter"),
|
||||
("stop", "stop"),
|
||||
("length", "length"),
|
||||
];
|
||||
|
||||
pub fn finish_reason_for(provider_reason: &str) -> &'static str {
|
||||
FINISH_REASONS
|
||||
.iter()
|
||||
.find(|(reason, _)| *reason == provider_reason)
|
||||
.map_or("stop", |(_, mapped)| *mapped)
|
||||
}
|
||||
|
||||
/// Python folds cache tokens into `prompt_tokens` and reports the split under
|
||||
/// `prompt_tokens_details`; mirror that so cost tracking agrees on both paths.
|
||||
pub fn usage_from_parts(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cache_read_tokens: u64,
|
||||
cache_creation_tokens: u64,
|
||||
) -> ChatCompletionsUsage {
|
||||
let prompt_tokens = input_tokens + cache_read_tokens + cache_creation_tokens;
|
||||
ChatCompletionsUsage {
|
||||
prompt_tokens,
|
||||
completion_tokens: output_tokens,
|
||||
total_tokens: prompt_tokens + output_tokens,
|
||||
prompt_tokens_details: PromptTokensDetails {
|
||||
cached_tokens: cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
text_tokens: input_tokens,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn unix_now() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |elapsed| elapsed.as_secs())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn maps_every_reason_the_route_can_observe() {
|
||||
assert_eq!(finish_reason_for("end_turn"), "stop");
|
||||
assert_eq!(finish_reason_for("stop_sequence"), "stop");
|
||||
assert_eq!(finish_reason_for("max_tokens"), "length");
|
||||
assert_eq!(finish_reason_for("refusal"), "content_filter");
|
||||
assert_eq!(finish_reason_for("guardrail_intervened"), "content_filter");
|
||||
// Converse emits these two, and folding them into `stop` would report a
|
||||
// filtered completion as a normal one.
|
||||
assert_eq!(finish_reason_for("content_filtered"), "content_filter");
|
||||
assert_eq!(finish_reason_for("content_filter"), "content_filter");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_an_unmapped_reason_to_stop_like_python() {
|
||||
// Python warns and falls back to `stop` for a reason its own map does
|
||||
// not carry, so only a reason absent from `_FINISH_REASON_MAP` belongs
|
||||
// here.
|
||||
assert_eq!(finish_reason_for("something_new"), "stop");
|
||||
assert_eq!(finish_reason_for(""), "stop");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_cache_tokens_into_prompt_tokens() {
|
||||
let usage = usage_from_parts(10, 4, 7, 3);
|
||||
assert_eq!(usage.prompt_tokens, 20);
|
||||
assert_eq!(usage.completion_tokens, 4);
|
||||
assert_eq!(usage.total_tokens, 24);
|
||||
assert_eq!(usage.prompt_tokens_details.cached_tokens, 7);
|
||||
assert_eq!(usage.prompt_tokens_details.cache_creation_tokens, 3);
|
||||
assert_eq!(usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_raw_input_tokens_when_no_cache_is_involved() {
|
||||
let usage = usage_from_parts(12, 5, 0, 0);
|
||||
assert_eq!(usage.prompt_tokens, 12);
|
||||
assert_eq!(usage.total_tokens, 17);
|
||||
assert_eq!(usage.prompt_tokens_details.text_tokens, 12);
|
||||
}
|
||||
}
|
||||
820
litellm-rust/crates/core/src/chat_completions/tests.rs
Normal file
820
litellm-rust/crates/core/src/chat_completions/tests.rs
Normal file
|
|
@ -0,0 +1,820 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::error::CoreError;
|
||||
|
||||
use super::prepare::prepare_chat_completions_call;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
use super::types::ChatCompletionsRequest;
|
||||
|
||||
fn request<'a>(
|
||||
model: &'a str,
|
||||
provider: Option<&'a str>,
|
||||
messages: Value,
|
||||
optional_params: Value,
|
||||
) -> ChatCompletionsRequest<'a> {
|
||||
ChatCompletionsRequest {
|
||||
model,
|
||||
messages,
|
||||
optional_params: match optional_params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
},
|
||||
api_key: Some("sk-test"),
|
||||
api_base: None,
|
||||
custom_llm_provider: provider,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
|
||||
/// carry resolved credentials), so unwrap the failure case by hand.
|
||||
fn decline(request: ChatCompletionsRequest<'_>) -> CoreError {
|
||||
match prepare_chat_completions_call(request) {
|
||||
Err(error) => error,
|
||||
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_the_provider_from_the_model_prefix() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
|
||||
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_an_explicit_provider_prefix_from_the_model() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adds_the_auth_and_default_headers() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.contains(&("x-api-key".to_string(), "sk-test".to_string()))
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
|
||||
);
|
||||
assert!(matches!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
|
||||
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
|
||||
// overwrites a forwarded one. Honouring the caller's would let whoever sends
|
||||
// the request choose the Anthropic principal it bills to.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"X-Api-Key".to_string(),
|
||||
json!("sk-caller"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
|
||||
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
|
||||
// for an OAuth token, so re-adding the key here would put the credential into
|
||||
// a header the host removed on purpose.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-token"),
|
||||
),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert!(
|
||||
!prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
|
||||
"the resolved key must not be applied over an OAuth bearer, got {:?}",
|
||||
prepared.upstream_headers
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
|
||||
// Only an OAuth bearer replaces the credential. Python sends the deployment's
|
||||
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
|
||||
// the mere presence of that header would drop the deployment's auth.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
("Authorization".to_string(), json!("Bearer unrelated")),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer unrelated"),
|
||||
"the unrelated authorization must survive, got {:?}",
|
||||
prepared.upstream_headers
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_an_unsupported_request_before_resolving_credentials() {
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true}),
|
||||
);
|
||||
call.api_key = None;
|
||||
// No api_key is set and no env is consulted: the gate must run first, so the
|
||||
// error is the decline rather than a missing-credential error.
|
||||
assert_eq!(decline(call), CoreError::Unsupported("streaming"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_unknown_provider() {
|
||||
assert_eq!(
|
||||
decline(request(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider("openai".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_a_model_with_no_resolvable_provider() {
|
||||
assert!(matches!(
|
||||
decline(request(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_empty_or_malformed_message_list() {
|
||||
assert_eq!(
|
||||
decline(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest("chat completions requires at least one message".to_string())
|
||||
);
|
||||
assert!(matches!(
|
||||
decline(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!("not a list"),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_string_extra_headers() {
|
||||
let mut call = request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
|
||||
assert_eq!(
|
||||
decline(call),
|
||||
CoreError::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn prepares_a_bedrock_call_without_resolving_credentials() {
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.api_key = None;
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert_eq!(
|
||||
prepared.url,
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "us-east-1".to_string()
|
||||
}
|
||||
);
|
||||
// SigV4 signs the serialized body, so prepare must not have added an
|
||||
// Authorization header; the handler does it.
|
||||
assert!(
|
||||
!prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
);
|
||||
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
|
||||
// Python signs only the AWS header set and reattaches the rest, so a header
|
||||
// the caller forwarded rides along without joining the canonical request.
|
||||
// Signing it makes Converse 403 on a deployment that works on Python.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
// A key would resolve to a bearer token and never reach the signer.
|
||||
call.api_key = None;
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"x-request-id".to_string(),
|
||||
json!("abc-123"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
|
||||
.await
|
||||
.expect("signs");
|
||||
|
||||
let authorization = signed
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.clone())
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
authorization.starts_with("AWS4-HMAC-SHA256"),
|
||||
"expected a SigV4 signature, got {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
// It still goes on the wire, it is just not part of the signature.
|
||||
assert!(
|
||||
signed
|
||||
.iter()
|
||||
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
|
||||
"forwarded header was dropped instead of reattached"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_header_the_signer_computes_declines_to_python() {
|
||||
// Reattaching the caller's copy next to the computed one puts the name on
|
||||
// the wire twice and Bedrock rejects the pair, so a request carrying one
|
||||
// has to go to Python instead of being signed here.
|
||||
for forwarded in [
|
||||
"Authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"Date",
|
||||
] {
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
call.api_key = None;
|
||||
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
|
||||
.await
|
||||
.expect_err("{forwarded} should decline instead of being signed");
|
||||
assert!(
|
||||
matches!(error, CoreError::Unsupported(_)),
|
||||
"{forwarded} declined as {error:?}, which the host would not fall back on"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
|
||||
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
|
||||
// once a bearer token resolves, so the deployment's identity wins on
|
||||
// Python. Keeping the caller's would authorize and bill the call as a
|
||||
// different principal, and only when the deployment carries `rust: true`.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer caller-supplied"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let authorizations: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
authorizations,
|
||||
vec!["Bearer sk-test"],
|
||||
"the deployment token must be the only authorization on the wire"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
|
||||
// The opposite precedence, and deliberate: Anthropic's own transform
|
||||
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
|
||||
// generalized into a rule that the configured key always wins.
|
||||
//
|
||||
// An OAuth bearer is the whole of that exception. This forwarded a plain
|
||||
// `x-api-key` until round 17, which read as the same claim and was not:
|
||||
// Python overwrites a forwarded `x-api-key` with the deployment's.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-forwarded"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-forwarded")
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
|
||||
// The configured bearer identity has its own account and quota boundary,
|
||||
// so a request carrying one must not be signed as whatever principal the
|
||||
// host's AWS credentials resolve to.
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::Bearer {
|
||||
token: "sk-test".to_string()
|
||||
}
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-test"),
|
||||
"prepare did not carry the bearer token"
|
||||
);
|
||||
}
|
||||
|
||||
fn decline_reason(
|
||||
model: &str,
|
||||
provider: Option<&str>,
|
||||
messages: Value,
|
||||
params: Value,
|
||||
) -> Option<&'static str> {
|
||||
let params = match params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
};
|
||||
super::chat_completions_decline_reason(model, provider, messages, ¶ms)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_accepts_what_prepare_accepts() {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_declines_without_resolving_credentials_or_calling_out() {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true}),
|
||||
),
|
||||
Some("streaming")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
),
|
||||
Some("provider is not on the rust chat completions path")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
),
|
||||
Some("provider is not on the rust chat completions path")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!("nope"),
|
||||
json!({})
|
||||
),
|
||||
Some("unreadable message list")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})),
|
||||
Some("empty message list")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
|
||||
// A gate that accepts what prepare then declines would make the host emit
|
||||
// its pre-call logging on a path that falls back, so pin the agreement.
|
||||
for (messages, params) in [
|
||||
(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 8}),
|
||||
),
|
||||
(
|
||||
json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]),
|
||||
json!({"temperature": 0.1}),
|
||||
),
|
||||
(
|
||||
json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]),
|
||||
json!({}),
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
messages.clone(),
|
||||
params.clone()
|
||||
),
|
||||
None,
|
||||
"gate declined {messages}"
|
||||
);
|
||||
prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
messages.clone(),
|
||||
params,
|
||||
))
|
||||
.unwrap_or_else(|error| panic!("prepare declined {messages}: {error}"));
|
||||
}
|
||||
}
|
||||
|
||||
mod round_trip {
|
||||
use super::*;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::chat_completions::chat_completions;
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
let header_end = loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break request.len();
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break position + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&request[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
while request.len().saturating_sub(header_end) < content_length {
|
||||
let n = socket.read(&mut buffer).await.expect("reads body");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
fn http_response(status: &str, body: &str) -> String {
|
||||
format!(
|
||||
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
)
|
||||
}
|
||||
|
||||
/// Serve one request from a stub upstream and hand back what it received.
|
||||
async fn serve_once(
|
||||
status: &'static str,
|
||||
body: &'static str,
|
||||
) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let port = listener.local_addr().expect("addr").port();
|
||||
let handle = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts");
|
||||
let received = read_http_request(&mut socket).await;
|
||||
socket
|
||||
.write_all(http_response(status, body).as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
socket.flush().await.expect("flushes");
|
||||
received
|
||||
});
|
||||
(format!("http://127.0.0.1:{port}/v1/messages"), handle)
|
||||
}
|
||||
|
||||
fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> {
|
||||
ChatCompletionsRequest {
|
||||
model: "anthropic/claude-sonnet-4-5",
|
||||
messages,
|
||||
optional_params: match params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
},
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(api_base),
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: Some(std::time::Duration::from_secs(10)),
|
||||
}
|
||||
}
|
||||
|
||||
const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_sends_the_translated_body_and_normalizes_the_response() {
|
||||
let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await;
|
||||
let response = chat_completions(call(
|
||||
&api_base,
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let received = handle.await.expect("server task");
|
||||
let sent: Value = serde_json::from_str(
|
||||
received
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("request has a body")
|
||||
.1,
|
||||
)
|
||||
.expect("body is json");
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
|
||||
);
|
||||
assert_eq!(
|
||||
sent["system"],
|
||||
json!([{"type": "text", "text": "be terse"}])
|
||||
);
|
||||
assert_eq!(sent["max_tokens"], json!(16));
|
||||
assert!(received.to_lowercase().contains("x-api-key: sk-test"));
|
||||
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello")
|
||||
);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_response_it_cannot_normalize_is_reported_as_already_sent() {
|
||||
// The provider was called and billed, so the host must not retry this
|
||||
// on its own path. `MissingField` here would read as a pre-send
|
||||
// decline and be retried; `InvalidResponse` cannot.
|
||||
const NO_USAGE: &str =
|
||||
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#;
|
||||
let (api_base, handle) = serve_once("200 OK", NO_USAGE).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() {
|
||||
const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#;
|
||||
let (api_base, handle) = serve_once("200 OK", TOOL_USE).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_upstream_error_status_keeps_its_code() {
|
||||
let (api_base, handle) =
|
||||
serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("upstream rejects");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::Http { status: 429, .. }),
|
||||
"expected a 429, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_connection_that_is_never_established_declines_instead_of_failing() {
|
||||
// Nothing was sent, so nothing was billed and the host can still serve
|
||||
// the request. Classing this with the post-send failures would turn a
|
||||
// recoverable fallback into a user-facing error on exactly the
|
||||
// deployments whose transport is configured only on the Python client.
|
||||
let port = {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
listener.local_addr().expect("has an address").port()
|
||||
// Dropped here, so the port is closed and the connect is refused.
|
||||
};
|
||||
let err = chat_completions(call(
|
||||
&format!("http://127.0.0.1:{port}/v1/messages"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("nothing is listening");
|
||||
assert!(
|
||||
matches!(err, CoreError::Connect(_)),
|
||||
"expected a pre-send connect failure, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
|
||||
use crate::chat_completions::handler::as_response_error;
|
||||
|
||||
for original in [
|
||||
CoreError::MissingField("usage"),
|
||||
CoreError::Unsupported("non-text response content block"),
|
||||
CoreError::InvalidRequest("whatever".to_string()),
|
||||
CoreError::Auth("whatever".to_string()),
|
||||
] {
|
||||
let label = format!("{original:?}");
|
||||
assert!(
|
||||
matches!(as_response_error(original), CoreError::InvalidResponse(_)),
|
||||
"{label} must not stay retryable once the provider has answered"
|
||||
);
|
||||
}
|
||||
// An upstream status is already unambiguous, so it survives intact.
|
||||
assert!(matches!(
|
||||
as_response_error(CoreError::Http {
|
||||
status: 500,
|
||||
body: "boom".to_string()
|
||||
}),
|
||||
CoreError::Http { status: 500, .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
155
litellm-rust/crates/core/src/chat_completions/transformation.rs
Normal file
155
litellm-rust/crates/core/src/chat_completions/transformation.rs
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use super::types::{
|
||||
ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
|
||||
/// How the upstream call is authenticated. API-key strategies are resolved in
|
||||
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ChatCompletionsAuth {
|
||||
Header { name: &'static str, value: String },
|
||||
Bearer { token: String },
|
||||
AwsSigV4 { region: String },
|
||||
}
|
||||
|
||||
/// Why a request cannot be served by the Rust path.
|
||||
///
|
||||
/// The core declines rather than guessing: the host turns this into a
|
||||
/// transparent fallback to the Python implementation, which covers the full
|
||||
/// surface. Acceptance is an allowlist, so a parameter or message shape the
|
||||
/// core has never seen declines by construction instead of being translated
|
||||
/// wrong.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct Unsupported(pub &'static str);
|
||||
|
||||
pub const STREAM_PARAM: &str = "stream";
|
||||
|
||||
/// Message fields that carry no meaning for the upstream body, so their
|
||||
/// presence does not make a request untranslatable.
|
||||
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
|
||||
|
||||
pub trait ChatCompletionsProviderConfig: Sync {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth>;
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[("content-type", "application/json")]
|
||||
}
|
||||
|
||||
/// Whether an auth header the caller already supplied is the credential this
|
||||
/// request should authenticate with, so the resolved one is not applied.
|
||||
///
|
||||
/// Defaults to false: the deployment's credential outranks anything
|
||||
/// forwarded, which is what every provider wants for its own auth header.
|
||||
/// A provider overrides this only for a scheme it hands off to entirely.
|
||||
fn defers_to_forwarded_auth(&self, _headers: &[(String, String)]) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Provider parameter names (post-mapping) the Rust path knows how to place
|
||||
/// in the upstream body. Anything outside this set declines the request.
|
||||
fn supported_params(&self) -> &'static [&'static str];
|
||||
|
||||
/// Parameters consumed as call configuration (credentials, endpoints)
|
||||
/// rather than placed in the body. Accepted, never serialized.
|
||||
fn config_params(&self) -> &'static [&'static str] {
|
||||
&[]
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(
|
||||
self.supported_params(),
|
||||
self.config_params(),
|
||||
optional_params,
|
||||
)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData>;
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse>;
|
||||
}
|
||||
|
||||
pub fn unsupported_param(
|
||||
supported: &'static [&'static str],
|
||||
config: &'static [&'static str],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
if optional_params
|
||||
.get(STREAM_PARAM)
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Some(Unsupported("streaming"));
|
||||
}
|
||||
optional_params
|
||||
.keys()
|
||||
.any(|key| {
|
||||
key != STREAM_PARAM
|
||||
&& !supported.contains(&key.as_str())
|
||||
&& !config.contains(&key.as_str())
|
||||
})
|
||||
.then_some(Unsupported("unrecognized request parameter"))
|
||||
}
|
||||
|
||||
/// Message shapes the core can translate faithfully: text content, either a
|
||||
/// plain string or a non-empty list of parts that are all
|
||||
/// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal
|
||||
/// parts decline so Python's fuller translation handles them.
|
||||
pub fn unsupported_message(message: &ChatMessage) -> Option<Unsupported> {
|
||||
if message
|
||||
.extra
|
||||
.keys()
|
||||
.any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str()))
|
||||
{
|
||||
return Some(Unsupported("unrecognized message field"));
|
||||
}
|
||||
if !matches!(message.role.as_str(), "system" | "user" | "assistant") {
|
||||
return Some(Unsupported("unrecognized message role"));
|
||||
}
|
||||
match &message.content {
|
||||
None => Some(Unsupported("message without content")),
|
||||
Some(ChatMessageContent::Text(_)) => None,
|
||||
Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => {
|
||||
Some(Unsupported("message without content"))
|
||||
}
|
||||
Some(ChatMessageContent::Parts(parts)) => parts
|
||||
.iter()
|
||||
.any(|part| {
|
||||
part.get("type").and_then(Value::as_str) != Some("text")
|
||||
|| part.get("text").and_then(Value::as_str).is_none()
|
||||
|| part.as_object().is_some_and(|object| object.len() != 2)
|
||||
})
|
||||
.then_some(Unsupported("non-text message content")),
|
||||
}
|
||||
}
|
||||
112
litellm-rust/crates/core/src/chat_completions/types.rs
Normal file
112
litellm-rust/crates/core/src/chat_completions/types.rs
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
|
||||
|
||||
/// A `/chat/completions` call as it crosses into the core.
|
||||
///
|
||||
/// `optional_params` arrives already mapped to the provider's own parameter
|
||||
/// names by the host, exactly as the messages route receives an already
|
||||
/// Anthropic-shaped body. The core owns the conversation translation, the
|
||||
/// provider call, and the response normalization.
|
||||
pub struct ChatCompletionsRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub messages: Value,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(super) struct ProviderChatCompletionsRequest {
|
||||
pub(super) model: String,
|
||||
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) auth: ChatCompletionsAuth,
|
||||
#[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))]
|
||||
pub(super) optional_params: Map<String, Value>,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
/// The provider-shaped request body a config produces. Named rather than a bare
|
||||
/// `Value` so the transform contract stays a typed one, mirroring
|
||||
/// [`crate::audio_transcription::types::AudioTranscriptionRequestData`].
|
||||
pub struct ProviderChatRequestData {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
/// The raw provider response body handed back to a config for normalization.
|
||||
pub struct ProviderChatResponseData {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatMessageContent {
|
||||
Text(String),
|
||||
Parts(Vec<Value>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<ChatMessageContent>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
/// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python
|
||||
/// path reports so cost tracking sees the same numbers on either path.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct PromptTokensDetails {
|
||||
pub cached_tokens: u64,
|
||||
pub cache_creation_tokens: u64,
|
||||
pub text_tokens: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsUsage {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub prompt_tokens_details: PromptTokensDetails,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsChoiceMessage {
|
||||
pub role: String,
|
||||
// Whether an empty turn is `None` or `""` is the provider's choice, not a
|
||||
// shared invariant: Anthropic's transform ends on `merged_text or None`
|
||||
// while Converse assigns the joined string unconditionally. Each config
|
||||
// mirrors its own, so keep this optional and serialize it even when None.
|
||||
pub content: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsChoice {
|
||||
pub index: u64,
|
||||
pub message: ChatCompletionsChoiceMessage,
|
||||
pub finish_reason: String,
|
||||
}
|
||||
|
||||
/// The normalized response handed back to the host.
|
||||
///
|
||||
/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the
|
||||
/// `ModelResponse` it already created, and echoing the provider's own id here
|
||||
/// would change it. Pinned by `response_carries_no_id` in `tests.rs`.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsResponse {
|
||||
pub created: u64,
|
||||
pub model: String,
|
||||
pub choices: Vec<ChatCompletionsChoice>,
|
||||
pub usage: ChatCompletionsUsage,
|
||||
}
|
||||
|
|
@ -12,8 +12,30 @@ pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
|
|||
|
||||
/// Max characters of an upstream error body echoed across the call boundary
|
||||
/// before truncation, so provider bodies are bounded and data-minimized.
|
||||
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
pub(crate) const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Provider name used for Anthropic Messages when a deployment's provider model
|
||||
/// does not carry an explicit provider prefix.
|
||||
pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
|
||||
|
||||
/// Prefix identifying an Anthropic OAuth token. Mirrors Python's
|
||||
/// `ANTHROPIC_OAUTH_TOKEN_PREFIX`, which is what makes `validate_environment`
|
||||
/// authenticate with `authorization` and drop `x-api-key` entirely.
|
||||
pub(crate) const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat";
|
||||
|
||||
/// Full-request timeout ceiling for chat completions provider calls, in
|
||||
/// seconds. Mirrors the Python chat completions default.
|
||||
pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Connect timeout for chat completions provider calls, in seconds.
|
||||
pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// `object` field every non-streaming chat completion response carries.
|
||||
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
|
||||
|
||||
/// Placeholder Python substitutes for empty or whitespace-only message text,
|
||||
/// which Anthropic and Bedrock both reject. Must match
|
||||
/// `_EMPTY_TEXT_PLACEHOLDER` in
|
||||
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
|
||||
pub const EMPTY_TEXT_PLACEHOLDER: &str =
|
||||
"[System: Empty message content sanitised to satisfy protocol]";
|
||||
|
|
|
|||
|
|
@ -23,8 +23,19 @@ pub enum CoreError {
|
|||
Http { status: u16, body: String },
|
||||
#[error("upstream network error: {0}")]
|
||||
Network(String),
|
||||
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
|
||||
/// before any byte of the request went out. Nothing was billed, so a host
|
||||
/// that keeps a reference implementation can serve the request itself.
|
||||
/// A timeout is deliberately not this, since the provider may have received
|
||||
/// and answered the request already.
|
||||
#[error("could not reach the provider: {0}")]
|
||||
Connect(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
/// The request is outside the surface this route covers in Rust. Hosts that
|
||||
/// keep a reference implementation treat this as "fall back", not "fail".
|
||||
#[error("unsupported by the rust path: {0}")]
|
||||
Unsupported(&'static str),
|
||||
}
|
||||
|
||||
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
||||
|
|
|
|||
112
litellm-rust/crates/core/src/http_utils.rs
Normal file
112
litellm-rust/crates/core/src/http_utils.rs
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
//! Header and upstream-body helpers shared by every route module.
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
|
||||
/// Bound an upstream error body before it crosses a host boundary, so provider
|
||||
/// bodies stay data-minimized.
|
||||
pub fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
pub fn string_headers(
|
||||
context: &'static str,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"{context} extra_headers.{key} must be a string, got {}",
|
||||
json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
if !name.eq_ignore_ascii_case("authorization") {
|
||||
return false;
|
||||
}
|
||||
let value = value.trim();
|
||||
value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("bearer ")
|
||||
&& !value[7..].trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn truncate_leaves_short_bodies_untouched() {
|
||||
assert_eq!(truncate_error_body("short"), "short");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_bounds_long_bodies_by_characters() {
|
||||
let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
assert_eq!(
|
||||
truncated.chars().count(),
|
||||
UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = Map::from_iter([("x-trace".to_string(), json!(7))]);
|
||||
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_lookup_is_case_insensitive() {
|
||||
let headers = vec![("X-Api-Key".to_string(), "k".to_string())];
|
||||
assert!(has_header(&headers, "x-api-key"));
|
||||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bearer_detection_requires_a_non_empty_token() {
|
||||
assert!(has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer abc".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer ".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Basic abc".to_string()
|
||||
)]));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,8 +1,10 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod caching;
|
||||
pub mod call_lifecycle;
|
||||
pub mod chat_completions;
|
||||
pub mod constants;
|
||||
pub mod error;
|
||||
pub mod http_utils;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod providers;
|
||||
|
|
|
|||
|
|
@ -1,19 +1,15 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::CoreResult;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body};
|
||||
|
||||
const HEADER_CONTEXT: &str = "messages";
|
||||
|
||||
pub(super) fn messages_provider_config(
|
||||
provider: &str,
|
||||
|
|
@ -28,37 +24,5 @@ pub(super) fn messages_provider_config(
|
|||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"messages extra_headers.{key} must be a string, got {}",
|
||||
json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
if !name.eq_ignore_ascii_case("authorization") {
|
||||
return false;
|
||||
}
|
||||
let value = value.trim();
|
||||
value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("bearer ")
|
||||
&& !value[7..].trim().is_empty()
|
||||
})
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,444 @@
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
fn params(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn transform(model: &str, msgs: Value, opts: Value) -> Value {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_request(model, messages(msgs), params(opts))
|
||||
.expect("request transforms")
|
||||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_response("claude-sonnet-4-5", ProviderChatResponseData { body })
|
||||
}
|
||||
|
||||
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_messages_body_python_builds() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"max_tokens": 128, "temperature": 0.2}),
|
||||
);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]}
|
||||
],
|
||||
"system": [{"type": "text", "text": "be terse"}],
|
||||
"max_tokens": 128,
|
||||
"temperature": 0.2
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn omits_system_when_no_system_message_is_present() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert!(body.get("system").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_turns_and_wraps_every_text_in_a_block() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "two"}]},
|
||||
{"role": "assistant", "content": "ack"}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "one"},
|
||||
{"type": "text", "text": "two"}
|
||||
]},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "ack"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn right_strips_a_trailing_assistant_prefill_like_python() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "Argentina "}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"][1]["content"][0]["text"],
|
||||
json!("Argentina")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_every_supported_param_through_untouched() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"max_tokens": 64,
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.9,
|
||||
"stop_sequences": ["STOP"]
|
||||
}),
|
||||
);
|
||||
assert_eq!(body["max_tokens"], json!(64));
|
||||
assert_eq!(body["temperature"], json!(0.1));
|
||||
assert_eq!(body["top_p"], json!(0.9));
|
||||
assert_eq!(body["stop_sequences"], json!(["STOP"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_top_k_because_python_gates_it_by_model_below_this_point() {
|
||||
// `temperature` and `top_p` arrive already resolved, because
|
||||
// `map_openai_params` applies `_apply_sampling_param` to them before the
|
||||
// gate runs. `top_k` bypasses that and is gated inside `transform_request`,
|
||||
// the function this route replaces, so forwarding it would send `top_k` to
|
||||
// a model that removed sampling params and take a 400 after the call, where
|
||||
// Python drops it and succeeds.
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"top_k": 40})
|
||||
),
|
||||
Some(Unsupported("unrecognized request parameter"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_streaming_before_anything_else() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true, "max_tokens": 16})
|
||||
),
|
||||
Some(Unsupported("streaming"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_an_explicit_stream_false() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": false, "max_tokens": 16})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_any_param_outside_the_allowlist() {
|
||||
for param in [
|
||||
json!({"tools": []}),
|
||||
json!({"tool_choice": {"type": "auto"}}),
|
||||
json!({"thinking": {"type": "enabled"}}),
|
||||
json!({"system": "injected"}),
|
||||
json!({"metadata": {"user_id": "u1"}}),
|
||||
json!({"output_config": {"effort": "high"}}),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
|
||||
Some(Unsupported("unrecognized request parameter")),
|
||||
"expected {param} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_tool_calls_tool_results_and_multimodal_content() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": null, "tool_calls": [
|
||||
{"id": "c1", "type": "function",
|
||||
"function": {"name": "f", "arguments": "{}"}}
|
||||
]}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("unrecognized message field"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "ok"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("unrecognized message field"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "https://x/y.png"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_message_whose_content_list_is_empty() {
|
||||
// An empty list passes every per-part check, so without this it would reach
|
||||
// the provider as an empty `content` array and fail after the call rather
|
||||
// than declining to Python before it.
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": []}]), json!({})),
|
||||
Some(Unsupported("message without content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
|
||||
json!({})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_conversation_that_does_not_open_on_a_user_turn() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "assistant", "content": "prefill"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("conversation does not open on a user turn"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_a_plain_text_conversation() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "again"}]}
|
||||
]),
|
||||
json!({"max_tokens": 16, "temperature": 0.5})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_a_text_response_into_openai_shape() {
|
||||
let response = transform_response(json!({
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5-20260101",
|
||||
"content": [{"type": "text", "text": "hello"}, {"type": "text", "text": " there"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": null,
|
||||
"usage": {"input_tokens": 11, "output_tokens": 4}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.model, "claude-sonnet-4-5-20260101");
|
||||
assert_eq!(response.choices.len(), 1);
|
||||
assert_eq!(response.choices[0].index, 0);
|
||||
assert_eq!(response.choices[0].message.role, "assistant");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello there")
|
||||
);
|
||||
assert_eq!(response.choices[0].finish_reason, "stop");
|
||||
assert_eq!(response.usage.prompt_tokens, 11);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_cache_tokens_into_prompt_tokens_like_python() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 2,
|
||||
"cache_read_input_tokens": 5,
|
||||
"cache_creation_input_tokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.prompt_tokens, 18);
|
||||
assert_eq!(response.usage.total_tokens, 20);
|
||||
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
|
||||
assert_eq!(
|
||||
response.usage.prompt_tokens_details.cache_creation_tokens,
|
||||
3
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_max_tokens_stop_reason_to_length() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "max_tokens",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].finish_reason, "length");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_refusal_returns_the_completion_python_returns() {
|
||||
// `refusal` is a stop_reason, not a content block type, so the content is
|
||||
// ordinary text and this normalizes rather than declining. Python maps it
|
||||
// to content_filter in _FINISH_REASON_MAP and returns the completion.
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "I can't help with that."}],
|
||||
"stop_reason": "refusal",
|
||||
"usage": {"input_tokens": 9, "output_tokens": 6}
|
||||
}))
|
||||
.expect("a refusal still transforms");
|
||||
assert_eq!(response.choices[0].finish_reason, "content_filter");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("I can't help with that.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_no_content_rather_than_an_empty_string() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].message.content, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_carries_no_id_so_python_keeps_its_chatcmpl_id() {
|
||||
let response = transform_response(json!({
|
||||
"id": "msg_should_not_leak",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
let value = serde_json::to_value(response).expect("serializable");
|
||||
assert!(
|
||||
value.get("id").is_none(),
|
||||
"the rust response must not carry an id, got {value}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_response_carrying_a_non_text_block() {
|
||||
let err = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}],
|
||||
"stop_reason": "tool_use",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect_err("non-text block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("messages response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"),
|
||||
CoreError::MissingField("content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "content": []})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"content": [], "usage": {}})).expect_err("no model"),
|
||||
CoreError::MissingField("model")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_the_messages_url_and_x_api_key_auth() {
|
||||
let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
|
||||
.expect("url builds"),
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
);
|
||||
assert_eq!(
|
||||
config
|
||||
.auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None)
|
||||
.expect("auth resolves"),
|
||||
ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
value: "sk-x".to_string()
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
config.default_headers(),
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,211 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::chat_completions::conversation::{Conversation, build_conversation};
|
||||
use crate::chat_completions::transformation::{
|
||||
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
|
||||
unsupported_param,
|
||||
};
|
||||
use crate::chat_completions::types::{
|
||||
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage,
|
||||
ProviderChatRequestData, ProviderChatResponseData,
|
||||
};
|
||||
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::providers::anthropic::messages::transformation::{
|
||||
complete_anthropic_url, resolve_anthropic_api_key,
|
||||
};
|
||||
|
||||
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
|
||||
|
||||
/// Anthropic parameter names, post `map_openai_params`, that the Rust path can
|
||||
/// place verbatim in the Messages body.
|
||||
///
|
||||
/// `top_k` is deliberately absent even though the Messages API takes it.
|
||||
/// `temperature` and `top_p` reach this gate already resolved, because
|
||||
/// `map_openai_params` runs first and applies `_apply_sampling_param` to them.
|
||||
/// `top_k` bypasses `map_openai_params` entirely, so Python applies that same
|
||||
/// per-model gate inside `transform_request`, the function this route replaces.
|
||||
/// Forwarding it would send `top_k` to a model that removed sampling params and
|
||||
/// take a 400 after the call, where Python drops it and succeeds.
|
||||
const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"];
|
||||
|
||||
pub struct AnthropicChatCompletionsConfig;
|
||||
|
||||
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig =
|
||||
AnthropicChatCompletionsConfig;
|
||||
|
||||
fn text_block(text: &str) -> Value {
|
||||
json!({"type": "text", "text": text})
|
||||
}
|
||||
|
||||
fn anthropic_body(model: &str, conversation: &Conversation, params: Map<String, Value>) -> Value {
|
||||
let messages: Vec<Value> = conversation
|
||||
.turns
|
||||
.iter()
|
||||
.map(|turn| {
|
||||
json!({
|
||||
"role": turn.role.as_str(),
|
||||
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
|
||||
|
||||
let body = Map::from_iter(
|
||||
[
|
||||
("model".to_string(), json!(model)),
|
||||
("messages".to_string(), json!(messages)),
|
||||
]
|
||||
.into_iter()
|
||||
// Python builds `{"model", "messages", **optional_params}` with
|
||||
// `system` already folded into optional_params, so a caller-supplied
|
||||
// key of the same name wins here too.
|
||||
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
|
||||
.chain(params),
|
||||
);
|
||||
Value::Object(body)
|
||||
}
|
||||
|
||||
impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_anthropic_url(api_base, env_lookup))
|
||||
}
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
Ok(ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
value: resolve_anthropic_api_key(api_key, env_lookup)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
}
|
||||
|
||||
/// An OAuth bearer is the whole credential: Python's `validate_environment`
|
||||
/// authenticates with it and drops `x-api-key` rather than resolving one, so
|
||||
/// the resolved key must not be applied over the top. Any other forwarded
|
||||
/// `authorization` is unrelated to this header and does not defer, which is
|
||||
/// also what Python does: it sends the deployment's `x-api-key` alongside.
|
||||
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
name.eq_ignore_ascii_case("authorization")
|
||||
&& value
|
||||
.strip_prefix("Bearer ")
|
||||
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
|
||||
})
|
||||
}
|
||||
|
||||
fn supported_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_PARAMS
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(SUPPORTED_PARAMS, &[], optional_params)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
// Anthropic rejects a request whose first turn is not a user turn.
|
||||
// Python only repairs that under `litellm.modify_params`, which the
|
||||
// core cannot observe, so decline instead of guessing.
|
||||
.or_else(|| {
|
||||
(!build_conversation(messages).opens_on_user_turn())
|
||||
.then_some(Unsupported("conversation does not open on a user turn"))
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: anthropic_body(model, &build_conversation(&messages), optional_params),
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
_model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("messages response is not an object".into())
|
||||
})?;
|
||||
|
||||
let content = body
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("content"))?;
|
||||
// The route declines tool and thinking requests, so a non-text block
|
||||
// means the response carries something this path never asked for.
|
||||
// Decline rather than silently dropping it; the host falls back.
|
||||
if content
|
||||
.iter()
|
||||
.any(|block| block.get("type").and_then(Value::as_str) != Some("text"))
|
||||
{
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str))
|
||||
.collect();
|
||||
|
||||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
created: unix_now(),
|
||||
model: body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("model"))?
|
||||
.to_string(),
|
||||
choices: vec![ChatCompletionsChoice {
|
||||
index: 0,
|
||||
message: ChatCompletionsChoiceMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: (!text.is_empty()).then_some(text),
|
||||
},
|
||||
finish_reason: finish_reason_for(
|
||||
body.get("stop_reason")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(""),
|
||||
)
|
||||
.to_string(),
|
||||
}],
|
||||
usage: usage_from_parts(
|
||||
field("input_tokens"),
|
||||
field("output_tokens"),
|
||||
field("cache_read_input_tokens"),
|
||||
field("cache_creation_input_tokens"),
|
||||
),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod chat_completions;
|
||||
pub mod messages;
|
||||
|
|
|
|||
|
|
@ -8,11 +8,8 @@ use crate::audio_transcription::types::{
|
|||
};
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
|
||||
use super::aws_base::AwsAuthConfig;
|
||||
use super::constants::{
|
||||
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE,
|
||||
DEFAULT_BEDROCK_REGION,
|
||||
};
|
||||
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
|
||||
|
||||
|
|
@ -21,64 +18,6 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
|
|||
|
||||
pub struct BedrockAudioTranscriptionConfig;
|
||||
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(3))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
|
||||
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
|
|
@ -203,32 +142,6 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
|
|||
|
|
@ -12,13 +12,15 @@ use aws_sigv4::http_request::{
|
|||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
|
||||
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT,
|
||||
AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE,
|
||||
DEFAULT_SESSION_NAME_PREFIX,
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
|
||||
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
|
||||
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
|
||||
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
|
||||
SIGV4_COMPUTED_HEADER_NAMES,
|
||||
};
|
||||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
|
|
@ -401,6 +403,33 @@ fn default_session_name() -> String {
|
|||
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
|
||||
}
|
||||
|
||||
/// The subset of `headers` SigV4 should cover.
|
||||
///
|
||||
/// Python signs only these and reattaches the rest afterwards, so a forwarded
|
||||
/// client header cannot change the canonical request and invalidate the
|
||||
/// signature. Signing everything instead makes the request 403 on a header the
|
||||
/// caller supplied, on a deployment that works on the Python path.
|
||||
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name = name.to_ascii_lowercase();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|
||||
|| name.starts_with("x-amz-")
|
||||
|| name.starts_with("x-amzn-")
|
||||
})
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether the signer produces `name` itself.
|
||||
///
|
||||
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
|
||||
/// the wire next to the computed one.
|
||||
pub fn is_sigv4_computed_header(name: &str) -> bool {
|
||||
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
|
||||
}
|
||||
|
||||
pub fn sign_bedrock_post(
|
||||
url: &str,
|
||||
body: &[u8],
|
||||
|
|
@ -441,6 +470,121 @@ pub fn sign_bedrock_post(
|
|||
.collect())
|
||||
}
|
||||
|
||||
/// Model-id and region parsing shared by every Bedrock route.
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
// Python splits the whole ARN and takes field 3, the region. Stripping
|
||||
// `arn:` first shifts every field down one, so the region is field 2
|
||||
// here; field 3 is the account id.
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(2))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Credentials a host resolved through its own chain and handed down verbatim.
|
||||
///
|
||||
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
|
||||
/// profiles, STS and boto sessions) passes the result here so the core signs
|
||||
/// with exactly those. Without this the core would re-derive from ambient
|
||||
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
|
||||
/// environment outranks explicit keys in [`classify_auth`] and the two sides
|
||||
/// would sign as different principals.
|
||||
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
};
|
||||
let access_key_id = value("aws_access_key_id")?;
|
||||
let secret_access_key = value("aws_secret_access_key")?;
|
||||
Some(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
value("aws_session_token").map(str::to_string),
|
||||
None,
|
||||
"litellm-host-supplied",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
@ -458,6 +602,18 @@ mod tests {
|
|||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
|
||||
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
|
||||
// takes field 3. Stripping `arn:` first shifts every field down one, so
|
||||
// the region is field 2 here. Taking field 3 after the strip returns
|
||||
// the account id, which is not a region at all.
|
||||
let (_, region) = bedrock_model_id_and_region(
|
||||
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
|
||||
);
|
||||
assert_eq!(region.as_deref(), Some("us-west-2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_preserves_python_precedence() {
|
||||
let config = AwsAuthConfig {
|
||||
|
|
@ -610,6 +766,52 @@ mod tests {
|
|||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
|
||||
// Python signs only the AWS header set, so a header a caller forwarded
|
||||
// cannot change the canonical request. Signing it instead makes the
|
||||
// request 403 the moment anything on the wire rewrites or drops it.
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-request-id".to_string(), "abc-123".to_string());
|
||||
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
|
||||
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
|
||||
let signable = aws_signature_headers(&headers);
|
||||
|
||||
assert!(!signable.contains_key("x-request-id"));
|
||||
assert!(!signable.contains_key("Accept-Encoding"));
|
||||
// The AWS-prefixed one is genuinely part of the signature.
|
||||
assert!(signable.contains_key("x-amzn-trace-id"));
|
||||
assert!(signable.contains_key("Content-Type"));
|
||||
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&signable,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
SystemTime::UNIX_EPOCH,
|
||||
)
|
||||
.expect("signs");
|
||||
let authorization = signed
|
||||
.get("Authorization")
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("accept-encoding"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_matches_botocore_golden_vector() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,580 @@
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
fn params(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn transform(msgs: Value, opts: Value) -> Value {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_request(
|
||||
"anthropic.claude-sonnet-4-5-v1:0",
|
||||
messages(msgs),
|
||||
params(opts),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response(
|
||||
"anthropic.claude-sonnet-4-5-v1:0",
|
||||
ProviderChatResponseData { body },
|
||||
)
|
||||
}
|
||||
|
||||
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_converse_body_python_builds() {
|
||||
let body = transform(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"maxTokens": 128, "temperature": 0.2}),
|
||||
);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"inferenceConfig": {"maxTokens": 128, "temperature": 0.2},
|
||||
"messages": [{"role": "user", "content": [{"text": "hi"}]}],
|
||||
"system": [{"text": "be terse"}]
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn always_emits_inference_config_even_when_empty() {
|
||||
let body = transform(json!([{"role": "user", "content": "hi"}]), json!({}));
|
||||
assert_eq!(body["inferenceConfig"], json!({}));
|
||||
assert!(body.get("system").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn places_only_inference_params_in_inference_config() {
|
||||
let body = transform(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 64,
|
||||
"temperature": 0.1,
|
||||
"topP": 0.9,
|
||||
"stopSequences": ["STOP"]
|
||||
}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["inferenceConfig"],
|
||||
json!({"maxTokens": 64, "temperature": 0.1, "topP": 0.9, "stopSequences": ["STOP"]})
|
||||
);
|
||||
assert!(body.get("additionalModelRequestFields").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_user_turns_into_one_message() {
|
||||
let body = transform(
|
||||
json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "two"}]},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "three"}
|
||||
]),
|
||||
json!({}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [{"text": "one"}, {"text": "two"}]},
|
||||
{"role": "assistant", "content": [{"text": "ack"}]},
|
||||
{"role": "user", "content": [{"text": "three"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_streaming() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true})
|
||||
),
|
||||
Some(Unsupported("streaming"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_top_k_because_python_routes_it_by_base_model() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"topK": 40})
|
||||
),
|
||||
Some(Unsupported("unrecognized request parameter"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_tools_and_other_params_outside_the_allowlist() {
|
||||
for param in [
|
||||
json!({"tools": []}),
|
||||
json!({"tool_choice": {"auto": {}}}),
|
||||
json!({"thinking": {"type": "enabled"}}),
|
||||
json!({"requestMetadata": {"k": "v"}}),
|
||||
json!({"outputConfig": {}}),
|
||||
json!({"_parallel_tool_use_config": {}}),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
|
||||
Some(Unsupported("unrecognized request parameter")),
|
||||
"expected {param} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_blank_text_rather_than_substituting_the_anthropic_placeholder() {
|
||||
for content in [
|
||||
json!(""),
|
||||
json!(" "),
|
||||
json!([{"type": "text", "text": " "}]),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("blank message text")),
|
||||
"expected blank content {content} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_message_whose_content_list_is_empty() {
|
||||
// The blank-text check scans parts, so an empty list clears it; Converse
|
||||
// rejects an empty `content` array, which is a decline the core owes the
|
||||
// host before the call rather than an error after it.
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": []}]), json!({})),
|
||||
Some(Unsupported("message without content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
|
||||
json!({})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_conversation_that_opens_or_closes_on_an_assistant_turn() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "assistant", "content": "prefill"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported(
|
||||
"conversation does not run user turn to user turn"
|
||||
))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "prefill"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported(
|
||||
"conversation does not run user turn to user turn"
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_a_user_to_user_text_conversation() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
{"role": "user", "content": "again"}
|
||||
]),
|
||||
json!({"maxTokens": 16})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_converse_url_from_the_region_in_the_model_id() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| {
|
||||
None
|
||||
})
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_the_region_env_then_the_default_region() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string());
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env)
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None)
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_an_explicit_runtime_endpoint_over_the_api_base() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"}));
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(
|
||||
Some("https://ignored.example"),
|
||||
"anthropic.claude-v2",
|
||||
&overrides,
|
||||
&|_| None
|
||||
)
|
||||
.expect("url builds"),
|
||||
"https://vpce.internal/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signs_with_sigv4_in_the_resolved_region() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.auth(
|
||||
None,
|
||||
"eu-central-1/anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
&|_| None
|
||||
)
|
||||
.expect("auth resolves"),
|
||||
ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "eu-central-1".to_string()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() {
|
||||
// Python's get_request_headers reads `api_key` as the Bedrock bearer token
|
||||
// and only falls back to the env when the caller passed none, so each case
|
||||
// pins one of its precedence rules. Signing as the host principal when a
|
||||
// bearer identity is configured would cross an account and quota boundary.
|
||||
let bedrock_env =
|
||||
|key: &str| (key == "AWS_BEARER_TOKEN_BEDROCK").then(|| "from-env".to_string());
|
||||
let no_env = |_: &str| None;
|
||||
let resolve = |api_key, env: &dyn Fn(&str) -> Option<String>| {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG
|
||||
.auth(
|
||||
api_key,
|
||||
"eu-central-1/anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
env,
|
||||
)
|
||||
.expect("auth resolves")
|
||||
};
|
||||
let bearer = |token: &str| ChatCompletionsAuth::Bearer {
|
||||
token: token.to_string(),
|
||||
};
|
||||
let sigv4 = ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "eu-central-1".to_string(),
|
||||
};
|
||||
|
||||
// A caller-supplied key is the bearer token, and outranks the env.
|
||||
assert_eq!(
|
||||
resolve(Some("bedrock-api-key"), &bedrock_env),
|
||||
bearer("bedrock-api-key")
|
||||
);
|
||||
// No key, so the env supplies it.
|
||||
assert_eq!(resolve(None, &bedrock_env), bearer("from-env"));
|
||||
// An empty key is not a bearer token, and deliberately does NOT reach for
|
||||
// the env, which is what Python's `is not None` check does.
|
||||
assert_eq!(resolve(Some(""), &bedrock_env), sigv4);
|
||||
// Whitespace is truthy in Python, so it stays a bearer token rather than
|
||||
// silently becoming a host-credentialed SigV4 request.
|
||||
assert_eq!(resolve(Some(" "), &no_env), bearer(" "));
|
||||
// Neither present, so SigV4 as before.
|
||||
assert_eq!(resolve(None, &no_env), sigv4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_a_converse_response_into_openai_shape() {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"role": "assistant", "content": [
|
||||
{"text": "hello"}, {"text": " there"}
|
||||
]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.model, "anthropic.claude-sonnet-4-5-v1:0");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello there")
|
||||
);
|
||||
assert_eq!(response.choices[0].finish_reason, "stop");
|
||||
assert_eq!(response.usage.prompt_tokens, 11);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_converse_stop_reasons_python_maps() {
|
||||
for (provider_reason, expected) in [
|
||||
("end_turn", "stop"),
|
||||
("stop_sequence", "stop"),
|
||||
("max_tokens", "length"),
|
||||
("guardrail_intervened", "content_filter"),
|
||||
// Converse emits this one, and Python's `_FINISH_REASON_MAP` carries
|
||||
// it. Folding it into `stop` reports a filtered completion as a normal
|
||||
// one to anything keying on the finish reason.
|
||||
("content_filtered", "content_filter"),
|
||||
("content_filter", "content_filter"),
|
||||
] {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": provider_reason,
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(
|
||||
response.choices[0].finish_reason, expected,
|
||||
"stopReason {provider_reason}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_an_empty_converse_answer_as_an_empty_string_not_null() {
|
||||
// Converse assigns the joined text unconditionally
|
||||
// (`chat_completion_message["content"] = content_str`), unlike Anthropic's
|
||||
// `merged_text or None`, so an empty answer is `""` on both paths. A caller
|
||||
// calling `.strip()` on it would break on the Rust path alone. Reachable
|
||||
// through a filtered or guardrail-intervened response.
|
||||
for content in [json!([]), json!([{"text": ""}])] {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": content}},
|
||||
"stopReason": "content_filtered",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 0}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].message.content, Some(String::new()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_the_total_tokens_converse_sent_rather_than_recomputing_them() {
|
||||
// Python reads `usage["totalTokens"]` straight through here, where Anthropic
|
||||
// has no such field and adds the two counts instead. The two agree while the
|
||||
// gate declines every cache_control request, so this is what keeps them
|
||||
// agreeing if that ever widens.
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 4, "cacheReadInputTokens": 7, "totalTokens": 14}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(
|
||||
response.usage.total_tokens, 14,
|
||||
"provider total was recomputed"
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens, 17);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_the_computed_total_when_converse_omits_it() {
|
||||
// Python raises a KeyError on a body with no `totalTokens`. Reporting a zero
|
||||
// instead would be a worse divergence than the one above, so the computed
|
||||
// total stands in.
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 4}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.total_tokens, 14);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_cache_control_message_so_widening_the_gate_is_a_red_test() {
|
||||
// Converse only reports cache token counts when the request carries a
|
||||
// cachePoint block, which is why the provider total and the computed one
|
||||
// cannot disagree today. This is the tripwire: whoever widens the gate to
|
||||
// admit prompt caching has to come back and re-check the usage mapping
|
||||
// rather than discovering a silent number change in production.
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_converse_cache_tokens_into_prompt_tokens() {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {
|
||||
"inputTokens": 10,
|
||||
"outputTokens": 2,
|
||||
"cacheReadInputTokens": 5,
|
||||
"cacheWriteInputTokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.prompt_tokens, 18);
|
||||
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
|
||||
assert_eq!(
|
||||
response.usage.prompt_tokens_details.cache_creation_tokens,
|
||||
3
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_response_carrying_a_tool_use_block() {
|
||||
let err = transform_response(json!({
|
||||
"output": {"message": {"content": [
|
||||
{"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}}
|
||||
]}},
|
||||
"stopReason": "tool_use",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1}
|
||||
}))
|
||||
.expect_err("tool use block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("converse response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"usage": {}})).expect_err("no output"),
|
||||
CoreError::MissingField("output.message.content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_aws_call_configuration_without_serializing_it() {
|
||||
let call_config = json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIA",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_session_token": "token",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_profile_name": "litellm-stage",
|
||||
"aws_role_name": "role",
|
||||
"aws_session_name": "session",
|
||||
"aws_web_identity_token": "wit",
|
||||
"aws_sts_endpoint": "https://sts.example",
|
||||
"aws_external_id": "ext",
|
||||
"aws_bedrock_runtime_endpoint": "https://vpce.internal"
|
||||
});
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
call_config.clone()
|
||||
),
|
||||
None
|
||||
);
|
||||
let body = transform(json!([{"role": "user", "content": "hi"}]), call_config);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"inferenceConfig": {"maxTokens": 16},
|
||||
"messages": [{"role": "user", "content": [{"text": "hi"}]}]
|
||||
}),
|
||||
"aws call configuration must not reach the Converse body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_a_complete_converse_url_untouched() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let already_built =
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse";
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(
|
||||
Some(already_built),
|
||||
"anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
&|_| None
|
||||
)
|
||||
.expect("url builds"),
|
||||
already_built,
|
||||
"a host that encoded the model id itself must not have it re-derived"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_supplied_credentials_outrank_ambient_profile_and_role_state() {
|
||||
use crate::providers::bedrock::aws_base::host_supplied_credentials;
|
||||
|
||||
let supplied = params(json!({
|
||||
"aws_access_key_id": "AKIAHOST",
|
||||
"aws_secret_access_key": "hostsecret",
|
||||
"aws_session_token": "hosttoken"
|
||||
}));
|
||||
let credentials = host_supplied_credentials(&supplied).expect("host credentials");
|
||||
assert_eq!(credentials.access_key_id(), "AKIAHOST");
|
||||
assert_eq!(credentials.secret_access_key(), "hostsecret");
|
||||
assert_eq!(credentials.session_token(), Some("hosttoken"));
|
||||
|
||||
// Without a full static pair there is nothing to honor, so the core falls
|
||||
// back to deriving credentials itself.
|
||||
assert!(host_supplied_credentials(¶ms(json!({"aws_access_key_id": "AKIA"}))).is_none());
|
||||
assert!(
|
||||
host_supplied_credentials(¶ms(
|
||||
json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"})
|
||||
))
|
||||
.is_none()
|
||||
);
|
||||
assert!(host_supplied_credentials(&Map::new()).is_none());
|
||||
}
|
||||
|
|
@ -0,0 +1,297 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
|
||||
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
|
||||
use crate::chat_completions::transformation::{
|
||||
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
|
||||
unsupported_param,
|
||||
};
|
||||
use crate::chat_completions::types::{
|
||||
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse,
|
||||
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
|
||||
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
|
||||
|
||||
/// Converse parameter names, post `map_openai_params`, that the Rust path can
|
||||
/// place verbatim in `inferenceConfig`.
|
||||
///
|
||||
/// `topK` is deliberately absent: Python routes it to
|
||||
/// `additionalModelRequestFields` for Anthropic base models and to
|
||||
/// `inferenceConfig` otherwise, and that branch reads the model catalog the
|
||||
/// core cannot see.
|
||||
const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"];
|
||||
|
||||
/// Params that belong in `inferenceConfig`, in the order Python's
|
||||
/// `AmazonConverseConfig` declares them, so bodies compare cleanly.
|
||||
const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS;
|
||||
|
||||
const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint";
|
||||
|
||||
/// AWS call configuration a host passes down: consumed for signing and endpoint
|
||||
/// resolution, never serialized into the Converse body.
|
||||
const CONFIG_PARAMS: &[&str] = &[
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
"aws_region_name",
|
||||
"aws_session_name",
|
||||
"aws_profile_name",
|
||||
"aws_role_name",
|
||||
"aws_web_identity_token",
|
||||
"aws_sts_endpoint",
|
||||
"aws_external_id",
|
||||
AWS_BEDROCK_RUNTIME_ENDPOINT,
|
||||
];
|
||||
|
||||
const CONVERSE_PATH_SUFFIX: &str = "/converse";
|
||||
|
||||
pub struct BedrockChatCompletionsConfig;
|
||||
|
||||
pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig =
|
||||
BedrockChatCompletionsConfig;
|
||||
|
||||
fn converse_body(conversation: &Conversation, params: &Map<String, Value>) -> Value {
|
||||
let messages: Vec<Value> = conversation
|
||||
.turns
|
||||
.iter()
|
||||
.map(|turn| {
|
||||
json!({
|
||||
"role": turn.role.as_str(),
|
||||
"content": turn.texts.iter().map(|text| json!({"text": text})).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| {
|
||||
params
|
||||
.get(*name)
|
||||
.map(|value| ((*name).to_string(), value.clone()))
|
||||
}));
|
||||
|
||||
let system: Vec<Value> = conversation
|
||||
.system
|
||||
.iter()
|
||||
.map(|text| json!({"text": text}))
|
||||
.collect();
|
||||
|
||||
Value::Object(Map::from_iter(
|
||||
[
|
||||
(
|
||||
"inferenceConfig".to_string(),
|
||||
Value::Object(inference_config),
|
||||
),
|
||||
("messages".to_string(), json!(messages)),
|
||||
]
|
||||
.into_iter()
|
||||
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))),
|
||||
))
|
||||
}
|
||||
|
||||
fn has_blank_text(message: &ChatMessage) -> bool {
|
||||
match &message.content {
|
||||
None => false,
|
||||
Some(ChatMessageContent::Text(text)) => text.trim().is_empty(),
|
||||
Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| {
|
||||
part.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(|text| text.trim().is_empty())
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let (model_id, model_region) = bedrock_model_id_and_region(model);
|
||||
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
|
||||
let endpoint = optional_params
|
||||
.get(AWS_BEDROCK_RUNTIME_ENDPOINT)
|
||||
.and_then(Value::as_str)
|
||||
.or(api_base)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion));
|
||||
let endpoint = endpoint.trim_end_matches('/');
|
||||
// A host that already built the full Converse URL (LiteLLM's Python
|
||||
// path encodes the model id itself) passes it through untouched, the
|
||||
// way the Anthropic config leaves a complete `/v1/messages` URL alone.
|
||||
if endpoint.ends_with(CONVERSE_PATH_SUFFIX) {
|
||||
return Ok(endpoint.to_string());
|
||||
}
|
||||
Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}"))
|
||||
}
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
// Python reads `api_key` as the Bedrock bearer token and consults the
|
||||
// env only when the caller passed none, so a caller-supplied empty key
|
||||
// falls through to SigV4 without reaching for the environment. An
|
||||
// all-whitespace token stays a bearer token here because Python sends
|
||||
// it too: treating it as absent would sign as the host principal
|
||||
// instead, which is the identity swap this branch exists to prevent.
|
||||
let bearer = match api_key {
|
||||
Some(key) => Some(key.to_string()),
|
||||
None => env_lookup(AWS_BEARER_TOKEN_BEDROCK),
|
||||
}
|
||||
.filter(|token| !token.is_empty());
|
||||
if let Some(token) = bearer {
|
||||
return Ok(ChatCompletionsAuth::Bearer { token });
|
||||
}
|
||||
let (_, model_region) = bedrock_model_id_and_region(model);
|
||||
Ok(ChatCompletionsAuth::AwsSigV4 {
|
||||
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
|
||||
})
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[("Content-Type", "application/json")]
|
||||
}
|
||||
|
||||
fn supported_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_PARAMS
|
||||
}
|
||||
|
||||
fn config_params(&self) -> &'static [&'static str] {
|
||||
CONFIG_PARAMS
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
// Python's Converse translation drops blank text blocks instead of
|
||||
// substituting the placeholder the shared conversation builder
|
||||
// applies, so decline blank text rather than diverge.
|
||||
.or_else(|| {
|
||||
messages
|
||||
.iter()
|
||||
.any(has_blank_text)
|
||||
.then_some(Unsupported("blank message text"))
|
||||
})
|
||||
// Converse has no assistant prefill: Python inserts a continue turn
|
||||
// when a conversation opens or closes on an assistant message, and
|
||||
// only under `litellm.modify_params`, which the core cannot see.
|
||||
// Declining both ends also keeps the shared builder's final
|
||||
// assistant right-strip (an Anthropic rule) unreachable here.
|
||||
.or_else(|| {
|
||||
let conversation = build_conversation(messages);
|
||||
let ends_on_assistant = conversation
|
||||
.turns
|
||||
.last()
|
||||
.is_some_and(|turn| turn.role == TurnRole::Assistant);
|
||||
(!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
|
||||
"conversation does not run user turn to user turn",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
_model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: converse_body(&build_conversation(&messages), &optional_params),
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("converse response is not an object".into())
|
||||
})?;
|
||||
|
||||
let content = body
|
||||
.get("output")
|
||||
.and_then(|output| output.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("output.message.content"))?;
|
||||
// The route declines tool requests, so anything other than a text block
|
||||
// is something this path never asked for. Decline; the host falls back.
|
||||
if content.iter().any(|block| {
|
||||
block
|
||||
.as_object()
|
||||
.is_none_or(|block| block.len() != 1 || !block.contains_key("text"))
|
||||
}) {
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str))
|
||||
.collect();
|
||||
|
||||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
let computed = usage_from_parts(
|
||||
field("inputTokens"),
|
||||
field("outputTokens"),
|
||||
field("cacheReadInputTokens"),
|
||||
field("cacheWriteInputTokens"),
|
||||
);
|
||||
// Converse reports `totalTokens` and Python passes it straight through,
|
||||
// where Anthropic has no such field and Python adds the two counts
|
||||
// instead, so only this provider overrides the computed total. Python
|
||||
// does a bare `usage["totalTokens"]` lookup, so a body without the key
|
||||
// raises there rather than reporting a zero; fall back to the computed
|
||||
// total, which is the closest thing to that without failing the call.
|
||||
let usage = ChatCompletionsUsage {
|
||||
total_tokens: usage
|
||||
.get("totalTokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(computed.total_tokens),
|
||||
..computed
|
||||
};
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
created: unix_now(),
|
||||
// Converse echoes no model id, so Python reports the requested one.
|
||||
model: model.to_string(),
|
||||
choices: vec![ChatCompletionsChoice {
|
||||
index: 0,
|
||||
message: ChatCompletionsChoiceMessage {
|
||||
role: "assistant".to_string(),
|
||||
// Converse assigns the joined string unconditionally, so an
|
||||
// empty response is `""` here and not `None` as it is on
|
||||
// Anthropic. A caller calling `.strip()` on it would break
|
||||
// on this path alone.
|
||||
content: Some(text),
|
||||
},
|
||||
finish_reason: finish_reason_for(
|
||||
body.get("stopReason").and_then(Value::as_str).unwrap_or(""),
|
||||
)
|
||||
.to_string(),
|
||||
}],
|
||||
usage,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
|
|
@ -11,6 +11,31 @@ pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
|
|||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
|
||||
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
|
||||
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
|
||||
/// Python's `_filter_headers_for_aws_signature` allowlist.
|
||||
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
|
||||
"host",
|
||||
"content-type",
|
||||
"date",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"x-amz-content-sha256",
|
||||
"x-amz-algorithm",
|
||||
"x-amz-credential",
|
||||
"x-amz-signedheaders",
|
||||
"x-amz-signature",
|
||||
];
|
||||
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
|
||||
/// which the reattach loop skips so a caller's copy cannot ride alongside the
|
||||
/// computed one.
|
||||
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
|
||||
"authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"date",
|
||||
];
|
||||
pub const BEDROCK_SERVICE: &str = "bedrock";
|
||||
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
|
||||
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
|
||||
|
|
|
|||
|
|
@ -5,4 +5,5 @@
|
|||
#[cfg(feature = "bedrock-auth")]
|
||||
pub mod audio_transcription;
|
||||
pub mod aws_base;
|
||||
pub mod chat_completions;
|
||||
mod constants;
|
||||
|
|
|
|||
|
|
@ -6,6 +6,10 @@ use litellm_ai_gateway::io::audio_transcription::{
|
|||
};
|
||||
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
|
||||
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
|
||||
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
|
||||
use litellm_core::chat_completions::{
|
||||
chat_completions as run_chat_completions, chat_completions_decline_reason,
|
||||
};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::messages::messages as run_messages;
|
||||
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
|
|
@ -16,6 +20,20 @@ use serde_json::{Map, Value};
|
|||
|
||||
mod gil;
|
||||
|
||||
pyo3::create_exception!(
|
||||
_native,
|
||||
RustBridgeDeclined,
|
||||
pyo3::exceptions::PyException,
|
||||
"The route declined before calling the provider, so the host may retry on its own path."
|
||||
);
|
||||
|
||||
pyo3::create_exception!(
|
||||
_native,
|
||||
RustUpstreamError,
|
||||
pyo3::exceptions::PyException,
|
||||
"The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
|
||||
);
|
||||
|
||||
type MarshaledOcrInputs = (
|
||||
Value,
|
||||
Option<Map<String, Value>>,
|
||||
|
|
@ -45,6 +63,15 @@ fn messages_response_to_py(
|
|||
json_to_py(py, value)
|
||||
}
|
||||
|
||||
fn chat_completions_response_to_py(
|
||||
py: Python<'_>,
|
||||
response: ChatCompletionsResponse,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let value =
|
||||
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
json_to_py(py, value)
|
||||
}
|
||||
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
|
|
@ -56,6 +83,33 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
|||
}
|
||||
}
|
||||
|
||||
/// Map a core error for a route whose host keeps a Python implementation.
|
||||
///
|
||||
/// The distinction the host needs is whether the provider was already called.
|
||||
/// Everything raised before the request goes out is safe for the host to retry
|
||||
/// on its own path; anything after it is not, because the provider has already
|
||||
/// done the work and billed for it.
|
||||
fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Unsupported(_)
|
||||
| CoreError::Auth(_)
|
||||
| CoreError::InvalidProvider(_)
|
||||
| CoreError::InvalidRequest(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_)
|
||||
| CoreError::Routing(_)
|
||||
// Nothing reached the provider, so serving it on Python cannot double
|
||||
// bill and is the only way the caller gets an answer at all.
|
||||
| CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
|
||||
CoreError::Http { status, body } => {
|
||||
RustUpstreamError::new_err((status, format!("{status}: {body}")))
|
||||
}
|
||||
CoreError::Network(message) | CoreError::InvalidResponse(message) => {
|
||||
RustUpstreamError::new_err((0u16, message))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_object_to_map(
|
||||
py: Python<'_>,
|
||||
name: &'static str,
|
||||
|
|
@ -430,6 +484,143 @@ fn amessages(
|
|||
})
|
||||
}
|
||||
|
||||
type MarshaledChatCompletionsInputs = (
|
||||
Value,
|
||||
Map<String, Value>,
|
||||
Option<Map<String, Value>>,
|
||||
Option<Duration>,
|
||||
);
|
||||
|
||||
fn marshal_chat_completions_inputs(
|
||||
py: Python<'_>,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<MarshaledChatCompletionsInputs> {
|
||||
let messages = py_to_json(py, messages.bind(py))?;
|
||||
if !messages.is_array() {
|
||||
return Err(PyValueError::new_err("messages must be a list"));
|
||||
}
|
||||
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
|
||||
let extra_headers = match extra_headers {
|
||||
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
|
||||
None => None,
|
||||
};
|
||||
Ok((
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
optional_timeout(timeout_seconds),
|
||||
))
|
||||
}
|
||||
|
||||
/// The decline reason for this request, or `None` when the Rust path accepts
|
||||
/// it. Resolves no credentials and performs no I/O, so a host can ask before
|
||||
/// committing to either path.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
|
||||
fn chat_completions_decline(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
custom_llm_provider: Option<String>,
|
||||
) -> PyResult<Option<String>> {
|
||||
let messages = py_to_json(py, messages.bind(py))?;
|
||||
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
|
||||
Ok(chat_completions_decline_reason(
|
||||
&model,
|
||||
custom_llm_provider.as_deref(),
|
||||
messages,
|
||||
&optional_params,
|
||||
)
|
||||
.map(str::to_string))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn chat_completions(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
|
||||
py,
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
let result = gil::release_gil(py, || {
|
||||
pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions(
|
||||
ChatCompletionsRequest {
|
||||
model: &model,
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: custom_llm_provider.as_deref(),
|
||||
extra_headers,
|
||||
timeout,
|
||||
},
|
||||
))
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(response) => chat_completions_response_to_py(py, response),
|
||||
Err(err) => Err(chat_completions_error_to_pyerr(err)),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn achat_completions(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
|
||||
py,
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let response = run_chat_completions(ChatCompletionsRequest {
|
||||
model: &model,
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: custom_llm_provider.as_deref(),
|
||||
extra_headers,
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.map_err(chat_completions_error_to_pyerr)?;
|
||||
|
||||
Python::attach(|py| chat_completions_response_to_py(py, response))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
|
|
@ -439,12 +630,18 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
|||
|
||||
#[pymodule]
|
||||
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let py = module.py();
|
||||
module.add_function(wrap_pyfunction!(ocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(aocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(transcription, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(atranscription, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(messages, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(amessages, module)?)?;
|
||||
module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
|
||||
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())?;
|
||||
module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(chat_completions, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(achat_completions, module)?)?;
|
||||
module.add_class::<ResponsesWebSocketConnection>()?;
|
||||
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
|
||||
Ok(())
|
||||
|
|
|
|||
|
|
@ -221,7 +221,7 @@ overwrite_user_with_key_hash: bool = (
|
|||
bedrock_request_metadata_fields: Optional[Sequence[str]] = (
|
||||
None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata`
|
||||
)
|
||||
store_audit_logs = False # Enterprise feature, allow users to see audit logs
|
||||
store_audit_logs: bool | None = None
|
||||
skip_system_message_in_guardrail: bool = False
|
||||
skip_tool_message_in_guardrail: bool = False
|
||||
### end of callbacks #############
|
||||
|
|
@ -453,6 +453,7 @@ max_end_user_budget_id: Optional[str] = None
|
|||
# backwards compatibility — arbitrary client-supplied identifiers still
|
||||
# pass through unchanged.
|
||||
validate_end_user_id_in_db: bool = False
|
||||
block_requests_for_models_without_pricing: bool = False
|
||||
disable_end_user_cost_tracking: Optional[bool] = None
|
||||
disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
|
||||
enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
|
||||
|
|
|
|||
|
|
@ -8,6 +8,12 @@ from logging import Formatter
|
|||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE,
|
||||
MAX_STRING_LENGTH_STDOUT_LOG,
|
||||
)
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int
|
||||
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, redact_structured_value
|
||||
|
|
@ -101,7 +107,7 @@ class SecretRedactionFilter(logging.Filter):
|
|||
# Redact exception tracebacks
|
||||
if record.exc_info and record.exc_info[1] is not None:
|
||||
try:
|
||||
record.exc_text = _redact_string(self._formatter.formatException(record.exc_info))
|
||||
record.exc_text = _redact_string(record.exc_text or self._formatter.formatException(record.exc_info))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
|
@ -116,6 +122,72 @@ class SecretRedactionFilter(logging.Filter):
|
|||
_secret_filter: Final = SecretRedactionFilter()
|
||||
|
||||
|
||||
def _get_max_string_length_stdout_log() -> int:
|
||||
"""Read the limit per record so a value loaded later via proxy config
|
||||
environment_variables is honored."""
|
||||
return get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", MAX_STRING_LENGTH_STDOUT_LOG)
|
||||
|
||||
|
||||
def _stdout_truncation_marker(skipped_chars: int) -> str:
|
||||
return (
|
||||
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. "
|
||||
f"{LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE}) ..."
|
||||
)
|
||||
|
||||
|
||||
def _truncate_for_stdout_log(text: str, limit: int) -> str:
|
||||
kept_chars: Final = limit - len(_stdout_truncation_marker(len(text)))
|
||||
if kept_chars <= 0:
|
||||
return text[:limit]
|
||||
head_chars: Final = kept_chars // 2
|
||||
tail_chars: Final = kept_chars - head_chars
|
||||
return f"{text[:head_chars]}{_stdout_truncation_marker(len(text) - kept_chars)}{text[-tail_chars:]}"
|
||||
|
||||
|
||||
class StdoutLogTruncationFilter(logging.Filter):
|
||||
"""Bounds how much of an oversized log line reaches stdout.
|
||||
|
||||
A provider error string can echo the whole request payload, so one failed agentic
|
||||
request writes hundreds of KB to stdout, repeatedly as the exception propagates from
|
||||
the router to the proxy handler and into its traceback, all inline on the event loop.
|
||||
|
||||
DEBUG records pass through untouched, since dumping full payloads is the point of
|
||||
`--detailed_debug`, and logging callbacks (OTEL, Datadog, etc.) don't run through
|
||||
logging filters at all, so they still get the untruncated error.
|
||||
"""
|
||||
|
||||
_formatter = logging.Formatter()
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if record.levelno < logging.INFO:
|
||||
return True
|
||||
|
||||
limit: Final = _get_max_string_length_stdout_log()
|
||||
if limit <= 0:
|
||||
return True
|
||||
|
||||
try:
|
||||
message: Final = record.getMessage()
|
||||
except (TypeError, ValueError):
|
||||
return True
|
||||
|
||||
if len(message) > limit:
|
||||
record.msg = _truncate_for_stdout_log(message, limit) # rebind-ok: the Filter interface mutates the record
|
||||
record.args = None # rebind-ok: args are consumed by the truncated message above
|
||||
|
||||
if isinstance(record.exc_info, tuple):
|
||||
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
|
||||
if len(exc_text) > limit:
|
||||
record.exc_text = _truncate_for_stdout_log( # rebind-ok: the Filter interface mutates the record
|
||||
exc_text, limit
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
_stdout_truncation_filter: Final = StdoutLogTruncationFilter()
|
||||
|
||||
|
||||
class CorrelationContextFilter(logging.Filter):
|
||||
"""Stamps each log record with the current request's trace_id and session_id from contextvars.
|
||||
|
||||
|
|
@ -301,6 +373,7 @@ def _setup_json_exception_handlers(formatter):
|
|||
error_handler: Final = logging.StreamHandler()
|
||||
error_handler.setFormatter(formatter)
|
||||
error_handler.addFilter(_secret_filter)
|
||||
error_handler.addFilter(_stdout_truncation_filter)
|
||||
error_handler.addFilter(_correlation_filter)
|
||||
|
||||
# Setup excepthook for uncaught exceptions
|
||||
|
|
@ -365,6 +438,12 @@ verbose_router_logger.addHandler(handler)
|
|||
verbose_proxy_logger.addHandler(handler)
|
||||
verbose_logger.addHandler(handler)
|
||||
|
||||
# Filters attached to the logger, not the handler, survive callers swapping in their own
|
||||
# handlers (JSON mode, uvicorn log config, a host app's root handler).
|
||||
verbose_router_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_proxy_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_logger.addFilter(_stdout_truncation_filter)
|
||||
|
||||
|
||||
def _suppress_loggers():
|
||||
"""Suppress noisy loggers at INFO level"""
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing import Final
|
|||
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
from redis.credentials import CredentialProvider
|
||||
|
||||
from litellm import get_secret, get_secret_str
|
||||
from litellm._redis_credential_provider import (
|
||||
|
|
@ -134,6 +135,7 @@ def _get_redis_cluster_kwargs(client=None):
|
|||
"ssl_check_hostname",
|
||||
"ssl_ca_certs",
|
||||
"redis_connect_func", # Needed for sync clusters and IAM detection
|
||||
"credential_provider",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_redis_ad_token",
|
||||
|
|
@ -549,14 +551,22 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
|
|||
return sentinel.master_for(service_name, **connection_kwargs)
|
||||
|
||||
|
||||
def _sentinel_auth_kwargs(connection_kwargs: dict, sentinel_password: str | None) -> dict:
|
||||
"""The Sentinel monitors are separate servers that authenticate with their own password, so the
|
||||
data node's credential provider never belongs on them: leaving it there makes redis-py send the
|
||||
data node's token to a monitor, which fails whether the monitor is unauthenticated or has its
|
||||
own password."""
|
||||
kept: Final = ((k, v) for k, v in connection_kwargs.items() if k != "credential_provider")
|
||||
return dict(kept, password=sentinel_password)
|
||||
|
||||
|
||||
def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
||||
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
|
||||
sentinel_password: Final = redis_kwargs.get("sentinel_password")
|
||||
service_name: Final = redis_kwargs.get("service_name")
|
||||
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
|
||||
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
|
||||
sentinel_kwargs: Final = dict(connection_kwargs)
|
||||
sentinel_kwargs["password"] = sentinel_password
|
||||
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
|
@ -574,6 +584,36 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
|||
return sentinel.master_for(service_name, **connection_kwargs)
|
||||
|
||||
|
||||
def _async_credential_provider(redis_connect_func: object | None) -> CredentialProvider | None:
|
||||
"""The Azure AD and GCP IAM connect funcs run their AUTH exchange with the blocking client
|
||||
API, so on an async connection their ``send_command``/``read_response`` calls return
|
||||
coroutines nobody awaits and every connect fails. Async paths authenticate through a
|
||||
``CredentialProvider`` instead, which redis-py consults per connection so the token stays
|
||||
fresh. Any other ``redis_connect_func`` is left where it is, since redis-py awaits it
|
||||
itself when it is a coroutine function."""
|
||||
gcp_service_account: Final = getattr(redis_connect_func, "_gcp_service_account", None)
|
||||
if gcp_service_account is not None:
|
||||
return GCPIAMCredentialProvider(gcp_service_account)
|
||||
|
||||
azure_credential: Final = getattr(redis_connect_func, "_azure_credential", None)
|
||||
if azure_credential is not None:
|
||||
return AzureADCredentialProvider(azure_credential, username=os.environ.get("REDIS_USERNAME") or None)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
|
||||
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
|
||||
which supersedes any static username or password redis-py would otherwise reject it with."""
|
||||
credential_provider: Final = _async_credential_provider(redis_kwargs.get("redis_connect_func"))
|
||||
if credential_provider is None:
|
||||
return redis_kwargs
|
||||
|
||||
superseded: Final = frozenset({"redis_connect_func", "username", "password"})
|
||||
kept: Final = ((k, v) for k, v in redis_kwargs.items() if k not in superseded)
|
||||
return dict(kept, credential_provider=credential_provider) # mutable-ok: the branches below mutate these kwargs
|
||||
|
||||
|
||||
def get_redis_client(**env_overrides):
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
|
||||
|
|
@ -600,7 +640,7 @@ def get_redis_async_client(
|
|||
connection_pool: async_redis.BlockingConnectionPool | None = None,
|
||||
**env_overrides,
|
||||
) -> async_redis.Redis | async_redis.RedisCluster:
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
from redis.cluster import ClusterNode
|
||||
|
|
@ -611,28 +651,12 @@ def get_redis_async_client(
|
|||
if arg in args:
|
||||
cluster_kwargs[arg] = redis_kwargs[arg]
|
||||
|
||||
# Handle GCP IAM authentication for async clusters
|
||||
redis_connect_func = cluster_kwargs.pop("redis_connect_func", None)
|
||||
|
||||
# Use a CredentialProvider so the IAM token is regenerated on every new
|
||||
# connection — mirrors the sync path where redis_connect_func is invoked
|
||||
# per connection. Without this, the token would expire after ~1 hour.
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
# Handle Azure AD authentication for async clusters via CredentialProvider
|
||||
# so the credential's internal cache + silent refresh runs per connection
|
||||
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
|
||||
new_startup_nodes: Final[list[ClusterNode]] = []
|
||||
|
||||
for item in redis_kwargs["startup_nodes"]:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
cluster_kwargs.pop("startup_nodes", None)
|
||||
cluster_kwargs.pop("redis_connect_func", None)
|
||||
|
||||
# Default to a periodic health check + TCP keepalive so a connection silently dropped
|
||||
# by a cluster restart (e.g. ElastiCache Serverless maintenance) is revalidated and
|
||||
|
|
@ -667,19 +691,6 @@ def get_redis_async_client(
|
|||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_async_redis_sentinel(redis_kwargs)
|
||||
|
||||
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
|
||||
# Redis client. The async client doesn't support redis_connect_func, but it
|
||||
# does honour credential_provider — which is called per connection, so the
|
||||
# underlying SDK can refresh tokens silently before they expire.
|
||||
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
|
||||
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
|
||||
|
||||
if connection_pool is not None:
|
||||
|
|
@ -693,7 +704,7 @@ def get_redis_async_client(
|
|||
def get_redis_connection_pool(
|
||||
**env_overrides,
|
||||
) -> async_redis.BlockingConnectionPool | None:
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
|
||||
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
|
|
@ -714,18 +725,6 @@ def get_redis_connection_pool(
|
|||
)
|
||||
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
|
||||
|
||||
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
|
||||
# connections re-fetch tokens via the SDK's internal cache + silent refresh
|
||||
# rather than reusing a single token captured at pool creation.
|
||||
redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None)
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
|
||||
if redis_kwargs.pop("ssl", None):
|
||||
redis_kwargs["connection_class"] = async_redis.SSLConnection
|
||||
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
|
||||
|
|
|
|||
|
|
@ -7,9 +7,10 @@ import hashlib
|
|||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Final, NamedTuple, cast
|
||||
from typing import Any, Final, NamedTuple, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import (
|
||||
|
|
@ -38,11 +39,59 @@ class WXORequestParams(NamedTuple):
|
|||
thread_id: str | None
|
||||
|
||||
|
||||
class WXOLitellmParams(TypedDict, total=False):
|
||||
"""litellm_params keys read when routing an A2A request to watsonx Orchestrate."""
|
||||
|
||||
cp4d_host: ReadOnly[str]
|
||||
instance_id: ReadOnly[str]
|
||||
wxo_agent_id: ReadOnly[str]
|
||||
api_key: ReadOnly[str]
|
||||
username: ReadOnly[str | None]
|
||||
auth_mode: ReadOnly[str]
|
||||
thread_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _IBMCloudTokenBody(TypedDict):
|
||||
"""Fields read from the IBM Cloud IAM token response."""
|
||||
|
||||
access_token: ReadOnly[str]
|
||||
expires_in: ReadOnly[NotRequired[int]]
|
||||
|
||||
|
||||
class _CP4DTokenBody(TypedDict):
|
||||
"""Fields read from the CP4D authorize response."""
|
||||
|
||||
token: ReadOnly[str]
|
||||
expiration: ReadOnly[NotRequired[float]]
|
||||
|
||||
|
||||
class _WXORun(TypedDict, total=False):
|
||||
"""Fields the handler reads from a WXO run object or run event."""
|
||||
|
||||
status: ReadOnly[str]
|
||||
run_id: ReadOnly[str]
|
||||
id: ReadOnly[str]
|
||||
|
||||
|
||||
class _SSELineSource(Protocol):
|
||||
def aiter_lines(self) -> AsyncIterator[str]: ...
|
||||
|
||||
|
||||
class _WXOView(TypedDict, total=False):
|
||||
"""Typed reads of otherwise untyped watsonx Orchestrate and httpx values."""
|
||||
|
||||
ibm_cloud_token: ReadOnly[_IBMCloudTokenBody]
|
||||
cp4d_token: ReadOnly[_CP4DTokenBody]
|
||||
run: ReadOnly[_WXORun]
|
||||
content_type: ReadOnly[str]
|
||||
sse_source: ReadOnly[_SSELineSource]
|
||||
|
||||
|
||||
class WatsonxOrchestrateHandler:
|
||||
@staticmethod
|
||||
def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler:
|
||||
return get_async_httpx_client(
|
||||
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
|
||||
llm_provider=httpxSpecialProvider.A2AProvider,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
|
|
@ -57,7 +106,7 @@ class WatsonxOrchestrateHandler:
|
|||
return hashlib.sha256(material.encode()).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int:
|
||||
def _cp4d_token_ttl_seconds(expiration: float, now_wall: float | None = None) -> int:
|
||||
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
|
||||
expires_at: Final = int(expiration)
|
||||
wall: Final = now_wall if now_wall is not None else time.time()
|
||||
|
|
@ -90,9 +139,9 @@ class WatsonxOrchestrateHandler:
|
|||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
token = str(payload["access_token"])
|
||||
ttl_s = int(payload.get("expires_in", 3600))
|
||||
iam_payload: Final[_WXOView] = {"ibm_cloud_token": response.json()}
|
||||
token = str(iam_payload["ibm_cloud_token"]["access_token"])
|
||||
ttl_s = int(iam_payload["ibm_cloud_token"].get("expires_in", 3600))
|
||||
else:
|
||||
if not username:
|
||||
raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'")
|
||||
|
|
@ -103,9 +152,9 @@ class WatsonxOrchestrateHandler:
|
|||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
token = str(payload["token"])
|
||||
expiration: Final = payload.get("expiration")
|
||||
cp4d_payload: Final[_WXOView] = {"cp4d_token": response.json()}
|
||||
token = str(cp4d_payload["cp4d_token"]["token"])
|
||||
expiration: Final = cp4d_payload["cp4d_token"].get("expiration")
|
||||
if expiration is None:
|
||||
ttl_s = 3600
|
||||
else:
|
||||
|
|
@ -118,6 +167,16 @@ class WatsonxOrchestrateHandler:
|
|||
del _token_cache[stale_key]
|
||||
return token
|
||||
|
||||
@staticmethod
|
||||
def _run_body(response: httpx.Response) -> _WXORun:
|
||||
view: Final[_WXOView] = {"run": response.json()}
|
||||
return view["run"]
|
||||
|
||||
@staticmethod
|
||||
def _decode_run_event(payload: str | bytes) -> _WXORun:
|
||||
view: Final[_WXOView] = {"run": json.loads(payload)}
|
||||
return view["run"]
|
||||
|
||||
@staticmethod
|
||||
async def _poll_run(
|
||||
base_url: str,
|
||||
|
|
@ -126,14 +185,14 @@ class WatsonxOrchestrateHandler:
|
|||
client: AsyncHTTPHandler,
|
||||
max_attempts: int = _MAX_POLL_ATTEMPTS,
|
||||
interval_s: float = _POLL_INTERVAL_S,
|
||||
) -> dict[str, Any]:
|
||||
) -> _WXORun:
|
||||
url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}"
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
await asyncio.sleep(interval_s)
|
||||
response = await client.get(url, headers=auth_headers)
|
||||
response.raise_for_status()
|
||||
result: dict[str, Any] = response.json()
|
||||
result = WatsonxOrchestrateHandler._run_body(response)
|
||||
status = result.get("status", "")
|
||||
verbose_logger.debug("WXO: Poll %s/%s run='%s' status='%s'", attempt + 1, max_attempts, run_id, status)
|
||||
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
|
||||
|
|
@ -145,11 +204,11 @@ class WatsonxOrchestrateHandler:
|
|||
|
||||
@staticmethod
|
||||
async def _get_successful_run_data(
|
||||
run_data: dict[str, Any],
|
||||
run_data: _WXORun,
|
||||
base_url: str,
|
||||
auth_headers: dict[str, str],
|
||||
client: AsyncHTTPHandler,
|
||||
) -> dict[str, Any]:
|
||||
) -> _WXORun:
|
||||
status = run_data.get("status", "")
|
||||
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
|
||||
run_id: Final = run_data.get("run_id") or run_data.get("id") or ""
|
||||
|
|
@ -170,15 +229,16 @@ class WatsonxOrchestrateHandler:
|
|||
|
||||
@staticmethod
|
||||
async def _accumulate_wxo_sse_text(response: Any) -> str:
|
||||
source: Final[_WXOView] = {"sse_source": response}
|
||||
accumulated_text = ""
|
||||
async for line in response.aiter_lines():
|
||||
async for line in source["sse_source"].aiter_lines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data_str = line[5:].strip()
|
||||
if not data_str or data_str == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
event = json.loads(data_str)
|
||||
event = WatsonxOrchestrateHandler._decode_run_event(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event)
|
||||
|
|
@ -187,7 +247,7 @@ class WatsonxOrchestrateHandler:
|
|||
return accumulated_text
|
||||
|
||||
@staticmethod
|
||||
def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
|
||||
def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams:
|
||||
cp4d_host: Final = litellm_params.get("cp4d_host") or ""
|
||||
instance_id: Final = litellm_params.get("instance_id") or ""
|
||||
wxo_agent_id: Final = litellm_params.get("wxo_agent_id") or ""
|
||||
|
|
@ -215,9 +275,9 @@ class WatsonxOrchestrateHandler:
|
|||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
params: dict[str, object],
|
||||
litellm_params: WXOLitellmParams,
|
||||
) -> dict[str, object]:
|
||||
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
|
||||
client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0)
|
||||
|
|
@ -246,7 +306,8 @@ class WatsonxOrchestrateHandler:
|
|||
headers=auth_headers,
|
||||
)
|
||||
run_response.raise_for_status()
|
||||
run_data: dict[str, Any] = run_response.json()
|
||||
started: Final[_WXOView] = {"run": run_response.json()}
|
||||
run_data: _WXORun = started["run"]
|
||||
|
||||
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
|
||||
run_data=run_data,
|
||||
|
|
@ -261,11 +322,11 @@ class WatsonxOrchestrateHandler:
|
|||
@staticmethod
|
||||
async def handle_streaming(
|
||||
request_id: str,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
params: dict[str, object],
|
||||
litellm_params: WXOLitellmParams,
|
||||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
|
||||
client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0)
|
||||
|
|
@ -316,10 +377,11 @@ class WatsonxOrchestrateHandler:
|
|||
yield chunk
|
||||
return
|
||||
|
||||
content_type: Final = response.headers.get("content-type", "").lower()
|
||||
header_view: Final[_WXOView] = {"content_type": response.headers.get("content-type", "")}
|
||||
content_type: Final = header_view["content_type"].lower()
|
||||
if "text/event-stream" not in content_type:
|
||||
response_body: Final = await response.aread()
|
||||
result = json.loads(response_body)
|
||||
result = WatsonxOrchestrateHandler._decode_run_event(response_body)
|
||||
result = await WatsonxOrchestrateHandler._get_successful_run_data(
|
||||
run_data=result,
|
||||
base_url=base_url,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
|
|
@ -87,7 +87,7 @@ async def _handle_completed_batch(
|
|||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_input_entries(file_content),
|
||||
entries=_iter_batch_output_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
|
|
@ -111,43 +111,91 @@ def _iter_successful_output_line_stats(
|
|||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
for entry in entries:
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
if stats is not None:
|
||||
yield stats
|
||||
|
||||
|
||||
def _safe_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
"""Return the stats for one batch output line, or None for a line that is
|
||||
unsuccessful or cannot be costed, so a single bad line never aborts the
|
||||
whole batch's cost accounting."""
|
||||
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
|
||||
try:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
return None
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
|
||||
verbose_logger.warning(
|
||||
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
|
||||
"is still billed. custom_id=%s error=%s",
|
||||
custom_id,
|
||||
str(e),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _compute_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats:
|
||||
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details: Final = parse_prompt_tokens_details(usage)
|
||||
raw_model: Final = response_body.get("model")
|
||||
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
return _BatchOutputLineStats(
|
||||
cost=_output_line_cost(
|
||||
response_body=response_body,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
response_model=response_model,
|
||||
model_info=model_info,
|
||||
),
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
||||
def _output_line_cost(
|
||||
response_body: Mapping[str, Any],
|
||||
usage: Usage,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
response_model: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> float:
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
continue
|
||||
response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details = parse_prompt_tokens_details(usage)
|
||||
raw_model = response_body.get("model")
|
||||
response_model = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
cost_model = model_name
|
||||
else:
|
||||
cost_model = response_model or model_name or ""
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
line_cost = prompt_cost + completion_cost
|
||||
else:
|
||||
line_cost = litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
yield _BatchOutputLineStats(
|
||||
cost=line_cost,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
|
||||
return litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
cost_model: Final = (
|
||||
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
|
||||
)
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
def _aggregate_batch_cost_usage_models(
|
||||
|
|
@ -338,9 +386,10 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
|
||||
def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
|
||||
"""
|
||||
Get the file content as a list of dictionaries from JSON Lines format
|
||||
Get the file content as a list of dictionaries from JSON Lines format,
|
||||
skipping malformed lines
|
||||
"""
|
||||
return list(_iter_batch_input_entries(file_content))
|
||||
return list(_iter_batch_output_entries(file_content))
|
||||
|
||||
|
||||
def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
||||
|
|
@ -361,15 +410,29 @@ def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
|||
yield line
|
||||
|
||||
|
||||
def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]:
|
||||
def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]:
|
||||
"""
|
||||
Yield parsed batch input JSONL entries one at a time without materializing the
|
||||
whole file as a list, so peak memory stays bounded. Raises on a malformed line;
|
||||
callers that must survive bad rows should iterate ``_iter_batch_input_lines``
|
||||
and parse per-row instead.
|
||||
Yield parsed batch output JSONL entries one at a time without materializing
|
||||
the whole file as a list, so peak memory stays bounded. A malformed or
|
||||
non-object line is skipped with a warning so one bad line never aborts the
|
||||
whole batch's cost accounting.
|
||||
"""
|
||||
for line in _iter_batch_input_lines(file_content):
|
||||
yield json.loads(line)
|
||||
entry = _parse_batch_output_line(line)
|
||||
if entry is not None:
|
||||
yield entry
|
||||
|
||||
|
||||
def _parse_batch_output_line(line: bytes) -> dict | None:
|
||||
try:
|
||||
parsed: Final = json.loads(line)
|
||||
except ValueError as e:
|
||||
verbose_logger.warning("skipping malformed batch output line: %s", str(e))
|
||||
return None
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
verbose_logger.warning("skipping non-object batch output line of type %s", type(parsed).__name__)
|
||||
return None
|
||||
|
||||
|
||||
# A batch request's input tokens scale roughly with its serialized size, so this
|
||||
|
|
@ -440,7 +503,9 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_provider: str = "openai") -> Usage:
|
||||
def _get_batch_job_usage_from_response_body(
|
||||
response_body: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Usage:
|
||||
"""
|
||||
Get the tokens of a batch job from the response body
|
||||
"""
|
||||
|
|
@ -472,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
|
|||
return usage
|
||||
|
||||
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> dict:
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
|
||||
"""
|
||||
Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
|
||||
|
||||
|
|
@ -482,7 +547,9 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> d
|
|||
return batch_results_line.get("result", None) or {}
|
||||
|
||||
|
||||
def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> Any:
|
||||
def _get_response_from_batch_job_output_file(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Any:
|
||||
"""
|
||||
Get the response from the batch job output file
|
||||
"""
|
||||
|
|
@ -495,7 +562,9 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
|
|||
return _response_body
|
||||
|
||||
|
||||
def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> bool:
|
||||
def _batch_response_was_successful(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the batch job response was successful
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from collections.abc import Sequence
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
|
@ -60,6 +61,13 @@ def resolve_embedding_max_input_tokens(
|
|||
return deployment_max_input_tokens
|
||||
|
||||
|
||||
def resolve_embedding_timeout(configured_timeout: float | None) -> float:
|
||||
"""Explicit cache setting first, else the short semantic-cache default."""
|
||||
if configured_timeout is not None:
|
||||
return configured_timeout
|
||||
return SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
|
||||
def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str:
|
||||
"""Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call."""
|
||||
if max_input_tokens is None:
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ class Cache:
|
|||
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
|
||||
qdrant_semantic_cache_vector_size: int | None = None,
|
||||
semantic_cache_embedding_max_input_tokens: int | None = None,
|
||||
semantic_cache_embedding_timeout: float | None = None,
|
||||
# GCP IAM authentication parameters
|
||||
gcp_service_account: str | None = None,
|
||||
gcp_ssl_ca_certs: str | None = None,
|
||||
|
|
@ -124,6 +125,7 @@ class Cache:
|
|||
qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic".
|
||||
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
|
||||
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
|
||||
semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
|
||||
|
||||
# Disk Cache Args
|
||||
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
|
||||
|
|
@ -195,6 +197,7 @@ class Cache:
|
|||
embedding_model=redis_semantic_cache_embedding_model,
|
||||
index_name=redis_semantic_cache_index_name,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
|
||||
|
|
@ -211,6 +214,7 @@ class Cache:
|
|||
index_name=valkey_semantic_cache_index_name,
|
||||
startup_nodes=redis_startup_nodes,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
|
||||
|
|
@ -223,6 +227,7 @@ class Cache:
|
|||
embedding_model=qdrant_semantic_cache_embedding_model,
|
||||
vector_size=qdrant_semantic_cache_vector_size,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
)
|
||||
elif type == LiteLLMCacheType.LOCAL:
|
||||
self.cache = InMemoryCache()
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import asyncio
|
|||
import datetime
|
||||
import inspect
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -106,7 +106,7 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
|
|||
return "choices" in cached_result
|
||||
|
||||
|
||||
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bool:
|
||||
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
|
||||
"""
|
||||
When stream=True, do not run success callbacks at cache-hit time.
|
||||
|
||||
|
|
@ -119,11 +119,21 @@ def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bo
|
|||
return kwargs.get("stream", False) is True
|
||||
|
||||
|
||||
def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
|
||||
"""Dump prompt token details to an opaque field mapping, tolerating non-pydantic stand-ins."""
|
||||
return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
|
||||
|
||||
|
||||
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
|
||||
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
|
||||
return request_kwargs.get("cache_key", None)
|
||||
|
||||
|
||||
class LLMCachingHandler:
|
||||
def __init__(
|
||||
self,
|
||||
original_function: Callable,
|
||||
request_kwargs: dict[str, Any],
|
||||
request_kwargs: dict[str, object],
|
||||
start_time: datetime.datetime,
|
||||
):
|
||||
from litellm.caching import DualCache, RedisCache
|
||||
|
|
@ -150,7 +160,7 @@ class LLMCachingHandler:
|
|||
start_time: datetime.datetime,
|
||||
call_type: str,
|
||||
kwargs: dict[str, Any],
|
||||
args: tuple[Any, ...] | None = None,
|
||||
args: tuple[object, ...] | None = None,
|
||||
) -> CachingHandlerResponse | None:
|
||||
"""
|
||||
Internal method to get from the cache.
|
||||
|
|
@ -289,7 +299,7 @@ class LLMCachingHandler:
|
|||
start_time: datetime.datetime,
|
||||
call_type: str,
|
||||
kwargs: dict[str, Any],
|
||||
args: tuple[Any, ...] | None = None,
|
||||
args: tuple[object, ...] | None = None,
|
||||
) -> CachingHandlerResponse:
|
||||
cached_result: Any | None = None
|
||||
|
||||
|
|
@ -366,7 +376,7 @@ class LLMCachingHandler:
|
|||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
|
||||
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, Any]) -> list[str]:
|
||||
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[str]:
|
||||
"""
|
||||
Handles the input of kwargs['input'] being a list or a string
|
||||
"""
|
||||
|
|
@ -548,8 +558,8 @@ class LLMCachingHandler:
|
|||
if details2 is None:
|
||||
return details1
|
||||
|
||||
dict1: Final = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
|
||||
dict2: Final = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
|
||||
dict1: Final = _prompt_tokens_details_as_mapping(details1)
|
||||
dict2: Final = _prompt_tokens_details_as_mapping(details2)
|
||||
|
||||
merged: Final[dict] = {}
|
||||
for key in set(dict1.keys()) | set(dict2.keys()):
|
||||
|
|
@ -671,7 +681,9 @@ class LLMCachingHandler:
|
|||
cache_hit=cache_hit,
|
||||
)
|
||||
|
||||
async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None:
|
||||
async def _retrieve_from_cache(
|
||||
self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...]
|
||||
) -> Any | None:
|
||||
"""
|
||||
Internal method to
|
||||
- get cache key
|
||||
|
|
@ -727,7 +739,8 @@ class LLMCachingHandler:
|
|||
cached_result = None
|
||||
else:
|
||||
request_kwargs: Final = new_kwargs.copy()
|
||||
request_cache_key: Final = request_kwargs.pop("cache_key", None)
|
||||
request_cache_key: Final = _request_cache_key(request_kwargs)
|
||||
request_kwargs.pop("cache_key", None)
|
||||
if litellm.cache._supports_async() is True:
|
||||
## check if dual cache is supported ##
|
||||
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
|
||||
|
|
@ -749,10 +762,10 @@ class LLMCachingHandler:
|
|||
self,
|
||||
cached_result: Any,
|
||||
call_type: str,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
model: str,
|
||||
args: tuple[Any, ...],
|
||||
args: tuple[object, ...],
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> (
|
||||
ModelResponse
|
||||
|
|
@ -948,7 +961,7 @@ class LLMCachingHandler:
|
|||
result: Any,
|
||||
original_function: Callable,
|
||||
kwargs: dict[str, Any],
|
||||
args: tuple[Any, ...] | None = None,
|
||||
args: tuple[object, ...] | None = None,
|
||||
):
|
||||
"""
|
||||
Internal method to check the type of the result & cache used and adds the result to the cache accordingly
|
||||
|
|
@ -1013,8 +1026,8 @@ class LLMCachingHandler:
|
|||
def sync_set_cache(
|
||||
self,
|
||||
result: Any,
|
||||
kwargs: dict[str, Any],
|
||||
args: tuple[Any, ...] | None = None,
|
||||
kwargs: dict[str, object],
|
||||
args: tuple[object, ...] | None = None,
|
||||
):
|
||||
"""
|
||||
Sync internal method to add the result to the cache
|
||||
|
|
@ -1204,8 +1217,8 @@ class LLMCachingHandler:
|
|||
|
||||
def convert_args_to_kwargs(
|
||||
original_function: Callable,
|
||||
args: tuple[Any, ...] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
args: tuple[object, ...] | None = None,
|
||||
) -> dict[str, object]:
|
||||
# Get the signature of the original function
|
||||
signature: Final = inspect.signature(original_function)
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue