Merge branch 'litellm_internal_staging' into litellm_cyberark_admin_ui_config

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-08-29 19:59:50 +00:00
commit ccef0345f6
976 changed files with 82886 additions and 8885 deletions

View file

@ -2421,45 +2421,6 @@ jobs:
- wait_for_service:
url: http://localhost:4000
timeout: "300"
# Add Ruby installation and testing before the existing Node.js and Python tests
- run:
name: Install Ruby and Bundler
command: |
# Clone RVM at pinned tag and verify the commit SHA matches the
# published tag before running its install script.
RVM_VERSION="1.29.12"
RVM_EXPECTED_SHA="6bfc9213c9d6914fe756f524eb034a403d51db81"
git clone --depth 1 --branch "$RVM_VERSION" https://github.com/rvm/rvm.git /tmp/rvm
RVM_ACTUAL_SHA="$(git -C /tmp/rvm rev-parse HEAD)"
if [ "$RVM_ACTUAL_SHA" != "$RVM_EXPECTED_SHA" ]; then
echo "RVM tag $RVM_VERSION resolved to $RVM_ACTUAL_SHA; expected $RVM_EXPECTED_SHA" >&2
exit 1
fi
# Import RVM signing keys (used by `rvm install` to verify Ruby tarballs)
gpg --keyserver hkp://keyserver.ubuntu.com --recv-keys 409B6B1796C275462A1703113804BB82D39DC0E3 7D2BAF1CF37B13E2069D6956105BD0E739499BDB
# Install RVM from the verified checkout. The install script
# sources `scripts/functions/installer` using paths relative to
# its own working directory, so it must be run from /tmp/rvm.
(cd /tmp/rvm && ./install --path "$HOME/.rvm")
source "$HOME/.rvm/scripts/rvm"
# Install Ruby 3.2.2 (RVM verifies the tarball PGP signature)
rvm install 3.2.2
rvm use 3.2.2 --default
# Install latest Bundler
gem install bundler
- run:
name: Run Ruby tests
command: |
source $HOME/.rvm/scripts/rvm
cd tests/pass_through_tests/ruby_passthrough_tests
bundle install
bundle exec rspec
no_output_timeout: 30m
# Install Node.js directly from nodejs.org with SHA256 verification,
# instead of piping NodeSource's setup_24.x apt-repo installer into
# sudo bash (which runs a mutable upstream script unattended).

View file

@ -83,6 +83,24 @@ jobs:
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Regenerate the lazy OpenAPI snapshot
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
- name: Fail if the lazy OpenAPI snapshot is stale
if: steps.changes.outputs.relevant == 'true'
run: |
if ! git diff --exit-code -- litellm/proxy/_lazy_openapi_snapshot.json; then
echo "::error file=litellm/proxy/_lazy_openapi_snapshot.json::The lazy OpenAPI snapshot is out of sync with the lazily loaded routes."
echo ""
echo "A lazily loaded route or model changed without regenerating the snapshot that /openapi.json serves for unloaded features."
echo "To fix, run from the repo root:"
echo " uv run python -m litellm.proxy._lazy_openapi_snapshot"
echo "then run npm run gen:api from ui/litellm-dashboard and commit both files."
exit 1
fi
echo "_lazy_openapi_snapshot.json is in sync with the lazily loaded routes."
- name: Set up Node.js
if: steps.changes.outputs.relevant == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0

View file

@ -0,0 +1,68 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: litellm_internal_staging
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base litellm_internal_staging
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

View file

@ -114,4 +114,4 @@ jobs:
- name: Audit provider endpoints against the schema
working-directory: terraform/provider
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json"
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json" -coverage-allowlist ./tools/endpointaudit/coverage_allowlist.txt

View file

@ -141,6 +141,7 @@ jobs:
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/list_api
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers

View file

@ -1,36 +1,36 @@
{
"reportAny": {
"limit": 18483
"limit": 17270
},
"reportArgumentType": {
"limit": 2564
"limit": 2539
},
"reportAssignmentType": {
"limit": 320
"limit": 319
},
"reportAttributeAccessIssue": {
"limit": 483
"limit": 480
},
"reportCallIssue": {
"limit": 113
"limit": 112
},
"reportConstantRedefinition": {
"limit": 40
},
"reportDeprecated": {
"limit": 213
"limit": 212
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 5960
"limit": 5486
},
"reportFunctionMemberAccess": {
"limit": 7
},
"reportGeneralTypeIssues": {
"limit": 154
"limit": 101
},
"reportIncompatibleMethodOverride": {
"limit": 56
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5659
"limit": 5658
},
"reportMissingTypeArgument": {
"limit": 15484
"limit": 15425
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1058
"limit": 1055
},
"reportOptionalOperand": {
"limit": 0
@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1810
"limit": 1808
},
"reportRedeclaration": {
"limit": 8
@ -93,25 +93,25 @@
"limit": 213
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
"limit": 25
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44528
"limit": 44526
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38804
"limit": 38721
},
"reportUnknownParameterType": {
"limit": 19829
"limit": 19778
},
"reportUnknownVariableType": {
"limit": 30355
"limit": 30290
},
"reportUnnecessaryCast": {
"limit": 117
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 833
"limit": 829
},
"reportUntypedBaseClass": {
"limit": 0
@ -135,12 +135,12 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 139
"limit": 138
},
"reportUnusedImport": {
"limit": 544
"limit": 543
},
"reportUnusedVariable": {
"limit": 145
"limit": 139
}
}

View file

@ -73,6 +73,11 @@ ARRAY_KEYS: dict[str, JsonSchema] = {
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
@ -215,6 +220,15 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"audio_transcription_config": STRING,
}

View file

@ -10,6 +10,11 @@
-- partitioned, so existing installs are unaffected until you run this.
--
-- IMPORTANT
-- * After partitioning, `prisma db push` (including the proxy's
-- --use_prisma_db_push startup mode) is NOT supported: it tries to rewrite
-- the primary key back to ("request_id"), which Postgres rejects on a
-- partitioned table. The proxy detects this and exits with guidance.
-- Use the default startup path (`prisma migrate deploy`) instead.
-- * Test on a staging copy first and take a backup.
-- * Postgres cannot convert a populated table to partitioned in place, so this
-- renames the old table aside and creates a fresh partitioned table.

View file

@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
GET - /audit - Get all audit logs
"""
from typing import Any, Dict, List, Optional
from typing import TYPE_CHECKING, Final, Optional
#### AUDIT LOGGING ####
from fastapi import APIRouter, Depends, HTTPException, Query
@ -18,11 +18,16 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import AuditLogRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
router = APIRouter()
def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]:
"""
Build an OR condition that matches a value inside a JSON column at the
given key, checking both before_value and updated_values.
@ -101,46 +106,37 @@ async def get_audit_logs(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
# Build filter conditions
where_conditions: Dict[str, Any] = {}
if changed_by:
where_conditions["changed_by"] = changed_by
if changed_by_api_key:
where_conditions["changed_by_api_key"] = changed_by_api_key
if action:
where_conditions["action"] = action
if table_name:
where_conditions["table_name"] = table_name
if object_id:
where_conditions["object_id"] = object_id
if start_date or end_date:
date_filter: Dict[str, Any] = {}
if start_date:
date_filter["gte"] = start_date
if end_date:
date_filter["lte"] = end_date
where_conditions["updated_at"] = date_filter
date_filter: Final[dict[str, str]] = {
**({"gte": start_date} if start_date else {}),
**({"lte": end_date} if end_date else {}),
}
# JSON field filters (PostgreSQL only) — each filter is AND'd with the
# others, but checks both before_value and updated_values internally (OR).
if object_team_id:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("team_id", object_team_id)
]
if object_key_hash:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("token", object_key_hash)
]
json_field_conditions: Final[list[dict[str, object]]] = [
*([_build_json_field_or_condition("team_id", object_team_id)] if object_team_id else []),
*([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []),
]
# Build sort conditions
order_by: Dict[str, Any] = {}
if sort_by and isinstance(sort_by, str):
order_by[sort_by] = sort_order
else:
order_by["updated_at"] = sort_order # Default sort by updated_at
# Build filter conditions
where_conditions: Final[dict[str, object]] = {
**({"changed_by": changed_by} if changed_by else {}),
**({"changed_by_api_key": changed_by_api_key} if changed_by_api_key else {}),
**({"action": action} if action else {}),
**({"table_name": table_name} if table_name else {}),
**({"object_id": object_id} if object_id else {}),
**({"updated_at": date_filter} if start_date or end_date else {}),
**({"AND": json_field_conditions} if json_field_conditions else {}),
}
order_by: Final[dict[str, str]] = (
{sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order}
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
# Get paginated results
audit_logs = await prisma_client.db.litellm_auditlog.find_many(
audit_logs: Final = await audit_log_table.find_many(
where=where_conditions,
order=order_by,
skip=(page - 1) * page_size,
@ -148,13 +144,14 @@ async def get_audit_logs(
)
# Get total count for pagination
total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions)
total_pages = -(-total_count // page_size) # Ceiling division
total_count: Final = await audit_log_table.count(where=where_conditions)
total_pages: Final = -(-total_count // page_size) # Ceiling division
# Return paginated response
return PaginatedAuditLogResponse(
audit_logs=[
AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs
AuditLogResponse.model_validate(audit_log.model_dump())
for audit_log in audit_logs
]
if audit_logs
else [],
@ -198,8 +195,10 @@ async def get_audit_log_by_id(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
# Get the audit log by ID
audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id})
audit_log: Final = await audit_log_table.find_unique(where={"id": id})
if audit_log is None:
raise HTTPException(
@ -207,4 +206,4 @@ async def get_audit_log_by_id(
)
# Convert to response model
return AuditLogResponse(**audit_log.model_dump())
return AuditLogResponse.model_validate(audit_log.model_dump())

View file

@ -1,6 +1,8 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by the get-responses call.
Cost tracking is handled by the get-responses call, which prices normally only because the
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
same route are non-inference and free.
"""
from datetime import datetime, timedelta, timezone
@ -9,12 +11,14 @@ from typing import TYPE_CHECKING, Dict, Optional, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -113,7 +117,8 @@ class CheckResponsesCost:
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is automatically tracked by the get-responses call
- Cost is tracked by the get-responses call, billed because the poll is stamped
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
- Mark responses in a terminal state as complete in the database
"""
try:
@ -153,6 +158,7 @@ class CheckResponsesCost:
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
# Add model information if available

View file

@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import (
CallTypes,
@ -222,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object=file_object,
model_mappings=model_mappings,
flat_model_file_ids=list(model_mappings.values()),
created_by=user_api_key_dict.user_id,
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
)
@ -238,7 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_file_id": file_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_file_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
}
@ -342,7 +343,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_object": file_object.model_dump_json(),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
@ -473,19 +474,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
page_size: Final = min(limit or 20, 100)
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
batches = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
matches: Final = await self._collect_listed_batches(
where_clause=where_clause,
after=after,
wanted=page_size + 1,
user_api_key_dict=user_api_key_dict,
)
return build_list_page(list(matches[:page_size]), has_more=len(matches) > page_size)
has_more = len(batches) > page_size
async def _collect_listed_batches(
self,
where_clause: Mapping[str, object],
after: Optional[str],
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
"""Read chunks newest-first until ``wanted`` batches survive parsing and
file-id resolution or the caller's rows run out, so a run of rows that will
not parse refills the page instead of emptying it. The first chunk is
``wanted`` rows, so a healthy page still costs one query; a scan that has to
continue widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` like ``afile_list``,
and every chunk advances the keyset cursor, so the walk ends once the
caller's rows are exhausted."""
matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks
cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row
chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk
while len(matches) < wanted:
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
)
matches = matches + await self._resolve_listed_rows(
rows=chunk, wanted=wanted - len(matches), user_api_key_dict=user_api_key_dict
)
if len(chunk) < chunk_size:
break
cursor_id = chunk[-1].unified_object_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return matches
async def _resolve_listed_rows(
self,
rows: "Sequence[PrismaManagedObjectRow]",
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
parsed_rows: Final = tuple(
(row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None
(row, batch_obj) for row in rows if (batch_obj := _parse_managed_batch_row(row)) is not None
)
unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
raw_file_ids=frozenset(
@ -496,19 +534,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
prisma_client=self.prisma_client,
)
resolved_batches: Final = [
await self._resolve_listed_batch(
resolved: Final[list[LiteLLMBatch]] = [] # mutable-ok: resolution stops as soon as the page is full
for row, batch_obj in parsed_rows:
if len(resolved) == wanted:
break
resolved_batch = await self._resolve_listed_batch(
row=row,
batch_obj=batch_obj,
unified_id_by_raw_id=unified_id_by_raw_id,
user_api_key_dict=user_api_key_dict,
)
for row, batch_obj in parsed_rows
]
return build_list_page(
[batch_obj for batch_obj in resolved_batches if batch_obj is not None],
has_more=has_more,
)
if resolved_batch is not None:
resolved.append(resolved_batch)
return tuple(resolved)
async def _resolve_listed_batch(
self,

View file

@ -12,9 +12,10 @@ Endpoints for /project operations
import json
from collections.abc import Sequence
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -26,37 +27,50 @@ from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma.actions import (
LiteLLM_ProjectTableActions,
LiteLLM_TeamTableActions,
LiteLLM_VerificationTokenActions,
)
from litellm import Router
router = APIRouter()
def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = prisma_client.db.litellm_teamtable
return team_table
_OBJECT_PERMISSION_PAYLOAD: Final = TypeAdapter(dict[str, object])
def _project_table(prisma_client: PrismaClient) -> "LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable]":
project_table: LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable] = (
prisma_client.db.litellm_projecttable
)
return project_table
def _team_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_TeamTable"]:
return TeamRepository(prisma_client).table
def _project_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_ProjectTable"]:
return ProjectRepository(prisma_client).table
def _verification_token_table(
prisma_client: PrismaClient,
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
verification_token_table: LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken] = (
prisma_client.db.litellm_verificationtoken
)
return verification_token_table
) -> TableActions["prisma_models.LiteLLM_VerificationToken"]:
return VerificationTokenRepository(prisma_client).table
def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]:
return BudgetRepository(prisma_client).table
def _object_permission_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]:
return ObjectPermissionRepository(prisma_client).table
def _user_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_UserTable"]:
return UserRepository(prisma_client).table
def _jsonified(prisma_client: PrismaClient, payload: dict[str, object]) -> dict[str, object]:
@ -205,6 +219,114 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
def _project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> tuple[str, ...]:
"""Entries project auth expands to many concrete models (`all-proxy-models`, `*` patterns,
access groups). The rate limiter looks quotas up by the exact requested model name, so a
quota keyed on one of these entries is never applied."""
return tuple(
model
for model in (models or ())
if model == SpecialModelNames.all_proxy_models.value or "*" in model or model in access_group_names
)
def _raise_on_project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> None:
expanding: Final = _project_models_expanding_at_request_time(models, access_group_names)
if not expanding:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
},
)
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
async def _create_budget_for_project(
data: NewProjectRequest,
user_id: str | None,
@ -219,7 +341,7 @@ async def _create_budget_for_project(
new_budget = _jsonified(prisma_client, budget_row.model_dump(exclude_none=True))
_budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create(
_budget: Final = await _budget_table(prisma_client).create(
data={
**new_budget,
"created_by": user_id or litellm_proxy_admin_name,
@ -242,10 +364,8 @@ async def _set_project_object_permission(
return None
if data.object_permission is not None:
created_object_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=data.object_permission.model_dump(exclude_none=True),
)
created_object_permission: Final = await _object_permission_table(prisma_client).create(
data=data.object_permission.model_dump(exclude_none=True),
)
del data.object_permission
return created_object_permission.object_permission_id
@ -352,7 +472,9 @@ async def new_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
)
@ -399,6 +521,10 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -470,10 +596,8 @@ async def new_project(
new_project_row = _remove_budget_fields_from_project_data(new_project_row)
verbose_proxy_logger.info(f"new_project_row: {json.dumps(new_project_row, indent=2)}")
response: prisma_models.LiteLLM_ProjectTable = await prisma_client.db.litellm_projecttable.create(
data={
**new_project_row, # type: ignore
},
response: Final = await _project_table(prisma_client).create(
data={**new_project_row},
include={"litellm_budget_table": True},
)
@ -538,7 +662,9 @@ async def update_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
user_api_key_cache,
@ -642,6 +768,12 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
)
# Prepare update data
update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"}))
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
@ -652,7 +784,7 @@ async def update_project(
if budget_updates and existing_project.budget_id:
# Update existing budget
await prisma_client.db.litellm_budgettable.update(
await _budget_table(prisma_client).update(
where={"budget_id": existing_project.budget_id},
data={
**budget_updates,
@ -667,18 +799,17 @@ async def update_project(
if "object_permission" in update_data:
object_permission_data = update_data.pop("object_permission")
if object_permission_data:
object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data)
if existing_project.object_permission_id:
# Update existing permission
await prisma_client.db.litellm_objectpermissiontable.update(
await _object_permission_table(prisma_client).update(
where={"object_permission_id": existing_project.object_permission_id},
data=object_permission_data,
data=object_permission_payload,
)
else:
# Create new permission
created_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=object_permission_data,
)
created_permission: Final = await _object_permission_table(prisma_client).create(
data=object_permission_payload,
)
update_data["object_permission_id"] = created_permission.object_permission_id
@ -694,7 +825,7 @@ async def update_project(
update_data = _remove_budget_fields_from_project_data(update_data)
# Update project
updated_project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.update(
updated_project: Final = await _project_table(prisma_client).update(
where={"project_id": data.project_id},
data=update_data,
include={"litellm_budget_table": True, "object_permission": True},
@ -934,7 +1065,7 @@ async def list_projects(
# Look up the user's team memberships via the reverse-index on
# LiteLLM_UserTable.teams (maintained by team_member_add alongside
# members_with_roles). This avoids a full scan of all team rows.
user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique(
user_record: Final = await _user_table(prisma_client).find_unique(
where={"user_id": user_api_key_dict.user_id},
)
user_team_ids: list[str] = user_record.teams if user_record is not None and user_record.teams else []

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.60"
version = "0.1.62"
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.60"
version = "0.1.62"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,18 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -0,0 +1,22 @@
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cost" DOUBLE PRECISION;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "shadow_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cache_hit" BOOLEAN NOT NULL DEFAULT false;
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ShadowEvalFunnel" (
"job_id" TEXT NOT NULL,
"not_sampled" INTEGER NOT NULL DEFAULT 0,
"unjudgeable" INTEGER NOT NULL DEFAULT 0,
"shed" INTEGER NOT NULL DEFAULT 0,
"withheld" INTEGER NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_ShadowEvalFunnel_pkey" PRIMARY KEY ("job_id")
);

View file

@ -754,6 +754,7 @@ model LiteLLM_DailyUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -789,6 +790,7 @@ model LiteLLM_DailyOrganizationSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -824,6 +826,7 @@ model LiteLLM_DailyEndUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -858,6 +861,7 @@ model LiteLLM_DailyAgentSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -892,6 +896,7 @@ model LiteLLM_DailyTeamSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -929,6 +934,7 @@ model LiteLLM_DailyTagSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -1527,12 +1533,27 @@ model LiteLLM_ShadowEvalAttempt {
confidence Float?
judge_cost Float @default(0)
shadow_cost Float @default(0)
real_cost Float? // NULL = row predates cost measurement; comparisons read only measured rows
real_classifier_cost Float @default(0)
shadow_classifier_cost Float @default(0)
real_cache_hit Boolean @default(false)
error String?
created_at DateTime @default(now())
@@index([job_id])
}
// Per-leg sampling funnel counters the attempt rows cannot derive: requests an
// admitting job saw but did not judge. attempted = the leg's attempt rows; the
// leg's eligible traffic = not_sampled + unjudgeable + shed + withheld + attempted.
model LiteLLM_ShadowEvalFunnel {
job_id String @id
not_sampled Int @default(0)
unjudgeable Int @default(0)
shed Int @default(0)
withheld Int @default(0)
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//

View file

@ -40,6 +40,65 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
)
_SPEND_LOGS_PK_CLAUSE_RE = re.compile(
r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
re.IGNORECASE,
)
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
"reconciles the database against schema.prisma, which declares the unpartitioned "
"primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
"constraint on partitioned table must include all partitioning columns. Start the "
"proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
"applies shipped migrations and leaves the partitioned primary key alone."
)
def _without_sql_comments(statement: str) -> str:
return "\n".join(
line
for line in statement.splitlines()
if line.strip() and not line.strip().startswith("--")
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
kept = tuple(
clause.strip()
for clause in statement[prefix_match.end():].split(",\n")
if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
)
if not kept:
return None
return statement[: prefix_match.end()] + ",\n".join(kept)
def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
"""Drop statements from a `prisma migrate diff` script that fight the
SpendLogs partitioning runbook (db_scripts/partition_spend_logs.sql): the
primary-key rewrite on "LiteLLM_SpendLogs", which Postgres rejects on a
partitioned table, and drops of runbook artifacts such as
"LiteLLM_SpendLogs_legacy"."""
kept = tuple(
filtered
for statement in diff_sql.split(";")
for bare in (_without_sql_comments(statement),)
if bare and not _SPEND_LOGS_ARTIFACT_DROP_RE.match(bare)
for filtered in (_without_spend_logs_pk_clauses(bare),)
if filtered is not None
)
return "".join(f"{statement};\n\n" for statement in kept)
def _migration_timestamp(name: str) -> int:
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
@ -355,7 +414,24 @@ class ProxyExtrasDBManager:
return
logger.info(f"Migration diff created at {diff_sql_path}")
if ProxyExtrasDBManager.spend_logs_is_partitioned():
filtered_sql = filter_partitioned_spend_logs_diff(
diff_sql_path.read_text()
)
diff_sql_path.write_text(filtered_sql)
logger.info(
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
"rewrite and partitioning artifacts from the drift script"
)
if not filtered_sql.strip():
logger.info("Drift script is empty after filtering; nothing to apply")
if not mark_all_applied:
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
return
# 2. Run prisma db execute to apply the migration
applied_ok = False
try:
logger.info("Running prisma db execute to apply the migration diff...")
result = subprocess.run(
@ -376,6 +452,7 @@ class ProxyExtrasDBManager:
)
logger.info(f"prisma db execute stdout: {result.stdout}")
logger.info("✅ Migration diff applied successfully")
applied_ok = True
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to apply migration diff: {e.stderr}")
except subprocess.TimeoutExpired:
@ -384,6 +461,16 @@ class ProxyExtrasDBManager:
# 3. Mark all migrations as applied
if not mark_all_applied:
return
if not applied_ok:
logger.warning(
"Drift script failed to apply; NOT marking migrations as "
"applied so a later migration run can retry them"
)
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
@staticmethod
def _mark_migrations_applied(migrations_dir: str) -> None:
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
logger.info(f"Resolving {len(migration_names)} migrations")
for migration_name in migration_names:
@ -410,6 +497,55 @@ class ProxyExtrasDBManager:
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
partitioned table in Prisma's target schema (the `schema` URL param,
falling back to Prisma's default target, public), i.e. the operator
ran db_scripts/partition_spend_logs.sql. Returns False when psycopg is
unavailable or the database cannot be reached, preserving the
pre-existing behavior in those cases."""
database_url = os.getenv("DATABASE_URL")
if not database_url:
return False
try:
import psycopg
except ImportError:
return False
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
try:
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
row = conn.execute(
"SELECT 1 "
"FROM pg_partitioned_table pt "
"JOIN pg_class c ON c.oid = pt.partrelid "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE c.relname = 'LiteLLM_SpendLogs' "
" AND n.nspname = %s",
(
ProxyExtrasDBManager._prisma_schema_param(database_url)
or "public",
),
).fetchone()
except (psycopg.OperationalError, psycopg.DatabaseError):
return False
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
None,
)
@staticmethod
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
@ -528,7 +664,8 @@ class ProxyExtrasDBManager:
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
# Preserve `prisma db push` path unchanged.
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
@ -972,6 +1109,8 @@ class ProxyExtrasDBManager:
)
raise
else:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.89"
version = "0.4.91"
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.89"
version = "0.4.91"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

340
litellm-rust/Cargo.lock generated
View file

@ -2,6 +2,36 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "aho-corasick"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
dependencies = [
"memchr",
]
[[package]]
name = "alloca"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4"
dependencies = [
"cc",
]
[[package]]
name = "anes"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
[[package]]
name = "anstyle"
version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
[[package]]
name = "arc-swap"
version = "1.9.2"
@ -506,6 +536,12 @@ dependencies = [
"either",
]
[[package]]
name = "cast"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
[[package]]
name = "cc"
version = "1.3.0"
@ -541,6 +577,58 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "ciborium"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
dependencies = [
"ciborium-io",
"ciborium-ll",
"serde",
]
[[package]]
name = "ciborium-io"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
[[package]]
name = "ciborium-ll"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
dependencies = [
"ciborium-io",
"half",
]
[[package]]
name = "clap"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca"
dependencies = [
"clap_builder",
]
[[package]]
name = "clap_builder"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889"
dependencies = [
"anstyle",
"clap_lex",
]
[[package]]
name = "clap_lex"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
[[package]]
name = "cmake"
version = "0.1.58"
@ -596,6 +684,72 @@ dependencies = [
"libc",
]
[[package]]
name = "criterion"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3"
dependencies = [
"alloca",
"anes",
"cast",
"ciborium",
"clap",
"criterion-plot",
"itertools",
"num-traits",
"oorandom",
"page_size",
"plotters",
"rayon",
"regex",
"serde",
"serde_json",
"tinytemplate",
"walkdir",
]
[[package]]
name = "criterion-plot"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
dependencies = [
"cast",
"itertools",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
[[package]]
name = "crunchy"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]]
name = "crypto-common"
version = "0.1.7"
@ -856,6 +1010,17 @@ dependencies = [
"tracing",
]
[[package]]
name = "half"
version = "2.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
dependencies = [
"cfg-if",
"crunchy",
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1179,6 +1344,15 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "itertools"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.18"
@ -1255,10 +1429,13 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"litellm-ai-gateway",
"litellm-core",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"serde",
"serde_json",
"tokio",
]
@ -1340,6 +1517,12 @@ version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "oorandom"
version = "11.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
[[package]]
name = "openssl-probe"
version = "0.2.1"
@ -1352,6 +1535,16 @@ version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e"
[[package]]
name = "page_size"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@ -1376,6 +1569,34 @@ version = "0.3.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e"
[[package]]
name = "plotters"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
dependencies = [
"num-traits",
"plotters-backend",
"plotters-svg",
"wasm-bindgen",
"web-sys",
]
[[package]]
name = "plotters-backend"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
[[package]]
name = "plotters-svg"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
dependencies = [
"plotters-backend",
]
[[package]]
name = "portable-atomic"
version = "1.14.0"
@ -1486,6 +1707,16 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "pythonize"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89"
dependencies = [
"pyo3",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -1613,12 +1844,61 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "rayon"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
dependencies = [
"either",
"rayon-core",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
dependencies = [
"crossbeam-deque",
"crossbeam-utils",
]
[[package]]
name = "regex"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "regex-automata"
version = "0.4.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
]
[[package]]
name = "regex-lite"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
[[package]]
name = "regex-syntax"
version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "reqwest"
version = "0.12.28"
@ -1774,6 +2054,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "same-file"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
dependencies = [
"winapi-util",
]
[[package]]
name = "schannel"
version = "0.1.29"
@ -2099,6 +2388,16 @@ dependencies = [
"zerovec",
]
[[package]]
name = "tinytemplate"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
dependencies = [
"serde",
"serde_json",
]
[[package]]
name = "tinyvec"
version = "1.12.0"
@ -2363,6 +2662,16 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "walkdir"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
dependencies = [
"same-file",
"winapi-util",
]
[[package]]
name = "want"
version = "0.3.1"
@ -2475,6 +2784,37 @@ dependencies = [
"rustls-pki-types",
]
[[package]]
name = "winapi"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
dependencies = [
"winapi-i686-pc-windows-gnu",
"winapi-x86_64-pc-windows-gnu",
]
[[package]]
name = "winapi-i686-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
[[package]]
name = "winapi-util"
version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "winapi-x86_64-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-link"
version = "0.2.1"

View file

@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.29.0"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }

View file

@ -9,10 +9,23 @@ repository.workspace = true
name = "_native"
crate-type = ["cdylib"]
[features]
default = ["extension-module"]
extension-module = ["pyo3/extension-module"]
[dependencies]
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
pyo3 = { workspace = true, features = ["extension-module"] }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
pythonize.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
[[bench]]
name = "serialization"
harness = false

View file

@ -0,0 +1,103 @@
use std::hint::black_box;
use std::time::Duration;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::{Value, json};
const PAYLOAD_SIZES: &[(&str, usize)] = &[
("1_KiB", 1024),
("64_KiB", 64 * 1024),
("1_MiB", 1024 * 1024),
("4_MiB", 4 * 1024 * 1024),
("16_MiB", 16 * 1024 * 1024),
];
fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value {
let json = py.import("json").expect("Python json module should import");
let encoded: String = json
.call_method1("dumps", (value,))
.expect("payload should serialize")
.extract()
.expect("json.dumps should return a string");
serde_json::from_str(&encoded).expect("serialized JSON should parse")
}
fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value {
pythonize::depythonize(value).expect("payload should depythonize")
}
fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
let json = py.import("json").expect("Python json module should import");
let encoded = serde_json::to_string(value).expect("response should serialize");
json.call_method1("loads", (encoded,))
.expect("serialized response should parse in Python")
.unbind()
}
fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
pythonize::pythonize(py, value)
.expect("response should pythonize")
.unbind()
}
fn serialization(c: &mut Criterion) {
Python::initialize();
Python::attach(|py| {
for &(label, payload_bytes) in PAYLOAD_SIZES {
let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes));
let document = PyDict::new(py);
document
.set_item("type", "image_url")
.expect("document type should be set");
document
.set_item("image_url", &data_uri)
.expect("document URL should be set");
let response = json!({
"pages": [{
"index": 0,
"markdown": "OCR text",
"images": [{"image_base64": data_uri}],
}],
"model": "mistral-ocr-latest",
"document_annotation": null,
"usage_info": {"pages_processed": 1},
"object": "ocr",
});
c.bench_with_input(
BenchmarkId::new("python_to_rust_json", label),
&document,
|b, document| {
b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any())))
},
);
c.bench_with_input(
BenchmarkId::new("python_to_rust_pythonize", label),
&document,
|b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_json", label),
&response,
|b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_pythonize", label),
&response,
|b, response| b.iter(|| pythonize_to_py(py, black_box(response))),
);
}
});
}
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(20)
.warm_up_time(Duration::from_secs(1))
.measurement_time(Duration::from_secs(4));
targets = serialization
}
criterion_main!(benches);

View file

@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict};
use serde_json::{Map, Value};
mod gil;
mod marshal;
use marshal::{from_py, to_py};
pyo3::create_exception!(
_native,
@ -41,35 +44,18 @@ type MarshaledOcrInputs = (
Option<Duration>,
);
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
let json = py.import("json")?;
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string()))
}
fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
let json = py.import("json")?;
let encoded =
serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?;
Ok(json.call_method1("loads", (encoded,))?.unbind())
}
fn messages_response_to_py(
py: Python<'_>,
response: AnthropicMessagesResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
to_py(py, &response)
}
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)
to_py(py, &response)
}
fn core_error_to_pyerr(err: CoreError) -> PyErr {
@ -116,7 +102,7 @@ fn optional_object_to_map(
value: Option<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
match value {
Some(value) => match py_to_json(py, value.bind(py))? {
Some(value) => match from_py(value.bind(py))? {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
},
@ -139,7 +125,7 @@ fn marshal_headers(
headers: Option<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => py_to_json(py, headers.bind(py))?,
Some(headers) => from_py(headers.bind(py))?,
None => Value::Object(Map::new()),
};
let Value::Object(headers) = value else {
@ -211,7 +197,7 @@ fn marshal_inputs(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
let document = py_to_json(py, document.bind(py))?;
let document = from_py(document.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -262,7 +248,7 @@ fn ocr(
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -307,7 +293,7 @@ fn aocr(
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -325,7 +311,7 @@ fn transcription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -351,7 +337,7 @@ fn transcription(
))
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -370,7 +356,7 @@ fn atranscription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -394,7 +380,7 @@ fn atranscription(
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -406,7 +392,7 @@ fn marshal_messages_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledMessagesInputs> {
let body = py_to_json(py, body.bind(py))?;
let body: Value = from_py(body.bind(py))?;
if !body.is_object() {
return Err(PyValueError::new_err("body must be a dict"));
}
@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledChatCompletionsInputs> {
let messages = py_to_json(py, messages.bind(py))?;
let messages: Value = from_py(messages.bind(py))?;
if !messages.is_array() {
return Err(PyValueError::new_err("messages must be a list"));
}
@ -527,7 +513,7 @@ fn chat_completions_decline(
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
let messages = py_to_json(py, messages.bind(py))?;
let messages = from_py(messages.bind(py))?;
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,

View file

@ -0,0 +1,20 @@
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;
pub fn from_py<T>(value: &Bound<'_, PyAny>) -> PyResult<T>
where
T: DeserializeOwned,
{
pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string()))
}
pub fn to_py<T>(py: Python<'_>, value: &T) -> PyResult<Py<PyAny>>
where
T: Serialize + ?Sized,
{
pythonize::pythonize(py, value)
.map(Bound::unbind)
.map_err(|error| PyValueError::new_err(error.to_string()))
}

View file

@ -0,0 +1,52 @@
use std::fs;
use std::path::{Path, PathBuf};
const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[
"py.import(\"json\")",
"pythonize::",
"serde_json::to_string",
"serde_json::from_str",
];
fn source_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("src")
}
fn rust_sources(directory: &Path) -> Vec<PathBuf> {
fs::read_dir(directory)
.expect("bridge source directory should be readable")
.map(|entry| {
entry
.expect("bridge source entry should be readable")
.path()
})
.flat_map(|path| {
if path.is_dir() {
rust_sources(&path)
} else if path.extension().is_some_and(|extension| extension == "rs") {
vec![path]
} else {
Vec::new()
}
})
.collect()
}
#[test]
fn serialization_is_centralized_in_marshal_module() {
let root = source_root();
for path in rust_sources(&root) {
if path == root.join("marshal.rs") {
continue;
}
let source = fs::read_to_string(&path).expect("bridge source should be readable");
for disallowed in DISALLOWED_OUTSIDE_MARSHAL {
assert!(
!source.contains(disallowed),
"{} bypasses the typed marshal module with `{disallowed}`",
path.display()
);
}
}
}

View file

@ -274,7 +274,6 @@ databricks_key: Optional[str] = None
openai_like_key: Optional[str] = None
azure_key: Optional[str] = None
anthropic_key: Optional[str] = None
autorouter_savings_baseline_model: Optional[str] = None
replicate_key: Optional[str] = None
bytez_key: Optional[str] = None
gdc_key: Optional[str] = None
@ -445,6 +444,7 @@ max_ui_session_budget: Optional[float] = (
1.0 # USD budget for each dashboard login session (playground, test connection)
)
internal_user_budget_duration: Optional[str] = None
budget_rollover: bool = False # carry spend beyond max_budget into the next window instead of zeroing it
tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None
max_end_user_budget: Optional[float] = None
max_end_user_budget_id: Optional[str] = None
@ -486,6 +486,7 @@ public_mcp_servers: Optional[List[str]] = None
public_mcp_hub_strict_whitelist: bool = True
public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
agent_search_embedding_model: Optional[str] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
# New format: { "displayName": { "url": "...", "index": 0 } }
# Old format: { "displayName": "url" } (for backward compatibility)

View file

@ -18,7 +18,10 @@ until they're actually needed.
import importlib
import sys
from collections.abc import Callable
from typing import Any, Final, cast
from types import ModuleType
from typing import TYPE_CHECKING, Any, Final, cast
from typing_extensions import ReadOnly, TypedDict
# Import all the data structures that define what can be lazy-loaded
# These are just lists of names and maps of where to find them
@ -53,6 +56,9 @@ from ._lazy_imports_registry import (
UTILS_NAMES,
)
if TYPE_CHECKING:
from tiktoken import Encoding
def get_litellm_globals() -> dict:
"""
@ -78,10 +84,10 @@ def _get_utils_globals() -> dict:
# They're separate from the main lazy import system because they have specific use cases
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
_default_encoding: Any | None = None
_default_encoding: "Encoding | None" = None
def _get_default_encoding() -> Any:
def _get_default_encoding() -> "Encoding":
"""
Lazily load and cache the default OpenAI encoding.
@ -100,10 +106,10 @@ def _get_default_encoding() -> Any:
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
_get_modified_max_tokens_func: Any | None = None
_get_modified_max_tokens_func: "Callable[..., int | None] | None" = None
def _get_modified_max_tokens() -> Any:
def _get_modified_max_tokens() -> "Callable[..., int | None]":
"""
Lazily load and cache the get_modified_max_tokens function.
@ -124,10 +130,10 @@ def _get_modified_max_tokens() -> Any:
# Lazy loader for token_counter to avoid importing token_counter module at module import time
_token_counter_new_func: Any | None = None
_token_counter_new_func: "Callable[..., int] | None" = None
def _get_token_counter_new() -> Any:
def _get_token_counter_new() -> "Callable[..., int]":
"""
Lazily load and cache the token_counter function (aliased as token_counter_new).
@ -154,10 +160,10 @@ def _get_token_counter_new() -> Any:
# This registry maps attribute names (like "ModelResponse") to handler functions
# It's built once the first time someone accesses a lazy-loaded attribute
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]:
"""
Build the registry that maps attribute names to their handler functions.
@ -206,7 +212,18 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
return _LAZY_IMPORT_REGISTRY
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any:
class _AttributeView(TypedDict):
"""Holds one module attribute so the lazily fetched value is read back as ``object``."""
value: ReadOnly[object]
def _module_attribute(module: ModuleType, attr_name: str) -> object:
attribute: Final[_AttributeView] = {"value": getattr(module, attr_name)}
return attribute["value"]
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object:
"""
Generic function that handles lazy importing for most attributes.
@ -255,7 +272,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value: Final = getattr(module, attr_name)
value: Final = _module_attribute(module, attr_name)
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
@ -272,62 +289,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# The registry (above) maps attribute names to these handler functions.
def _lazy_import_utils(name: str) -> Any:
def _lazy_import_utils(name: str) -> object:
"""Handler for utils module attributes (ModelResponse, token_counter, etc.)"""
return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils")
def _lazy_import_cost_calculator(name: str) -> Any:
def _lazy_import_cost_calculator(name: str) -> object:
"""Handler for cost calculator functions (completion_cost, cost_per_token, etc.)"""
return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator")
def _lazy_import_token_counter(name: str) -> Any:
def _lazy_import_token_counter(name: str) -> object:
"""Handler for token counter utilities"""
return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter")
def _lazy_import_bedrock_types(name: str) -> Any:
def _lazy_import_bedrock_types(name: str) -> object:
"""Handler for Bedrock type aliases"""
return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types")
def _lazy_import_types_utils(name: str) -> Any:
def _lazy_import_types_utils(name: str) -> object:
"""Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)"""
return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils")
def _lazy_import_caching(name: str) -> Any:
def _lazy_import_caching(name: str) -> object:
"""Handler for caching classes (Cache, DualCache, RedisCache, etc.)"""
return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching")
def _lazy_import_dotprompt(name: str) -> Any:
def _lazy_import_dotprompt(name: str) -> object:
"""Handler for dotprompt integration globals"""
return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt")
def _lazy_import_types(name: str) -> Any:
def _lazy_import_types(name: str) -> object:
"""Handler for type classes (GuardrailItem, etc.)"""
return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types")
def _lazy_import_llm_configs(name: str) -> Any:
def _lazy_import_llm_configs(name: str) -> object:
"""Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)"""
return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config")
def _lazy_import_litellm_logging(name: str) -> Any:
def _lazy_import_litellm_logging(name: str) -> object:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
def _lazy_import_llm_provider_logic(name: str) -> Any:
def _lazy_import_llm_provider_logic(name: str) -> object:
"""Handler for LLM provider logic functions (get_llm_provider, etc.)"""
return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
def _lazy_import_utils_module(name: str) -> Any:
def _lazy_import_utils_module(name: str) -> object:
"""
Handler for utils module lazy imports.
@ -355,7 +372,7 @@ def _lazy_import_utils_module(name: str) -> Any:
module = importlib.import_module(module_path)
# Get the actual attribute from the module
value: Final = getattr(module, attr_name)
value: Final = _module_attribute(module, attr_name)
# Cache it so we don't have to import again next time
_globals[name] = value
@ -370,7 +387,7 @@ def _lazy_import_utils_module(name: str) -> Any:
# These handlers have custom logic that doesn't fit the generic pattern
def _lazy_import_llm_client_cache(name: str) -> Any:
def _lazy_import_llm_client_cache(name: str) -> object:
"""
Handler for LLM client cache - has special logic for singleton instance.
@ -386,8 +403,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
return _globals[name]
# Import the class
module: Final = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache: Final = getattr(module, "LLMClientCache")
from litellm.caching.llm_caching_handler import LLMClientCache
# If they want the class itself, return it
if name == "LLMClientCache":
@ -403,7 +419,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}")
def _lazy_import_http_handlers(name: str) -> Any:
def _lazy_import_http_handlers(name: str) -> object:
"""
Handler for HTTP clients - has special logic for creating client instances.

View file

@ -5,7 +5,7 @@ import os
import sys
from datetime import datetime
from logging import Formatter
from typing import Any, Final
from typing import Any, Final, TextIO
import litellm
from litellm.constants import (
@ -234,11 +234,65 @@ class CorrelationContextFilter(logging.Filter):
_correlation_filter: Final = CorrelationContextFilter()
json_logs = bool(os.getenv("JSON_LOGS", False))
_LOG_FORMAT_PREFIX: Final = "%(asctime)s - %(name)s:%(levelname)s"
_LOG_FORMAT_SUFFIX: Final = ": %(filename)s:%(lineno)s - %(message)s"
_PLAIN_LOG_FORMAT: Final = _LOG_FORMAT_PREFIX + _LOG_FORMAT_SUFFIX
_COLOR_LOG_FORMAT: Final = f"\033[92m{_LOG_FORMAT_PREFIX}\033[0m{_LOG_FORMAT_SUFFIX}"
def _stream_is_tty(stream: TextIO | None) -> bool:
"""True when the stream is an open interactive terminal; never raises.
A stream can be None (pythonw/embedded interpreters), lack isatty entirely
(GUI log-redirect shims), or be closed; import must survive all three.
"""
try:
return stream is not None and stream.isatty()
except (AttributeError, ValueError):
return False
def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str:
"""The plain-text log format, colorized only when both streams are an interactive terminal.
Honors the NO_COLOR convention from no-color.org: color is disabled when
NO_COLOR is present with a non-empty value.
"""
if os.environ.get("NO_COLOR"):
return _PLAIN_LOG_FORMAT
return _COLOR_LOG_FORMAT if _stream_is_tty(stdout) and _stream_is_tty(stderr) else _PLAIN_LOG_FORMAT
class LevelRoutingStreamHandler(logging.StreamHandler):
"""Writes records below WARNING to stdout and WARNING and above to stderr.
Collectors that derive severity from the stream report every stderr line as an error.
"""
def emit(self, record: logging.LogRecord) -> None:
preferred: Final = sys.stdout if record.levelno < logging.WARNING else sys.stderr
if preferred is None or getattr(preferred, "closed", False):
self.stream = sys.stderr # rebind-ok: fall back to the pre-fix stream rather than raising per record
else:
self.stream = preferred # rebind-ok: StreamHandler.emit writes self.stream under the handler lock
super().emit(record)
def _parse_json_logs_env(value: str | None) -> bool:
"""Strict opt-in parse for the JSON_LOGS env var: only "true" (any case) enables JSON logs.
Matches the reader in litellm-proxy-extras/_logging.py. The previous
bool(os.getenv(...)) treated any non-empty value, including "false" and "0",
as enabled.
"""
return (value or "").lower() == "true"
json_logs: Final = _parse_json_logs_env(os.getenv("JSON_LOGS"))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: Final[str] = getattr(logging, log_level.upper())
handler: Final = logging.StreamHandler()
handler: Final = LevelRoutingStreamHandler()
handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)
handler.addFilter(_correlation_filter)
@ -447,7 +501,7 @@ if json_logs:
_setup_json_exception_handlers(JsonFormatter())
else:
formatter: Final = CorrelationPlainFormatter(
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
_plain_log_format(sys.stdout, sys.stderr),
datefmt="%H:%M:%S",
)
@ -628,7 +682,7 @@ def _turn_on_json():
- Adds a JSON formatter to all loggers
"""
handler: Final = logging.StreamHandler()
handler: Final = LevelRoutingStreamHandler()
handler.setFormatter(JsonFormatter())
_initialize_loggers_with_handler(handler)
# Set up exception handlers

View file

@ -17,11 +17,27 @@ A2A Streaming Events:
- Artifact update (kind: "artifact-update") - Content/artifact delivery
"""
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import TYPE_CHECKING, Final
from uuid import uuid4
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
_STR_KEY_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_object_mapping(value: object) -> Mapping[str, object]:
try:
return _STR_KEY_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return {}
class A2AStreamingContext:
@ -30,7 +46,7 @@ class A2AStreamingContext:
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: dict[str, Any]):
def __init__(self, request_id: str, input_message: Mapping[str, JsonValue]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
@ -46,44 +62,46 @@ class A2ACompletionBridgeTransformation:
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
def _text_from_a2a_part(part: JsonValue) -> str | None:
if not isinstance(part, dict):
return None
text: Final = part.get("text")
if text is None:
return None
if part.get("kind") not in (None, "", "text"):
return None
return str(text)
@staticmethod
def _extract_text_from_a2a_parts(parts: Sequence[JsonValue]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: Final[list[str]] = []
for part in parts:
if not isinstance(part, dict):
continue
kind = part.get("kind")
text = part.get("text")
if text is None:
continue
if kind in (None, "", "text"):
content_parts.append(str(text))
return "\n".join(content_parts)
extracted: Final = (A2ACompletionBridgeTransformation._text_from_a2a_part(part) for part in parts)
return "\n".join(text for text in extracted if text is not None)
@staticmethod
def get_forward_metadata(
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
a2a_message: Mapping[str, JsonValue],
params: Mapping[str, JsonValue] | None = None,
) -> Mapping[str, JsonValue] | None:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message see ``apply_forward_metadata_to_completion_params``.
"""
merged: Final[dict[str, Any]] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
params_metadata: Final = params.get("metadata") if params else None
message_metadata: Final = a2a_message.get("metadata")
if isinstance(message_metadata, dict):
merged.update(message_metadata)
merged: Final[dict[str, JsonValue]] = {
**(params_metadata if isinstance(params_metadata, dict) else {}),
**(message_metadata if isinstance(message_metadata, dict) else {}),
}
return merged or None
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: dict[str, Any],
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
completion_params: MutableMapping[str, object],
a2a_message: Mapping[str, JsonValue],
params: Mapping[str, JsonValue] | None = None,
) -> None:
"""
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
@ -97,24 +115,20 @@ class A2ACompletionBridgeTransformation:
if not forward_metadata:
return
extra_body = completion_params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
extra_body: Final = _as_object_mapping(completion_params.get("extra_body"))
# Layer client-supplied A2A metadata under any agent-owner-configured
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata: Final = extra_body.get("metadata")
existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
existing_dict: Final = _as_object_mapping(extra_body.get("metadata"))
merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict}
completion_params["extra_body"] = {**extra_body, "metadata": merged_metadata}
verbose_logger.debug("A2A -> completion forward metadata keys=%s", list(forward_metadata.keys()))
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: dict[str, Any],
) -> list[dict[str, Any]]:
a2a_message: Mapping[str, JsonValue],
) -> list[dict[str, object]]:
"""
Transform an A2A message to OpenAI message format.
@ -125,25 +139,19 @@ class A2ACompletionBridgeTransformation:
List of OpenAI-format messages
"""
role: Final = a2a_message.get("role", "user")
parts = a2a_message.get("parts", [])
raw_parts: Final = a2a_message.get("parts", [])
# Map A2A roles to OpenAI roles
openai_role = role
if role == "user":
openai_role = "user"
elif role == "assistant":
openai_role = "assistant"
elif role == "system":
openai_role = "system"
if not isinstance(parts, list):
parts = []
openai_role: Final = (
"user" if role == "user" else "assistant" if role == "assistant" else "system" if role == "system" else role
)
parts: Final = raw_parts if isinstance(raw_parts, list) else []
content: Final = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts)
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content}
openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content}
verbose_logger.debug(
"A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content)
@ -151,11 +159,20 @@ class A2ACompletionBridgeTransformation:
return [openai_message]
@staticmethod
def _extract_response_content(response: "ModelResponse | CustomStreamWrapper") -> str:
if not isinstance(response, ModelResponse) or not response.choices:
return ""
choice: Final = response.choices[0]
if not choice.message:
return ""
return choice.message.content or ""
@staticmethod
def openai_response_to_a2a_response(
response: Any,
response: "ModelResponse | CustomStreamWrapper",
request_id: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
@ -166,12 +183,7 @@ class A2ACompletionBridgeTransformation:
Returns:
A2A SendMessageResponse dict
"""
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice: Final = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
content: Final = A2ACompletionBridgeTransformation._extract_response_content(response)
# Build A2A message
a2a_message: Final = {
@ -182,7 +194,7 @@ class A2ACompletionBridgeTransformation:
}
# Build A2A response
a2a_response: Final = {
a2a_response: Final[dict[str, object]] = {
"jsonrpc": "2.0",
"id": request_id,
"result": a2a_message,
@ -200,7 +212,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create the initial task event with status 'submitted'.
@ -235,7 +247,7 @@ class A2ACompletionBridgeTransformation:
state: str,
final: bool = False,
message_text: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create a status update event.
@ -245,7 +257,7 @@ class A2ACompletionBridgeTransformation:
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Final[dict[str, Any]] = {
status: Final[dict[str, object]] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
@ -277,7 +289,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create an artifact update event with content.

View file

@ -86,7 +86,7 @@ A2ACardResolver: Final = LiteLLMA2ACardResolver
def _set_usage_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
prompt_tokens: int,
completion_tokens: int,
) -> None:
@ -99,7 +99,7 @@ def _set_usage_on_logging_obj(
completion_tokens: Number of output tokens
"""
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
usage: Final = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@ -109,7 +109,7 @@ def _set_usage_on_logging_obj(
def _set_agent_id_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
agent_id: str | None,
) -> None:
"""
@ -123,7 +123,7 @@ def _set_agent_id_on_logging_obj(
return
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
# Set agent_id directly on model_call_details (same pattern as custom_llm_provider)
litellm_logging_obj.model_call_details["agent_id"] = agent_id
@ -132,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output
def _set_litellm_params_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> None:
"""
@ -144,18 +144,22 @@ def _set_litellm_params_on_logging_obj(
context, so merge the pricing keys in rather than replacing the dict.
"""
logging_obj: Final = kwargs.get("litellm_logging_obj")
if logging_obj is None:
if not isinstance(logging_obj, Logging):
return
cost_params = {key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None}
cost_params: Final = {
key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None
}
if not cost_params:
return
existing: Final = logging_obj.model_call_details.get("litellm_params") or {}
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
logging_obj.model_call_details["litellm_params"] = {
**(logging_obj.model_call_details.get("litellm_params") or {}),
**cost_params,
}
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str:
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str:
"""
Extract agent info and set model/custom_llm_provider for cost tracking.
@ -175,7 +179,7 @@ def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) ->
# Set on litellm_logging_obj if available (for standard logging payload)
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
litellm_logging_obj.model = model
litellm_logging_obj.custom_llm_provider = custom_llm_provider
litellm_logging_obj.model_call_details["model"] = model
@ -498,7 +502,7 @@ async def asend_message(
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
# Calculate token usage from request and response
response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True)
response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True)
(
prompt_tokens,
completion_tokens,

View file

@ -390,7 +390,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
] = "openai",
logging_obj: Any | None = None,
logging_obj: LiteLLMLoggingObj | None = None,
):
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:

View file

@ -59,9 +59,11 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
AllMessageValues,
ChatCompletionFileObject,
ChatCompletionImageObject,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolReferenceObject,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import Choices
@ -175,6 +177,16 @@ def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Li
return "length"
def _input_file_from_file_value(file_value: object) -> dict[str, object]:
if not isinstance(file_value, dict):
return {"type": "input_file"}
file_dict: Final = cast("dict[str, object]", file_value) # cast-ok: runtime dict checked
return {
"type": "input_file",
**{key: file_dict[key] for key in ("file_id", "file_data", "filename") if key in file_dict},
}
def _incomplete_reason_from_response_payload(response_payload: object) -> str | None:
if not isinstance(response_payload, Mapping):
return None
@ -957,7 +969,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
content: str
| list[object]
| Iterable[
Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]
Union[
"OpenAIMessageContentListBlock",
"ChatCompletionThinkingBlock",
"ChatCompletionRedactedThinkingBlock",
"ChatCompletionToolReferenceObject",
]
]
| None,
role: str,
@ -1006,17 +1023,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
result.append(converted)
verbose_logger.debug("Chat provider: image -> %s", converted)
elif item_type == "file":
# Map Chat Completion file to Responses API input_file
# {"type": "file", "file": {"file_data": "...", "filename": "..."}}
# -> {"type": "input_file", "file_data": "...", "filename": "..."}
file_data = item.get("file", {})
converted = {"type": "input_file"}
if isinstance(file_data, dict):
for key in ["file_id", "file_data", "filename"]:
if key in file_data:
converted[key] = file_data[key]
converted = _input_file_from_file_value(
cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked
)
result.append(converted)
verbose_logger.debug("Chat provider: file -> %s", converted)
elif item_type == "tool_reference":
verbose_logger.debug(
"Chat provider: tool_reference has no responses API equivalent; skipped"
)
elif item_type in [
"input_text",
"input_image",

View file

@ -296,6 +296,9 @@ GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
)
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES: Final = 500_000
PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096
PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
@ -626,6 +629,15 @@ LITELLM_CHAT_PROVIDERS: Final = [
"amazon_nova",
]
# Resolving these providers runs an OAuth device flow (their provider info IS the login), so any
# metadata or capability lookup against them can block for minutes waiting on a human.
PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
{
"github_copilot",
"chatgpt",
}
)
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
"openai",
"azure",
@ -1364,8 +1376,6 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
AUTO_ROUTED_REQUEST_METADATA_KEY: Final = "_auto_routed_request"
ROUTER_MODEL_NAME_RESPONSE_FIELD: Final = "router_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
@ -1475,6 +1485,12 @@ LITELLM_PROXY_MASTER_KEY_ALIAS: Final = "litellm_proxy_master_key"
# ``ProxyLogging._handle_logging_proxy_only_error``.
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: Final = "litellm_no_upstream_llm_call"
# Key/team metadata fields naming the OTel Resource ``service.name``, highest
# precedence first. Shared between the OTel v2 tenant router (which reads them
# out of ``user_api_key_auth_metadata``) and proxy request setup (which re-applies
# the key's values after the team metadata merge so a key outranks its team).
OTEL_SERVICE_NAME_METADATA_KEYS: Final = ("otel_service_name_override", "otel_service_name")
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED: Final = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS: Final = int(
@ -1648,6 +1664,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"enable_anthropic_prompt_caching",
"anthropic_prompt_caching_ttl",
"max_ui_session_budget",
"budget_rollover",
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
@ -1814,6 +1831,43 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
# A retrieved response replays the usage of the call that created it, so pricing these
# read/management routes like inference bills the same tokens twice.
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
{
"get_responses",
"aget_responses",
"delete_responses",
"adelete_responses",
"cancel_responses",
"acancel_responses",
"list_input_items",
"alist_input_items",
"vector_store_create",
"avector_store_create",
"vector_store_retrieve",
"avector_store_retrieve",
"vector_store_list",
"avector_store_list",
"vector_store_update",
"avector_store_update",
"vector_store_delete",
"avector_store_delete",
"vector_store_file_create",
"avector_store_file_create",
"vector_store_file_list",
"avector_store_file_list",
"vector_store_file_retrieve",
"avector_store_file_retrieve",
"vector_store_file_content",
"avector_store_file_content",
"vector_store_file_update",
"avector_store_file_update",
"vector_store_file_delete",
"avector_store_file_delete",
}
)
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
# spend under the table's composite unique constraint.

View file

@ -2,7 +2,7 @@
## File for 'response_cost' calculation in Logging
import logging
import time
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Literal, cast
@ -76,7 +76,10 @@ from litellm.llms.perplexity.cost_calculator import (
from litellm.llms.tencent.cost_calculator import (
cost_per_token as tencent_cost_per_token,
)
from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
from litellm.llms.together_ai.cost_calculator import (
get_model_params_and_category,
has_together_registry_pricing,
)
from litellm.llms.vertex_ai.cost_calculator import (
cost_per_character as google_cost_per_character,
)
@ -557,9 +560,10 @@ def cost_per_token(
)
elif call_type == "atranscription" or call_type == "transcription":
if _transcription_usage_has_token_details(usage_block):
return openai_cost_per_token(
return generic_cost_per_token(
model=model_without_prefix,
usage=usage_block,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
data_residency=data_residency,
)
@ -735,6 +739,13 @@ def _get_provider_for_cost_calc(
return custom_llm_provider
def _get_hidden_str_for_cost_calc(hidden_params: object, key: str) -> str | None:
if not isinstance(hidden_params, Mapping):
return None
value: Final[object] = hidden_params.get(key)
return value if isinstance(value, str) and value else None
def _select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -751,7 +762,6 @@ def _select_model_name_for_cost_calc(
"""
return_model: str | None = None
region_name: str | None = None
custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider)
completion_response_model: str | None = None
@ -761,6 +771,14 @@ def _select_model_name_for_cost_calc(
elif isinstance(completion_response, dict):
completion_response_model = completion_response.get("model", None)
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None)
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
explicit_pricing: Final = custom_pricing is True or base_model is not None
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
region_name: Final = (
_get_hidden_str_for_cost_calc(hidden_params, "region_name")
if not explicit_pricing and priced_from_response
else None
)
if custom_pricing is True:
if router_model_id is not None and router_model_id in litellm.model_cost:
@ -776,14 +794,12 @@ def _select_model_name_for_cost_calc(
else:
return_model = model
elif base_model is not None:
return_model = base_model
elif base_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else provider_response_model
elif completion_response_model is None and hidden_params is not None:
if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0:
return_model = hidden_params.get("model", model)
elif hidden_params is not None and hidden_params.get("region_name", None) is not None:
region_name = hidden_params.get("region_name", None)
if return_model is None and completion_response_model is not None:
return_model = completion_response_model
@ -1568,10 +1584,9 @@ def completion_cost(
return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj)
# Calculate cost based on prompt_tokens, completion_tokens
if "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai":
# together ai prices based on size of llm
# get_model_params_and_category takes a model name and returns the category of LLM size it is in model_prices_and_context_window.json
if (
"togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai"
) and not has_together_registry_pricing(model, litellm.model_cost):
model = get_model_params_and_category(model, call_type=CallTypes(call_type))
# replicate llms are calculate based on time for request running

View file

@ -7,6 +7,7 @@ import base64
import os
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from importlib import metadata
from typing import Any, Final, TypeVar
import httpx
@ -21,6 +22,18 @@ try:
streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
except ImportError:
pass
MCP_STREAMABLE_HTTP_REQUIREMENT: Final = "mcp>=1.28.1"
def missing_streamable_http_client_error() -> ImportError:
return ImportError(
f"MCP streamable HTTP transport requires {MCP_STREAMABLE_HTTP_REQUIREMENT}, but the installed "
f"mcp {metadata.version('mcp')} does not provide streamable_http_client. "
"Fix with: pip install 'litellm[mcp]' (or upgrade mcp directly: pip install -U mcp)"
)
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
@ -43,6 +56,9 @@ from litellm.types.mcp import (
MCPStdioConfig,
MCPTransport,
MCPTransportType,
credential_redirect_hook,
has_header,
without_header,
)
@ -260,6 +276,7 @@ class MCPClient:
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: str | dict[str, str] | None = None,
auth_header_name: str | None = None,
timeout: float | None = None,
stdio_config: MCPStdioConfig | None = None,
extra_headers: dict[str, str] | None = None,
@ -275,6 +292,11 @@ class MCPClient:
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout if timeout is not None else MCP_CLIENT_TIMEOUT
self._mcp_auth_value: str | dict[str, str] | None = None
# The one place this client decides which header its credential occupies: the operator's
# configured slot on the v1 path, or the slot the v2 resolver's auth object already owns.
# Every consumer reads this rather than re-deriving it, since each re-derivation so far
# picked up a different bug.
self._credential_slot: str | None = auth_header_name or getattr(resolved_auth, "header_name", None)
self.stdio_config: MCPStdioConfig | None = stdio_config
self.extra_headers: dict[str, str] | None = extra_headers
self.ssl_verify: VerifyTypes | None = ssl_verify
@ -323,7 +345,7 @@ class MCPClient:
)
# HTTP transport (default)
if streamable_http_client is None:
raise ImportError("streamable_http_client is not available. Please install mcp with HTTP support.")
raise missing_streamable_http_client_error()
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
@ -488,26 +510,33 @@ class MCPClient:
else:
self._mcp_auth_value = mcp_auth_value
def _header_slot(self, default: str) -> str:
return self._credential_slot or default
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers: Final = {}
if self._mcp_auth_value:
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
static_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {static_bearer}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
headers[self._header_slot("Authorization")] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
headers[self._header_slot("X-API-Key")] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
# This auth type means the caller owns the whole header value.
headers["Authorization"] = self._mcp_auth_value
headers[self._header_slot("Authorization")] = self._mcp_auth_value
elif self.auth_type == MCPAuth.oauth2:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
oauth2_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {oauth2_bearer}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {strip_auth_scheme(self._mcp_auth_value, 'token')}"
scheme_token: Final = strip_auth_scheme(self._mcp_auth_value, "token")
headers[self._header_slot("Authorization")] = f"token {scheme_token}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
exchanged_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {exchanged_bearer}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
@ -515,7 +544,14 @@ class MCPClient:
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
# update the headers with the extra headers
if self.extra_headers:
headers.update(self.extra_headers)
# Mirrors _resolve_v2_auth: when the operator named a slot for the credential the
# gateway resolved, no injected header may shadow it, case-insensitively, since HTTP
# header names are. Without a configured slot the old precedence stands unchanged.
slot: Final = self._credential_slot
injected: Final = (
without_header(self.extra_headers, slot) if slot and has_header(headers, slot) else self.extra_headers
)
headers.update(injected or {})
return _strip_header_whitespace(headers)
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
@ -543,12 +579,14 @@ class MCPClient:
# SigV4 aws_auth. Both are None for the common case — no behavior change.
fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth
effective_auth: Final = auth if auth is not None else fallback_auth
guard: Final = credential_redirect_hook(self.server_url, self._credential_slot)
return httpx.AsyncClient(
headers=headers,
timeout=timeout,
auth=effective_auth,
verify=ssl_config,
follow_redirects=True,
event_hooks={"request": [guard]} if guard else {},
)
return factory

View file

@ -2,6 +2,7 @@ import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
@ -65,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
custom_llm_provider: str,
hidden_params: dict[str, Any] | None = None,
):
self.litellm_logging_obj = litellm_logging_obj
@ -72,6 +74,10 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
self.start_time = datetime.now()
self.collected_chunks: list[bytes] = []
self.model = model
self.custom_llm_provider = custom_llm_provider
self.endpoint_type: Final = (
EndpointType.GEMINI if custom_llm_provider == litellm.LlmProviders.GEMINI.value else EndpointType.VERTEX_AI
)
self._hidden_params: dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
@ -89,7 +95,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route="/v1/generateContent",
request_body=self.request_body or {},
endpoint_type=EndpointType.VERTEX_AI,
endpoint_type=self.endpoint_type,
start_time=self.start_time,
raw_bytes=self.collected_chunks,
end_time=end_time,
@ -118,13 +124,13 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
self.response = response
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.iter_lines()
@ -169,13 +175,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
self.response = response
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.aiter_lines()

View file

@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_proxy_logger
from .ms_teams import MS_TEAMS_ALERTING_DESTINATION, build_ms_teams_payload
if TYPE_CHECKING:
from .slack_alerting import SlackAlerting as _SlackAlerting
@ -62,14 +64,17 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count)
if count > 1:
payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}"
request_body: Final = (
build_ms_teams_payload(payload["text"]) if item.get("format") == MS_TEAMS_ALERTING_DESTINATION else payload
)
response: Final = await slackAlertingInstance.async_http_handler.post(
url=item["url"],
headers=item["headers"],
data=json.dumps(payload),
data=json.dumps(request_body),
)
if response.status_code != 200:
verbose_proxy_logger.debug("Error sending slack alert to url=%s. Error=%s", item["url"], response.text)
verbose_proxy_logger.debug("Error sending alert to url=%s. Error=%s", item["url"], response.text)
except Exception as e:
verbose_proxy_logger.debug("Error sending slack alert: %s", e)
verbose_proxy_logger.debug("Error sending alert: %s", e)
finally:
_print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance)

View file

@ -0,0 +1,75 @@
"""Microsoft Teams alert delivery helpers.
Teams incoming webhooks (Workflows and legacy connectors) accept an Adaptive
Card wrapped in a message attachment, so alert text is delivered as a single
wrapped TextBlock.
"""
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from typing_extensions import ReadOnly, TypedDict
from litellm.types.integrations.slack_alerting import AlertType
MS_TEAMS_WEBHOOK_URL_ENV: Final = "MS_TEAMS_WEBHOOK_URL"
MS_TEAMS_ALERTING_DESTINATION: Final = "ms_teams"
MS_TEAMS_ALERT_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"Content-type": "application/json"})
class MSTeamsTextBlock(TypedDict):
type: ReadOnly[str]
text: ReadOnly[str]
wrap: ReadOnly[bool]
class MSTeamsAdaptiveCard(TypedDict):
type: ReadOnly[str]
version: ReadOnly[str]
body: ReadOnly[tuple[MSTeamsTextBlock, ...]]
class MSTeamsAttachment(TypedDict):
contentType: ReadOnly[str]
content: ReadOnly[MSTeamsAdaptiveCard]
class MSTeamsMessage(TypedDict):
type: ReadOnly[str]
attachments: ReadOnly[tuple[MSTeamsAttachment, ...]]
class MSTeamsAlertText(TypedDict):
text: ReadOnly[str]
class MSTeamsQueueItem(TypedDict):
url: ReadOnly[str]
headers: ReadOnly[Mapping[str, str]]
payload: ReadOnly[MSTeamsAlertText]
alert_type: ReadOnly[AlertType]
format: ReadOnly[str]
def get_ms_teams_webhook_url() -> str | None:
return os.getenv(MS_TEAMS_WEBHOOK_URL_ENV)
def build_ms_teams_payload(text: str) -> MSTeamsMessage:
return MSTeamsMessage(
type="message",
attachments=(
MSTeamsAttachment(
contentType="application/vnd.microsoft.card.adaptive",
content=MSTeamsAdaptiveCard(
type="AdaptiveCard",
version="1.4",
body=(MSTeamsTextBlock(type="TextBlock", text=text, wrap=True),),
),
),
),
)

View file

@ -57,6 +57,13 @@ from litellm.types.proxy.model_deprecation import (
from ..email_templates.templates import *
from .batching_handler import send_to_webhook, squash_payloads
from .ms_teams import (
MS_TEAMS_ALERT_HEADERS,
MS_TEAMS_ALERTING_DESTINATION,
MSTeamsAlertText,
MSTeamsQueueItem,
get_ms_teams_webhook_url,
)
from .utils import process_slack_alerting_variables
if TYPE_CHECKING:
@ -1431,13 +1438,43 @@ Model Info:
# only send budget alerts over Email
await self.send_email_alert_using_smtp(webhook_event=user_info, alert_type=alert_type)
if "slack" not in self.alerting:
send_to_slack: Final = "slack" in self.alerting
send_to_ms_teams: Final = MS_TEAMS_ALERTING_DESTINATION in self.alerting
if not send_to_slack and not send_to_ms_teams:
return
if alert_type not in self.alert_types:
return
from datetime import datetime
current_time: Final = datetime.now().strftime("%H:%M:%S")
_proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
alert_type_name: Final = getattr(alert_type, "name", alert_type)
alert_type_formatted: Final = f"Alert type: `{alert_type_name}`"
if alert_type == "daily_reports" or alert_type == "new_model_added":
formatted_message = alert_type_formatted + message
else:
formatted_message = (
f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
)
if kwargs:
for key, value in kwargs.items():
formatted_message += f"\n\n{key}: `{value}`\n\n"
if alerting_metadata:
for key, value in alerting_metadata.items():
formatted_message += f"\n\n*Alerting Metadata*: \n{key}: `{value}`\n\n"
if _proxy_base_url is not None:
formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
if send_to_ms_teams:
self._enqueue_ms_teams_alert(formatted_message=formatted_message, alert_type=alert_type)
if not send_to_slack:
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
return
# Check if digest mode is enabled for this alert type
alert_type_name_str: Final = getattr(alert_type, "value", str(alert_type))
_atc: Final = self.alert_type_config.get(alert_type_name_str)
@ -1473,28 +1510,6 @@ Model Info:
)
return # Suppress immediate alert; will be emitted by _flush_digest_buckets
# Get the current timestamp
current_time: Final = datetime.now().strftime("%H:%M:%S")
_proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
# Use .name if it's an enum, otherwise use as is
alert_type_name: Final = getattr(alert_type, "name", alert_type)
alert_type_formatted: Final = f"Alert type: `{alert_type_name}`"
if alert_type == "daily_reports" or alert_type == "new_model_added":
formatted_message = alert_type_formatted + message
else:
formatted_message = (
f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
)
if kwargs:
for key, value in kwargs.items():
formatted_message += f"\n\n{key}: `{value}`\n\n"
if alerting_metadata:
for key, value in alerting_metadata.items():
formatted_message += f"\n\n*Alerting Metadata*: \n{key}: `{value}`\n\n"
if _proxy_base_url is not None:
formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
# check if we find the slack webhook url in self.alert_to_webhook_url
if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url:
slack_webhook_url: str | list[str] | None = self.alert_to_webhook_url[alert_type]
@ -1531,6 +1546,24 @@ Model Info:
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
def _enqueue_ms_teams_alert(self, formatted_message: str, alert_type: AlertType) -> None:
ms_teams_webhook_url: Final = get_ms_teams_webhook_url()
if ms_teams_webhook_url is None:
verbose_proxy_logger.error(
"MS Teams alerting is enabled but MS_TEAMS_WEBHOOK_URL is not set. Dropping alert type=%s",
alert_type,
)
return
payload: Final[MSTeamsAlertText] = {"text": formatted_message}
item: Final[MSTeamsQueueItem] = {
"url": ms_teams_webhook_url,
"headers": MS_TEAMS_ALERT_HEADERS,
"payload": payload,
"alert_type": alert_type,
"format": MS_TEAMS_ALERTING_DESTINATION,
}
self.log_queue.append(item)
async def async_send_batch(self):
if not self.log_queue:
return

View file

@ -24,6 +24,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
CacheControlInjectionPoint,
CacheControlMessageInjectionPoint,
)
@ -185,7 +187,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
reserved_blocks: Final = (
1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
processed_messages = self._apply_message_injections(
points=applied_message_points,
messages=processed_messages,
@ -194,7 +196,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages) > breakpoints_before
and AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) > breakpoints_before
):
non_default_params.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
@ -236,7 +238,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return provider
@staticmethod
def _count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
def count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
system_blocks: Final = (
sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
)
@ -258,7 +260,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
``max_blocks`` is reached. Injection points are honored in config order,
so earlier points win when slots are scarce.
"""
used_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints(messages)
used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
limit_reached = False
for point in points:
@ -376,7 +378,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# 2. list of objects - only apply to last item per Anthropic spec
elif isinstance(message_content, list):
if len(message_content) > 0 and isinstance(message_content[-1], dict):
message_content[-1]["cache_control"] = control
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
return message
@staticmethod
@ -454,8 +456,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
)
max_blocks: Final = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks
message_blocks: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
system_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints((), processed_system)
message_blocks: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
system_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints((), processed_system)
if system_points and processed_system is not None and message_blocks + system_blocks < max_blocks:
system_already_has_cc: Final = isinstance(processed_system, list) and any(
@ -589,7 +591,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > 0:
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
return any(
@ -749,6 +751,64 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if points:
non_default_params["cache_control_injection_points"] = points
@staticmethod
def record_gateway_injection(
request_kwargs: Mapping[str, object],
added: int,
) -> None:
"""Name the deployment whose payload the gateway, not the client, put breakpoints on.
Spend accounting only asks whether litellm acted, so what it needs is which
deployment, not a count. Recording that is what makes the mark attempt-scoped: the
metadata bucket is one dict shared by every retry, failover and fallback of a
request, and ``litellm_call_id`` is shared with it, so anything request-scoped
written by one attempt is read by all of them and each boundary would have to
remember to strip it. The deployment is the part that actually changes when the
request moves, so a leg that injected nothing is never credited for one that did.
It also makes a zero delta (hook re-entry) and a negative one (a prompt manager
replacing the messages) harmless, since neither rewrites an earlier mark.
A pass that runs before a deployment is chosen, which is what the proxy does for
prompt templates, injects into the payload every leg goes on to send, so it marks
the request for all of them rather than for one.
Only what this pass actually placed counts. A ``tool_config`` point is placed by
the Bedrock converse transform, and only when the request carries tools, so the
presence of one here says nothing about whether a breakpoint reaches the wire;
claiming it marked three request shapes out of four that inject nothing. Missing
that Bedrock credit is the fail-closed direction, and the alternative is a
provider transform that carries spend-attribution state.
Reads whichever bucket the request actually carries rather than asking the shared
name resolver, which answers on key presence: ``litellm_params`` declares
``litellm_metadata`` as None on every request, so the resolver names a bucket that
is not there and the mark is dropped.
Never CREATES the bucket. The proxy seeds it on every request and is the marker's
only reader, so a request without one is a bare SDK call nothing would consume it
from. Creating it would also add a key to a dict call sites splat as ``**kwargs``,
and on the Responses API ``metadata`` is both this bucket's default name and an
explicit parameter, so the splat collides with the caller's own value.
"""
if added <= 0:
return
bucket: Final = next(
(
candidate
for candidate in (request_kwargs.get("litellm_metadata"), request_kwargs.get("metadata"))
if isinstance(candidate, dict)
),
None,
)
if bucket is not None:
model_info: Final = request_kwargs.get("model_info")
bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = (
model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT)
if isinstance(model_info, dict)
else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT
)
@staticmethod
def maybe_inject_cache_control(
messages: list[dict],
@ -798,17 +858,18 @@ class AnthropicCacheControlHook(CustomPromptManagement):
openai_dialect: Final = AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(
model, custom_llm_provider, api_base, kwargs.get("prompt_cache_options")
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system)
messages, system, remaining = AnthropicCacheControlHook.apply_to_anthropic_messages_request(
messages=messages,
system=system,
injection_points=injection_points,
openai_dialect=openai_dialect,
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > breakpoints_before
):
breakpoints_added: Final = (
AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) - breakpoints_before
)
AnthropicCacheControlHook.record_gateway_injection(kwargs, breakpoints_added)
if openai_dialect and breakpoints_added > 0:
kwargs.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
if remaining:
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining)

View file

@ -4,11 +4,38 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
import base64
import urllib.parse
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, TypedDict
from typing_extensions import NotRequired, ReadOnly
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class BitBucketSrcEntry(TypedDict):
path: ReadOnly[NotRequired[str]]
type: ReadOnly[NotRequired[str]]
class BitBucketSrcListing(TypedDict):
values: ReadOnly[NotRequired[list[BitBucketSrcEntry]]]
class BitBucketBranch(TypedDict):
name: ReadOnly[NotRequired[str]]
type: ReadOnly[NotRequired[str]]
class BitBucketBranchListing(TypedDict):
values: ReadOnly[NotRequired[list[BitBucketBranch]]]
class BitBucketFileMetadata(TypedDict):
content_type: ReadOnly[str | None]
content_length: ReadOnly[str | None]
last_modified: ReadOnly[str | None]
def _sanitize_file_path(file_path: str) -> str:
"""Reject path traversal and URL-encode each path segment."""
if "#" in file_path or "?" in file_path:
@ -31,7 +58,7 @@ class BitBucketClient:
- Branch-specific file fetching
"""
def __init__(self, config: dict[str, Any]):
def __init__(self, config: Mapping[str, object]):
"""
Initialize the BitBucket client.
@ -135,8 +162,8 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
files: Final = []
data: Final[BitBucketSrcListing] = response.json()
files: Final[list[str]] = []
for item in data.get("values", []):
if item.get("type") == "commit_file":
@ -162,7 +189,7 @@ class BitBucketClient:
else:
raise Exception(f"Error listing files in '{directory_path}': {e}")
def get_repository_info(self) -> dict[str, Any]:
def get_repository_info(self) -> Mapping[str, object]:
"""
Get information about the repository.
@ -191,7 +218,7 @@ class BitBucketClient:
except Exception:
return False
def get_branches(self) -> list[dict[str, Any]]:
def get_branches(self) -> list[BitBucketBranch]:
"""
Get list of branches in the repository.
@ -204,12 +231,12 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
data: Final[BitBucketBranchListing] = response.json()
return data.get("values", [])
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str) -> dict[str, Any] | None:
def get_file_metadata(self, file_path: str) -> BitBucketFileMetadata | None:
"""
Get metadata about a file (size, last modified, etc.).

View file

@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
import time
import uuid
from typing import Any, ClassVar, Final, cast
from typing import TYPE_CHECKING, Any, ClassVar, Final, cast
from litellm._logging import verbose_logger
from litellm.compression import compress
@ -22,6 +22,9 @@ from litellm.types.integrations.custom_logger import (
)
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
_CACHE_TTL_SECONDS: Final = 15 * 60
@ -222,7 +225,7 @@ class CompressionInterceptionLogger(CustomLogger):
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj | None",
stream: bool,
kwargs: dict,
) -> AgenticLoopPlan:

View file

@ -45,7 +45,7 @@ class CustomBatchLogger(CustomLogger):
super().__init__(**kwargs)
async def periodic_flush(self):
async def periodic_flush(self) -> None:
while True:
await asyncio.sleep(self.flush_interval)
verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval)

View file

@ -60,6 +60,12 @@ _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
DEFAULT_ADVISORY_MESSAGE: Final = (
"The user's latest message was flagged for {reason} by a content safety "
"guardrail. This may be a false positive. Use your judgment: respond "
"helpfully if the request is legitimate, or decline if it is not."
)
_guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar(
"litellm_guardrail_self_recorded", default=False
)
@ -158,6 +164,7 @@ class CustomGuardrail(CustomLogger):
sensitive_data_route_to_model: str | None = None,
sticky_session_routing: bool = True,
run_in_parallel: bool = False,
scan_raw_request: bool = False,
only_scan_new_messages: bool = False,
**kwargs,
):
@ -180,6 +187,13 @@ class CustomGuardrail(CustomLogger):
run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with
other opted-in guardrails of the same hook. Only safe for block-only guardrails that
do not mutate the request or response.
scan_raw_request: When True, this pre_call guardrail always evaluates the request as it
was before any guardrail in this hook ran, regardless of where it's declared in the
guardrails list -- so an earlier guardrail that masks/rewrites content (e.g. PII
redaction) can never hide a violation from this one. Only safe for block-only
guardrails: any data this guardrail returns is discarded, matching run_in_parallel's
contract, since applying its mutations on top of a stale snapshot would silently
undo whatever later guardrails already did to the live request.
"""
self.guardrail_name = guardrail_name
self.supported_event_hooks = supported_event_hooks
@ -195,6 +209,7 @@ class CustomGuardrail(CustomLogger):
self.sensitive_data_route_to_model: str | None = sensitive_data_route_to_model
self.sticky_session_routing: bool = sticky_session_routing
self.run_in_parallel: bool = run_in_parallel
self.scan_raw_request: bool = scan_raw_request
self.only_scan_new_messages: bool = only_scan_new_messages
if supported_event_hooks:
@ -281,6 +296,82 @@ class CustomGuardrail(CustomLogger):
original_response=original_response,
)
def inject_advisory_message(
self,
data: dict[str, Any], # mutable-ok: caller's dict is mutated in place, matching mark_pre_call_hook_ran
message: str,
) -> bool:
"""
Append an advisory system message to the request in place, so the LLM
itself can weigh a possible false-positive guardrail flag rather than
the request being hard-blocked or silently allowed.
Unlike raise_passthrough_exception, this does NOT short-circuit the LLM
call; the request proceeds normally with the extra message appended.
Guardrails should call this from on_flagged handling analogous to how
passthrough-supporting guardrails call raise_passthrough_exception.
Args:
data: The request data dictionary, mutated in place to append the
advisory message to its "messages" list and/or "input"/
"instructions" text.
message: The formatted advisory message to append as a system message.
Returns:
True if the advisory was actually written somewhere the model will
see it. False if ``data["input"]`` is a structured Responses-API
list (not a plain string) -- the Responses API reads only
``input``, so appending to ``messages`` would be inert regardless
of whether a ``messages`` list also happens to be present, and
there is no field this helper can safely append into. The caller
must treat this like any other case where the mitigation can't
land and degrade to blocking instead of silently letting the
flagged request through unmodified.
"""
advisory_message: Final = {"role": "system", "content": message} # mutable-ok: plain dict for live request
existing_messages: Final = data.get("messages")
existing_input: Final = data.get("input")
existing_instructions: Final = data.get("instructions")
if isinstance(existing_instructions, str):
# Responses API "instructions" is the privileged, developer-set
# system-level field the model treats as authoritative -- unlike
# "input", which the caller controls and could use to tell the
# model to disregard a trailing warning. Prefer it over "input"
# whenever present.
if isinstance(existing_messages, list):
messages_with_instructions_note: Final = [ # mutable-ok: fresh list
*existing_messages,
advisory_message,
]
data["messages"] = messages_with_instructions_note # rebind-ok: mutates caller's dict by design
data["instructions"] = f"{existing_instructions}\n\n{message}" # rebind-ok: mutates caller's dict by design
return True
if isinstance(existing_input, str):
# A plain-string "input" doesn't rule out "messages" also being a
# real, read field (e.g. a chat-completions call carrying a stray
# "input"), so write to both when both are present.
if isinstance(existing_messages, list):
messages_with_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
data["messages"] = messages_with_input_note # rebind-ok: mutates caller's dict by design
# The Responses API reads "input", not "messages" -- appending only to
# "messages" would leave the advisory unreachable for that endpoint.
data["input"] = f"{existing_input}\n\n{message}" # rebind-ok: mutates caller's dict by design
return True
if existing_input is not None:
# existing_input is a structured (non-string) Responses-API item
# list. That endpoint reads only "input", so appending to
# "messages" -- even if "messages" also happens to be present --
# would never reach the model. Leave data untouched and report
# non-delivery so the caller degrades to blocking.
return False
if isinstance(existing_messages, list):
messages_without_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
data["messages"] = messages_without_input_note # rebind-ok: mutates caller's dict by design
return True
sole_message: Final = [advisory_message] # mutable-ok: plain list for the live JSON request
data["messages"] = sole_message # rebind-ok: mutates caller's dict by design
return True
def raise_sensitive_data_route_exception(
self,
route_to_model: str,

View file

@ -209,6 +209,8 @@ class DotpromptManager(CustomPromptManagement):
prompt_spec=prompt_spec,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
async def async_get_chat_completion_prompt(

View file

@ -149,7 +149,6 @@ class PromptManager:
)
self.prompts[template_id] = template
except Exception:
# Optional: print(f"Error loading prompt from JSON: {template_id}")
pass
def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate:

View file

@ -9,9 +9,12 @@ Flow:
from __future__ import annotations
import gzip
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, Protocol
from urllib.parse import urlparse
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -28,6 +31,34 @@ _MAVVRIK_ALLOWED_SUFFIXES: Final = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app
_GCS_CHUNK_SIZE: Final = 8 * 1024 * 1024 # 8 MB
class MavvrikRegisterBody(TypedDict):
metricsMarker: ReadOnly[NotRequired[int | str]]
class MavvrikUploadUrlBody(TypedDict):
url: ReadOnly[NotRequired[str]]
class _RegisterResponse(Protocol):
def json(self) -> MavvrikRegisterBody: ...
class _UploadUrlResponse(Protocol):
def json(self) -> MavvrikUploadUrlBody: ...
def _register_body(response: _RegisterResponse) -> MavvrikRegisterBody:
return response.json()
def _upload_url_body(response: _UploadUrlResponse) -> MavvrikUploadUrlBody:
return response.json()
def _header_value(headers: Mapping[str, str], name: str) -> str | None:
return headers.get(name)
def _validate_api_endpoint(api_endpoint: str) -> None:
if not api_endpoint.startswith("https://"):
raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL")
@ -56,12 +87,12 @@ class FocusMavvrikDestination(FocusDestination):
self,
*,
prefix: str,
config: dict[str, Any] | None = None,
config: Mapping[str, str] | None = None,
) -> None:
config = config or {}
api_key: Final = config.get("api_key")
api_endpoint: Final = config.get("api_endpoint")
connection_id: Final = config.get("connection_id")
resolved_config: Final[Mapping[str, str]] = config or {}
api_key: Final = resolved_config.get("api_key")
api_endpoint: Final = resolved_config.get("api_endpoint")
connection_id: Final = resolved_config.get("connection_id")
if not api_key:
raise ValueError(
@ -100,7 +131,7 @@ class FocusMavvrikDestination(FocusDestination):
def _auth_headers(self) -> dict[str, str]:
return {"Content-Type": "application/json", "x-api-key": self.api_key}
async def _ensure_registered(self) -> int | None:
async def _ensure_registered(self) -> int | str | None:
"""POST agent endpoint to register/initialize the connector (once per instance).
Returns metricsMarker from the Mavvrik response the last date index
@ -127,7 +158,7 @@ class FocusMavvrikDestination(FocusDestination):
if resp.status_code >= 400:
raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}")
self._registered = True
metrics_marker: Final = resp.json().get("metricsMarker", 0)
metrics_marker: Final = _register_body(resp).get("metricsMarker", 0)
verbose_logger.debug(
"Mavvrik FOCUS destination: connector registered (metricsMarker=%s)",
metrics_marker,
@ -148,7 +179,7 @@ class FocusMavvrikDestination(FocusDestination):
raise RuntimeError(
f"Mavvrik FOCUS destination: failed to get signed URL ({resp.status_code}): {resp.text[:200]}"
)
signed_url: Final = resp.json().get("url")
signed_url: Final = _upload_url_body(resp).get("url")
if not signed_url:
raise RuntimeError(f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}")
_validate_gcs_url(signed_url, "signed URL")
@ -190,7 +221,7 @@ class FocusMavvrikDestination(FocusDestination):
f"Mavvrik FOCUS destination: GCS session init failed ({init_resp.status_code}): {init_resp.text[:400]}"
)
session_uri: Final = init_resp.headers.get("Location")
session_uri: Final = _header_value(init_resp.headers, "Location")
if not session_uri:
raise RuntimeError("Mavvrik FOCUS destination: GCS session init missing Location header")
_validate_gcs_url(session_uri, "session URI")
@ -264,7 +295,7 @@ class FocusMavvrikDestination(FocusDestination):
)
verbose_logger.debug("Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch)
async def get_metrics_marker(self) -> int | None:
async def get_metrics_marker(self) -> int | str | None:
"""Register with Mavvrik and return the current metricsMarker.
Always calls the Mavvrik register API unlike deliver() which skips
@ -287,7 +318,7 @@ class FocusMavvrikDestination(FocusDestination):
if resp.status_code >= 400:
raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}")
self._registered = True
metrics_marker: Final = resp.json().get("metricsMarker", 0)
metrics_marker: Final = _register_body(resp).get("metricsMarker", 0)
verbose_logger.debug("Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker)
return metrics_marker

View file

@ -416,17 +416,8 @@ class GenericPromptManager(CustomPromptManagement):
tools=tools,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=(
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
if prompt_spec
else False
),
ignore_prompt_manager_optional_params=(
ignore_prompt_manager_optional_params
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
if prompt_spec
else False
),
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
def get_chat_completion_prompt(
@ -457,17 +448,8 @@ class GenericPromptManager(CustomPromptManagement):
prompt_spec=prompt_spec,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=(
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
if prompt_spec
else False
),
ignore_prompt_manager_optional_params=(
ignore_prompt_manager_optional_params
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
if prompt_spec
else False
),
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
def clear_cache(self) -> None:

View file

@ -5,6 +5,7 @@ import os
import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
@ -137,6 +138,16 @@ def resolve_langfuse_credentials(
return public_key, secret_key, resolved_host
@lru_cache(maxsize=8)
def _warn_invalid_deployment_environment(raw_value: str, error: str) -> None:
verbose_logger.warning(
"Ignoring invalid LANGFUSE_TRACING_ENVIRONMENT=%r for the langfuse callback: %s. "
"Traces will be sent to Langfuse's default environment.",
raw_value,
error,
)
class LangFuseLogger:
# Class variables or attributes
def __init__(
@ -165,9 +176,11 @@ class LangFuseLogger:
# add http:// if unset, assume communicating over private network - e.g. render
self.langfuse_host = "http://" + self.langfuse_host
_env_override: Final = str(langfuse_environment).strip() if langfuse_environment is not None else None
self.langfuse_environment = _env_override or os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
if self.langfuse_environment:
validate_langfuse_environment_value(self.langfuse_environment)
if _env_override:
validate_langfuse_environment_value(_env_override)
self.langfuse_environment: str | None = _env_override
else:
self.langfuse_environment = self.resolve_deployment_environment()
self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")
self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(flush_interval)
@ -953,6 +966,20 @@ class LangFuseLogger:
verbose_logger.warning("Failed to apply masking function: %s. Returning original data.", e)
return data
@staticmethod
def resolve_deployment_environment() -> str | None:
"""Resolve LANGFUSE_TRACING_ENVIRONMENT: stripped value, "default" plus a warning when invalid, None when unset."""
raw: Final = os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
if not raw:
return None
value: Final = raw.strip()
try:
validate_langfuse_environment_value(value)
except ValueError as e:
_warn_invalid_deployment_environment(raw, str(e))
return "default"
return value
@staticmethod
def _get_langfuse_flush_interval(flush_interval: int) -> int:
"""

View file

@ -1,5 +1,3 @@
import os
"""
This file contains the LangFuseHandler class
@ -8,6 +6,7 @@ Used to get the LangFuseLogger for a given request
Handles Key/Team Based Langfuse Logging
"""
import os
from typing import TYPE_CHECKING, Any, Final
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
@ -157,7 +156,11 @@ class LangFuseHandler:
if raw is None:
return None
value = str(raw).strip()
if not value or value == os.getenv("LANGFUSE_TRACING_ENVIRONMENT"):
if (
not value
or value == os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
or value == LangFuseLogger.resolve_deployment_environment()
):
return None
return value

View file

@ -2,6 +2,7 @@
Call Hook for LiteLLM Proxy which allows Langfuse prompt management.
"""
import inspect
import os
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
@ -109,6 +110,9 @@ def langfuse_client_init(
cert=os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate),
)
if "environment" in inspect.signature(Langfuse.__init__).parameters:
parameters["environment"] = LangFuseLogger.resolve_deployment_environment()
client: Final = Langfuse(**parameters)
return client

View file

@ -0,0 +1,395 @@
"""
New Relic Metric API Integration - sends per-team cost/usage metrics to /metric/v1
NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-api/introduction-metric-api/
`async_log_success_event` / `async_log_failure_event` queue one record per request;
at flush the queue is aggregated by (team, model group, model, provider, status)
into count/summary metrics. `interval.ms` is the real window between flushes,
computed at flush time.
Team-scoped by construction: the ingest key is injected explicitly and there is
deliberately no environment-variable fallback, so a team's metrics are never sent
with the proxy operator's credentials (mirrors ``allow_env_credentials=False`` on
the Datadog team logger).
Error policy on flush: 4xx drops the batch (a retry would fail identically; 403
is a permanent credential failure), 5xx/network re-queues capped at
``max_queue_size`` records with the oldest dropped.
For batching specific details see CustomBatchLogger class
"""
import asyncio
import gzip
import time
import traceback
from collections.abc import Mapping
from math import ceil
from types import MappingProxyType
from typing import Final
from httpx import HTTPStatusError, Response
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.newrelic import (
NEWRELIC_DEFAULT_REGION,
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN,
NEWRELIC_METRIC_COMPLETION_TOKENS,
NEWRELIC_METRIC_COST_USD,
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TOTAL_TOKENS,
NEWRELIC_METRICS_MAX_BATCH_SIZE,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
NewRelicCountMetric,
NewRelicMetric,
NewRelicMetricCommon,
NewRelicMetricEnvelope,
NewRelicMetricRecord,
NewRelicSummaryMetric,
NewRelicSummaryValue,
)
from litellm.types.utils import StandardLoggingPayload
# 408 (request timeout) and 429 (rate limit) are transient client errors the
# Metric API expects a retry on, unlike 400/403 which a retry would only repeat.
_RETRYABLE_CLIENT_STATUSES: Final = frozenset({408, 429})
def resolve_newrelic_metric_endpoint(newrelic_region: str | None) -> str:
if not newrelic_region:
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
endpoint: Final = NEWRELIC_METRIC_ENDPOINT_BY_REGION.get(newrelic_region.lower())
if endpoint is None:
verbose_logger.warning(
"New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.",
newrelic_region,
", ".join(sorted(NEWRELIC_METRIC_ENDPOINT_BY_REGION)),
)
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
return endpoint
def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) -> NewRelicMetricRecord:
metadata: Final = standard_logging_object.get("metadata")
team_id: Final = ((metadata.get("user_api_key_team_id") or metadata.get("team_id")) if metadata else None) or ""
team_alias: Final = (
(metadata.get("user_api_key_team_alias") or metadata.get("team_alias")) if metadata else None
) or ""
return NewRelicMetricRecord(
team_id=team_id,
team_alias=team_alias,
model_group=standard_logging_object.get("model_group") or "",
model=standard_logging_object.get("model") or "",
custom_llm_provider=standard_logging_object.get("custom_llm_provider") or "",
status=str(standard_logging_object.get("status") or "success"),
response_cost=float(standard_logging_object.get("response_cost") or 0.0),
prompt_tokens=int(standard_logging_object.get("prompt_tokens") or 0),
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
)
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
first: Final = bucket_records[0]
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
for key, value in (
("team_id", first.team_id),
("team_alias", first.team_alias),
("model_group", first.model_group),
("model", first.model),
("custom_llm_provider", first.custom_llm_provider),
("status", first.status),
)
if value
}
durations: Final = tuple(record.duration_ms for record in bucket_records)
counts: Final[tuple[tuple[str, float], ...]] = (
(NEWRELIC_METRIC_REQUESTS, float(len(bucket_records))),
(NEWRELIC_METRIC_COST_USD, sum(record.response_cost for record in bucket_records)),
(NEWRELIC_METRIC_PROMPT_TOKENS, float(sum(record.prompt_tokens for record in bucket_records))),
(NEWRELIC_METRIC_COMPLETION_TOKENS, float(sum(record.completion_tokens for record in bucket_records))),
(NEWRELIC_METRIC_TOTAL_TOKENS, float(sum(record.total_tokens for record in bucket_records))),
)
count_metrics: Final[tuple[NewRelicMetric, ...]] = tuple(
NewRelicCountMetric(name=name, type="count", value=value, attributes=attributes) for name, value in counts
)
summary_metric: Final = NewRelicSummaryMetric(
name=NEWRELIC_METRIC_REQUEST_DURATION_MS,
type="summary",
value=NewRelicSummaryValue(
count=len(durations),
sum=sum(durations),
min=min(durations),
max=max(durations),
),
attributes=attributes,
)
return (*count_metrics, summary_metric)
def build_metric_payload(
records: tuple[NewRelicMetricRecord, ...],
*,
window_start: float,
now: float,
) -> tuple[NewRelicMetricEnvelope, ...]:
"""Aggregates records into one Metric API envelope for the flush window."""
interval_ms: Final = max(1, int((now - window_start) * 1000))
bucket_keys: Final = tuple(dict.fromkeys(record.bucket_key for record in records))
metrics: Final = tuple(
metric
for key in bucket_keys
for metric in _bucket_metrics(tuple(record for record in records if record.bucket_key == key))
)
common: Final[NewRelicMetricCommon] = {
"timestamp": int(window_start * 1000),
"interval.ms": interval_ms,
}
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
class NewRelicMetricsLogger(CustomBatchLogger):
def __init__(
self,
newrelic_api_key: str,
newrelic_region: str | None = None,
) -> None:
if not newrelic_api_key:
raise ValueError(
"newrelic_api_key is required for NewRelicMetricsLogger; "
"team-scoped metrics never fall back to environment credentials"
)
self.newrelic_api_key: Final = newrelic_api_key
self.metric_api_url: Final = resolve_newrelic_metric_endpoint(newrelic_region)
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self._stopped: bool = False
self._drain_lock = asyncio.Lock()
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
super().__init__(
flush_lock=self.flush_lock,
batch_size=NEWRELIC_METRICS_MAX_BATCH_SIZE,
max_queue_size=NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
)
def stop(self) -> None:
"""Ends the periodic flush loop; called on DynamicLoggingCache eviction.
Schedules one final drain of anything still queued, so eviction never
silently discards records. Guarded so it can never raise into the
cache's eviction path.
"""
self._stopped = True
try:
asyncio.get_running_loop().create_task(self._final_drain())
except Exception: # noqa: BLE001 # no running loop / shutdown; the periodic loop's final drain still runs
verbose_logger.debug("New Relic Metrics: could not schedule final drain on stop()", exc_info=True)
async def _drain_with_retry(self) -> None:
"""Deliver everything queued on a stopped logger, or drop it with a log.
A stopped logger has no periodic loop left, so every post-stop path
funnels through here. ``_drain_lock`` serializes drains: a callback that
appends and starts its own drain queues behind the running one instead
of racing it. Each pass attempts the whole current queue in
``batch_size`` chunks, unlike the periodic path it does not stop at the
first failing chunk, so a persistently failing head never starves the
tail. Only after ``_MAX_DRAIN_PASSES`` against a permanently failing
destination is the remainder dropped, and then only the records that were
queued when this drain began, so every dropped record got the full retry
budget: a record a callback appended mid-drain is not in that snapshot,
so it is left for its own serialized drain rather than dropped after
fewer attempts, and is never stranded.
"""
async with self._drain_lock:
attempted: Final = tuple(self.log_queue)
for _pass in range(NEWRELIC_METRICS_MAX_DRAIN_PASSES):
await self._drain_flush_once()
if not self.log_queue:
return
if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1:
await asyncio.sleep(2**_pass)
async with self.flush_lock:
tried_ids: Final = frozenset(id(record) for record in attempted)
survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids)
dropped: Final = len(self.log_queue) - len(survivors)
if dropped:
verbose_logger.warning(
"New Relic Metrics: dropping %s records after %s drain passes",
dropped,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
)
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
async def _drain_flush_once(self) -> None:
"""Attempt every queued record once, in ``batch_size`` chunks, without
stopping at the first failing chunk so a persistently failing head does
not starve the tail (the periodic ``flush_queue`` deliberately stops
instead). Takes the queue under ``flush_lock`` and re-queues only the
chunks a 5xx/network error left undelivered, so records a concurrent
request appends during the sends survive for the next pass."""
async with self.flush_lock:
pending: Final = tuple(self.log_queue)
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
del self.log_queue[:]
if not pending:
return
chunks: Final = tuple(
pending[start : start + self.batch_size] for start in range(0, len(pending), self.batch_size)
)
delivered: Final = tuple([await self._classify_and_send(chunk, window_start) for chunk in chunks])
failed: Final = tuple(record for chunk, ok in zip(chunks, delivered) for record in (() if ok else chunk))
if failed:
self._requeue(failed)
async def _final_drain(self) -> None:
await self._drain_with_retry()
async def periodic_flush(self) -> None:
while not self._stopped:
await asyncio.sleep(self.flush_interval)
if self._stopped:
break
await self.flush_queue()
await self._final_drain()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def _log_async_event(self, standard_logging_object: StandardLoggingPayload | None) -> None:
if standard_logging_object is None:
raise ValueError("standard_logging_object not found in kwargs")
self.log_queue.append(_metric_record_from_payload(standard_logging_object))
if self._stopped:
# A stopped logger has no periodic loop left; an in-flight callback
# that appends after the eviction drain delivers its own record.
await self._drain_with_retry()
return
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def flush_queue(self) -> None:
async with self.flush_lock:
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
queued: Final = len(self.log_queue)
if not queued:
return
verbose_logger.debug("New Relic Metrics: Flushing %s queued records", queued)
# Bounded by what is queued now: records appended mid-flush belong to
# the next window, and looping until empty would never end under load.
for _chunk in range(ceil(queued / self.batch_size)):
if not await self.async_send_batch(window_start=window_start):
return
async def async_send_batch(self, window_start: float | None = None) -> bool:
"""Sends the oldest ``batch_size`` records only, so a queue grown past that
by re-queues cannot breach the Metric API data point cap in one request.
Returns False once a chunk fails and is re-queued, so the caller stops."""
if not self.log_queue:
return False
batch_to_send: Final[tuple[NewRelicMetricRecord, ...]] = tuple(self.log_queue[: self.batch_size])
del self.log_queue[: len(batch_to_send)]
delivered: Final = await self._classify_and_send(
batch_to_send, window_start if window_start is not None else self.last_flush_time
)
if not delivered:
self._requeue(batch_to_send)
return delivered
async def _classify_and_send(self, batch: tuple[NewRelicMetricRecord, ...], window_start: float) -> bool:
"""Send one chunk and classify the outcome, never touching the queue.
Returns True when the batch is done with (delivered on any 2xx, or a 4xx
a retry would only repeat, 403 being a permanent bad-key rejection), and
False when a 5xx or network error means the caller should re-queue it.
``AsyncHTTPHandler.post`` raises ``HTTPStatusError`` on any non-2xx, so a
4xx never returns a response here; the status is read off the raised
error to keep the client-error path (drop) distinct from 5xx (retry)."""
payload: Final = build_metric_payload(records=batch, window_start=window_start, now=time.time())
try:
status = (
await self.async_send_compressed_data(payload)
).status_code # rebind-ok: reassigned from the raised HTTPStatusError below
except HTTPStatusError as e:
status = e.response.status_code
except Exception as e: # noqa: BLE001 # transport/network failure re-queues the batch
verbose_logger.warning(
"New Relic Metrics: network error sending %s records, will retry - %s",
len(batch),
e,
)
return False
if 200 <= status < 300:
return True
if 400 <= status < 500 and status not in _RETRYABLE_CLIENT_STATUSES:
verbose_logger.warning(
"New Relic Metrics: %s from Metric API%s, dropping %s records.",
status,
" (permanent credential failure: invalid or revoked team ingest key)" if status == 403 else "",
len(batch),
)
return True
verbose_logger.warning(
"New Relic Metrics: %s from Metric API, will retry %s records",
status,
len(batch),
)
return False
def _requeue(self, batch: tuple[NewRelicMetricRecord, ...]) -> None:
"""Prepends ``batch`` in place (never by assignment: records appended by
concurrent requests during the flush await must survive), keeping
chronological order so the cap drops the oldest records first."""
self.log_queue[:0] = batch
overflow: Final = len(self.log_queue) - self.max_queue_size
if overflow > 0:
del self.log_queue[:overflow]
verbose_logger.warning(
"New Relic Metrics: retry queue exceeded max_queue_size=%s; dropped %s oldest records.",
self.max_queue_size,
overflow,
)
async def async_send_compressed_data(self, payload: tuple[NewRelicMetricEnvelope, ...]) -> Response:
compressed_data: Final = gzip.compress(safe_dumps(payload).encode("utf-8"))
headers: Final[Mapping[str, str]] = MappingProxyType(
{
"Content-Type": "application/json",
"Content-Encoding": "gzip",
"Api-Key": self.newrelic_api_key,
}
)
return await self.async_client.post(
url=self.metric_api_url,
data=compressed_data,
headers=headers,
)

View file

@ -0,0 +1,90 @@
"""
New Relic Team Handler
Used to get the NewRelicMetricsLogger for a given request.
Handles Key/Team Based New Relic metrics, following the same pattern as DataDogHandler.
"""
from typing import TYPE_CHECKING, Final
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
from .newrelic_metrics import NewRelicMetricsLogger
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
class NewRelicLoggingConfig(TypedDict):
newrelic_api_key: ReadOnly[str | None]
newrelic_region: ReadOnly[str | None]
class NewRelicHandler:
@staticmethod
def get_newrelic_logger_for_request(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
"""
Get a team-scoped NewRelicMetricsLogger for a given request.
Resolves and caches per-team NewRelicMetricsLogger instances using
DynamicLoggingCache, keyed by the team's New Relic credentials. Each unique
set of credentials gets its own logger instance with its own batch/flush loop.
Note: This handler is only called when a team-scoped newrelic_api_key is
present. The trace logger for the ``newrelic`` callback (OTel v2 / legacy
agent) is managed separately by _init_custom_logger_compatible_class via
_in_memory_loggers.
"""
_credentials: Final = NewRelicHandler.get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
temp_newrelic_logger = in_memory_dynamic_logger_cache.get_cache(
credentials=_credentials, service_name="newrelic"
)
if temp_newrelic_logger is None:
temp_newrelic_logger = NewRelicHandler._create_newrelic_logger_from_credentials(
credentials=_credentials,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
return temp_newrelic_logger
@staticmethod
def _create_newrelic_logger_from_credentials(
credentials: NewRelicLoggingConfig,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
newrelic_logger: Final = NewRelicMetricsLogger(
newrelic_api_key=credentials.get("newrelic_api_key") or "",
newrelic_region=credentials.get("newrelic_region"),
)
in_memory_dynamic_logger_cache.set_cache(
credentials=credentials,
service_name="newrelic",
logging_obj=newrelic_logger,
)
verbose_logger.debug("New Relic: Created and cached new NewRelicMetricsLogger for team-scoped credentials")
return newrelic_logger
@staticmethod
def get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> NewRelicLoggingConfig:
return NewRelicLoggingConfig(
newrelic_api_key=standard_callback_dynamic_params.get("newrelic_api_key"),
newrelic_region=standard_callback_dynamic_params.get("newrelic_region"),
)
@staticmethod
def _dynamic_newrelic_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
return standard_callback_dynamic_params.get("newrelic_api_key") is not None

View file

@ -22,6 +22,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
)
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
from litellm.integrations.otel.model.semconv import Metric
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.service_tier_utils import (
@ -1643,7 +1644,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
if self._operation_duration_histogram:
self._operation_duration_histogram.record(duration_s, attributes=common_attrs)
if response_obj and (usage := response_obj.get("usage")) and self._token_usage_histogram:
if (
self._token_usage_histogram
and response_obj
and not is_unbilled_non_inference_call_from_params(kwargs.get("call_type"), params, response_obj)
and (usage := response_obj.get("usage"))
):
in_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"}
out_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"}
self._token_usage_histogram.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
@ -1719,6 +1725,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
if not self._time_per_output_token_histogram:
return
if is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj
):
return
# Get completion tokens from response_obj
completion_tokens = None
if response_obj and (usage := response_obj.get("usage")):
@ -2049,6 +2060,26 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# serialise to JSON once so set_attribute never coerces.
guardrail_span.set_attribute("guardrail_violation_categories", safe_dumps(violation_categories))
# Billable usage counters and USD cost stamped by the provider hook
# (e.g. Azure Prompt Shield text records, Bedrock policy units).
guardrail_usage = guardrail_information.get("guardrail_usage")
if guardrail_usage is not None:
guardrail_span.set_attribute("guardrail_usage", safe_dumps(guardrail_usage))
guardrail_cost = guardrail_information.get("guardrail_cost")
if guardrail_cost is not None:
self.safe_set_attribute(
span=guardrail_span,
key="guardrail_cost",
value=guardrail_cost,
)
guardrail_cost_in_spend = guardrail_information.get("guardrail_cost_in_spend")
if isinstance(guardrail_cost_in_spend, bool):
self.safe_set_attribute(
span=guardrail_span,
key="guardrail_cost_in_spend",
value=guardrail_cost_in_spend,
)
self._set_team_attributes_from_kwargs(guardrail_span, kwargs)
guardrail_span.end(end_time=self._to_ns(end_time_datetime))
@ -2468,7 +2499,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload)
usage: Final = response_obj and response_obj.get("usage")
usage: Final = (
response_obj.get("usage")
if response_obj
and not is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), litellm_params, response_obj
)
else None
)
if usage:
self.safe_set_attribute(
span=span,

View file

@ -146,7 +146,7 @@ class SpanEmitter:
For callers that own and manage their own span lifecycle. ``tracer``
overrides the bound tracer for this span only, used for per-request
multi-tenant credential routing. ``links`` records related-but-not-parent
spans (e.g. the transport span of an MCP message, per MCP semconv).
spans (e.g. the trace context an MCP client propagated in ``params._meta``).
"""
return (tracer or self._tracer).start_span(
name,
@ -196,8 +196,8 @@ class SpanEmitter:
Return the span, or ``None`` if it was deduplicated away. ``tracer``
overrides the bound tracer for this span, used for per-request routing.
``links`` records related-but-not-parent spans (the transport span of an
MCP message).
``links`` records related-but-not-parent spans (e.g. the trace context an
MCP client propagated in ``params._meta``).
"""
# LLM-call and MCP tool-call spans carry a dedup key (their request's
# call id), so a sync+async double-firing coalesces. ``isinstance`` narrows

View file

@ -390,10 +390,10 @@ class OpenTelemetryV2(CustomLogger):
MCP tool calls reach the success/failure callbacks like any other request
(with ``call_type`` ``call_mcp_tool``), but they are not LLM calls and have
no ``pre_call`` carrier so they get their own CLIENT span here. Per the MCP
semconv it parents to the trace context the client propagated in
``params._meta`` (or starts a new root) and links the transport span, rather
than nesting under the HTTP/session span. Returns whether it handled the
no ``pre_call`` carrier so they get their own CLIENT span here. It nests
under the transport span of the request carrying this message, and trace
context the client propagated in ``params._meta`` is recorded as a span
link (see ``resolve_mcp_span_context``). Returns whether it handled the
event, so the caller skips the LLM-call path. The whole span is emitted at
once (there is no boundary to open it at), deduped on the call id.
"""
@ -436,9 +436,9 @@ class OpenTelemetryV2(CustomLogger):
Like a tool call, listing reaches the success/failure callbacks (here with
``call_type`` ``list_mcp_tools``) with no ``pre_call`` carrier, so it gets its
own CLIENT span. Per the MCP semconv it parents to the ``params._meta`` trace
context (or starts a new root) and links the transport span, rather than
nesting under the HTTP/session span. Returns whether it handled the event so
own CLIENT span, nested under the transport span of the request carrying
this message with any ``params._meta`` trace context recorded as a span
link (see ``resolve_mcp_span_context``). Returns whether it handled the event so
the caller skips the LLM-call path.
"""
raw_payload: Final = kwargs.get("standard_logging_object")

View file

@ -136,6 +136,9 @@ class GenAIMapper:
LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id,
LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template,
LiteLLM.GUARDRAIL_DETECTION_METHOD: lambda d: d.detection_method,
LiteLLM.GUARDRAIL_USAGE: lambda d: d.usage_json,
LiteLLM.GUARDRAIL_COST: lambda d: d.cost,
LiteLLM.GUARDRAIL_COST_IN_SPEND: lambda d: d.cost_in_spend,
}
_SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = {

View file

@ -190,6 +190,15 @@ class GuardrailSpanData:
guardrail_id: str | None = None
policy_template: str | None = None
detection_method: str | None = None
# Provider-reported billable usage counters (JSON-serialized) and the USD cost
# priced from them by the provider hook (``guardrail_usage`` /
# ``guardrail_cost`` on ``StandardLoggingGuardrailInformation``).
usage_json: str | None = None
cost: float | None = None
# Whether ``cost`` participates in the request's billed spend (absent means
# billed, the default; False means report-only). Mirrors
# ``guardrail_cost_in_spend`` so trace consumers can avoid double-counting.
cost_in_spend: bool | None = None
# Set when the guardrail intervened/blocked or failed, so the emitter marks
# the span ERROR — a blocking guardrail is an error outcome for that span.
error: SpanError | None = None
@ -209,6 +218,8 @@ class GuardrailSpanData:
get: Final = cast(Mapping[str, object], entry).get
status: Final = as_str(get("guardrail_status"))
response: Final = get("guardrail_response")
usage: Final = get("guardrail_usage")
in_spend: Final = get("guardrail_cost_in_spend")
error: Final = (
SpanError(error_type=status, message=as_str(get("guardrail_action")))
if status in cls._ERROR_STATUSES
@ -231,6 +242,9 @@ class GuardrailSpanData:
guardrail_id=as_str(get("guardrail_id")),
policy_template=as_str(get("policy_template")),
detection_method=as_str(get("detection_method")),
usage_json=_json_or_none(usage) if usage is not None else None,
cost=as_float(get("guardrail_cost")),
cost_in_spend=in_spend if isinstance(in_spend, bool) else None,
error=error,
)

View file

@ -32,6 +32,7 @@ class GenAIOperation(str, Enum):
EXECUTE_TOOL = "execute_tool" # MCP tool-call spans
LITELLM_VECTOR_STORE_MANAGEMENT = "litellm.vector_store_management"
LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management"
LITELLM_RESPONSES_MANAGEMENT = "litellm.responses_management"
LITELLM_MODERATION = "litellm.moderation"
@ -307,6 +308,15 @@ class LiteLLM:
GUARDRAIL_ID: Final = "litellm.guardrail.id"
GUARDRAIL_POLICY_TEMPLATE: Final = "litellm.guardrail.policy_template"
GUARDRAIL_DETECTION_METHOD: Final = "litellm.guardrail.detection_method"
# Provider-reported billable usage counters, JSON-serialized into one value.
GUARDRAIL_USAGE: Final = "litellm.guardrail.usage"
# Numeric USD cost of the guardrail invocation; lives under the litellm.cost.*
# namespace (COST_PREFIX) beside the LLM call's litellm.cost.total.
GUARDRAIL_COST: Final = "litellm.cost.guardrail"
# Whether litellm.cost.guardrail is already inside litellm.cost.total (True,
# the billed default) or reported alongside it (False) — without this a trace
# consumer cannot tell whether adding the two double-counts.
GUARDRAIL_COST_IN_SPEND: Final = "litellm.guardrail.cost_in_spend"
SERVICE_NAME: Final = "litellm.service.name"
SERVICE_CALL_TYPE: Final = "litellm.service.call_type"
PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms"
@ -374,6 +384,14 @@ _OPERATION_BY_CALL_TYPE: Final[dict[str, GenAIOperation]] = {
"aembedding": GenAIOperation.EMBEDDINGS,
"responses": GenAIOperation.CHAT,
"aresponses": GenAIOperation.CHAT,
"get_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"aget_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"delete_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"adelete_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"cancel_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"acancel_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"list_input_items": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"alist_input_items": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
"image_generation": GenAIOperation.GENERATE_CONTENT,
"aimage_generation": GenAIOperation.GENERATE_CONTENT,
"moderation": GenAIOperation.LITELLM_MODERATION,

View file

@ -10,6 +10,8 @@ Canonical hierarchy::
DB_CALL (CLIENT) # its key/user/team lookups nest here
GUARDRAIL (INTERNAL) # request-lifecycle hook, sibling of LLM_CALL
LLM_CALL (CLIENT)
MCP_TOOL_CALL (CLIENT) # nests under the POST carrying the message
MCP_LIST_TOOLS (CLIENT) # (client-propagated context is a span link)
DB_CALL (CLIENT) # e.g. the spend-log write
Guardrails parent to PROXY_REQUEST, not LLM_CALL: pre/during/post-call guardrail
@ -18,14 +20,14 @@ before the LLM call even starts), so a guardrail is a sibling of the LLM call,
not a child of it. The emitter parents every span to the ambient OTel context
(the active server span), which matches this.
MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) have two shapes, chosen at emit
time by :func:`resolve_mcp_span_context`. When the client propagates trace context
in ``params._meta`` MCP and the HTTP transport are independent contexts per the
OTel GenAI MCP semconv, so the span parents to that propagated context and records
the ``PROXY_REQUEST`` transport span as a span *link*, never a parent the shape
this registry's ``parent=None, links=PROXY_REQUEST`` entry encodes. When nothing is
propagated (the common case) the span nests under the transport span of the request
carrying that message, so the tool call stays in one trace.
MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) are parented at emit time by
:func:`resolve_mcp_span_context`: they nest under the ``PROXY_REQUEST`` transport
span of the request carrying that message, so the tool call stays in one trace.
Trace context the client propagated in ``params._meta`` (SEP-414) is recorded as
a span *link*, never the parent a remote parent would root the span in a trace
whose root never reaches the gateway's tracing backend. Links always target that
remote client context, never a registry role, so ``SpanSpec`` declares no link
field; the concrete transport parent is resolved per message at emit time.
Not every service call becomes a span :func:`span_role_for_service` decides:
@ -85,25 +87,19 @@ class SpanSpec:
role: SpanRole
kind: LiteLLMSpanKind
parent: SpanRole | None
links: SpanRole | None = None
SPAN_REGISTRY: Final[dict[SpanRole, SpanSpec]] = {
SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None),
SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST),
# The proxy is an MCP client to the upstream server, so MCP spans are CLIENT
# spans. With trace context propagated in ``params._meta``, MCP and the HTTP
# transport are independent contexts (OTel GenAI MCP semconv): the span parents
# to the propagated context and records the PROXY_REQUEST transport span as a
# span *link*, never a parent — the shape ``parent=None, links=PROXY_REQUEST``
# encodes. With nothing propagated, ``resolve_mcp_span_context`` nests the span
# under that message's transport span instead, keeping the call in one trace.
SpanRole.MCP_TOOL_CALL: SpanSpec(
SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST
),
SpanRole.MCP_LIST_TOOLS: SpanSpec(
SpanRole.MCP_LIST_TOOLS, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST
),
# spans. ``resolve_mcp_span_context`` nests them under the PROXY_REQUEST
# transport span of the request carrying that message (resolved per message at
# emit time), keeping the call in one trace. Trace context the client
# propagated in ``params._meta`` becomes a span *link* to that remote context,
# which is not a registry role, so ``SpanSpec`` has no link field.
SpanRole.MCP_TOOL_CALL: SpanSpec(SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST),
SpanRole.MCP_LIST_TOOLS: SpanSpec(SpanRole.MCP_LIST_TOOLS, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST),
SpanRole.GUARDRAIL: SpanSpec(SpanRole.GUARDRAIL, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST),
SpanRole.DB_CALL: SpanSpec(SpanRole.DB_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST),
SpanRole.SERVICE: SpanSpec(SpanRole.SERVICE, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST),
@ -209,8 +205,8 @@ def service_span_name(data: "ServiceSpanData") -> str:
def root_roles() -> list[SpanRole]:
"""Roles with no in-process parent. They start a new trace unless they adopt a
remote parent (e.g. an MCP span joining the client's propagated context)."""
"""Roles with no in-process parent, i.e. they start a new trace (only the
instrumentor-owned ``PROXY_REQUEST`` server span today)."""
return [role for role, spec in SPAN_REGISTRY.items() if spec.parent is None]
@ -227,8 +223,6 @@ def validate_registry(
raise ValueError(f"SPAN_REGISTRY[{role}] has mismatched role {spec.role}")
if spec.parent is not None and spec.parent not in reg:
raise ValueError(f"span role {role} declares unknown parent {spec.parent}")
if spec.links is not None and spec.links not in reg:
raise ValueError(f"span role {role} declares unknown link target {spec.links}")
missing: Final = [role for role in SpanRole if role not in reg]
if missing:
raise ValueError(f"SPAN_REGISTRY is missing roles: {missing}")

View file

@ -57,8 +57,8 @@ def request_root_span() -> "Span | None":
# The W3C trace-context carrier (``traceparent``/``tracestate``/``baggage``) the
# MCP client propagated in the current request's ``params._meta``. The MCP gateway
# sets it per message so the MCP span can parent to the client's span rather than
# to the transport. A ``ContextVar`` because, like the root-span anchor, it must
# sets it per message so the MCP span can record the client's span as a span
# link. A ``ContextVar`` because, like the root-span anchor, it must
# ride the request task and be readable by the inline success-logging callback.
_mcp_message_trace_carrier: Final["ContextVar[Mapping[str, str] | None]"] = ContextVar(
"litellm_otel_mcp_message_trace_carrier", default=None
@ -148,10 +148,10 @@ def _mcp_transport_span_context() -> "SpanContext | None":
Prefers the transport the gateway published for this specific message; falls
back to the ambient request anchor for paths that emit an MCP span on the
request task itself (the REST MCP endpoints, the SDK). Parenting and linking
only need the immutable context, and unlike ``mcp_message_transport_span`` they
stay correct against a transport that has already finished, so this does not
require the span to still be recording.
request task itself (the REST MCP endpoints). Parenting needs only the
immutable context, and unlike ``mcp_message_transport_span`` it stays correct
against a transport that has already finished, so this does not require the
span to still be recording.
"""
published: Final = _mcp_message_transport_span.get()
if published is not None:
@ -222,25 +222,31 @@ def resolve_mcp_span_context(
) -> "tuple[Context, tuple[Link, ...]]":
"""Parent context + links for an MCP message span.
The span always nests under the transport span of the request carrying this
message, so a tool call and the ``POST`` that carried it stay in one trace.
The transport comes from :func:`_mcp_transport_span_context`, which is the
*current message's* POST rather than whatever request happened to open the
session, so a long-lived session does not glue every message under its first
request.
When the client propagates W3C trace context in the request's ``params._meta``
(SEP-414), MCP and the underlying transport are independent lifecycles one
streamable-HTTP session multiplexes many messages, and the client's own span is
the truthful parent. So, per the OTel GenAI MCP semconv:
(SEP-414), that remote context is recorded as a span *link*, never the parent.
The OTel GenAI MCP semconv prefers the inverse (remote parent, transport link),
but the gateway's tracing backend only ever receives the gateway's half of such
a trace: parenting into the client's trace id roots the span in a trace whose
root span never reaches the backend, so the span is unreachable from the trace
view and the transport transaction shows a dangling link (observed with
clients that propagate synthetic trace ids). Anchoring to the gateway's own
request and linking the client's context keeps every trace renderable while
preserving the client-side correlation.
* parent to the trace context the client propagated (a *remote* parent), and
* record the transport span as a *link*, never the parent.
Almost no client implements SEP-414 yet, so in practice nothing is propagated.
Rooting the span there splits a single tool call into two disconnected traces
joined only by a link, which is how it surfaces in APM: the ``POST`` transaction
and the ``tools/call`` span share no trace. With no remote parent to honor,
parent to the transport span of the request carrying this message instead, so
the call stays in one trace; no link is added since the transport is now the
real parent. The transport comes from :func:`_mcp_transport_span_context`, which
is the *current message's* POST rather than whatever request happened to open
the session, so a long-lived session does not glue every message under its
first request. With neither a remote parent nor a transport the returned context
carries no span and the span legitimately starts its own root trace.
With no transport at all the span starts its own root trace, still carrying
the link the client context is only ever a link, so this event keeps one
shape everywhere. Both returned contexts are built on an explicitly empty
base, so ambient (stale session) state can never leak in, and the span
inherits the transport's sampling decision exactly like every other
request-level span a client's sampled flag neither forces nor suppresses
recording.
Only trace context (``traceparent``/``tracestate``) is extracted, never the
client's W3C Baggage: ``params._meta`` is caller-controlled, and the otel
@ -251,13 +257,12 @@ def resolve_mcp_span_context(
never fall through to the ambient (stale session) span.
"""
source: Final = carrier if carrier is not None else _mcp_message_trace_carrier.get()
parent: Final = _PROPAGATOR.extract(dict(source or {}), context=Context())
propagated: Final = get_current_span(_PROPAGATOR.extract(dict(source or {}), context=Context()))
links: Final = (Link(propagated.get_span_context()),) if is_recordable_span(propagated) else ()
transport: Final = _mcp_transport_span_context()
if is_recordable_span(get_current_span(parent)):
return parent, (Link(transport),) if transport is not None else ()
if transport is not None:
return context_from_span(NonRecordingSpan(transport)), ()
return parent, ()
if transport is None:
return Context(), links
return context_from_span(NonRecordingSpan(transport), context=Context()), links
def is_recordable_span(obj: object) -> bool:

View file

@ -32,6 +32,7 @@ from litellm.integrations.otel.model.semconv import (
resolve_provider,
)
from litellm.integrations.otel.model.utils import to_seconds
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -198,16 +199,21 @@ class GenAIMetricRecorder:
) -> None:
common_attrs: Final = self._filter_attributes(self._bounded_attributes(kwargs))
duration_s: Final = (end_time - start_time).total_seconds()
usage_is_replayed: Final = is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj
)
self._metrics.operation_duration.record(duration_s, attributes=common_attrs)
self._record_token_usage(response_obj, common_attrs)
if not usage_is_replayed:
self._record_token_usage(response_obj, common_attrs)
cost: Final = kwargs.get("response_cost")
if cost:
self._metrics.token_cost.record(cost, attributes=common_attrs)
self._record_time_to_first_token(kwargs, common_attrs)
self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs)
if not usage_is_replayed:
self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs)
self._record_response_duration(kwargs, end_time, common_attrs)
def record_failure(

View file

@ -2,12 +2,13 @@
When a request carries team/key vendor credentials in
``standard_callback_dynamic_params``, or the key/team config resolved at auth
names a destination project, its spans must export through a
``TracerProvider`` whose OTLP headers carry those credentials / that project.
``TenantTracerCache`` builds and caches one provider per distinct
(credentials, project) pair, and otherwise hands back the logger's default
tracer. This lets a single logger fan requests out to many tenants without
needing a logger per tenant.
names a destination project or a service name, its spans must export through a
``TracerProvider`` whose OTLP headers carry those credentials / that project,
or whose Resource carries that ``service.name``. ``TenantTracerCache`` builds
and caches one provider per distinct (credentials, project, service name)
tuple, and otherwise hands back the logger's default tracer. This lets a
single logger fan requests out to many tenants without needing a logger per
tenant.
"""
import threading
@ -22,6 +23,7 @@ from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import Tracer
from litellm._logging import verbose_logger
from litellm.constants import OTEL_SERVICE_NAME_METADATA_KEYS
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.plumbing.providers import (
build_tracer_provider,
@ -65,8 +67,30 @@ _MAX_RETIRED_PROVIDERS: Final = 64
_HeaderItems: TypeAlias = tuple[tuple[str, str], ...]
_RouteKey: TypeAlias = tuple[_HeaderItems, _HeaderItems, str | None, str | None]
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
#: Key/team config fields naming the Resource ``service.name``, highest
#: precedence first. Read only from ``user_api_key_auth_metadata`` (the config
#: the proxy resolved at auth), never from client-supplied request metadata:
#: the service name picks the dataset/service traces land in (Honeycomb routes
#: datasets by it), so a caller must not be able to choose one.
_SERVICE_NAME_KEYS: Final = OTEL_SERVICE_NAME_METADATA_KEYS
def tenant_service_name(auth_metadata: Mapping[str, str] | None) -> str | None:
"""The per-request ``service.name`` override for this key/team, if any.
``None`` keeps the env-configured default (``OTEL_SERVICE_NAME``).
"""
if not auth_metadata:
return None
return next(
(stripped for key in _SERVICE_NAME_KEYS if (stripped := (auth_metadata.get(key) or "").strip())),
None,
)
def _shutdown_provider(provider: TracerProvider) -> None:
"""Flush + stop an evicted provider's processors (reclaims their threads).
@ -116,7 +140,7 @@ class TenantRoute:
class TenantTracerCache:
"""Credential/project-scoped ``TracerProvider`` cache keyed by the routing headers."""
"""Tenant-scoped ``TracerProvider`` cache keyed by routing headers and service name."""
def __init__(
self,
@ -131,7 +155,7 @@ class TenantTracerCache:
# thread-pool workers concurrently with the event loop, so cache
# updates, span counts, and retirement must be atomic.
self._lock: Final = threading.Lock()
self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems, str | None], TracerProvider] = (
self._providers: OrderedDict[_RouteKey, TracerProvider] = (
OrderedDict() # mutable-ok: bounded LRU; eviction needs in-place ordered mutation
)
self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state
@ -172,10 +196,11 @@ class TenantTracerCache:
) -> TenantRoute:
"""Return the tracer (and trace-detachment flag) for this request.
Use ``default`` unless the request's dynamic credentials or its key/team
project require a scoped tracer, in which case build (or reuse) one. The
cache is a bounded LRU: the least-recently-used provider is flushed and
shut down on overflow so its exporter threads don't accumulate.
Use ``default`` unless the request's dynamic credentials, its key/team
project, or its key/team service name require a scoped tracer, in
which case build (or reuse) one. The cache is a bounded LRU: the
least-recently-used provider is flushed and shut down on overflow so
its exporter threads don't accumulate.
A routed provider is returned already held its open-span count is
incremented in the same critical section as the cache update so a
@ -184,7 +209,8 @@ class TenantTracerCache:
"""
credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
project_headers: Final = self._project_headers(auth_metadata)
if not credential_headers and not project_headers:
service_name: Final = tenant_service_name(auth_metadata)
if not credential_headers and not project_headers and service_name is None:
return TenantRoute(tracer=default, detached=False)
# A fixed per-integration region endpoint (New Relic us/eu), never a
# caller-supplied host; ``None`` keeps the preset's own endpoint.
@ -193,9 +219,12 @@ class TenantTracerCache:
tuple(sorted(credential_headers.items())),
tuple(sorted(project_headers.items())),
endpoint,
service_name,
)
with self._lock:
provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers, endpoint)
provider: Final = self._cached_provider_locked(
cache_key, credential_headers, project_headers, endpoint, service_name
)
self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1
evicted: Final = self._evicted_on_overflow_locked()
if evicted is not None:
@ -208,16 +237,19 @@ class TenantTracerCache:
def _cached_provider_locked(
self,
cache_key: tuple[_HeaderItems, _HeaderItems, str | None],
cache_key: _RouteKey,
credential_headers: Mapping[str, str],
project_headers: Mapping[str, str],
endpoint: str | None,
service_name: str | None,
) -> TracerProvider:
cached: Final = self._providers.get(cache_key)
if cached is not None:
self._providers.move_to_end(cache_key)
return cached
built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers, endpoint))
built: Final = build_tracer_provider(
self._routed_config(credential_headers, project_headers, endpoint, service_name)
)
self._providers[cache_key] = built
return built
@ -267,6 +299,7 @@ class TenantTracerCache:
credential_headers: Mapping[str, str],
project_headers: Mapping[str, str],
endpoint: str | None = None,
service_name: str | None = None,
) -> OpenTelemetryV2Config:
"""Clone the config, rewriting headers on the callback's own exporter.
@ -285,7 +318,10 @@ class TenantTracerCache:
self._routed_exporter(spec, credential_headers, project_headers, endpoint)
for spec in self._config.exporters
]
return self._config.model_copy(update={"exporters": exporters})
update: Final = (
{"exporters": exporters} if service_name is None else {"exporters": exporters, "service_name": service_name}
)
return self._config.model_copy(update=update)
def _routed_exporter(
self,

View file

@ -7,6 +7,9 @@ import time
from datetime import datetime, timedelta
from typing import Final
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm import get_secret
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
@ -18,10 +21,32 @@ PROMETHEUS_URL: Final[str | None] = get_secret("PROMETHEUS_URL")
PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTED_INSTANCE")
async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
_RAW_JSON_PAYLOAD: Final = TypeAdapter(object)
class PrometheusRangeSample(BaseModel):
"""One ``matrix`` series of the Prometheus HTTP query API."""
metric: dict[str, object]
values: list[tuple[float, str]]
class PrometheusQueryData(BaseModel):
result: list[PrometheusRangeSample]
class PrometheusQueryResponse(BaseModel):
data: PrometheusQueryData
class PrometheusDailySpend(TypedDict):
date: ReadOnly[str]
spend: ReadOnly[float]
async def get_metric_from_prometheus(
metric_name: str,
):
) -> list[PrometheusRangeSample]:
# Get the start of the current day in Unix timestamp
if PROMETHEUS_URL is None:
raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env")
@ -31,13 +56,13 @@ async def get_metric_from_prometheus(
response: Final = await async_http_handler.get(
f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now}
) # End of the day
_json_response: Final = response.json()
_json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json())
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result
return results
async def get_fallback_metric_from_prometheus():
async def get_fallback_metric_from_prometheus() -> str:
"""
Gets fallback metrics from prometheus for the last 24 hours
"""
@ -55,17 +80,17 @@ async def get_fallback_metric_from_prometheus():
verbose_logger.debug("response json %s", response_json)
for result in response_json:
verbose_logger.debug("result= %s", result)
metric = result["metric"]
metric_values = result["values"]
metric_labels = result.metric
metric_values = result.values
most_recent_value = metric_values[0]
if PROMETHEUS_SELECTED_INSTANCE is not None:
if metric.get("instance") != PROMETHEUS_SELECTED_INSTANCE:
if metric_labels.get("instance") != PROMETHEUS_SELECTED_INSTANCE:
continue
value = int(float(most_recent_value[1])) # Convert value to integer
primary_model = metric.get("primary_model", "Unknown")
fallback_model = metric.get("fallback_model", "Unknown")
primary_model = metric_labels.get("primary_model", "Unknown")
fallback_model = metric_labels.get("fallback_model", "Unknown")
response_message += f"`{value} successful fallback requests` with primary model=`{primary_model}` -> fallback model=`{fallback_model}`"
response_message += "\n"
verbose_logger.debug("response message %s", response_message)
@ -96,7 +121,7 @@ def _quote_promql_string_literal(value: str) -> str:
return json.dumps(value, ensure_ascii=False)
async def get_daily_spend_from_prometheus(api_key: str | None):
async def get_daily_spend_from_prometheus(api_key: str | None) -> list[PrometheusDailySpend]:
"""
Expected Response Format:
[
@ -133,17 +158,16 @@ async def get_daily_spend_from_prometheus(api_key: str | None):
}
response: Final = await async_http_handler.get(url, params=params)
_json_response: Final = response.json()
_json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json())
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
formatted_results: Final = []
for result in results:
metric_data = result["values"]
for timestamp, value in metric_data:
# Convert timestamp to ISO 8601 string with UTC offset
date = datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00"
spend = float(value)
formatted_results.append({"date": date, "spend": spend})
results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result
formatted_results: Final[list[PrometheusDailySpend]] = [
{
"date": datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00",
"spend": float(value),
}
for result in results
for timestamp, value in result.values
]
return formatted_results

View file

@ -19,6 +19,19 @@ class PromptManagementClient(TypedDict):
completed_messages: list[AllMessageValues] | None
def resolve_prompt_manager_ignore_flags(
prompt_spec: PromptSpec | None,
ignore_prompt_manager_model: bool | None,
ignore_prompt_manager_optional_params: bool | None,
) -> tuple[bool, bool]:
spec_params: Final = prompt_spec.litellm_params if prompt_spec is not None else None
return (
bool(ignore_prompt_manager_model) or bool(spec_params is not None and spec_params.ignore_prompt_manager_model),
bool(ignore_prompt_manager_optional_params)
or bool(spec_params is not None and spec_params.ignore_prompt_manager_optional_params),
)
class PromptManagementBase(ABC):
@property
@abstractmethod
@ -182,13 +195,18 @@ class PromptManagementBase(ABC):
prompt_version=prompt_version,
)
resolved_ignore_model, resolved_ignore_optional_params = resolve_prompt_manager_ignore_flags(
prompt_spec=prompt_spec,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
return self.post_compile_prompt_processing(
prompt_template=prompt_template,
messages=messages,
non_default_params=non_default_params,
model=model,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
ignore_prompt_manager_model=resolved_ignore_model,
ignore_prompt_manager_optional_params=resolved_ignore_optional_params,
)
async def async_get_chat_completion_prompt(
@ -224,11 +242,16 @@ class PromptManagementBase(ABC):
prompt_version=prompt_version,
)
resolved_ignore_model, resolved_ignore_optional_params = resolve_prompt_manager_ignore_flags(
prompt_spec=prompt_spec,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
return self.post_compile_prompt_processing(
prompt_template=prompt_template,
messages=messages,
non_default_params=non_default_params,
model=model,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
ignore_prompt_manager_model=resolved_ignore_model,
ignore_prompt_manager_optional_params=resolved_ignore_optional_params,
)

View file

@ -17,6 +17,7 @@ import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@ -222,7 +223,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://"
return f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{encoded_key}"
return f"{self.s3_endpoint_url}/{self.s3_bucket_name}/{encoded_key}"
return f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{encoded_key}"
return (
f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}."
f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}"
)
def _sse_headers(self) -> Mapping[str, str]:
candidates: Final = {

View file

@ -38,6 +38,7 @@ from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalD
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.db.shadow_eval_funnel import ShadowEvalFunnelStage
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
from litellm.types.utils import StandardLoggingPayload
@ -386,6 +387,13 @@ def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> s
)
def _leg_eval_spend(sums: Mapping[str, object]) -> float:
return sum(
float(raw) if isinstance(raw := sums.get(column), (int, float)) else 0.0
for column in ("judge_cost", "shadow_cost", "shadow_classifier_cost")
)
def _job_spend_counter_key(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
@ -412,6 +420,15 @@ async def _add_job_spend_to_counter(counter_key: str, cost: float) -> None:
verbose_logger.warning("shadow_eval: spend counter increment failed for %s: %s", counter_key, e)
def _record_funnel_event(job_id: str, stage: "ShadowEvalFunnelStage") -> None:
try:
from litellm.proxy.db.shadow_eval_funnel import record_shadow_eval_funnel_event
record_shadow_eval_funnel_event(job_id, stage)
except Exception as e: # noqa: BLE001 # coverage stats are advisory; sampling must proceed
verbose_logger.debug("shadow_eval: funnel increment failed for %s: %s", job_id, e)
async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
"""Whether the shadowed key or its team is over budget, decided by the same owners
the request path uses, so counter keys and thresholds can never drift from auth's.
@ -452,6 +469,14 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
return False
def _forwarded_team_id(metadata: Mapping[str, object]) -> str | None:
"""The shadowed key's team, the identity the judge call already carries in its metadata
and the router already selects deployments with. Read here too so the arm choice, which
happens before the router sees the call, is made under the same team."""
team_id: Final = metadata.get("user_api_key_team_id")
return team_id if isinstance(team_id, str) and team_id else None
def _routing_decision(metadata: Mapping[str, object]) -> Mapping[str, object]:
"""The routing decision a pre-routing strategy wrote to a call's metadata, empty when
a plain model served it. Read off the sampled request for the control arm, and off the
@ -466,6 +491,13 @@ def _routed_tier(metadata: Mapping[str, object]) -> str | None:
return str(raw) if raw is not None else None
def _decision_classifier_cost(metadata: Mapping[str, object]) -> float:
"""What the arm's own routing decision says its classifier call billed: the money a
completion cost alone omits, and 0 for a plain model that never classifies."""
raw: Final = _routing_decision(metadata).get("classifier_cost")
return float(raw) if isinstance(raw, (int, float)) else 0.0
def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool:
"""Whether the router under evaluation served this request, which is what decides
the direction it belongs to. A forward job skips its own router's traffic, since
@ -481,6 +513,7 @@ class _CallFailure:
error: str
cost: float = 0.0
classifier_cost: float = 0.0
@dataclass(frozen=True, slots=True)
@ -491,6 +524,7 @@ class _ShadowResponse:
model: str
tier: str | None
cost: float
classifier_cost: float
@dataclass(frozen=True, slots=True)
@ -567,6 +601,7 @@ class ShadowEvalLogger(CustomLogger):
jobs_cache: InMemoryCache | None = None,
job_spend_reader: Callable[[str, float, float], Awaitable[float]] | None = None,
job_spend_writer: Callable[[str, float], Awaitable[None]] | None = None,
funnel_recorder: Callable[[str, "ShadowEvalFunnelStage"], None] | None = None,
) -> None:
"""Providers are callables so the proxy's lazily-initialized globals are resolved
at call time, not at logger construction. The spend reader and writer wrap the
@ -576,6 +611,7 @@ class ShadowEvalLogger(CustomLogger):
self._jobs_cache = jobs_cache or _jobs_cache
self._read_job_spend = job_spend_reader or _job_spend_from_counter
self._write_job_spend = job_spend_writer or _add_job_spend_to_counter
self._record_funnel = funnel_recorder or _record_funnel_event
self._inflight_shadow_tasks: int = 0
# Starts per job since the last cache fill, never decremented within a
# generation; the refill absorbs written rows and resets.
@ -602,7 +638,8 @@ class ShadowEvalLogger(CustomLogger):
await prisma.db.litellm_shadowevalattempt.group_by(
by=["job_id"],
count=True,
sum={"judge_cost": True, "shadow_cost": True}, # mutable-ok: Prisma aggregate spec
# mutable-ok: Prisma aggregate spec
sum={"judge_cost": True, "shadow_cost": True, "shadow_classifier_cost": True},
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
)
if records
@ -611,8 +648,7 @@ class ShadowEvalLogger(CustomLogger):
attempt_stats: Final = { # mutable-ok: frozen snapshot of the grouped read
str(row["job_id"]): (
int(row["_count"]["_all"]),
float((row["_sum"] or {}).get("judge_cost") or 0.0)
+ float((row["_sum"] or {}).get("shadow_cost") or 0.0),
_leg_eval_spend(row["_sum"] or _EMPTY_METADATA),
)
for row in grouped or []
}
@ -638,6 +674,32 @@ class ShadowEvalLogger(CustomLogger):
#### hook ####
def _sampled_jobs(
self,
active_jobs: Sequence[ActiveShadowEvalJob],
request_metadata: Mapping[str, object],
request_id: str,
) -> tuple[ActiveShadowEvalJob, ...]:
"""The jobs that sample this request. A key can hold one job per direction, and a
request routed by one job's router while bypassing the other's qualifies for both;
each is separately budgeted, so both fire. An admitting job that loses the sampling
dice is counted, so results can weigh judged rows against the traffic they stand for."""
eligible: list[ActiveShadowEvalJob] = [] # mutable-ok: bucketed per-job admission
now: Final = datetime.now(timezone.utc)
for job in active_jobs:
if (
now >= job.ends_at
or job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns
or (job.max_budget is not None and job.spend >= job.max_budget)
or _request_was_routed_by(request_metadata, job.router_name) != (job.direction == "reverse")
):
continue
if not _sample_hits(request_id, job.id, job.shadow_percentage):
self._record_funnel(job.id, "not_sampled")
continue
eligible.append(job)
return tuple(eligible)
async def async_log_success_event(
self,
kwargs: Mapping[str, object],
@ -669,18 +731,8 @@ class ShadowEvalLogger(CustomLogger):
return # only surfaces this table can normalize are comparable; unknown types fail closed
if ops.wire_params and _request_mutating_guardrail_ran(request_metadata):
return # the wire-body snapshot predates the rewrite; replaying it would resurrect stripped content
# A key can hold one job per direction, and a request routed by one job's
# router while bypassing the other's qualifies for both. Each is separately
# budgeted, so both fire; the request is normalized once, and only when at
# least one job sampled it.
eligible: Final = tuple(
job
for job in (await self._active_jobs()).get(str(api_key_hash), ())
if datetime.now(timezone.utc) < job.ends_at
and job.attempts + self._job_starts.get(job.id, 0) < job.max_turns
and (job.max_budget is None or job.spend < job.max_budget)
and _sample_hits(request_id, job.id, job.shadow_percentage)
and _request_was_routed_by(request_metadata, job.router_name) == (job.direction == "reverse")
eligible: Final = self._sampled_jobs(
(await self._active_jobs()).get(str(api_key_hash), ()), request_metadata, request_id
)
if not eligible:
return
@ -691,12 +743,18 @@ class ShadowEvalLogger(CustomLogger):
response_obj,
)
if sample is None:
for job in eligible:
self._record_funnel(job.id, "unjudgeable")
return
messages, shadow_params, real_text = sample
control_tier: Final = _routed_tier(request_metadata)
real_cost: Final = float(payload.get("response_cost") or 0.0)
real_cache_hit: Final = payload.get("cache_hit") is True
real_classifier_cost: Final = _decision_classifier_cost(request_metadata)
for job in eligible:
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
return
self._record_funnel(job.id, "shed")
continue
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1
self._inflight_shadow_tasks += 1
asyncio.create_task(
@ -706,6 +764,9 @@ class ShadowEvalLogger(CustomLogger):
messages=messages,
real_text=real_text,
real_model=payload.get("model") or "",
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
control_tier=control_tier,
shadow_params=shadow_params,
parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot
@ -726,37 +787,66 @@ class ShadowEvalLogger(CustomLogger):
messages: Sequence[Mapping[str, object]],
real_text: str,
real_model: str,
real_cost: float,
real_classifier_cost: float,
real_cache_hit: bool,
control_tier: str | None,
shadow_params: Mapping[str, object],
parent_metadata: Mapping[str, object],
) -> None:
"""Budget gate -> shadow call -> blind judge -> one attempt row. The prisma gate
sits above the dispatch so no provider spend happens without a place to record
the outcome, and the budget read lives here rather than in the success hook."""
"""Budget gate -> shadow call -> blind judge -> one attempt row, and every exit
in exactly one coverage bucket: the gates that decline to spend on an admitted
sample (no DB to record into, an over-budget key, an unverifiable or exhausted
eval budget) count it withheld, so eligible traffic still reconciles as
not_sampled + unjudgeable + shed + withheld + attempt rows. The prisma gate sits
above the dispatch so no provider spend happens without a place to record the
outcome, and the budget read lives here rather than in the success hook."""
prisma: Final = self._prisma_provider()
try:
if prisma is None:
self._record_funnel(job.id, "withheld")
return
if await _key_or_team_is_over_budget(parent_metadata):
self._record_funnel(job.id, "withheld")
return
if job.max_budget is not None:
try:
spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget)
except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it
verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e)
self._record_funnel(job.id, "withheld")
return
if spend >= job.max_budget:
self._record_funnel(job.id, "withheld")
return
shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata)
except Exception as e: # noqa: BLE001 # detached task: nothing billed yet, record and never raise
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
await self._record_attempt(
prisma, job, request_id, control_tier, outcome="error", error=f"pipeline error: {e}"
prisma,
job,
request_id,
control_tier,
outcome="error",
error=f"pipeline error: {e}",
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
)
return
if isinstance(shadow, _CallFailure):
await self._record_attempt(
prisma, job, request_id, control_tier, outcome="error", error=shadow.error, shadow_cost=shadow.cost
prisma,
job,
request_id,
control_tier,
outcome="error",
error=shadow.error,
shadow_cost=shadow.cost,
shadow_classifier_cost=shadow.classifier_cost,
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
)
return
# From here the shadow call has billed, so every exit records its cost.
@ -779,6 +869,10 @@ class ShadowEvalLogger(CustomLogger):
shadow=shadow,
judge_cost=verdict.cost,
shadow_cost=shadow.cost,
shadow_classifier_cost=shadow.classifier_cost,
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
)
return
await self._record_attempt(
@ -792,6 +886,10 @@ class ShadowEvalLogger(CustomLogger):
confidence=verdict.confidence,
judge_cost=verdict.cost,
shadow_cost=shadow.cost,
shadow_classifier_cost=shadow.classifier_cost,
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
)
except Exception as e: # noqa: BLE001 # detached task: the shadow call billed, record its cost, never raise
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
@ -804,6 +902,10 @@ class ShadowEvalLogger(CustomLogger):
error=f"pipeline error: {e}",
shadow=shadow,
shadow_cost=shadow.cost,
shadow_classifier_cost=shadow.classifier_cost,
real_cost=real_cost,
real_classifier_cost=real_classifier_cost,
real_cache_hit=real_cache_hit,
)
async def _record_attempt(
@ -814,15 +916,20 @@ class ShadowEvalLogger(CustomLogger):
control_tier: str | None,
*,
outcome: str,
real_cost: float,
real_classifier_cost: float,
real_cache_hit: bool,
shadow: _ShadowResponse | None = None,
real_model: str = "",
confidence: float | None = None,
judge_cost: float = 0.0,
shadow_cost: float = 0.0,
shadow_classifier_cost: float = 0.0,
error: str | None = None,
) -> None:
if judge_cost + shadow_cost > 0:
await self._write_job_spend(_job_spend_counter_key(job.id), judge_cost + shadow_cost)
eval_spend: Final = judge_cost + shadow_cost + shadow_classifier_cost
if eval_spend > 0:
await self._write_job_spend(_job_spend_counter_key(job.id), eval_spend)
if prisma is None:
return
try:
@ -837,6 +944,10 @@ class ShadowEvalLogger(CustomLogger):
"confidence": confidence,
"judge_cost": judge_cost,
"shadow_cost": shadow_cost,
"shadow_classifier_cost": shadow_classifier_cost,
"real_cost": real_cost,
"real_classifier_cost": real_classifier_cost,
"real_cache_hit": real_cache_hit,
"error": error[:_MAX_ERROR_CHARS] if error else None,
}
)
@ -873,15 +984,23 @@ class ShadowEvalLogger(CustomLogger):
)
except Exception as e: # noqa: BLE001 # provider errors become error rows, not crashes
verbose_logger.debug("shadow_eval: router call failed: %s", e)
return _CallFailure(f"shadow router call failed: {_failure_detail(e)}")
return _CallFailure(
f"shadow router call failed: {_failure_detail(e)}",
classifier_cost=_decision_classifier_cost(shadow_metadata),
)
text: Final = _chat_final_text(response)
if not text:
return _CallFailure("shadow router returned an empty response", cost=_call_cost(response))
return _CallFailure(
"shadow router returned an empty response",
cost=_call_cost(response),
classifier_cost=_decision_classifier_cost(shadow_metadata),
)
return _ShadowResponse(
text=text,
model=str(getattr(response, "model", None) or _routing_decision(shadow_metadata).get("routed_model") or ""),
tier=_routed_tier(shadow_metadata),
cost=_call_cost(response),
classifier_cost=_decision_classifier_cost(shadow_metadata),
)
async def _call_judge(
@ -915,6 +1034,7 @@ class ShadowEvalLogger(CustomLogger):
self._router_provider(),
judge_model,
judge_messages, # pyright: ignore[reportArgumentType] # plain SDK message dicts
team_id=_forwarded_team_id(parent_metadata),
temperature=0,
max_tokens=JUDGE_MAX_OUTPUT_TOKENS,
response_format=PAIRWISE_JUDGE_RESPONSE_FORMAT,

View file

@ -0,0 +1,279 @@
"""Provider-agnostic SRT/WebVTT subtitle synthesis from timestamped transcription tokens."""
import unicodedata
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import accumulate, groupby
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
CUE_MAX_CHARS: Final = 84
CUE_MAX_DURATION_MS: Final = 7000
CUE_GAP_MS: Final = 700
SRT_RESPONSE_FORMAT: Final = "srt"
VTT_RESPONSE_FORMAT: Final = "vtt"
SUBTITLE_RESPONSE_FORMATS: Final = frozenset((SRT_RESPONSE_FORMAT, VTT_RESPONSE_FORMAT))
_SENTENCE_END_CHARS: Final = (".", "!", "?", "", "", "", "؟", "۔", "", "", "։", "")
_CJK_RANGES: Final = (
(0x3400, 0x4DBF),
(0x4E00, 0x9FFF),
(0xF900, 0xFAFF),
(0x3040, 0x309F),
(0x30A0, 0x30FF),
(0x31F0, 0x31FF),
)
_CJK_NO_BREAK_BEFORE: Final = "、。,.!?:;・ー…」』)〉》】〕"
_CJK_NO_BREAK_AFTER: Final = "「『(〈《【〔"
@dataclass(frozen=True, slots=True)
class SubtitleToken:
text: str
start_ms: int | None = None
end_ms: int | None = None
speaker: str | int | None = None
@dataclass(frozen=True, slots=True)
class SubtitleCue:
start_ms: int
end_ms: int
text: str
@dataclass(frozen=True, slots=True)
class _Word:
text: str
start_ms: int | None
end_ms: int | None
speaker: str | int | None
def _is_cjk(ch: str) -> bool:
cp: Final = ord(ch)
return any(lo <= cp <= hi for lo, hi in _CJK_RANGES)
def _is_cjk_word_boundary(prev_ch: str, next_ch: str) -> bool:
if not (_is_cjk(prev_ch) or _is_cjk(next_ch)):
return False
return next_ch not in _CJK_NO_BREAK_BEFORE and prev_ch not in _CJK_NO_BREAK_AFTER
def _text_width(text: str) -> int:
return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text)
def _starts_new_word(prev: SubtitleToken, token: SubtitleToken) -> bool:
prev_last: Final = prev.text[-1:]
first: Final = token.text[0]
return (
first.isspace()
or prev_last.isspace()
or token.speaker != prev.speaker
or _is_cjk_word_boundary(prev_last, first)
)
def _build_word(group: Sequence[SubtitleToken]) -> _Word:
return _Word(
text="".join(t.text for t in group),
start_ms=next((t.start_ms for t in group if t.start_ms is not None), None),
end_ms=next((t.end_ms for t in reversed(group) if t.end_ms is not None), None),
speaker=group[0].speaker,
)
def _merge_tokens_into_words(tokens: Sequence[SubtitleToken]) -> tuple[_Word, ...]:
"""
Merge subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words.
A token starts a new word when its text begins with whitespace, when the
previous token's text ends with whitespace, when the speaker changes, or
at a CJK character boundary (CJK scripts carry no spaces, so without this
an entire utterance would fuse into a single unbreakable "word"; CJK
punctuation stays attached to the preceding character per kinsoku rules).
Each word carries the first/last available timestamps of its tokens.
"""
kept: Final = tuple(t for t in tokens if t.text != "")
starts: Final = tuple(i for i, t in enumerate(kept) if i == 0 or _starts_new_word(kept[i - 1], t))
return tuple(_build_word(kept[begin:end]) for begin, end in zip(starts, (*starts[1:], len(kept))))
def _cue_start(ws: Sequence[_Word]) -> int | None:
return next((w.start_ms for w in ws if w.start_ms is not None), None)
def _cue_end(ws: Sequence[_Word]) -> int | None:
return next((w.end_ms for w in reversed(ws) if w.end_ms is not None), _cue_start(ws))
def _cue_text(ws: Sequence[_Word]) -> str:
return "".join(w.text for w in ws).strip()
def _should_break(cue: Sequence[_Word], word: _Word) -> bool:
speaker_changed: Final = word.speaker is not None and any(
w.speaker is not None and w.speaker != word.speaker for w in cue
)
cue_start: Final = _cue_start(cue)
cue_end: Final = _cue_end(cue)
gap_exceeded: Final = word.start_ms is not None and cue_end is not None and (word.start_ms - cue_end) >= CUE_GAP_MS
chars_exceeded: Final = _text_width(_cue_text(cue)) + _text_width(word.text) > CUE_MAX_CHARS
word_end: Final = word.end_ms if word.end_ms is not None else word.start_ms
duration_exceeded: Final = (
word_end is not None and cue_start is not None and (word_end - cue_start) > CUE_MAX_DURATION_MS
)
return speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded
def _cue_start_indices(words: Sequence[_Word]) -> tuple[int, ...]:
def next_start(start: int, index: int) -> int:
if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS):
return index
if _should_break(words[start:index], words[index]):
return index
return start
if not words:
return ()
return tuple(start for start, _ in groupby(accumulate(range(1, len(words)), next_start, initial=0)))
def _build_cue(ws: Sequence[_Word]) -> SubtitleCue | None:
text: Final = _cue_text(ws)
start: Final = _cue_start(ws)
if not text or start is None:
return None
end: Final = _cue_end(ws)
return SubtitleCue(start_ms=start, end_ms=end if end is not None else start, text=text)
def group_subtitle_tokens_into_cues(tokens: Sequence[SubtitleToken]) -> tuple[SubtitleCue, ...]:
"""
Group transcription tokens into subtitle cues aligned to the actual speech.
Cues only ever break at word boundaries (tokens may be subwords, so they
are first merged into words). A new cue starts when:
- the speaker changes (if diarization is on),
- a silence gap of at least CUE_GAP_MS separates two words, so
subtitles never bridge pauses in speech,
- adding the next word would exceed CUE_MAX_CHARS of display width
(~two subtitle lines; East-Asian wide characters count double), or
- adding the next word would make the cue span more than
CUE_MAX_DURATION_MS.
A cue also ends after sentence-final punctuation, which keeps cue breaks
at natural seams. Cue timestamps come straight from token timestamps;
words without timestamps stay attached to the surrounding cue, and a cue
whose words carry no timestamps at all is dropped.
"""
words: Final = _merge_tokens_into_words(tokens)
starts: Final = _cue_start_indices(words)
return tuple(
cue
for begin, end in zip(starts, (*starts[1:], len(words)))
if (cue := _build_cue(words[begin:end])) is not None
)
def _format_timestamp(total_ms: int, millis_separator: str) -> str:
clamped: Final = max(total_ms, 0)
hours, hour_remainder = divmod(clamped, 3_600_000)
minutes, minute_remainder = divmod(hour_remainder, 60_000)
seconds, millis = divmod(minute_remainder, 1_000)
return f"{hours:02d}:{minutes:02d}:{seconds:02d}{millis_separator}{millis:03d}"
def _render_srt(cues: Sequence[SubtitleCue]) -> str:
lines: Final = tuple(
line
for index, cue in enumerate(cues, start=1)
for line in (
str(index),
f"{_format_timestamp(cue.start_ms, ',')} --> {_format_timestamp(cue.end_ms, ',')}",
cue.text,
"",
)
)
return "\n".join(lines)
def _render_vtt(cues: Sequence[SubtitleCue]) -> str:
cue_lines: Final = tuple(
line
for cue in cues
for line in (
f"{_format_timestamp(cue.start_ms, '.')} --> {_format_timestamp(cue.end_ms, '.')}",
cue.text,
"",
)
)
return "\n".join(("WEBVTT", "", *cue_lines))
def render_subtitle_tokens_as_srt(tokens: Sequence[SubtitleToken]) -> str:
"""Render tokens as an SRT document; empty string when no token has timestamp data."""
cues: Final = group_subtitle_tokens_into_cues(tokens)
if not cues:
return ""
return _render_srt(cues)
def render_subtitle_tokens_as_vtt(tokens: Sequence[SubtitleToken]) -> str:
"""Render tokens as a WebVTT document; the WEBVTT header is emitted even without cues."""
return _render_vtt(group_subtitle_tokens_into_cues(tokens))
class TranscriptionWordTiming(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
word: str = ""
start: float | None = None
end: float | None = None
speaker: str | None = None
_WORD_TIMINGS_ADAPTER: Final = TypeAdapter(tuple[TranscriptionWordTiming, ...])
def _seconds_to_ms(seconds: float | None) -> int | None:
if seconds is None:
return None
return round(seconds * 1000)
def _word_to_subtitle_token(word: TranscriptionWordTiming) -> SubtitleToken:
return SubtitleToken(
text=f"{word.word} ",
start_ms=_seconds_to_ms(word.start),
end_ms=_seconds_to_ms(word.end),
speaker=word.speaker,
)
def _parse_word_timings(words: object) -> tuple[TranscriptionWordTiming, ...]:
try:
return _WORD_TIMINGS_ADAPTER.validate_python(words)
except ValidationError:
return ()
def synthesize_subtitle_document(words: object, response_format: str) -> str | None:
"""
Build an SRT/VTT document from OpenAI verbose_json-style word dicts
(word/start/end in float seconds, optional speaker). Returns None when the
format is not a subtitle format or the words carry no usable timestamps.
"""
if response_format not in SUBTITLE_RESPONSE_FORMATS:
return None
tokens: Final = tuple(_word_to_subtitle_token(word) for word in _parse_word_timings(words))
cues: Final = group_subtitle_tokens_into_cues(tokens)
if not cues:
return None
return _render_srt(cues) if response_format == SRT_RESPONSE_FORMAT else _render_vtt(cues)

View file

@ -0,0 +1,55 @@
import re
from types import MappingProxyType
from typing import Final, NamedTuple
class AwsPartition(NamedTuple):
partition: str
dns_suffix: str
_COMMERCIAL_PARTITION: Final = AwsPartition(partition="aws", dns_suffix="amazonaws.com")
_PARTITIONS_BY_REGION_PREFIX: Final = MappingProxyType(
{
"cn-": AwsPartition(partition="aws-cn", dns_suffix="amazonaws.com.cn"),
"us-gov-": AwsPartition(partition="aws-us-gov", dns_suffix="amazonaws.com"),
"us-isob-": AwsPartition(partition="aws-iso-b", dns_suffix="sc2s.sgov.gov"),
"us-isof-": AwsPartition(partition="aws-iso-f", dns_suffix="csp.hci.ic.gov"),
"us-iso-": AwsPartition(partition="aws-iso", dns_suffix="c2s.ic.gov"),
"eu-isoe-": AwsPartition(partition="aws-iso-e", dns_suffix="cloud.adc-e.uk"),
}
)
_BEDROCK_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:bedrock")
_BEDROCK_ARN_PREFIX_PATTERN: Final = re.compile(r"\Aarn:aws(?:-[a-z0-9-]+)?:bedrock:")
_AWS_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:")
def get_aws_partition(aws_region_name: str | None) -> AwsPartition:
if not aws_region_name:
return _COMMERCIAL_PARTITION
return next(
(partition for prefix, partition in _PARTITIONS_BY_REGION_PREFIX.items() if aws_region_name.startswith(prefix)),
_COMMERCIAL_PARTITION,
)
def get_aws_dns_suffix(aws_region_name: str | None) -> str:
return get_aws_partition(aws_region_name).dns_suffix
def get_aws_arn_prefix(aws_region_name: str | None) -> str:
return f"arn:{get_aws_partition(aws_region_name).partition}:"
def contains_bedrock_arn(value: str) -> bool:
return _BEDROCK_ARN_PATTERN.search(value) is not None
def is_bedrock_arn(value: str) -> bool:
return _BEDROCK_ARN_PREFIX_PATTERN.match(value) is not None
def contains_aws_arn(value: str) -> bool:
return _AWS_ARN_PATTERN.search(value) is not None

View file

@ -454,6 +454,62 @@ def safe_deep_copy(data):
return new_data
def independent_snapshot(
data: dict, # mutable-ok: caller-defined request-payload shape
) -> dict: # mutable-ok: caller-defined request-payload shape
"""
A copy of ``data`` whose top-level keys are deep-copied independently
where possible -- always attempted, regardless of
``litellm.safe_memory_mode``. Unlike ``safe_deep_copy``, which can return
the *original* object outright under that mode (defeating any isolation
guarantee for every key, not just the ones that need it), this never
skips copying wholesale.
Real proxy requests carry ``data["litellm_logging_obj"]`` (a ``Logging``
instance nesting a live OTel span with a real lock) by the time
``pre_call_hook`` runs, which can never be deep-copied. Any individual
key that fails to deep-copy falls back to sharing its original
reference, same crash tolerance as ``safe_deep_copy``'s own per-key
fallback; callers needing true isolation (e.g. a guardrail's
``scan_raw_request`` snapshot) only depend on the keys that are plain,
cleanly-copyable structures (``messages``/``input``,
``metadata``/``litellm_metadata``).
"""
sanitized: Final = {
key: (
{ # mutable-ok: same request-payload shape as data
inner_key: ("placeholder" if inner_key == "litellm_parent_otel_span" else inner_value)
for inner_key, inner_value in value.items()
}
if key in ("metadata", "litellm_metadata") and isinstance(value, dict)
else value
)
for key, value in data.items()
}
def _copied_value(key: str, sanitized_value: object) -> object:
try:
copied_value: Final = copy.deepcopy(sanitized_value)
except Exception: # noqa: BLE001 # any unpicklable value falls back to the original reference for this key only
return data.get(key)
original_value: Final = data.get(key)
if (
key in ("metadata", "litellm_metadata")
and isinstance(copied_value, dict)
and isinstance(original_value, dict)
and "litellm_parent_otel_span" in original_value
):
return { # mutable-ok: same request-payload shape as data
**copied_value,
"litellm_parent_otel_span": original_value["litellm_parent_otel_span"],
}
return copied_value
return { # mutable-ok: same request-payload shape as data
key: _copied_value(key, value) for key, value in sanitized.items()
}
def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
"""
Recursively filter out Exception objects and callable objects from dicts/lists.

View file

@ -2222,6 +2222,8 @@ def _map_exception_by_status(
status_code: Final = original_exception.status_code if hasattr(original_exception, "status_code") else None
if not isinstance(status_code, int) or status_code < 400:
return
if getattr(original_exception, "status_code_is_synthesized", False):
return
message: Final = f"{exception_provider} - {error_str}"
response: Final = original_exception.response if hasattr(original_exception, "response") else None
match status_code:
@ -2341,6 +2343,7 @@ def exception_type(
litellm_response_headers: Final = _get_response_headers(original_exception=original_exception)
try:
error_str = redact_string(str(original_exception)) if _ENABLE_SECRET_REDACTION else str(original_exception)
extra_information = ""
if model or custom_llm_provider:
if hasattr(original_exception, "message"):
error_str = (
@ -2357,7 +2360,6 @@ def exception_type(
# Common Extra information needed for all providers
# We pass num retries, api_base, vertex_deployment etc to the exception here
################################################################################
extra_information = ""
try:
_api_base: Final = litellm.get_api_base(model=model, optional_params=extra_kwargs)
messages: Final = litellm.get_first_chars_messages(kwargs=completion_kwargs)

View file

@ -2,7 +2,7 @@ from typing import Final, cast
from urllib.parse import urlparse
import litellm
from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH
from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, REPLICATE_MODEL_NAME_WITH_ID_LENGTH
from litellm.litellm_core_utils.fallback_generalizations import (
match_routing_generalization,
)
@ -127,6 +127,18 @@ def handle_anthropic_text_model_custom_llm_provider(
return model, custom_llm_provider
def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
"""The authenticating provider this pair already names, or None.
get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their
provider info includes the key it unlocks. For a metadata question that flow is pure hazard,
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
adopt the declaration instead of resolving.
"""
declared: Final = custom_llm_provider or model.split("/", 1)[0]
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
def get_llm_provider(
model: str,
custom_llm_provider: str | None = None,

View file

@ -2,6 +2,7 @@ from typing import Final, Literal
import litellm
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
from litellm.types.utils import LlmProviders, LlmProvidersSet
@ -30,6 +31,10 @@ def get_supported_openai_params(
- List if custom_llm_provider is mapped
- None if unmapped
"""
if not custom_llm_provider:
custom_llm_provider = declared_authenticating_provider(
model
) # rebind-ok: resolving would run the provider's OAuth flow
if not custom_llm_provider:
try:
custom_llm_provider = litellm.get_llm_provider(model=model)[1]

View file

@ -20,8 +20,8 @@ from __future__ import annotations
from collections.abc import Mapping
from typing import Final
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.types.utils import InternalCallOrigin
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
@ -45,6 +45,60 @@ budget-checked like the request that spawned it. Everything else on the parent's
be a lie on a sub-call that runs after it returned."""
def is_background_response(response: object) -> bool:
"""Whether a retrieved object is a response created with ``background=true``.
Such a create returns ``status="queued"`` and no usage at all, so nothing has billed the
job by the time anyone reads it back. Accepts the response as a mapping or a model,
because the callers hold it in both shapes.
"""
if isinstance(response, Mapping):
return response.get("background") is True
return getattr(response, "background", None) is True
def is_unbilled_non_inference_call(
call_type: str | None,
metadata: Mapping[str, object] | None,
response: object,
) -> bool:
"""A read/management route priced at zero, because the usage it reports belongs to the
call that created the object it just read.
Retrieving a background response is the exception, and the enterprise cost poller's read
is the same exception seen from the other side: that job's create billed nothing, so its
retrieval is the only place the spend is ever visible. Pricing those at zero would lose
the spend rather than deduplicate it.
"""
if call_type not in NON_INFERENCE_CALL_TYPES:
return False
if is_background_response(response):
return False
if metadata is None:
return True
return metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
def is_unbilled_non_inference_call_from_params(
call_type: str | None,
litellm_params: Mapping[str, object] | None,
response: object,
) -> bool:
""":func:`is_unbilled_non_inference_call` for callers holding raw ``litellm_params``.
The call-type membership test runs first so that inference traffic, which is every
request in a normal workload, never pays for the metadata merge behind it.
"""
if call_type not in NON_INFERENCE_CALL_TYPES:
return False
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
metadata: Final = (
StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) if litellm_params is not None else None
)
return is_unbilled_non_inference_call(call_type, metadata, response)
def sanitize_user_api_key_auth(auth: object) -> object:
"""Copy of the auth object with its budget reservation removed; the cost callback
falls back to reading the reservation from inside the auth object."""

View file

@ -64,6 +64,7 @@ from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
cost_breakdown_with_guardrail,
guardrail_information_cost,
@ -612,37 +613,60 @@ class Logging(LiteLLMLoggingBaseClass):
processed_list: Final[list[str | Callable | CustomLogger]] = []
for callback in callback_list:
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
# For callbacks that support team-scoped credentials (e.g. datadog),
# pass only the relevant dynamic params as custom_logger_init_args.
_custom_logger_init_args: dict | None = None
if callback == "datadog":
# dd_* params are blocked from standard_callback_dynamic_params
# (request-level security); only the proxy-stamped team/key
# callback vars are admin-configured and trusted.
_custom_logger_init_args = {k: v for k, v in self._trusted_callback_vars if k.startswith("dd_")}
callback_class = _init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is not None:
processed_list.append(callback_class)
for callback_instance in self._resolve_dynamic_callback_string(callback):
processed_list.append(callback_instance)
# If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks
if dynamic_callbacks_type == "success":
if self.dynamic_async_success_callbacks is None:
self.dynamic_async_success_callbacks = []
self.dynamic_async_success_callbacks.append(callback_class)
self.dynamic_async_success_callbacks.append(callback_instance)
elif dynamic_callbacks_type == "failure":
if self.dynamic_async_failure_callbacks is None:
self.dynamic_async_failure_callbacks = []
self.dynamic_async_failure_callbacks.append(callback_class)
self.dynamic_async_failure_callbacks.append(callback_instance)
else:
processed_list.append(callback)
return processed_list
def _resolve_dynamic_callback_string(self, callback: str) -> "tuple[CustomLogger, ...]":
"""
Resolve a known callback name to the logger instance(s) it dispatches to.
For callbacks that support team-scoped credentials (datadog, newrelic),
only the proxy-stamped team/key callback vars are passed as
custom_logger_init_args: dd_*/newrelic_* params are blocked from
standard_callback_dynamic_params (request-level security), so the
trusted-vars channel is the only way credentials reach a per-team logger.
"""
_trusted_var_prefix: Final = "dd_" if callback == "datadog" else "newrelic_" if callback == "newrelic" else None
_custom_logger_init_args: Final[dict | None] = (
{k: v for k, v in self._trusted_callback_vars if k.startswith(_trusted_var_prefix)}
if _trusted_var_prefix is not None
else None
)
callback_class: Final = _init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is None:
return ()
# With team creds, "newrelic" resolves to the per-team METRICS logger;
# resolve the name again without creds so the trace logger (OTel v2 /
# legacy agent) keeps receiving this request.
_newrelic_trace_class: Final = (
_init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
if callback == "newrelic" and _custom_logger_init_args and _custom_logger_init_args.get("newrelic_api_key")
else None
)
if _newrelic_trace_class is not None and _newrelic_trace_class is not callback_class:
return (callback_class, _newrelic_trace_class)
return (callback_class,)
def initialize_standard_callback_dynamic_params(self, kwargs: dict | None = None) -> StandardCallbackDynamicParams:
"""
Initialize the standard callback dynamic params from the kwargs
@ -864,7 +888,10 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_management_logger: CustomLogger | None = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs
) -> tuple[str, list[AllMessageValues], dict]:
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
custom_logger: Final = prompt_management_logger or self.get_custom_logger_for_prompt_management(
model=model,
non_default_params=non_default_params,
@ -874,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if custom_logger:
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
(
model,
messages,
@ -889,6 +917,11 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_label=prompt_label,
prompt_version=prompt_version,
)
if request_kwargs is not None:
AnthropicCacheControlHook.record_gateway_injection(
request_kwargs,
AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before,
)
self.messages = messages
return model, messages, non_default_params
@ -904,7 +937,10 @@ class Logging(LiteLLMLoggingBaseClass):
tools: list[dict] | None = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs
) -> tuple[str, list[AllMessageValues], dict]:
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
custom_logger: Final = prompt_management_logger or self.get_custom_logger_for_prompt_management(
model=model,
tools=tools,
@ -915,6 +951,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if custom_logger:
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
(
model,
messages,
@ -932,6 +969,11 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_label=prompt_label,
prompt_version=prompt_version,
)
if request_kwargs is not None:
AnthropicCacheControlHook.record_gateway_injection(
request_kwargs,
AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before,
)
self.messages = messages
return model, messages, non_default_params
@ -1586,6 +1628,11 @@ class Logging(LiteLLMLoggingBaseClass):
if cache_hit is True:
return 0.0
if is_unbilled_non_inference_call(
self.call_type, StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params), result
):
return 0.0
transformed_result: Final = self._generate_content_result_as_model_response(result)
if transformed_result is not None:
result = transformed_result
@ -4636,6 +4683,19 @@ def _init_custom_logger_compatible_class(
_in_memory_loggers.append(gitlab_logger)
return gitlab_logger
elif logging_integration == "newrelic":
if custom_logger_init_args.get("newrelic_api_key"):
# Team-scoped credentials: per-team METRICS logger, isolated per
# credential set via DynamicLoggingCache. The trace logger for
# this name stays on the global path below.
from litellm.integrations.newrelic.newrelic_team_handler import (
NewRelicHandler,
)
return NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=custom_logger_init_args,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
_v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
if _v2 is not None:
return _v2
@ -5057,7 +5117,7 @@ class StandardLoggingPayloadSetup:
return messages
@staticmethod
def merge_litellm_metadata(litellm_params: dict) -> dict:
def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict:
"""
Merge both litellm_metadata and metadata from litellm_params.
@ -5819,7 +5879,7 @@ def get_standard_logging_object_payload(
cache_hit: Final = kwargs.get("cache_hit", False)
# Extract usage as a plain dict, avoiding Pydantic round-trip
raw_usage_dict: Final = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj=response_obj,
response_obj=None if is_unbilled_non_inference_call(call_type, metadata, response_obj) else response_obj,
combined_usage_object=cast(Usage | None, kwargs.get("combined_usage_object")),
)
usage_dict: Final = (
@ -5998,7 +6058,7 @@ def get_standard_logging_object_payload(
prompt_tokens=usage_dict.get("prompt_tokens", 0),
completion_tokens=usage_dict.get("completion_tokens", 0),
request_tags=request_tags,
end_user=end_user_id or "",
end_user=end_user_id,
api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "",
model_group=_model_group,
model_id=_model_id,

View file

@ -21,11 +21,13 @@ class GuardrailCostEntry(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
guardrail_cost: float | None = None
# ``bool | None`` because the TypedDict sanctions None; None means "not set"
# and keeps the default billed behavior, so a None-carrying entry must not
# fail union validation and silently zero a sibling entry's real cost.
guardrail_cost_in_spend: bool | None = True
GuardrailInformationShape = tuple[GuardrailCostEntry, ...] | GuardrailCostEntry | None
_GUARDRAIL_INFORMATION_ADAPTER: Final[TypeAdapter[GuardrailInformationShape]] = TypeAdapter(GuardrailInformationShape)
_GUARDRAIL_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailCostEntry]] = TypeAdapter(GuardrailCostEntry)
def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None:
@ -47,23 +49,55 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str
return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items())
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"
def azure_prompt_shield_guardrail_cost(
usage_units: Mapping[str, int],
cost_tier: str | None,
price_per_1000_text_records: float | None,
) -> float | None:
"""USD cost of an Azure Prompt Shield invocation from its text-record count.
Returns 0.0 on the free tier, ``text_records * price / 1000`` when a price is
configured, and None when pricing is not configured (usage-only tracking).
"""
if cost_tier == "free":
return 0.0
if price_per_1000_text_records is None:
return None
return usage_units.get(AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT, 0) * price_per_1000_text_records / 1000.0
def _billable_entry_cost(entry: GuardrailCostEntry) -> float:
if entry.guardrail_cost_in_spend is False:
return 0.0
cost: Final = entry.guardrail_cost
if cost is None or not math.isfinite(cost) or cost <= 0.0:
return 0.0
return cost
def guardrail_information_cost(guardrail_information: object) -> float:
def _validated_entry_cost(raw: object) -> float:
"""Billable cost of one raw ``guardrail_information`` entry.
Validated per entry so one malformed entry (e.g. a custom hook stamping a
non-boolean ``guardrail_cost_in_spend``) prices to 0.0 by itself instead of
failing a whole-payload validation and silently zeroing a sibling entry's
real billable cost."""
try:
parsed: Final = _GUARDRAIL_INFORMATION_ADAPTER.validate_python(guardrail_information)
except ValidationError:
return _billable_entry_cost(_GUARDRAIL_COST_ENTRY_ADAPTER.validate_python(raw))
except ValidationError as e:
verbose_logger.warning("Ignoring malformed guardrail_information entry for guardrail cost: %s", e)
return 0.0
if parsed is None:
def guardrail_information_cost(guardrail_information: object) -> float:
if guardrail_information is None:
return 0.0
if isinstance(parsed, GuardrailCostEntry):
return _billable_entry_cost(parsed)
return sum(_billable_entry_cost(entry) for entry in parsed)
if isinstance(guardrail_information, (list, tuple)):
return sum(_validated_entry_cost(entry) for entry in guardrail_information)
return _validated_entry_cost(guardrail_information)
def cost_breakdown_with_guardrail(cost_breakdown: CostBreakdown | None, guardrail_cost: float) -> CostBreakdown | None:

View file

@ -7,7 +7,9 @@ from typing import Any, Final, Literal
import litellm
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
from litellm.litellm_core_utils.llm_cost_calc.utils import (
get_web_search_requests_from_usage,
)
from litellm.types.llms.openai import (
FileSearchTool,
ResponsesAPIResponse,
@ -368,7 +370,7 @@ class StandardBuiltInToolCostTracking:
get_anthropic_web_search_requests_from_response,
)
if usage is not None and (_get_web_search_requests(getattr(usage, "server_tool_use", None)) is not None):
if usage is not None and (get_web_search_requests_from_usage(usage) is not None):
return usage
web_search_requests: Final = get_anthropic_web_search_requests_from_response(response_object)
if web_search_requests is None:
@ -416,7 +418,7 @@ class StandardBuiltInToolCostTracking:
# Anthropic Claude (direct API and Vertex AI) uses server_tool_use.web_search_requests.
# Without this check, Claude ModelResponse always falls through to return False
# and _handle_web_search_cost() is never called.
if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
if get_web_search_requests_from_usage(usage) is not None:
return True
# xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
# answer with no url_citation annotations has no other chat-path signal
@ -429,16 +431,12 @@ class StandardBuiltInToolCostTracking:
response_object=response_object, output_type="web_search_call"
)
elif usage is not None:
if (
hasattr(usage, "server_tool_use")
and _get_web_search_requests(usage.server_tool_use) is not None
or (
hasattr(usage, "prompt_tokens_details")
and usage.prompt_tokens_details is not None
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
and hasattr(usage.prompt_tokens_details, "web_search_requests")
and usage.prompt_tokens_details.web_search_requests is not None
)
if get_web_search_requests_from_usage(usage) is not None or (
hasattr(usage, "prompt_tokens_details")
and usage.prompt_tokens_details is not None
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
and hasattr(usage.prompt_tokens_details, "web_search_requests")
and usage.prompt_tokens_details.web_search_requests is not None
):
return True
if _usage_reports_server_side_web_search_calls(usage):

View file

@ -1,6 +1,7 @@
# What is this?
## Helper utilities for cost_per_token()
import re
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
@ -72,7 +73,20 @@ def _get_token_detail_value(details: object, key: str) -> int | None:
return value if isinstance(value, int) else None
def _get_web_search_requests(server_tool_use: Any) -> int | None:
_IMAGE_SIZE_PATTERN: Final = re.compile(r"\d+(?:x|-x-)\d+")
def _requested_image_param(optional_params: Mapping[str, object] | None, key: str) -> str | None:
value: Final = None if optional_params is None else optional_params.get(key)
return value if isinstance(value, str) else None
def _requested_image_size(optional_params: Mapping[str, object] | None) -> str | None:
value: Final = _requested_image_param(optional_params, "size")
return value if value is not None and _IMAGE_SIZE_PATTERN.fullmatch(value) else None
def get_web_search_requests(server_tool_use: Any) -> int | None:
"""
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
@ -92,6 +106,16 @@ def _get_web_search_requests(server_tool_use: Any) -> int | None:
return getattr(server_tool_use, "web_search_requests", None)
def get_web_search_requests_from_usage(usage: Usage) -> int | None:
"""Read ``web_search_requests`` from a ``Usage``'s ``server_tool_use``.
``Usage`` deletes unset optional fields from ``__dict__`` (see
``SafeAttributeModel``), so direct attribute access can raise
``AttributeError``; ``getattr`` with a default is required here.
"""
return get_web_search_requests(getattr(usage, "server_tool_use", None))
def _is_above_128k(tokens: float) -> bool:
if tokens > 128000:
return True
@ -889,11 +913,22 @@ def generic_cost_per_token(
total_details: Final = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
has_double_counting: Final = (cache_hit > 0 or cache_creation > 0) and total_details > usage.prompt_tokens
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
if has_double_counting:
# cached and per-modality counts are both subsets of prompt_tokens and may overlap, so a
# modality can only bill what the cache did not already cover or the overlap is billed twice
uncached_budget: Final = max(usage.prompt_tokens - cache_hit - cache_creation, 0)
billable_audio: Final = min(audio_tokens, uncached_budget)
billable_image: Final = min(image_tokens, uncached_budget - billable_audio)
billable_video: Final = min(video_tokens, uncached_budget - billable_audio - billable_image)
prompt_tokens_details["audio_tokens"] = billable_audio
prompt_tokens_details["image_tokens"] = billable_image
prompt_tokens_details["video_tokens"] = billable_video
prompt_tokens_details["text_tokens"] = uncached_budget - billable_audio - billable_image - billable_video
elif text_tokens == 0 and prompt_tokens_details["image_count"] == 0:
# Clamp to zero: inconsistent streaming usage
text_tokens = max(text_tokens, 0)
prompt_tokens_details["text_tokens"] = text_tokens
prompt_tokens_details["text_tokens"] = max(
usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
)
(
prompt_base_cost,
@ -1063,15 +1098,17 @@ def get_token_type_cost_breakdown(
reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
# Reasoning is billed at the selected tier's reasoning rate for tiered models,
# else at the explicit per-reasoning-token rate when the model defines one,
# otherwise at the standard output-token rate - this mirrors how the total
# completion cost is computed, so the breakdown can never diverge from it.
# else at the service-tier-aware per-reasoning-token rate - this mirrors how the
# total completion cost is computed, so the breakdown can never diverge from it.
tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage)
flat_reasoning_rate: Final = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None)
reasoning_rate: Final = (
tiered_reasoning_rate
if tiered_reasoning_rate is not None
else (flat_reasoning_rate if flat_reasoning_rate is not None else completion_base_cost)
else _resolve_reasoning_token_cost(
model_info=model_info,
service_tier=service_tier,
completion_base_cost=completion_base_cost,
)
)
reasoning_cost = float(reasoning_tokens) * reasoning_rate
@ -1288,12 +1325,13 @@ class CostCalculatorUtils:
cost_calculator as vertex_ai_image_cost_calculator,
)
if size is None:
size = completion_response.size or "1024-x-1024"
if quality is None:
quality = completion_response.quality or "standard"
if n is None:
n = len(completion_response.data) if completion_response.data else 0
resolved_size: Final = (
size or completion_response.size or _requested_image_size(optional_params) or "1024-x-1024"
)
resolved_quality: Final = (
quality or completion_response.quality or _requested_image_param(optional_params, "quality") or "standard"
)
resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0)
if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value:
if isinstance(completion_response, ImageResponse):
@ -1305,7 +1343,7 @@ class CostCalculatorUtils:
if isinstance(completion_response, ImageResponse):
return bedrock_image_cost_calculator(
model=model,
size=size,
size=resolved_size,
image_response=completion_response,
optional_params=optional_params,
)
@ -1401,19 +1439,19 @@ class CostCalculatorUtils:
# Fall through to default for DALL-E models
return default_image_cost_calculator(
model=model,
quality=quality,
quality=resolved_quality,
custom_llm_provider=custom_llm_provider,
n=n,
size=size,
n=resolved_n,
size=resolved_size,
optional_params=optional_params,
)
else:
return default_image_cost_calculator(
model=model,
quality=quality,
quality=resolved_quality,
custom_llm_provider=custom_llm_provider,
n=n,
size=size,
n=resolved_n,
size=resolved_size,
optional_params=optional_params,
)
return 0.0

View file

@ -4,7 +4,9 @@ from __future__ import annotations
import json
import re
from typing import TYPE_CHECKING, Final
from dataclasses import dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Final, Literal
import litellm
@ -56,17 +58,62 @@ def extract_text_from_content(content: object) -> str:
return ""
def router_resolves_model(router: Router | None, model: str) -> bool:
"""Whether the model name resolves through the proxy's router (configured deployment
or model-group alias), the same check the judge dispatch itself makes, so start-time
validation cannot accept a name the call path then fails on."""
return router is not None and bool(model in router.model_group_alias or router.get_model_list(model_name=model))
@lru_cache(maxsize=512)
def _provider_qualified(model: str) -> str | None:
"""`model` in the one spelling litellm itself resolves it to, or None if it maps to no
provider.
A deployment may be configured as `openai/gpt-4o` and a judge given as `gpt-4o`; both
reach the same model, so an identity that keeps them apart reports two models where
there is one. None is a different answer from "unchanged": a name that is already
provider-qualified normalises to itself, and reading that as a failure would call every
correctly-spelled public model unresolvable.
"""
try:
stripped, provider, _, _ = litellm.get_llm_provider(model=model)
except Exception: # noqa: BLE001 # an unmapped name has no provider, which is the answer
return None
return f"{provider}/{stripped}" if provider and stripped else None
@dataclass(frozen=True, slots=True)
class JudgeTarget:
"""Where a call to one model name goes for one caller, and what answers it.
The single answer to that question: the resolvability gate, the judge-vs-candidate
gate and the dispatch all read it, so none of them can decide it differently. Splitting
it is what let start-time validation accept a team's own model while dispatch sent the
literal name to the SDK.
"""
via: Literal["router", "sdk", "nothing"]
models: frozenset[str]
def judge_target(router: Router | None, model: str, team_id: str | None = None) -> JudgeTarget:
"""Resolve `model` the way a call from `team_id` would be.
Three outcomes and no others: the router serves it (a deployment, a team-public name,
an alias, a routing group or a wildcard, exactly the channels `get_model_list`
composes); the SDK serves it because litellm recognises the provider; or nothing does,
which is the only case a caller may refuse on.
`team_id` is part of the question, not a refinement of it. A team-public name resolves
only for its own team and a team's own deployment resolves for nobody else, so asking
without it answers for a caller who does not exist.
"""
served: Final = router.resolved_litellm_models(model, team_id=team_id) if router is not None else ()
if served:
return JudgeTarget("router", frozenset(_provider_qualified(m) or m for m in served))
qualified: Final = _provider_qualified(model)
return JudgeTarget("sdk", frozenset({qualified})) if qualified is not None else JudgeTarget("nothing", frozenset())
async def judge_acompletion(
router: Router | None,
judge_model: str,
messages: list[AllMessageValues], # mutable-ok: the SDK acompletion signature takes a list
team_id: str | None = None,
**params: object,
) -> ModelResponse:
"""Dispatch a judge call through the proxy's router when the judge model is a
@ -74,9 +121,13 @@ async def judge_acompletion(
provider-qualified public names. The router path never retries or falls back:
a failed judge call is the caller's counted failure, not a spend multiplier.
Sampling preferences are advisory: models that removed sampling params (e.g.
claude-sonnet-5) drop them instead of rejecting the judge call."""
if router_resolves_model(router, judge_model):
return await router.acompletion( # pyright: ignore[reportOptionalMemberAccess] # router_resolves_model implies router is not None
claude-sonnet-5) drop them instead of rejecting the judge call.
The arm is chosen by `judge_target` under the caller's own team, the same call
start-time validation makes, so a judge a team can reach cannot be validated as a
deployment and then dispatched as a public name the SDK has never heard of."""
if judge_target(router, judge_model, team_id).via == "router":
return await router.acompletion( # pyright: ignore[reportOptionalMemberAccess] # a router target implies router is not None
model=judge_model,
messages=messages,
num_retries=0,

View file

@ -2,6 +2,7 @@ from collections.abc import Mapping
from typing import Final
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
def _form_field_value(value: object) -> str:
@ -13,18 +14,31 @@ def _form_field_value(value: object) -> str:
def _flatten_form_field(key: str, value: object) -> tuple[tuple[str, str], ...]:
if isinstance(value, Mapping):
return tuple(
item for subkey, subvalue in value.items() for item in _flatten_form_field(f"{key}[{subkey}]", subvalue)
)
if isinstance(value, (list, tuple)):
return tuple(item for entry in value for item in _flatten_form_field(f"{key}[]", entry))
if value is None:
return ()
serialized: Final = _form_field_value(value)
if not serialized:
return ()
return ((key, serialized),)
pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names
list[tuple[str, object, int]]
] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names
(key, value, 0)
]
flat_fields: Final[list[tuple[str, str]]] = [] # mutable-ok: local accumulator
while pending_fields:
current_key, current_value, depth = pending_fields.pop()
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError("form field nesting exceeds max depth")
if isinstance(current_value, Mapping):
pending_fields.extend(
(f"{current_key}[{subkey}]", subvalue, depth + 1)
for subkey, subvalue in reversed(tuple(current_value.items()))
)
continue
if isinstance(current_value, (list, tuple)):
pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value)))
continue
if current_value is None:
continue
serialized = _form_field_value(current_value)
if serialized:
flat_fields.append((current_key, serialized))
return tuple(flat_fields)
def _is_form_scalar(value: object) -> bool:
@ -32,23 +46,36 @@ def _is_form_scalar(value: object) -> bool:
def _flatten_form_data_field(key: str, value: object) -> tuple[tuple[str, str | tuple[str, ...]], ...]:
if isinstance(value, Mapping):
return tuple(
item
for subkey, subvalue in value.items()
for item in _flatten_form_data_field(f"{key}[{subkey}]", subvalue)
)
if isinstance(value, (list, tuple)):
if all(_is_form_scalar(entry) for entry in value):
serialized_fields: Final = tuple(field for entry in value if (field := _form_field_value(entry)))
return ((key, serialized_fields),) if serialized_fields else ()
return tuple(item for entry in value for item in _flatten_form_data_field(f"{key}[]", entry))
if value is None:
return ()
serialized: Final = _form_field_value(value)
if not serialized:
return ()
return ((key, serialized),)
pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names
list[tuple[str, object, int]]
] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names
(key, value, 0)
]
flat_fields: Final[list[tuple[str, str | tuple[str, ...]]]] = [] # mutable-ok: local accumulator
while pending_fields:
current_key, current_value, depth = pending_fields.pop()
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError("form field nesting exceeds max depth")
if isinstance(current_value, Mapping):
pending_fields.extend(
(f"{current_key}[{subkey}]", subvalue, depth + 1)
for subkey, subvalue in reversed(tuple(current_value.items()))
)
continue
if isinstance(current_value, (list, tuple)):
if all(_is_form_scalar(entry) for entry in current_value):
serialized_fields = tuple(field for entry in current_value if (field := _form_field_value(entry)))
if serialized_fields:
flat_fields.append((current_key, serialized_fields))
continue
pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value)))
continue
if current_value is None:
continue
serialized = _form_field_value(current_value)
if serialized:
flat_fields.append((current_key, serialized))
return tuple(flat_fields)
def flatten_form_field_values(*sources: Mapping[str, object] | None) -> tuple[tuple[str, str | tuple[str, ...]], ...]:

View file

@ -178,7 +178,7 @@ def update_response_metadata(
- response._hidden_params["litellm_overhead_time_ms"]
- response.response_time_ms
"""
if result is None:
if result is None or not hasattr(result, "_hidden_params"):
return
metadata: Final = ResponseMetadata(result)

View file

@ -2,9 +2,21 @@
Utility functions for ModelResponse and ModelResponseStream objects.
"""
from typing import Any, Final
from collections.abc import Mapping
from typing import Final
from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream
from typing_extensions import ReadOnly, TypedDict
from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream, StreamingChoices
class _AttributeView(TypedDict):
value: ReadOnly[object]
def _attribute_of(source: object, name: str) -> object:
attribute: Final[_AttributeView] = {"value": getattr(source, name)}
return attribute["value"]
def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
@ -40,10 +52,10 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
return False
# Check model_extra for dynamically added fields (this is where Pydantic stores them)
if hasattr(model_response, "model_extra") and model_response.model_extra:
for extra_field_name, extra_field_value in model_response.model_extra.items():
if _has_meaningful_content(extra_field_value):
return False
stream_extra_fields: Final[Mapping[str, object]] = model_response.model_extra or {}
for extra_field_value in stream_extra_fields.values():
if _has_meaningful_content(extra_field_value):
return False
# Check for any non-base fields that are set
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
@ -57,7 +69,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
continue
# Check if any other field has meaningful content
model_response_value = getattr(model_response, model_response_field, None)
model_response_value: object = getattr(model_response, model_response_field, None)
if _has_meaningful_content(model_response_value):
return False
@ -71,7 +83,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool:
return True
def _has_meaningful_content(value: Any) -> bool:
def _has_meaningful_content(value: object) -> bool:
"""
Check if a value contains meaningful content.
@ -102,7 +114,7 @@ def _has_meaningful_content(value: Any) -> bool:
return True
def _is_choice_non_empty(choice: Any) -> bool:
def _is_choice_non_empty(choice: StreamingChoices) -> bool:
"""
Deep check if a choice contains any meaningful content.
@ -113,41 +125,41 @@ def _is_choice_non_empty(choice: Any) -> bool:
bool: True if the choice has meaningful content, False otherwise
"""
# Check finish_reason
if hasattr(choice, "finish_reason") and choice.finish_reason is not None:
if getattr(choice, "finish_reason", None) is not None:
return True
# Check logprobs
if hasattr(choice, "logprobs") and choice.logprobs is not None:
if getattr(choice, "logprobs", None) is not None:
return True
# Check enhancements (if present)
if hasattr(choice, "enhancements") and choice.enhancements is not None:
if getattr(choice, "enhancements", None) is not None:
return True
# Deep check delta object
if hasattr(choice, "delta") and choice.delta is not None:
if _is_delta_non_empty(choice.delta):
return True
choice_delta: Final[Delta | None] = getattr(choice, "delta", None)
if choice_delta is not None and _is_delta_non_empty(choice_delta):
return True
# Check model_extra for dynamically added fields on the choice
if hasattr(choice, "model_extra") and choice.model_extra:
for extra_field_name, extra_field_value in choice.model_extra.items():
# Skip certain structural fields that are just default/None placeholders
if extra_field_name == "index" and extra_field_value == 0:
continue
if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None:
continue
if extra_field_name == "delta":
continue
if _has_meaningful_content(extra_field_value):
return True
choice_extra_fields: Final[Mapping[str, object]] = choice.model_extra or {}
for extra_field_name, extra_field_value in choice_extra_fields.items():
# Skip certain structural fields that are just default/None placeholders
if extra_field_name == "index" and extra_field_value == 0:
continue
if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None:
continue
if extra_field_name == "delta":
continue
if _has_meaningful_content(extra_field_value):
return True
# Check for any other non-standard fields on the choice
for attr_name in dir(choice):
# Skip private attributes, methods, and known empty fields
if (
attr_name.startswith("_")
or callable(getattr(choice, attr_name))
or callable(_attribute_of(choice, attr_name))
or attr_name.startswith("model_")
or attr_name
in {
@ -160,8 +172,8 @@ def _is_choice_non_empty(choice: Any) -> bool:
):
continue
attr_value = getattr(choice, attr_name, None)
if _has_meaningful_content(attr_value):
choice_attr_value: object = getattr(choice, attr_name, None)
if _has_meaningful_content(choice_attr_value):
return True
return False
@ -178,20 +190,20 @@ def _is_delta_non_empty(delta: Delta) -> bool:
bool: True if the delta has meaningful content, False otherwise
"""
# Check model_extra for dynamically added fields (this is where Pydantic stores them)
if hasattr(delta, "model_extra") and delta.model_extra:
for extra_field_name, extra_field_value in delta.model_extra.items():
# Even structural fields are meaningful if they have actual content
if _has_meaningful_content(extra_field_value):
return True
delta_extra_fields: Final[Mapping[str, object]] = delta.model_extra or {}
for extra_field_value in delta_extra_fields.values():
# Even structural fields are meaningful if they have actual content
if _has_meaningful_content(extra_field_value):
return True
# Check all regular attributes of the delta object
for attr_name in dir(delta):
# Skip private attributes, methods, and Pydantic-specific fields
if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"):
if attr_name.startswith("_") or callable(_attribute_of(delta, attr_name)) or attr_name.startswith("model_"):
continue
attr_value = getattr(delta, attr_name, None)
if _has_meaningful_content(attr_value):
delta_attr_value: object = getattr(delta, attr_name, None)
if _has_meaningful_content(delta_attr_value):
return True
return False

View file

@ -1747,6 +1747,46 @@ def hoist_images_from_tool_messages(
]
def _is_tool_reference_part(part: object) -> bool:
return isinstance(part, dict) and part.get("type") == "tool_reference"
def _tool_message_carries_tool_reference(message: AllMessageValues) -> bool:
if message.get("role") != "tool":
return False
content = message.get("content")
return isinstance(content, list) and any(_is_tool_reference_part(part) for part in content)
def _drop_tool_reference_parts(message: AllMessageValues) -> AllMessageValues:
if not _tool_message_carries_tool_reference(message):
return message
content = cast(list, message.get("content")) # cast-ok: shape checked by _tool_message_carries_tool_reference
remaining_parts = [ # mutable-ok: tool message content must stay a json list
part for part in content if not _is_tool_reference_part(part)
]
new_content = remaining_parts if remaining_parts else ""
rewritten = {**message, "content": new_content} # mutable-ok: chat messages are plain json dicts
return cast(AllMessageValues, rewritten) # cast-ok: dict spread keeps keys like cache_control
def drop_tool_reference_parts_from_tool_messages(
messages: list[AllMessageValues], # mutable-ok: message pipelines type messages as mutable lists
) -> list[AllMessageValues]: # mutable-ok: message pipelines type messages as mutable lists
"""
Remove tool_reference content parts from role:"tool" messages.
The OpenAI chat spec only accepts text in tool messages, so a tool_reference
part carried through the Anthropic adapter makes strict providers reject the
request. The reference names an already-declared tool rather than carrying
content, so it is dropped; a reference-only result keeps its tool message with
empty text so the preceding tool_call stays answered.
"""
if not any(_tool_message_carries_tool_reference(message) for message in messages):
return messages
return [_drop_tool_reference_parts(message) for message in messages] # mutable-ok: pipelines mutate message lists
def _attempt_json_repair(s: str) -> Any | None:
"""
Attempt to repair truncated JSON produced by LLM tool calls.

View file

@ -1412,7 +1412,7 @@ def convert_to_gemini_tool_call_result(
)
except Exception as e:
verbose_logger.warning("Failed to process image in tool response: %s", e)
elif content_type in ("file", "input_file"):
elif content_type in ("file", "input_file"): # pyright: ignore[reportUnnecessaryContains] # loose runtime dict
# Extract file for inline_data (for tool results with PDF, audio, video, etc.)
file_data = content.get("file_data", "")
if not file_data:
@ -1564,14 +1564,23 @@ def convert_to_anthropic_tool_result(
}
"""
anthropic_content: (
str | list[AnthropicMessagesToolResultContent | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam]
str
| list[
AnthropicMessagesToolResultContent
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
| ToolReference
]
) = ""
if isinstance(message["content"], str):
anthropic_content = message["content"]
elif isinstance(message["content"], list):
content_list: Final = message["content"]
anthropic_content_list: list[
AnthropicMessagesToolResultContent | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam
AnthropicMessagesToolResultContent
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
| ToolReference
] = []
for content in content_list:
if content["type"] == "text":
@ -1614,6 +1623,8 @@ def convert_to_anthropic_tool_result(
original_content_element=content,
)
anthropic_content_list.append(cast(AnthropicMessagesImageParam, _anthropic_image_param))
elif content["type"] == "tool_reference":
anthropic_content_list.append(ToolReference(type="tool_reference", tool_name=content["tool_name"]))
elif content["type"] == "file":
file_content = cast(ChatCompletionFileObject, content)
_file_block = anthropic_process_openai_file_message(file_content)

View file

@ -330,6 +330,24 @@ class RealTimeStreaming:
except (AttributeError, TypeError):
pass
def _flush_unbilled_transcription_usage(self) -> None:
if self.provider_config is None:
return
usage: Final = self.provider_config.unbilled_usage_on_session_close(self.model)
if usage is None:
return
flush_event: Final = (
cast( # cast-ok: usage-only partial event, the same shape _capture_transcription_usage logs
OpenAIRealtimeEvents,
{
"type": "conversation.item.input_audio_transcription.completed",
"usage": usage,
},
)
)
self.store_message(flush_event)
self._capture_transcription_usage(flush_event)
def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtimeEvents) -> None:
"""Extract function_call items from response.done events for spend logging."""
try:
@ -955,6 +973,7 @@ class RealTimeStreaming:
transcript = event.get("transcript", "")
self._collect_user_input_from_backend_event(cast(dict, event))
self.store_message(event_str)
self._capture_transcription_usage(event)
await self._send_event_to_client(event, event_str)
blocked = await self.run_realtime_guardrails(
cast(str, transcript),
@ -1068,6 +1087,7 @@ class RealTimeStreaming:
except Exception as e:
verbose_logger.exception("Error in backend to client send messages: %s", e)
finally:
self._flush_unbilled_transcription_usage()
await self.log_messages()
@staticmethod

View file

@ -13,6 +13,7 @@ import json
from typing import Any, Final
import litellm
from litellm._logging import verbose_logger
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
from ...caching import InMemoryCache
@ -46,6 +47,15 @@ class LangfuseInMemoryCache(InMemoryCache):
_created_langfuse_logger.Langfuse.flush()
_created_langfuse_logger.Langfuse.shutdown()
# Loggers with a periodic flush task (e.g. NewRelicMetricsLogger) expose
# stop() so eviction actually ends the task instead of leaking it.
_evicted_stop: Final = getattr(self.cache_dict[key], "stop", None)
if callable(_evicted_stop):
try:
_evicted_stop()
except Exception: # noqa: BLE001 # a failing stop() must not block eviction
verbose_logger.debug("DynamicLoggingCache: stop() raised during eviction", exc_info=True)
#########################################################
# Call parent class to remove key from cache
#########################################################

View file

@ -239,6 +239,22 @@ class ChunkProcessor:
model_response._hidden_params = chunk.get("_hidden_params", {})
return model_response
@staticmethod
def _get_provider_response_model(
chunks: Sequence["_BaseChunk"],
first_chunk_model: str,
) -> str | None:
models: Final = tuple(
model
for chunk in chunks
if isinstance((hidden_params := chunk.get("_hidden_params")), Mapping)
if isinstance((model := hidden_params.get("provider_response_model")), str) and model
)
return next(
(model for model in models if model != first_chunk_model),
models[0] if models else None,
)
@staticmethod
def apply_provider_assembled_streaming_metadata(
response: ModelResponse,
@ -360,6 +376,15 @@ class ChunkProcessor:
)
response = self.update_model_response_with_hidden_params(model_response=response, chunk=chunk)
provider_response_model: Final = self._get_provider_response_model(
chunks,
first_chunk_model,
)
if provider_response_model is not None:
response._hidden_params = dict( # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter
response._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params getter
provider_response_model=provider_response_model,
)
return response
@staticmethod

View file

@ -8,11 +8,12 @@ import time
import traceback
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
import anyio
import httpx
from pydantic import BaseModel
from pydantic import BaseModel, ValidationError
from typing_extensions import NotRequired, TypedDict
import litellm
@ -182,6 +183,48 @@ class _VertexChunkLike(Protocol):
candidates: Sequence[_VertexCandidateLike]
class _ParsedChunkHiddenParams(BaseModel):
provider_specific_fields: Mapping[str, object] | None = None
def _provider_response_model(chunk: object) -> str | None:
model: Final[object] = chunk.get("model") if isinstance(chunk, Mapping) else getattr(chunk, "model", None)
return model if isinstance(model, str) and model else None
def _parsed_provider_hidden_params(hidden: object) -> _ParsedChunkHiddenParams | None:
if not isinstance(hidden, dict):
return None
try:
return _ParsedChunkHiddenParams.model_validate(hidden)
except ValidationError:
return None
def _provider_hidden_params(
chunk: object,
provider_response_model: str | None,
) -> Mapping[str, object] | None:
hidden: Final[object] = getattr(chunk, "_hidden_params", None)
parsed: Final = _parsed_provider_hidden_params(hidden)
provider_specific_fields: Final[object | None] = (
dict(parsed.provider_specific_fields) # mutable-ok: stream assembly merges provider metadata into this dict
if parsed is not None and parsed.provider_specific_fields
else None
)
params: Final[Mapping[str, object]] = MappingProxyType(
{
key: value
for key, value in (
("provider_response_model", provider_response_model),
("provider_specific_fields", provider_specific_fields),
)
if value is not None
}
)
return params or None
class CustomStreamWrapper:
def __init__(
self,
@ -211,6 +254,7 @@ class CustomStreamWrapper:
self.thinking_content = ""
self.system_fingerprint: str | None = None
self._provider_response_model: str | None = None
self.received_finish_reason: str | None = None
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
self.special_tokens = [
@ -801,7 +845,9 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def model_response_creator(self, chunk: dict | None = None, hidden_params: dict | None = None):
def model_response_creator(
self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None
) -> ModelResponseStream:
_model: Final = self._cached_model_name
_logging_obj_llm_provider: Final = self._cached_logging_llm_provider
@ -1504,7 +1550,12 @@ class CustomStreamWrapper:
def chunk_creator(self, chunk: Any):
if hasattr(chunk, "id"):
self.response_id = chunk.id
model_response = self.model_response_creator()
provider_response_model: Final = _provider_response_model(chunk)
if provider_response_model is not None:
self._provider_response_model = provider_response_model
model_response = self.model_response_creator(
hidden_params=_provider_hidden_params(chunk, self._provider_response_model)
)
response_obj: dict[str, Any] = {}
try:
# return this for all models
@ -2318,6 +2369,7 @@ class CustomStreamWrapper:
partial_response: Final = litellm.stream_chunk_builder(
chunks=self.chunks,
messages=self.messages if isinstance(self.messages, list) else None,
logging_obj=self.logging_obj,
)
if partial_response is None:
return

View file

@ -3,7 +3,7 @@
import base64
import io
import struct
from collections.abc import Callable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any, Final, Literal, cast
import tiktoken
@ -25,14 +25,21 @@ from litellm.litellm_core_utils.default_encoding import encoding as default_enco
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.types.llms.anthropic import (
AnthropicContentParamSource,
AnthropicContentParamSourceFileId,
AnthropicContentParamSourceUrl,
AnthropicMessagesDocumentParam,
AnthropicMessagesImageParam,
AnthropicMessagesTextParam,
AnthropicMessagesToolResultParam,
AnthropicMessagesToolUseParam,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionDocumentObject,
ChatCompletionNamedToolChoiceParam,
ChatCompletionToolParam,
OpenAIMessageContent,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import Message, SelectTokenizerResponse
@ -346,7 +353,7 @@ def token_counter(
model="",
custom_tokenizer: dict | SelectTokenizerResponse | None = None,
text: str | list[str] | None = None,
messages: list[AllMessageValues | Message] | None = None,
messages: Sequence[AllMessageValues | Message] | None = None,
count_response_tokens: bool | None = False,
tools: list[ChatCompletionToolParam] | None = None,
tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
@ -646,6 +653,46 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type:
return expected_cls
def _anthropic_image_source_data(
source: AnthropicContentParamSource | AnthropicContentParamSourceUrl | AnthropicContentParamSourceFileId,
) -> str:
if source["type"] == "base64":
data: Final = source.get("data")
if not data:
return ""
media_type: Final = source.get("media_type") or "image/png"
return f"data:{media_type};base64,{data}"
if source["type"] == "url":
return source.get("url") or ""
return ""
def _count_document_tokens(
document: ChatCompletionDocumentObject | AnthropicMessagesDocumentParam,
count_function: TokenCounterFunction,
use_default_image_token_count: bool,
default_token_count: int | None,
) -> int:
source: Final = document["source"]
metadata_tokens: Final = sum(
count_function(text) for text in (document.get("title"), document.get("context")) if text
)
if source["type"] == "text":
return metadata_tokens + count_function(source["data"])
if source["type"] == "content":
content: Final = source["content"]
if isinstance(content, str):
return metadata_tokens + count_function(content)
return metadata_tokens + _count_content_list(
count_function, content, use_default_image_token_count, default_token_count
)
return metadata_tokens + calculate_img_tokens(
data=_anthropic_image_source_data(source),
mode="auto",
use_default_image_token_count=use_default_image_token_count,
)
def _count_anthropic_content(
content: Mapping[str, Any],
count_function: TokenCounterFunction,
@ -697,13 +744,17 @@ def _count_anthropic_content(
def _count_content_list(
count_function: TokenCounterFunction,
content_list: OpenAIMessageContent,
content_list: str
| Iterable[
OpenAIMessageContentListBlock
| AnthropicMessagesTextParam
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
],
use_default_image_token_count: bool,
default_token_count: int | None,
) -> int:
"""
Recursively count tokens from a list of content blocks.
"""
"""Recursively count tokens from a list of content blocks."""
try:
num_tokens = 0
for c in content_list:
@ -714,6 +765,19 @@ def _count_content_list(
elif c["type"] == "image_url":
image_url = c.get("image_url")
num_tokens += _count_image_tokens(image_url, use_default_image_token_count)
elif c["type"] == "image":
num_tokens += calculate_img_tokens(
data=_anthropic_image_source_data(c["source"]),
mode="auto",
use_default_image_token_count=use_default_image_token_count,
)
elif c["type"] == "document":
num_tokens += _count_document_tokens(
c,
count_function,
use_default_image_token_count,
default_token_count,
)
elif c["type"] in ("tool_use", "tool_result"):
num_tokens += _count_anthropic_content(
c,
@ -742,7 +806,8 @@ def _count_content_list(
content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__
raise ValueError(
f"Invalid content item type: {content_type}. "
f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking, tool_reference)."
f"Expected str or dict with 'type' field "
f"(text, image_url, image, document, tool_use, tool_result, thinking, tool_reference)."
)
return num_tokens
except Exception as e:

View file

@ -21,13 +21,61 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
import socket
from ipaddress import ip_address, ip_network
from typing import Any, Final
from typing import Any, Final, Protocol
from urllib.parse import quote, urlparse, urlunparse
import httpx
from typing_extensions import ReadOnly, TypedDict
import litellm
_SockAddr = tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]
class _LocationHeaderView(TypedDict):
location: ReadOnly[object]
class _ResponseView(TypedDict):
response: ReadOnly[httpx.Response]
class _UrlFetcher(Protocol):
"""The slice of ``httpx.Client`` / ``HTTPHandler`` that ``safe_get`` drives."""
def get(
self,
url: str,
*,
headers: dict[str, str] | None = None,
follow_redirects: bool = False,
) -> httpx.Response: ...
class _AsyncUrlFetcher(Protocol):
"""The slice of ``httpx.AsyncClient`` / ``AsyncHTTPHandler`` that ``async_safe_get`` drives."""
async def get(
self,
url: str,
*,
headers: dict[str, str] | None = None,
follow_redirects: bool = False,
) -> httpx.Response: ...
class _FetcherView(TypedDict):
fetcher: ReadOnly[_UrlFetcher]
class _AsyncFetcherView(TypedDict):
fetcher: ReadOnly[_AsyncUrlFetcher]
class _CallerHeadersView(TypedDict):
headers: ReadOnly[dict[str, str]]
# Globally-routable IPs that are cloud-internal. Everything else
# non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented by
# Python's ``ipaddress`` module). This list only holds IPs that are
@ -44,7 +92,7 @@ class SSRFError(ValueError):
"""Raised when a URL targets a blocked network."""
def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str:
def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str:
"""Percent-encode one user-controlled URL path segment.
``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986
@ -64,7 +112,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -
return quote(value_str, safe="")
def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str:
def encode_url_path_segments(value: object, *, field_name: str = "path") -> str:
"""Percent-encode a user-controlled URL path made of multiple segments.
Empty segments are rejected, so leading, trailing, or consecutive slashes
@ -77,11 +125,7 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str:
if value_str == "":
raise ValueError(f"{field_name} is required")
encoded_segments: Final = []
for segment in value_str.split("/"):
encoded_segments.append(encode_url_path_segment(segment, field_name=field_name))
return "/".join(encoded_segments)
return "/".join(encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/"))
def _is_blocked_ip(addr: str) -> bool:
@ -202,7 +246,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str:
return f"{bracketed}:{port}"
def _sockaddr_host(sockaddr: Any) -> str:
def _sockaddr_host(sockaddr: _SockAddr) -> str:
"""Return the host element of a ``getaddrinfo`` sockaddr as ``str``.
``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs
@ -285,8 +329,8 @@ def validate_url(url: str) -> tuple[str, str]:
raise SSRFError(f"No addresses found for '{hostname}'")
if not is_allowlisted:
for family, type_, proto, canonname, sockaddr in addrinfo:
resolved_ip = _sockaddr_host(sockaddr)
for addrinfo_entry in addrinfo:
resolved_ip = _sockaddr_host(addrinfo_entry[4])
if _is_blocked_ip(resolved_ip):
raise SSRFError(
f"URL targets a blocked address ({resolved_ip}). "
@ -363,9 +407,10 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None:
_MAX_REDIRECTS: Final = 10
def _extract_redirect_url(response: Any, request_url: str) -> str:
def _extract_redirect_url(response: httpx.Response, request_url: str) -> str:
"""Extract and resolve the redirect target from a response's Location header."""
location: Final = response.headers.get("location")
header_view: Final[_LocationHeaderView] = {"location": response.headers.get("location")}
location: Final = header_view["location"]
if not isinstance(location, str) or not location:
raise SSRFError("Redirect response has no Location header")
# Resolve relative URLs against the request URL
@ -393,14 +438,17 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
"""
if not getattr(litellm, "user_url_validation", True):
kwargs.setdefault("follow_redirects", True)
return client.get(url, **kwargs)
unvalidated: Final[_ResponseView] = {"response": client.get(url, **kwargs)}
return unvalidated["response"]
fetcher_view: Final[_FetcherView] = {"fetcher": client}
fetcher: Final = fetcher_view["fetcher"]
kwargs.pop("follow_redirects", None)
caller_headers: Final = kwargs.pop("headers", {})
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
response = client.get(
response = fetcher.get(
validated_url,
headers={**caller_headers, "Host": original_host},
headers={**headers_view["headers"], "Host": original_host},
follow_redirects=False,
**kwargs,
)
@ -416,14 +464,17 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
"""Async version of safe_get."""
if not getattr(litellm, "user_url_validation", True):
kwargs.setdefault("follow_redirects", True)
return await client.get(url, **kwargs)
unvalidated: Final[_ResponseView] = {"response": await client.get(url, **kwargs)}
return unvalidated["response"]
fetcher_view: Final[_AsyncFetcherView] = {"fetcher": client}
fetcher: Final = fetcher_view["fetcher"]
kwargs.pop("follow_redirects", None)
caller_headers: Final = kwargs.pop("headers", {})
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
response = await client.get(
response = await fetcher.get(
validated_url,
headers={**caller_headers, "Host": original_host},
headers={**headers_view["headers"], "Host": original_host},
follow_redirects=False,
**kwargs,
)

View file

@ -4,7 +4,7 @@ A2A Protocol Transformation for LiteLLM
import uuid
from collections.abc import Iterator
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -20,6 +20,11 @@ from ..common_utils import (
)
from .streaming_iterator import A2AModelResponseIterator
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class A2AConfig(BaseConfig):
"""
@ -246,12 +251,12 @@ class A2AConfig(BaseConfig):
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:

View file

@ -14,6 +14,8 @@ from litellm.types.llms.openai import (
from litellm.types.utils import ImageObject, ImageResponse
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
@ -169,7 +171,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:

View file

@ -16,6 +16,8 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, ModelResponse
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
@ -66,7 +68,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:

View file

@ -2,7 +2,7 @@
Translate from OpenAI's `/v1/chat/completions` to Amazon Nova's `/v1/chat/completions`
"""
from typing import Any, Final
from typing import TYPE_CHECKING, Final
import httpx
@ -16,6 +16,9 @@ from litellm.types.utils import ModelResponse
from ...openai_like.chat.transformation import OpenAILikeChatConfig
if TYPE_CHECKING:
import tiktoken
class AmazonNovaChatConfig(OpenAILikeChatConfig):
max_completion_tokens: int | None = None
@ -83,7 +86,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:

View file

@ -12,6 +12,8 @@ from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest
from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LoggingClass = LiteLLMLoggingObj
@ -261,7 +263,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:

View file

@ -24,10 +24,12 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
is_provider_native_tool_dict,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
anthropic_tool_names,
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
@ -360,7 +362,13 @@ class AnthropicMessagesHandler(BaseTranslation):
structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices]
tools_to_check: Final[list[ChatCompletionToolParam]] = (
[] if scan_only_tool_results else chat_completion_compatible_request.get("tools", [])
[]
if scan_only_tool_results
else [
tool
for tool in chat_completion_compatible_request.get("tools", [])
if not is_provider_native_tool_dict(tool)
]
)
# Step 1: Extract all text content and images
@ -419,7 +427,10 @@ class AnthropicMessagesHandler(BaseTranslation):
tool_name=anthropic_tool_name,
)
if scan_only_tool_results
else anthropic_tools
else [
*(tool for tool in data.get("tools") or [] if is_provider_native_tool_dict(tool)),
*anthropic_tools,
]
)
guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages")
@ -677,12 +688,9 @@ class AnthropicMessagesHandler(BaseTranslation):
)
def extract_request_tool_names(self, data: dict) -> list[str]:
"""Extract tool names from Anthropic messages request (tools[].name)."""
names: Final[list[str]] = []
for tool in data.get("tools") or []:
if isinstance(tool, dict) and tool.get("name"):
names.append(str(tool["name"]))
return names
"""Extract every tool name in an Anthropic messages request: tools[].name, plus
tools[].function.name for OpenAI-format tools the bridge forwards verbatim."""
return [name for tool in data.get("tools") or [] for name in anthropic_tool_names(tool)]
@classmethod
def _extract_input_text_and_images(

View file

@ -92,6 +92,8 @@ from ..common_utils import (
)
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LoggingClass = LiteLLMLoggingObj
@ -2575,7 +2577,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:

View file

@ -974,19 +974,25 @@ def strip_advisor_blocks_from_messages(messages: list[Any], replace_with_text: b
return messages
def is_anthropic_invalid_thinking_signature_error(error_text: str) -> bool:
def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool:
"""
Detect Anthropic 400 errors caused by missing or invalid thinking signatures.
Detect Anthropic 400 errors caused by invalid thinking blocks in replayed
history: a missing or invalid signature, or a block with empty thinking text.
Known error formats:
{"message":"messages.2.content.0.thinking.signature.str: Input should be a valid string"}
messages.N.content.M.thinking.signature.str: Input should be a valid string
messages.N.content.M: Invalid `signature` in `thinking` block
messages.N.content.M.thinking: each thinking block must contain thinking
"""
if not error_text:
return False
lower: Final = error_text.lower()
return "thinking" in lower and "signature" in lower and ("invalid" in lower or "valid string" in lower)
if "thinking" not in lower:
return False
if "signature" in lower and ("invalid" in lower or "valid string" in lower):
return True
return "must contain thinking" in lower
def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[Any]:
@ -1028,22 +1034,29 @@ def strip_thinking_blocks_from_anthropic_messages_request_dict(
data.pop("thinking", None)
def strip_empty_text_blocks_from_anthropic_messages(
def strip_empty_content_blocks_from_anthropic_messages(
messages: list[Any],
) -> list[Any]:
"""
Return a new message list with empty or whitespace-only ``{"type": "text"}``
content blocks removed.
and ``{"type": "thinking"}`` content blocks removed.
Anthropic's API rejects requests containing such blocks with
``"messages: text content blocks must be non-empty"``, but assistant
messages from Anthropic routinely arrive with ``{"type": "text", "text": ""}``
alongside ``tool_use`` blocks (see anthropics/anthropic-sdk-python#461).
``"messages: text content blocks must be non-empty"`` and
``"messages.N.content.M.thinking: each thinking block must contain
thinking"`` respectively. Assistant messages routinely arrive with
``{"type": "text", "text": ""}`` alongside ``tool_use`` blocks (see
anthropics/anthropic-sdk-python#461), and a turn served by a
non-Anthropic reasoning model through the /v1/messages bridge can carry
``{"type": "thinking", "thinking": ""}`` when the model produced no
reasoning text (e.g. it went straight to parallel tool calls).
Multi-turn tool-use clients (e.g. Claude Code) loop these prior responses
back as conversation history, which then causes the next request to 400
on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already
handles this in ``anthropic_messages_pt``; this helper provides the
equivalent guarantee for the native Anthropic Messages path.
``redacted_thinking`` blocks are never touched: they carry opaque
``data`` instead of thinking text.
Messages whose content is a list and becomes empty after stripping are
omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`.
@ -1056,7 +1069,7 @@ def strip_empty_text_blocks_from_anthropic_messages(
out.append(m)
continue
content = m["content"]
filtered = [b for b in content if not _is_empty_text_block(b)]
filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)]
if len(filtered) == len(content):
out.append(m)
elif filtered:
@ -1071,6 +1084,21 @@ def _is_empty_text_block(block: Any) -> bool:
return not isinstance(text, str) or not text.strip()
def is_empty_thinking_block(block: object) -> bool:
"""
True for a ``{"type": "thinking"}`` content block whose thinking text is
missing, not a string, or empty/whitespace-only after ``.strip()``.
Anthropic rejects such blocks with ``"each thinking block must contain
thinking"`` (whitespace-only included, verified live), regardless of any
signature they carry. ``redacted_thinking`` blocks are a different type
and always return False.
"""
if not isinstance(block, dict) or block.get("type") != "thinking":
return False
thinking: Final = block.get("thinking")
return not isinstance(thinking, str) or not thinking.strip()
def normalize_anthropic_tool_use_id(raw_id: str) -> str:
"""
Normalize a tool_use / tool_result id for Anthropic's ``^[a-zA-Z0-9_-]+$``

View file

@ -7,7 +7,7 @@ Litellm provider slug: `anthropic_text/<model_name>`
import json
import time
from collections.abc import AsyncIterator, Iterator
from typing import Final
from typing import TYPE_CHECKING, Final
import httpx
@ -32,6 +32,9 @@ from litellm.types.utils import (
Usage,
)
if TYPE_CHECKING:
import tiktoken
class AnthropicTextError(BaseLLMException):
def __init__(self, status_code, message):
@ -182,7 +185,7 @@ class AnthropicTextConfig(BaseConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: str,
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
@ -202,9 +205,10 @@ class AnthropicTextConfig(BaseConfig):
model_response.choices[0].finish_reason = completion_response["stop_reason"]
## CALCULATING USAGE
prompt_tokens: Final = len(encoding.encode(prompt)) ##[TODO] use the anthropic tokenizer here
tokenizer: Final = encoding if encoding is not None else litellm.encoding
prompt_tokens: Final = len(tokenizer.encode(prompt)) ##[TODO] use the anthropic tokenizer here
completion_tokens: Final = len(
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
tokenizer.encode(model_response["choices"][0]["message"].get("content", ""))
) ##[TODO] use the anthropic tokenizer here
model_response.created = int(time.time())

View file

@ -8,9 +8,9 @@ from typing import TYPE_CHECKING, Final, Optional
from pydantic import BaseModel, ValidationError
from litellm.litellm_core_utils.llm_cost_calc.utils import (
_get_web_search_requests,
generic_cost_per_token,
get_provider_specific_geo_multiplier,
get_web_search_requests_from_usage,
)
if TYPE_CHECKING:
@ -104,7 +104,7 @@ def get_cost_for_anthropic_web_search(
if usage is None:
return 0.0
web_search_requests: Final = _get_web_search_requests(getattr(usage, "server_tool_use", None))
web_search_requests: Final = get_web_search_requests_from_usage(usage)
if web_search_requests is None:
return 0.0

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