mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge branch 'litellm_internal_staging' into litellm_fix_bedrock_buffered_responses_stream
This commit is contained in:
commit
24c5846c75
1288 changed files with 37656 additions and 9736 deletions
22
.github/workflows/codspeed.yml
vendored
22
.github/workflows/codspeed.yml
vendored
|
|
@ -12,6 +12,7 @@ on:
|
|||
- "uv.lock"
|
||||
- ".github/workflows/codspeed.yml"
|
||||
- ".github/actions/setup-uv-with-retries/**"
|
||||
- ".github/actions/cache-cargo-build/**"
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
|
@ -23,6 +24,7 @@ on:
|
|||
- "uv.lock"
|
||||
- ".github/workflows/codspeed.yml"
|
||||
- ".github/actions/setup-uv-with-retries/**"
|
||||
- ".github/actions/cache-cargo-build/**"
|
||||
# Allow CodSpeed to trigger backtest performance analysis
|
||||
# in order to generate initial data
|
||||
workflow_dispatch:
|
||||
|
|
@ -55,6 +57,26 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
# Build the wheel and resolve every dependency outside the CodSpeed
|
||||
# runner: the same maturin build took 42 minutes inside `codspeed run`
|
||||
# versus under 3 minutes as a plain step (LIT-6183)
|
||||
- name: Build environment
|
||||
run: >
|
||||
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
|
||||
uv run --frozen --no-default-groups
|
||||
--with pytest==8.3.5
|
||||
--with pytest-codspeed==4.3.0
|
||||
--with "mcp>=1.26.0,<2.0"
|
||||
--with "a2a-sdk>=1.1.0,<2.0"
|
||||
pytest
|
||||
-p pytest_codspeed.plugin
|
||||
tests/benchmarks/
|
||||
--codspeed
|
||||
--collect-only -q
|
||||
|
||||
- name: Run benchmarks
|
||||
uses: CodSpeedHQ/action@1c8ae4843586d3ba879736b7f6b7b0c990757fab # v4.12.1
|
||||
with:
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 18483
|
||||
"limit": 16171
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2557
|
||||
"limit": 2226
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -12,25 +12,25 @@
|
|||
"limit": 480
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"limit": 113
|
||||
"limit": 112
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"limit": 40
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"limit": 213
|
||||
"limit": 211
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 5960
|
||||
"limit": 5199
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
},
|
||||
"reportGeneralTypeIssues": {
|
||||
"limit": 105
|
||||
"limit": 101
|
||||
},
|
||||
"reportIncompatibleMethodOverride": {
|
||||
"limit": 56
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5659
|
||||
"limit": 5611
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15482
|
||||
"limit": 15350
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1058
|
||||
"limit": 0
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -90,40 +90,40 @@
|
|||
"limit": 8
|
||||
},
|
||||
"reportReturnType": {
|
||||
"limit": 213
|
||||
"limit": 181
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 26
|
||||
"limit": 24
|
||||
},
|
||||
"reportUndefinedVariable": {
|
||||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44526
|
||||
"limit": 44368
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38779
|
||||
"limit": 38468
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19827
|
||||
"limit": 19665
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30348
|
||||
"limit": 30066
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
"limit": 111
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 697
|
||||
"limit": 695
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 831
|
||||
"limit": 828
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
@ -138,9 +138,9 @@
|
|||
"limit": 138
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 544
|
||||
"limit": 543
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 145
|
||||
"limit": 139
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ flag_management:
|
|||
carryforward: false
|
||||
- name: proxy-db-schema-migration
|
||||
carryforward: false
|
||||
- name: circleci
|
||||
carryforward: false
|
||||
|
||||
component_management:
|
||||
individual_components:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
#### 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.
|
||||
|
|
@ -53,33 +58,33 @@ async def get_audit_logs(
|
|||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(10, ge=1, le=100),
|
||||
# Filter parameters
|
||||
changed_by: Optional[str] = Query(
|
||||
changed_by: str | None = Query(
|
||||
None, description="Filter by user or system that performed the action"
|
||||
),
|
||||
changed_by_api_key: Optional[str] = Query(
|
||||
changed_by_api_key: str | None = Query(
|
||||
None, description="Filter by API key hash that performed the action"
|
||||
),
|
||||
action: Optional[str] = Query(
|
||||
action: str | None = Query(
|
||||
None, description="Filter by action type (create, update, delete)"
|
||||
),
|
||||
table_name: Optional[str] = Query(
|
||||
table_name: str | None = Query(
|
||||
None, description="Filter by table name that was modified"
|
||||
),
|
||||
object_id: Optional[str] = Query(
|
||||
object_id: str | None = Query(
|
||||
None, description="Filter by ID of the object that was modified"
|
||||
),
|
||||
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
|
||||
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
|
||||
object_team_id: Optional[str] = Query(
|
||||
start_date: str | None = Query(None, description="Filter logs after this date"),
|
||||
end_date: str | None = Query(None, description="Filter logs before this date"),
|
||||
object_team_id: str | None = Query(
|
||||
None,
|
||||
description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
),
|
||||
object_key_hash: Optional[str] = Query(
|
||||
object_key_hash: str | None = Query(
|
||||
None,
|
||||
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
),
|
||||
# Sorting parameters
|
||||
sort_by: Optional[str] = Query(
|
||||
sort_by: str | None = Query(
|
||||
None,
|
||||
description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')",
|
||||
),
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -2,9 +2,10 @@
|
|||
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
|
||||
"""
|
||||
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -626,6 +627,7 @@ class CheckBatchCost:
|
|||
later poll.
|
||||
"""
|
||||
from litellm.batches.batch_utils import (
|
||||
count_error_file_failed_requests,
|
||||
_get_file_content_as_dictionary,
|
||||
calculate_batch_cost_and_usage,
|
||||
)
|
||||
|
|
@ -761,16 +763,33 @@ class CheckBatchCost:
|
|||
model_id=model_id,
|
||||
deployment_model=litellm_model_name,
|
||||
)
|
||||
batch_cost, batch_usage, batch_models = (
|
||||
await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=file_content_as_dict,
|
||||
custom_llm_provider=llm_provider, # type: ignore
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info,
|
||||
batch_file_provider: Final = cast(
|
||||
Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], llm_provider
|
||||
)
|
||||
output_file_result: Final = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=file_content_as_dict,
|
||||
custom_llm_provider=batch_file_provider,
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info,
|
||||
)
|
||||
error_file_failed_requests: Final = await count_error_file_failed_requests(
|
||||
response,
|
||||
custom_llm_provider=batch_file_provider,
|
||||
litellm_params={
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
},
|
||||
)
|
||||
batch_result: Final = (
|
||||
output_file_result
|
||||
if not error_file_failed_requests
|
||||
else dataclasses_replace(
|
||||
output_file_result,
|
||||
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
|
||||
)
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=batch_models[0],
|
||||
model=batch_result.models[0],
|
||||
messages=[{"role": "user", "content": "<retrieve_batch>"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
|
|
@ -802,9 +821,11 @@ class CheckBatchCost:
|
|||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
batch_cost=batch_result.cost,
|
||||
batch_usage=batch_result.usage,
|
||||
batch_models=batch_result.models,
|
||||
batch_successful_requests=batch_result.successful_requests,
|
||||
batch_failed_requests=batch_result.failed_requests,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from collections.abc import Sequence
|
|||
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,39 +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]:
|
||||
|
|
@ -329,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,
|
||||
|
|
@ -352,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
|
||||
|
|
@ -586,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},
|
||||
)
|
||||
|
||||
|
|
@ -776,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,
|
||||
|
|
@ -791,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
|
||||
|
||||
|
|
@ -818,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},
|
||||
|
|
@ -1058,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 []
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.61"
|
||||
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.61"
|
||||
version = "0.1.62"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -5,6 +5,41 @@
|
|||
{{- $gatewayPort := .Values.gateway.service.port -}}
|
||||
{{- $backendPort := .Values.backend.service.port -}}
|
||||
{{- $uiPort := .Values.ui.service.port -}}
|
||||
{{/*
|
||||
Backends addressable from ingress.extraPaths, keyed by the `service` field.
|
||||
*/}}
|
||||
{{- $extraPathBackends := dict
|
||||
"gateway" (dict "name" $gatewayName "port" $gatewayPort)
|
||||
"backend" (dict "name" $backendName "port" $backendPort)
|
||||
"ui" (dict "name" $uiName "port" $uiPort)
|
||||
-}}
|
||||
{{/*
|
||||
UI paths (Next.js static export).
|
||||
|
||||
/ui/* is where the SPA serves its login + dashboard routes (e.g. /ui/login).
|
||||
Without it, /ui/* falls into the catch-all → backend → 404.
|
||||
|
||||
The App Router (output: "export", basePath: "") emits the RSC/flight payload
|
||||
for every route as a ROOT-level <route>.txt (/index.txt, /teams.txt,
|
||||
/__next._tree.txt, ...). The client router fetches these on every soft
|
||||
navigation / prefetch as <route>.txt?_rsc=<hash> (the query string is
|
||||
irrelevant to path matching). They are not under /ui, /_next, or
|
||||
/litellm-asset-prefix, so without /*.txt they fall to the backend catch-all
|
||||
→ 404 → client-side navigation never settles and the login flow spins in an
|
||||
infinite redirect loop (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt
|
||||
from the export; the rule only routes the request to it. Needs an ingress
|
||||
controller whose ImplementationSpecific path is a wildcard pattern
|
||||
(AWS ALB: `*` = 0+ chars); this chart targets the AWS Load Balancer
|
||||
Controller.
|
||||
*/}}
|
||||
{{- $uiPaths := list
|
||||
(dict "path" "/" "pathType" "Exact")
|
||||
(dict "path" "/favicon.ico" "pathType" "Exact")
|
||||
(dict "path" "/litellm-asset-prefix" "pathType" "Prefix")
|
||||
(dict "path" "/_next" "pathType" "Prefix")
|
||||
(dict "path" "/ui" "pathType" "Prefix")
|
||||
(dict "path" "/*.txt" "pathType" "ImplementationSpecific")
|
||||
-}}
|
||||
{{/*
|
||||
Gateway data-plane prefixes — must mirror gateway/routes/allowlist.py.
|
||||
Versioned paths are listed explicitly to avoid routing management routes
|
||||
|
|
@ -39,6 +74,21 @@
|
|||
routes at startup -> 404. So /test is rendered as a standalone Exact path
|
||||
and /test/* falls through to the backend catch-all.
|
||||
*/}}
|
||||
{{/*
|
||||
Every "<path>|<pathType>" this template renders on its own. An
|
||||
ingress.extraPaths entry that repeats one of these is rejected: duplicates
|
||||
in a single rule are resolved by position or by controller-specific tie
|
||||
breaking, so the operator entry could take over a built-in route (an entry
|
||||
at "/" Prefix would swallow the whole backend management API) instead of
|
||||
adding to it.
|
||||
*/}}
|
||||
{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
|
||||
{{- range $uiPaths }}
|
||||
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path .pathType) }}
|
||||
{{- end }}
|
||||
{{- range $gatewayPrefixes }}
|
||||
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|Prefix" .) }}
|
||||
{{- end }}
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
|
|
@ -64,65 +114,15 @@ spec:
|
|||
http:
|
||||
paths:
|
||||
# --- UI (Next.js static export) ---
|
||||
- path: /
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
- path: /favicon.ico
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
- path: /litellm-asset-prefix
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
- path: /_next
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
# /ui/* is where the Next.js SPA serves its login + dashboard
|
||||
# routes (e.g. /ui/login). Without this, /ui/* falls into the
|
||||
# catch-all → backend → 404.
|
||||
- path: /ui
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
# Next.js App Router (output: "export", basePath: "") emits the
|
||||
# RSC/flight payload for every route as a ROOT-level <route>.txt
|
||||
# (/index.txt, /teams.txt, /__next._tree.txt, ...). The client
|
||||
# router fetches these on every soft navigation / prefetch as
|
||||
# <route>.txt?_rsc=<hash> (the query string is irrelevant to path
|
||||
# matching). They are not under /ui, /_next, or
|
||||
# /litellm-asset-prefix, so without this rule they fall to the
|
||||
# backend catch-all → 404 → client-side navigation never settles
|
||||
# and the login flow spins in an infinite redirect loop
|
||||
# (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt from the
|
||||
# export; this rule only routes the request to it. Needs an
|
||||
# ingress controller whose ImplementationSpecific path is a
|
||||
# wildcard pattern (AWS ALB: `*` = 0+ chars); this chart targets
|
||||
# the AWS Load Balancer Controller.
|
||||
- path: /*.txt
|
||||
pathType: ImplementationSpecific
|
||||
{{- range $uiPaths }}
|
||||
- path: {{ .path }}
|
||||
pathType: {{ .pathType }}
|
||||
backend:
|
||||
service:
|
||||
name: {{ $uiName }}
|
||||
port:
|
||||
number: {{ $uiPort }}
|
||||
{{- end }}
|
||||
# --- Gateway data plane ---
|
||||
# Exact /test only (see the $gatewayPrefixes comment above);
|
||||
# /test/* MCP management endpoints fall to the backend catch-all.
|
||||
|
|
@ -142,6 +142,46 @@ spec:
|
|||
port:
|
||||
number: {{ $gatewayPort }}
|
||||
{{- end }}
|
||||
{{- /*
|
||||
--- Operator-supplied extra paths (ingress.extraPaths) ---
|
||||
Rendered after every built-in path so an entry can never take
|
||||
precedence over a default, and before the backend catch-all.
|
||||
Position only decides the match on controllers that honour manifest
|
||||
order: the AWS Load Balancer Controller this chart targets sorts
|
||||
Exact paths first and Prefix paths longest-first, but keeps
|
||||
ImplementationSpecific paths in manifest order, which is what the
|
||||
/*.txt rule above already depends on.
|
||||
*/}}
|
||||
{{- range $idx, $extra := .Values.ingress.extraPaths }}
|
||||
{{- if not (kindIs "map" $extra) }}
|
||||
{{- fail (printf "ingress.extraPaths[%d]: each entry must be a mapping with a 'path' key" $idx) }}
|
||||
{{- end }}
|
||||
{{- if not $extra.path }}
|
||||
{{- fail (printf "ingress.extraPaths[%d]: 'path' is required" $idx) }}
|
||||
{{- end }}
|
||||
{{- $service := $extra.service | default "gateway" }}
|
||||
{{- $target := get $extraPathBackends $service }}
|
||||
{{- if not $target }}
|
||||
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown service %q, expected one of backend, gateway, ui" $idx $extra.path $service) }}
|
||||
{{- end }}
|
||||
{{- $pathType := $extra.pathType | default "Prefix" }}
|
||||
{{- if not (has $pathType (list "Prefix" "Exact" "ImplementationSpecific")) }}
|
||||
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $pathType) }}
|
||||
{{- end }}
|
||||
{{- if eq $extra.path "/" }}
|
||||
{{- fail (printf "ingress.extraPaths[%d]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture" $idx) }}
|
||||
{{- end }}
|
||||
{{- if has (printf "%s|%s" $extra.path $pathType) $builtinPathKeys }}
|
||||
{{- fail (printf "ingress.extraPaths[%d]: path %s with pathType %s is already routed by this chart, and a duplicate would take it over rather than add to it" $idx $extra.path $pathType) }}
|
||||
{{- end }}
|
||||
- path: {{ $extra.path | quote }}
|
||||
pathType: {{ $pathType }}
|
||||
backend:
|
||||
service:
|
||||
name: {{ $target.name }}
|
||||
port:
|
||||
number: {{ $target.port }}
|
||||
{{- end }}
|
||||
# --- Catch-all → backend (management API: /key/*, /user/*, /team/*, ...) ---
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
|
|
|
|||
317
helm/litellm/tests/ingress_extra_paths_tests.yaml
Normal file
317
helm/litellm/tests/ingress_extra_paths_tests.yaml
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
suite: test ingress.extraPaths
|
||||
templates:
|
||||
- ingress.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: renders nothing extra between the built-in gateway prefixes and the backend catch-all when unset
|
||||
set:
|
||||
ingress.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-1]
|
||||
value:
|
||||
path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-backend
|
||||
port:
|
||||
number: 4001
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-2]
|
||||
value:
|
||||
path: /metrics
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
|
||||
- it: routes an extra path to the gateway by default, immediately before the backend catch-all
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-2]
|
||||
value:
|
||||
path: /watsonx
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-1]
|
||||
value:
|
||||
path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-backend
|
||||
port:
|
||||
number: 4001
|
||||
|
||||
- it: keeps every built-in path when extra paths are supplied
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.rules[0].http.paths
|
||||
content:
|
||||
path: /
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-ui
|
||||
port:
|
||||
number: 3000
|
||||
- contains:
|
||||
path: spec.rules[0].http.paths
|
||||
content:
|
||||
path: /ui
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-ui
|
||||
port:
|
||||
number: 3000
|
||||
- contains:
|
||||
path: spec.rules[0].http.paths
|
||||
content:
|
||||
path: /test
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
- contains:
|
||||
path: spec.rules[0].http.paths
|
||||
content:
|
||||
path: /v1/chat
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
- contains:
|
||||
path: spec.rules[0].http.paths
|
||||
content:
|
||||
path: /vertex_ai
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
|
||||
- it: renders every entry in order and honours the service and pathType selectors
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
service: gateway
|
||||
- path: /my-passthrough
|
||||
pathType: Exact
|
||||
service: backend
|
||||
- path: /brand.txt
|
||||
pathType: ImplementationSpecific
|
||||
service: ui
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-4]
|
||||
value:
|
||||
path: /watsonx
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-gateway
|
||||
port:
|
||||
number: 4000
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-3]
|
||||
value:
|
||||
path: /my-passthrough
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-backend
|
||||
port:
|
||||
number: 4001
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-2]
|
||||
value:
|
||||
path: /brand.txt
|
||||
pathType: ImplementationSpecific
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-ui
|
||||
port:
|
||||
number: 3000
|
||||
|
||||
- it: addresses the component services by their configured ports
|
||||
set:
|
||||
ingress.enabled: true
|
||||
gateway.service.port: 8000
|
||||
backend.service.port: 8001
|
||||
ui.service.port: 8080
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
- path: /my-passthrough
|
||||
service: backend
|
||||
- path: /brand.txt
|
||||
service: ui
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-4].backend.service.port.number
|
||||
value: 8000
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-3].backend.service.port.number
|
||||
value: 8001
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-2].backend.service.port.number
|
||||
value: 8080
|
||||
|
||||
- it: rejects an entry naming a service the chart does not deploy
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
service: proxy
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: 'ingress.extraPaths[0] (path /watsonx): unknown service "proxy", expected one of backend, gateway, ui'
|
||||
|
||||
- it: rejects an entry whose pathType is not a kubernetes pathType
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /watsonx
|
||||
pathType: prefix
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: 'ingress.extraPaths[0] (path /watsonx): unknown pathType "prefix", expected one of Exact, ImplementationSpecific, Prefix'
|
||||
|
||||
- it: rejects an entry with no path
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- service: gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: 'path' is required"
|
||||
|
||||
|
||||
- it: rejects a root entry that would take over the backend catch-all
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /
|
||||
service: gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
|
||||
|
||||
- it: rejects a root entry that would take over the UI root
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /
|
||||
pathType: Exact
|
||||
service: gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
|
||||
|
||||
# A root ImplementationSpecific entry duplicates no built-in pair, so the
|
||||
# duplicate check alone would admit it. It is still dead: the built-in
|
||||
# Exact / sorts ahead of it on the AWS Load Balancer Controller and claims
|
||||
# the only request its pattern matches, so it renders and never routes.
|
||||
- it: rejects a root entry that would render but never match
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /
|
||||
pathType: ImplementationSpecific
|
||||
service: gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
|
||||
|
||||
- it: rejects an entry that would take over a UI prefix
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /ui
|
||||
service: gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path /ui with pathType Prefix is already routed by this chart, and a duplicate would take it over rather than add to it"
|
||||
|
||||
- it: rejects an entry that would take over the UI RSC payload rule
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /*.txt
|
||||
pathType: ImplementationSpecific
|
||||
service: backend
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path /*.txt with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it"
|
||||
|
||||
- it: rejects an entry that would take over a gateway data-plane prefix
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /v1/chat
|
||||
service: backend
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path /v1/chat with pathType Prefix is already routed by this chart, and a duplicate would take it over rather than add to it"
|
||||
|
||||
- it: rejects an entry that would take over the exact /test route
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /test
|
||||
pathType: Exact
|
||||
service: backend
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
|
||||
|
||||
- it: allows a built-in path under a different pathType, which is a distinct rule
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- path: /ui
|
||||
pathType: Exact
|
||||
service: ui
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[-2]
|
||||
value:
|
||||
path: /ui
|
||||
pathType: Exact
|
||||
backend:
|
||||
service:
|
||||
name: RELEASE-NAME-litellm-ui
|
||||
port:
|
||||
number: 3000
|
||||
|
||||
- it: rejects a bare string entry instead of failing on template internals
|
||||
set:
|
||||
ingress.enabled: true
|
||||
ingress.extraPaths:
|
||||
- /watsonx
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ingress.extraPaths[0]: each entry must be a mapping with a 'path' key"
|
||||
|
|
@ -13,6 +13,27 @@ ingress:
|
|||
annotations: {}
|
||||
host: "" # optional; if set, becomes the rule's host
|
||||
tls: []
|
||||
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
|
||||
# UI / gateway / backend path is still rendered, these entries are placed
|
||||
# after them and before the backend catch-all, and an entry that repeats a
|
||||
# path the chart already routes is rejected at render time rather than
|
||||
# silently taking it over.
|
||||
#
|
||||
# The chart's built-in gateway prefix list is a snapshot of the data-plane
|
||||
# surface at release time. Use extraPaths for passthrough routes it does not
|
||||
# cover: a provider prefix added upstream after this chart version, or a
|
||||
# custom general_settings.pass_through_endpoints route.
|
||||
#
|
||||
# path required; the HTTP path to route
|
||||
# service which component serves it: gateway (default), backend, or ui
|
||||
# pathType Prefix (default), Exact, or ImplementationSpecific
|
||||
#
|
||||
# The target component only answers paths its own route allowlist keeps, so
|
||||
# a path here still has to be one that component serves.
|
||||
extraPaths: []
|
||||
# - path: /watsonx
|
||||
# pathType: Prefix
|
||||
# service: gateway
|
||||
|
||||
# Per-component ServiceAccounts for gateway, backend, and ui.
|
||||
#
|
||||
|
|
|
|||
|
|
@ -0,0 +1,12 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" (
|
||||
"entity_type" TEXT NOT NULL,
|
||||
"entity_id" TEXT NOT NULL,
|
||||
"window_duration" TEXT NOT NULL,
|
||||
"window_start" TIMESTAMP(3) NOT NULL,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration")
|
||||
);
|
||||
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_ModelAccessGroupBudgetTable" (
|
||||
"access_group_name" TEXT NOT NULL,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"budget_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_pkey" PRIMARY KEY ("access_group_name")
|
||||
);
|
||||
|
||||
-- AddForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_ModelAccessGroupBudgetTable" ADD CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
|
@ -29,6 +29,7 @@ model LiteLLM_BudgetTable {
|
|||
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
|
||||
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
|
||||
tags LiteLLM_TagTable[] // multiple tags can have the same budget
|
||||
model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget
|
||||
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
|
||||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
}
|
||||
|
|
@ -585,6 +586,20 @@ model LiteLLM_EndUserTable {
|
|||
blocked Boolean @default(false)
|
||||
}
|
||||
|
||||
// Budget and shared spend for a model access group. The groups themselves are not rows anywhere:
|
||||
// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here
|
||||
// exists only once someone gives that group a budget.
|
||||
model LiteLLM_ModelAccessGroupBudgetTable {
|
||||
access_group_name String @id
|
||||
spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String?
|
||||
}
|
||||
|
||||
// Track tags with budgets and spend
|
||||
model LiteLLM_TagTable {
|
||||
tag_name String @id
|
||||
|
|
@ -649,6 +664,18 @@ model LiteLLM_SpendLogs {
|
|||
@@index([session_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
entity_type String
|
||||
entity_id String
|
||||
window_duration String
|
||||
window_start DateTime
|
||||
spend Float @default(0.0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@id([entity_type, entity_id, window_duration])
|
||||
}
|
||||
|
||||
// View spend, model, api_key per request
|
||||
model LiteLLM_ErrorLogs {
|
||||
request_id String @id @default(uuid())
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.90"
|
||||
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.90"
|
||||
version = "0.4.91"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
340
litellm-rust/Cargo.lock
generated
340
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
103
litellm-rust/crates/python-bridge/benches/serialization.rs
Normal file
103
litellm-rust/crates/python-bridge/benches/serialization.rs
Normal 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);
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
20
litellm-rust/crates/python-bridge/src/marshal.rs
Normal file
20
litellm-rust/crates/python-bridge/src/marshal.rs
Normal 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()))
|
||||
}
|
||||
52
litellm-rust/crates/python-bridge/tests/marshal_boundary.rs
Normal file
52
litellm-rust/crates/python-bridge/tests/marshal_boundary.rs
Normal 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()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -17,8 +17,11 @@ until they're actually needed.
|
|||
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Final, cast
|
||||
from collections.abc import Callable, Mapping
|
||||
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,8 +56,12 @@ from ._lazy_imports_registry import (
|
|||
UTILS_NAMES,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
from tiktoken import Encoding
|
||||
|
||||
def get_litellm_globals() -> dict:
|
||||
|
||||
def get_litellm_globals() -> dict[str, object]:
|
||||
"""
|
||||
Get the globals dictionary of the litellm module.
|
||||
|
||||
|
|
@ -64,7 +71,7 @@ def get_litellm_globals() -> dict:
|
|||
return sys.modules["litellm"].__dict__
|
||||
|
||||
|
||||
def _get_utils_globals() -> dict:
|
||||
def _get_utils_globals() -> dict[str, object]:
|
||||
"""
|
||||
Get the globals dictionary of the utils module.
|
||||
|
||||
|
|
@ -74,14 +81,19 @@ def _get_utils_globals() -> dict:
|
|||
return sys.modules["litellm.utils"].__dict__
|
||||
|
||||
|
||||
def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None":
|
||||
"""Read the configured `litellm.request_timeout` used for the module level http clients."""
|
||||
return litellm_globals.get("request_timeout")
|
||||
|
||||
|
||||
# These are special lazy loaders for things that are used internally
|
||||
# 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 +112,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 +136,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 +166,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 +218,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 +278,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 +295,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 +378,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 +393,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 +409,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 +425,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.
|
||||
|
||||
|
|
@ -419,8 +441,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
|
|||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
|
||||
# Get timeout from module config (if set)
|
||||
timeout = _globals.get("request_timeout")
|
||||
params: Final = {"timeout": timeout, "client_alias": "module level aclient"}
|
||||
async_timeout: Final = _get_module_level_client_timeout(_globals)
|
||||
params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"}
|
||||
|
||||
# Create the client instance
|
||||
provider_id: Final = cast(Any, "litellm_module_level_client")
|
||||
|
|
@ -437,8 +459,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
|
|||
# Create a sync HTTP client
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
timeout = _globals.get("request_timeout")
|
||||
sync_client: Final = HTTPHandler(timeout=timeout)
|
||||
sync_timeout: Final = _get_module_level_client_timeout(_globals)
|
||||
sync_client: Final = HTTPHandler(timeout=sync_timeout)
|
||||
|
||||
# Cache it
|
||||
_globals["module_level_client"] = sync_client
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from enum import Enum
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import litellm
|
||||
|
|
@ -12,12 +14,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage
|
|||
from litellm.utils import token_counter
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchCostUsageResult:
|
||||
"""Aggregate cost, usage, and per-line pass/fail counts for a completed batch."""
|
||||
|
||||
cost: float
|
||||
usage: Usage
|
||||
models: list[str]
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
|
||||
|
||||
async def calculate_batch_cost_and_usage(
|
||||
file_content_dictionary: list[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, Usage, list[str]]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Calculate the cost and usage of a batch.
|
||||
|
||||
|
|
@ -32,8 +45,7 @@ async def calculate_batch_cost_and_usage(
|
|||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
return batch_cost, batch_usage, [model_name]
|
||||
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=file_content_dictionary,
|
||||
|
|
@ -49,7 +61,7 @@ async def _handle_completed_batch(
|
|||
model_name: str | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, Usage, list[str]]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""Fetch a completed batch's output file and aggregate its cost, usage, and
|
||||
models in a single pass over the JSONL lines, so the parsed file content is
|
||||
never materialized in memory.
|
||||
|
|
@ -72,27 +84,49 @@ async def _handle_completed_batch(
|
|||
# The generic retrieval helper keeps raising for callers that explicitly ask
|
||||
# for a missing output file.
|
||||
if batch.output_file_id is None:
|
||||
return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), []
|
||||
return BatchCostUsageResult(
|
||||
cost=0.0,
|
||||
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
|
||||
models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str]
|
||||
successful_requests=0,
|
||||
failed_requests=await count_error_file_failed_requests(
|
||||
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
|
||||
),
|
||||
)
|
||||
|
||||
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
|
||||
|
||||
if (
|
||||
custom_llm_provider == "vertex_ai"
|
||||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
_get_file_content_as_dictionary(file_content), model_name
|
||||
)
|
||||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_output_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
error_file_failed_requests: Final = await count_error_file_failed_requests(
|
||||
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
output_file_result: Final = (
|
||||
calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name)
|
||||
if (
|
||||
custom_llm_provider == "vertex_ai"
|
||||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
)
|
||||
else _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_output_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
)
|
||||
)
|
||||
|
||||
if not error_file_failed_requests:
|
||||
return output_file_result
|
||||
return dataclasses_replace(
|
||||
output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests
|
||||
)
|
||||
|
||||
|
||||
class _LineOutcome(Enum):
|
||||
"""A batch output line that yielded no billable stats."""
|
||||
|
||||
PROVIDER_FAILED = "provider_failed"
|
||||
UNCOSTABLE = "uncostable"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BatchOutputLineStats:
|
||||
|
|
@ -102,19 +136,27 @@ class _BatchOutputLineStats:
|
|||
total_tokens: int
|
||||
cache_read_tokens: int
|
||||
cache_creation_tokens: int
|
||||
reasoning_tokens: int
|
||||
model: str | None
|
||||
|
||||
|
||||
def _iter_successful_output_line_stats(
|
||||
def _classify_output_line_stats(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
|
||||
"""Classify every output line in a single pass, so counting failures never needs
|
||||
a second read of a potentially huge output file. A line the provider reported as
|
||||
failed yields ``PROVIDER_FAILED``; a successful line litellm could not price
|
||||
yields ``UNCOSTABLE`` and still counts as a successful request billed at $0, so
|
||||
the counts stay reconcilable with the provider's own ``request_counts``."""
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
yield _LineOutcome.PROVIDER_FAILED
|
||||
continue
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
if stats is not None:
|
||||
yield stats
|
||||
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
|
||||
|
||||
|
||||
def _safe_output_line_stats(
|
||||
|
|
@ -123,13 +165,11 @@ def _safe_output_line_stats(
|
|||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
"""Return the stats for one batch output line, or None for a line that is
|
||||
unsuccessful or cannot be costed, so a single bad line never aborts the
|
||||
whole batch's cost accounting."""
|
||||
"""Return the stats for one provider-successful batch output line, or None when
|
||||
it cannot be costed, so a single bad line never aborts the whole batch's cost
|
||||
accounting."""
|
||||
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
|
||||
try:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
return None
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
|
||||
verbose_logger.warning(
|
||||
|
|
@ -152,6 +192,7 @@ def _compute_output_line_stats(
|
|||
prompt_details: Final = parse_prompt_tokens_details(usage)
|
||||
raw_model: Final = response_body.get("model")
|
||||
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
completion_details: Final = usage.completion_tokens_details
|
||||
return _BatchOutputLineStats(
|
||||
cost=_output_line_cost(
|
||||
response_body=response_body,
|
||||
|
|
@ -166,6 +207,7 @@ def _compute_output_line_stats(
|
|||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
reasoning_tokens=(completion_details.reasoning_tokens if completion_details else None) or 0,
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
|
@ -203,10 +245,14 @@ def _aggregate_batch_cost_usage_models(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, Usage, list[str]]:
|
||||
"""Aggregate cost, usage, and models from batch output entries in a single
|
||||
pass, holding one small stats record per line instead of the parsed file."""
|
||||
line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
) -> BatchCostUsageResult:
|
||||
"""Aggregate cost, usage, models, and pass/fail counts from batch output
|
||||
entries in a single pass, holding one small stats record per line instead
|
||||
of the parsed file."""
|
||||
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
|
||||
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
|
||||
successful_requests: Final = len(all_results) - failed_requests
|
||||
|
||||
cache_token_params: Final = {
|
||||
key: tokens
|
||||
|
|
@ -220,18 +266,32 @@ def _aggregate_batch_cost_usage_models(
|
|||
total_tokens=sum(stats.total_tokens for stats in line_stats),
|
||||
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
|
||||
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
|
||||
reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats),
|
||||
**cache_token_params,
|
||||
)
|
||||
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
|
||||
total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
|
||||
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
|
||||
return total_cost, batch_usage, batch_models
|
||||
verbose_logger.debug(
|
||||
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
|
||||
total_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
successful_requests,
|
||||
failed_requests,
|
||||
)
|
||||
return BatchCostUsageResult(
|
||||
cost=total_cost,
|
||||
usage=batch_usage,
|
||||
models=batch_models,
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
)
|
||||
|
||||
|
||||
def calculate_vertex_ai_batch_cost_and_usage(
|
||||
vertex_ai_batch_responses: list[dict],
|
||||
model_name: str | None = None,
|
||||
) -> tuple[float, Usage]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Calculate both cost and usage from raw Vertex AI batch responses.
|
||||
|
||||
|
|
@ -242,6 +302,10 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}}
|
||||
|
||||
usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
|
||||
|
||||
A row with no ``response`` is counted as failed - the same signal already
|
||||
used to skip it from cost/usage aggregation, since Vertex batch prediction
|
||||
output doesn't establish a distinct error shape in this (non-default) path.
|
||||
"""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
|
|
@ -249,12 +313,16 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
total_tokens = 0
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
|
||||
failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
|
||||
actual_model_name: Final = model_name or "gemini-2.0-flash-001"
|
||||
|
||||
for response in vertex_ai_batch_responses:
|
||||
response_body = response.get("response")
|
||||
if response_body is None:
|
||||
failed_requests += 1
|
||||
continue
|
||||
successful_requests += 1
|
||||
|
||||
usage_metadata = response_body.get("usageMetadata", {})
|
||||
_prompt = usage_metadata.get("promptTokenCount", 0) or 0
|
||||
|
|
@ -282,17 +350,25 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
total_tokens += _total
|
||||
|
||||
verbose_logger.info(
|
||||
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
|
||||
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
|
||||
total_cost,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens,
|
||||
successful_requests,
|
||||
failed_requests,
|
||||
)
|
||||
|
||||
return total_cost, Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
return BatchCostUsageResult(
|
||||
cost=total_cost,
|
||||
usage=Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
),
|
||||
models=[actual_model_name],
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -322,6 +398,36 @@ def _provider_output_file_id(output_file_id: str) -> str:
|
|||
return extracted
|
||||
|
||||
|
||||
async def _fetch_batch_managed_file_content(
|
||||
file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
litellm_params: dict | None = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
Fetch a batch's output or error file and return its raw JSONL bytes.
|
||||
|
||||
Args:
|
||||
file_id: The provider or unified (litellm-managed) file id to fetch
|
||||
custom_llm_provider: The LLM provider
|
||||
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
|
||||
Required for Azure and other providers that need authentication
|
||||
"""
|
||||
from litellm.files.main import afile_content
|
||||
|
||||
# Build kwargs for afile_content with credentials from litellm_params
|
||||
file_content_kwargs: Final = {
|
||||
"file_id": _provider_output_file_id(file_id),
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
|
||||
# Extract and add credentials for file access
|
||||
credentials: Final = _extract_file_access_credentials(litellm_params)
|
||||
file_content_kwargs.update(credentials)
|
||||
|
||||
_file_content: Final = await afile_content(**file_content_kwargs)
|
||||
return _file_content.content
|
||||
|
||||
|
||||
async def _fetch_batch_output_file_content(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
|
|
@ -336,25 +442,36 @@ async def _fetch_batch_output_file_content(
|
|||
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
|
||||
Required for Azure and other providers that need authentication
|
||||
"""
|
||||
from litellm.files.main import afile_content
|
||||
|
||||
if batch.output_file_id is None:
|
||||
raise ValueError("Output file id is None cannot retrieve file content")
|
||||
|
||||
file_id: Final = _provider_output_file_id(batch.output_file_id)
|
||||
return await _fetch_batch_managed_file_content(
|
||||
batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
# Build kwargs for afile_content with credentials from litellm_params
|
||||
file_content_kwargs: Final = {
|
||||
"file_id": file_id,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
|
||||
# Extract and add credentials for file access
|
||||
credentials: Final = _extract_file_access_credentials(litellm_params)
|
||||
file_content_kwargs.update(credentials)
|
||||
async def count_error_file_failed_requests(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
litellm_params: dict | None,
|
||||
) -> int:
|
||||
"""Count failed requests reported only in the batch's separate error file.
|
||||
|
||||
_file_content: Final = await afile_content(**file_content_kwargs)
|
||||
return _file_content.content
|
||||
OpenAI-shaped batch providers write successful lines to ``output_file_id``
|
||||
and per-request failures (e.g. a rejected param) to a distinct
|
||||
``error_file_id`` - they never appear in the output file at all, so
|
||||
counting failures from the output file alone silently undercounts them.
|
||||
"""
|
||||
if batch.error_file_id is None:
|
||||
return 0
|
||||
try:
|
||||
error_file_content = await _fetch_batch_managed_file_content(
|
||||
batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch
|
||||
verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e)
|
||||
return 0
|
||||
return sum(1 for _ in _iter_batch_input_lines(error_file_content))
|
||||
|
||||
|
||||
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic.
|
|||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -64,7 +65,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
async_client: AsyncRedis | None = None,
|
||||
embedding_max_input_tokens: int | None = None,
|
||||
embedding_timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
**kwargs: object,
|
||||
):
|
||||
if similarity_threshold is None:
|
||||
raise ValueError("similarity_threshold must be provided, passed None")
|
||||
|
|
@ -87,11 +88,13 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
self.key_prefix = f"{self.index_name}:"
|
||||
self._index_dim: int | None = None
|
||||
|
||||
resolved_url = None
|
||||
if sync_client is None or async_client is None:
|
||||
resolved_url = redis_url or self._build_valkey_url(host, port, password, ssl)
|
||||
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
|
||||
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
|
||||
if sync_client is not None and async_client is not None:
|
||||
self.sync_client = sync_client
|
||||
self.async_client = async_client
|
||||
else:
|
||||
resolved_url: Final = redis_url or self._build_valkey_url(host, port, password, ssl)
|
||||
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
|
||||
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
|
||||
|
||||
print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}")
|
||||
|
||||
|
|
@ -118,7 +121,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
return hashlib.sha256(str(key).encode("utf-8")).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def _embedding_to_bytes(embedding: list[float]) -> bytes:
|
||||
def _embedding_to_bytes(embedding: Sequence[float]) -> bytes:
|
||||
return pack_vector(embedding)
|
||||
|
||||
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
|
||||
|
|
@ -192,7 +195,9 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
def _doc_key(self, key: str) -> str:
|
||||
return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}"
|
||||
|
||||
def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict:
|
||||
def _doc_mapping(
|
||||
self, key: str, prompt: str, value_str: str, embedding: Sequence[float]
|
||||
) -> Mapping[str | bytes, str | bytes]:
|
||||
return {
|
||||
self.CACHE_KEY_FIELD_NAME: self._scope_tag(key),
|
||||
self.PROMPT_FIELD_NAME: prompt,
|
||||
|
|
@ -208,30 +213,49 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
)
|
||||
return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2)
|
||||
|
||||
async def _async_search(self, key: str, embedding: Sequence[float]) -> object:
|
||||
"""Run the KNN query on the async client, stopping the untyped search surface here."""
|
||||
return await self.async_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None:
|
||||
docs: Final = getattr(search_result, "docs", [])
|
||||
def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None:
|
||||
docs: Final[Sequence[object]] = getattr(search_result, "docs", [])
|
||||
if not docs:
|
||||
return None
|
||||
doc: Final = docs[0]
|
||||
response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
|
||||
distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
|
||||
return _ValkeyCacheHit(
|
||||
response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)),
|
||||
distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)),
|
||||
response=str(response_field),
|
||||
distance=float(distance_field),
|
||||
)
|
||||
|
||||
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any:
|
||||
@staticmethod
|
||||
def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None:
|
||||
"""Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``."""
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
|
||||
@staticmethod
|
||||
def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""The request metadata forwarded to the embedding call."""
|
||||
return kwargs.get("metadata")
|
||||
|
||||
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object:
|
||||
if hit is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
similarity: Final = 1 - hit.distance
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
self._record_similarity(kwargs, similarity)
|
||||
|
||||
if similarity < self.similarity_threshold:
|
||||
return None
|
||||
return self._get_cache_logic(cached_response=hit.response)
|
||||
|
||||
def set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
|
|
@ -250,12 +274,12 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache set_cache: {e}")
|
||||
|
||||
def get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
def get_cache(self, key: str, **kwargs: object) -> object:
|
||||
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
embedding: Final = self._get_embedding(prompt)
|
||||
|
|
@ -263,14 +287,14 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
|
||||
search_result: Final = self.sync_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
|
||||
)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache get_cache: {e}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
|
||||
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
|
|
@ -278,7 +302,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
doc_key: Final = self._doc_key(key)
|
||||
|
|
@ -289,31 +313,28 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}")
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
||||
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
search_result: Final = await self.async_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
)
|
||||
search_result: Final[object] = await self._async_search(key, embedding)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
|
||||
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None:
|
||||
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None:
|
||||
try:
|
||||
await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list])
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}")
|
||||
|
||||
async def _index_info(self) -> dict:
|
||||
async def _index_info(self) -> Mapping[str, object]:
|
||||
return await self.async_client.ft(self.index_name).info()
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECO
|
|||
DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
|
||||
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
|
||||
DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
|
||||
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
|
||||
|
||||
# Maximum wall-clock seconds a streaming response is allowed to run.
|
||||
# Streams exceeding this duration are terminated with a Timeout error.
|
||||
|
|
@ -288,6 +289,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update
|
|||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer"
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
|
||||
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
|
||||
LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
|
||||
|
|
@ -1681,6 +1683,7 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16
|
|||
# Ceilings on the cached auth registries; larger tables fall back to per-row lookups
|
||||
# instead of holding an unbounded id set in every worker.
|
||||
TAG_REGISTRY_MAX_SIZE: Final = 5000
|
||||
MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000
|
||||
END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000
|
||||
# How long a failed registry load is remembered as "unusable", so a degraded Postgres
|
||||
# is not re-scanned on every request on top of the per-id lookups it falls back to.
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import json
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, TypeAlias, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
|
||||
|
|
@ -11,7 +12,6 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionImageObject,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionSystemMessage,
|
||||
ChatCompletionTextObject,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
|
|
@ -23,35 +23,63 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
AdapterCompletionStreamWrapper,
|
||||
ChatCompletionDeltaCustomToolCall,
|
||||
ChatCompletionDeltaToolCall,
|
||||
ChatCompletionMessageCustomToolCall,
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
Delta,
|
||||
Function,
|
||||
Message,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
class _GenAITextPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
_JsonDict: TypeAlias = dict[str, object]
|
||||
_JsonDictList: TypeAlias = list[_JsonDict]
|
||||
|
||||
|
||||
class _GenAISystemInstruction(TypedDict, total=False):
|
||||
parts: ReadOnly[list[_GenAITextPart]]
|
||||
class _ToolCallAccumulator(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
arguments: ReadOnly[str]
|
||||
|
||||
|
||||
class _GenAIFunctionCall(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
args: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _GenAIPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
functionCall: ReadOnly[dict[str, object]]
|
||||
functionCall: ReadOnly[_GenAIFunctionCall]
|
||||
|
||||
|
||||
class _GenAIFunctionResponse(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
response: ReadOnly[object]
|
||||
|
||||
|
||||
class _GenAIRequestFunctionCall(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
args: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _GenAIContentPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
inline_data: ReadOnly[Mapping[str, str]]
|
||||
functionResponse: ReadOnly[_GenAIFunctionResponse]
|
||||
functionCall: ReadOnly[_GenAIRequestFunctionCall]
|
||||
|
||||
|
||||
class _GenAIFunctionDeclaration(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
description: ReadOnly[str]
|
||||
parametersJsonSchema: ReadOnly[dict[str, object]]
|
||||
parametersJsonSchema: ReadOnly[object]
|
||||
|
||||
|
||||
class _GenAITool(TypedDict, total=False):
|
||||
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
|
||||
functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]]
|
||||
|
||||
|
||||
class _GenAIFunctionCallingConfig(TypedDict, total=False):
|
||||
|
|
@ -62,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False):
|
|||
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
|
||||
|
||||
|
||||
def _decode_tool_call_arguments(raw_arguments: str) -> object:
|
||||
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
|
||||
return json.loads(raw_arguments)
|
||||
class _GenAISystemInstruction(TypedDict, total=False):
|
||||
parts: ReadOnly[Sequence[Mapping[str, str]]]
|
||||
|
||||
|
||||
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
||||
|
|
@ -74,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
"""
|
||||
|
||||
sent_first_chunk: bool = False
|
||||
# State tracking for accumulating partial tool calls
|
||||
accumulated_tool_calls: dict[int, dict[str, str]]
|
||||
_parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
|
||||
|
||||
def __init__(self, completion_stream: object):
|
||||
self.sent_first_chunk = False
|
||||
self.accumulated_tool_calls = {}
|
||||
# State tracking for accumulating partial tool calls
|
||||
self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
|
||||
self._returned_response = False
|
||||
super().__init__(completion_stream)
|
||||
|
||||
|
|
@ -124,7 +154,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# After the stream is exhausted, check for any remaining accumulated tool calls
|
||||
if self.accumulated_tool_calls:
|
||||
try:
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
for (
|
||||
tool_call_index,
|
||||
tool_call_data,
|
||||
|
|
@ -132,7 +162,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
try:
|
||||
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
|
||||
# We default to an empty JSON object in this case.
|
||||
parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
|
||||
parsed_args: Mapping[str, object] = self._parse_accumulated_args(
|
||||
tool_call_data["arguments"] or "{}"
|
||||
)
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call_data["name"] or "undefined_tool_name",
|
||||
|
|
@ -149,7 +181,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
tool_call_data["arguments"],
|
||||
)
|
||||
if parts:
|
||||
final_chunk: Final[dict[str, object]] = {
|
||||
final_chunk: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -211,14 +243,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
class GoogleGenAIAdapter:
|
||||
"""Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
|
||||
|
||||
_parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def translate_generate_content_to_completion(
|
||||
self,
|
||||
model: str,
|
||||
contents: list[dict[str, Any]] | dict[str, Any],
|
||||
config: dict[str, Any] | None = None,
|
||||
contents: _JsonDictList | _JsonDict,
|
||||
config: Mapping[str, object] | None = None,
|
||||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
**kwargs,
|
||||
) -> dict[str, Any]:
|
||||
|
|
@ -250,7 +284,7 @@ class GoogleGenAIAdapter:
|
|||
messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction)
|
||||
|
||||
# Create base request as dict (which is compatible with ChatCompletionRequest)
|
||||
completion_request: Final[ChatCompletionRequest] = {
|
||||
completion_request: Final[_JsonDict] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
|
|
@ -312,9 +346,9 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _add_generic_litellm_params_to_request(
|
||||
self,
|
||||
completion_request_dict: dict[str, object],
|
||||
completion_request_dict: _JsonDict,
|
||||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
) -> dict[str, object]:
|
||||
) -> _JsonDict:
|
||||
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
|
||||
|
||||
Args:
|
||||
|
|
@ -326,7 +360,7 @@ class GoogleGenAIAdapter:
|
|||
"""
|
||||
allowed_fields: Final = GenericLiteLLMParams.model_fields.keys()
|
||||
if litellm_params:
|
||||
litellm_dict: Final = litellm_params.model_dump(exclude_none=True)
|
||||
litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True)
|
||||
for key, value in litellm_dict.items():
|
||||
if key in allowed_fields:
|
||||
completion_request_dict[key] = value
|
||||
|
|
@ -346,12 +380,12 @@ class GoogleGenAIAdapter:
|
|||
tools: Sequence[_GenAITool],
|
||||
) -> list[ChatCompletionToolParam]:
|
||||
"""Transform Google GenAI tools to OpenAI tools format"""
|
||||
openai_tools: Final[list[dict[str, object]]] = []
|
||||
openai_tools: Final = list[_JsonDict]()
|
||||
|
||||
for tool in tools:
|
||||
if "functionDeclarations" in tool:
|
||||
for func_decl in tool["functionDeclarations"]:
|
||||
function_chunk: dict[str, object] = {
|
||||
function_chunk: _JsonDict = {
|
||||
"name": func_decl.get("name", ""),
|
||||
}
|
||||
|
||||
|
|
@ -360,7 +394,7 @@ class GoogleGenAIAdapter:
|
|||
if "parametersJsonSchema" in func_decl:
|
||||
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
|
||||
|
||||
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
|
||||
openai_tool: _JsonDict = {"type": "function", "function": function_chunk}
|
||||
openai_tools.append(openai_tool)
|
||||
|
||||
# normalize the tool schemas
|
||||
|
|
@ -391,13 +425,13 @@ class GoogleGenAIAdapter:
|
|||
|
||||
# Handle system instruction
|
||||
if system_instruction:
|
||||
system_parts: Final = system_instruction.get("parts", [])
|
||||
system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", [])
|
||||
if system_parts and "text" in system_parts[0]:
|
||||
messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]))
|
||||
|
||||
for content in contents:
|
||||
role = content.get("role", "user")
|
||||
parts = content.get("parts", [])
|
||||
parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", [])
|
||||
|
||||
if role == "user":
|
||||
# Handle user messages with potential function responses
|
||||
|
|
@ -500,7 +534,7 @@ class GoogleGenAIAdapter:
|
|||
def translate_completion_to_generate_content(
|
||||
self,
|
||||
response: ModelResponse,
|
||||
) -> dict[str, object]:
|
||||
) -> _JsonDict:
|
||||
"""
|
||||
Transform litellm completion response to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -523,13 +557,13 @@ class GoogleGenAIAdapter:
|
|||
parts = self._transform_openai_message_to_google_genai_parts(choice.message)
|
||||
else:
|
||||
# Fallback for generic choice objects
|
||||
message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get(
|
||||
"content", ""
|
||||
)
|
||||
message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr(
|
||||
choice, "delta", _EMPTY_STR_MAPPING
|
||||
).get("content", "")
|
||||
parts = [{"text": message_content}] if message_content else []
|
||||
|
||||
# Create Google GenAI format response
|
||||
generate_content_response: Final[dict[str, object]] = {
|
||||
generate_content_response: Final[_JsonDict] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -563,7 +597,7 @@ class GoogleGenAIAdapter:
|
|||
self,
|
||||
response: ModelResponse | ModelResponseStream,
|
||||
wrapper: GoogleGenAIStreamWrapper,
|
||||
) -> dict[str, object] | None:
|
||||
) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Transform streaming litellm completion chunk to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -590,7 +624,7 @@ class GoogleGenAIAdapter:
|
|||
finish_reason: str | None = getattr(choice, "finish_reason", None)
|
||||
else:
|
||||
# Fallback for generic choice objects
|
||||
message_content: Final = getattr(choice, "delta", {}).get("content", "")
|
||||
message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "")
|
||||
parts = [{"text": message_content}] if message_content else []
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
|
||||
|
|
@ -599,7 +633,7 @@ class GoogleGenAIAdapter:
|
|||
return None
|
||||
|
||||
# Create Google GenAI streaming format response
|
||||
streaming_chunk: Final[dict[str, object]] = {
|
||||
streaming_chunk: Final[_JsonDict] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -635,10 +669,10 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _transform_openai_message_to_google_genai_parts(
|
||||
self,
|
||||
message: Any,
|
||||
) -> list[_GenAIPart]:
|
||||
message: Message,
|
||||
) -> Sequence[_GenAIPart]:
|
||||
"""Transform OpenAI message to Google GenAI parts format"""
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
|
||||
# Add text content if present
|
||||
if hasattr(message, "content") and message.content:
|
||||
|
|
@ -646,20 +680,22 @@ class GoogleGenAIAdapter:
|
|||
|
||||
# Add tool calls if present
|
||||
if hasattr(message, "tool_calls") and message.tool_calls:
|
||||
for tool_call in message.tool_calls:
|
||||
if hasattr(tool_call, "function") and tool_call.function:
|
||||
tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = (
|
||||
message.tool_calls
|
||||
)
|
||||
for tool_call in tool_calls:
|
||||
function: Function | None = getattr(tool_call, "function", None)
|
||||
if function:
|
||||
try:
|
||||
args = (
|
||||
_decode_tool_call_arguments(tool_call.function.arguments)
|
||||
if tool_call.function.arguments
|
||||
else {}
|
||||
args: Mapping[str, object] = (
|
||||
self._parse_tool_call_args(function.arguments) if function.arguments else {}
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call.function.name or "undefined_tool_name",
|
||||
"name": function.name or "undefined_tool_name",
|
||||
"args": args,
|
||||
}
|
||||
}
|
||||
|
|
@ -668,21 +704,23 @@ class GoogleGenAIAdapter:
|
|||
return parts if parts else [{"text": ""}]
|
||||
|
||||
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
|
||||
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
|
||||
) -> list[_GenAIPart]:
|
||||
self, delta: Delta, wrapper: GoogleGenAIStreamWrapper
|
||||
) -> Sequence[_GenAIPart]:
|
||||
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
|
||||
|
||||
# 1. Initialize wrapper state if it doesn't exist
|
||||
if not hasattr(wrapper, "accumulated_tool_calls"):
|
||||
wrapper.accumulated_tool_calls = {}
|
||||
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
|
||||
if hasattr(delta, "content") and delta.content:
|
||||
parts.append({"text": delta.content})
|
||||
|
||||
# 2. Ensure tool_calls is iterable
|
||||
tool_calls: Final = delta.tool_calls or []
|
||||
tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = (
|
||||
delta.tool_calls or []
|
||||
)
|
||||
|
||||
for tool_call in tool_calls:
|
||||
if not hasattr(tool_call, "function"):
|
||||
|
|
@ -701,19 +739,20 @@ class GoogleGenAIAdapter:
|
|||
}
|
||||
|
||||
# Accumulate name and arguments
|
||||
function_name = getattr(tool_call.function, "name", None)
|
||||
args_chunk = getattr(tool_call.function, "arguments", None)
|
||||
delta_function: Function | None = getattr(tool_call, "function", None)
|
||||
function_name: str | None = getattr(delta_function, "name", None)
|
||||
args_chunk: str | None = getattr(delta_function, "arguments", None)
|
||||
|
||||
# Optimization: Skip chunks that have no new data
|
||||
if not function_name and not args_chunk:
|
||||
verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index)
|
||||
continue
|
||||
|
||||
if function_name:
|
||||
wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name
|
||||
|
||||
if args_chunk:
|
||||
wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk
|
||||
previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index]
|
||||
wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator(
|
||||
name=function_name or previous_data["name"],
|
||||
arguments=previous_data["arguments"] + (args_chunk or ""),
|
||||
)
|
||||
|
||||
# Attempt to parse and emit a complete tool call
|
||||
accumulated_data = wrapper.accumulated_tool_calls[tool_call_index]
|
||||
|
|
@ -723,7 +762,7 @@ class GoogleGenAIAdapter:
|
|||
# 5. Attempt to parse arguments even if name hasn't arrived.
|
||||
try:
|
||||
# Attempt to parse the accumulated arguments string
|
||||
parsed_args = _decode_tool_call_arguments(accumulated_args)
|
||||
parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args)
|
||||
|
||||
# If parsing succeeds, but we don't have a name yet, wait.
|
||||
# The part will be created by a later chunk that brings the name.
|
||||
|
|
@ -757,7 +796,7 @@ class GoogleGenAIAdapter:
|
|||
|
||||
return mapping.get(finish_reason, "STOP")
|
||||
|
||||
def _map_usage(self, usage: Usage | None) -> dict[str, int]:
|
||||
def _map_usage(self, usage: object) -> Mapping[str, int]:
|
||||
"""Map OpenAI usage to Google GenAI usage format"""
|
||||
return {
|
||||
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,
|
||||
|
|
|
|||
|
|
@ -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,16 +162,12 @@ class BitBucketClient:
|
|||
response: Final = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
|
||||
data: Final = response.json()
|
||||
files: Final = []
|
||||
|
||||
for item in data.get("values", []):
|
||||
if item.get("type") == "commit_file":
|
||||
file_path = item.get("path", "")
|
||||
if file_path.endswith(file_extension):
|
||||
files.append(file_path)
|
||||
|
||||
return files
|
||||
data: Final[BitBucketSrcListing] = response.json()
|
||||
return [
|
||||
file_path
|
||||
for item in data.get("values", [])
|
||||
if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension)
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an HTTP error
|
||||
|
|
@ -162,7 +185,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 +214,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 +227,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.).
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ 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 collections.abc import Mapping, Sequence
|
||||
from typing import Any, ClassVar, Final, Protocol, cast
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.compression import compress
|
||||
|
|
@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
|
|||
_CACHE_TTL_SECONDS: Final = 15 * 60
|
||||
|
||||
|
||||
class _AgenticLoopParams(TypedDict, total=False):
|
||||
"""The ``agentic_loop_params`` entry the agentic loop driver records on the logging object."""
|
||||
|
||||
model: ReadOnly[str]
|
||||
|
||||
|
||||
class _AgenticLoopLoggingObj(Protocol):
|
||||
"""Logging object view exposing the untyped call details this handler reads."""
|
||||
|
||||
@property
|
||||
def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ...
|
||||
|
||||
|
||||
def _compression_savings_from_counts(
|
||||
original_tokens: object, compressed_tokens: object
|
||||
) -> CompressionSavingsMetadata | None:
|
||||
|
|
@ -80,7 +96,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
compression_trigger: int = 200_000,
|
||||
compression_target: int | None = None,
|
||||
embedding_model: str | None = None,
|
||||
embedding_model_params: dict[str, Any] | None = None,
|
||||
embedding_model_params: dict[str, object] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
|
|
@ -103,7 +119,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: dict[str, Any],
|
||||
callback_specific_params: dict[str, Any],
|
||||
callback_specific_params: Mapping[str, object],
|
||||
) -> "CompressionInterceptionLogger":
|
||||
compression_params: CompressionInterceptionConfig = {}
|
||||
if "compression_interception_params" in litellm_settings:
|
||||
|
|
@ -117,7 +133,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
)
|
||||
return CompressionInterceptionLogger.from_config_yaml(compression_params)
|
||||
|
||||
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, Any], call_type: CallTypes | None
|
||||
) -> dict[str, object] | None:
|
||||
if not self.enabled:
|
||||
return None
|
||||
if call_type is not None and call_type != CallTypes.anthropic_messages:
|
||||
|
|
@ -147,7 +165,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
cache: Final = cast(dict[str, str], compressed.get("cache", {}))
|
||||
skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason"))
|
||||
compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", []))
|
||||
compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", []))
|
||||
|
||||
# Only mutate kwargs when compression actually produced a result.
|
||||
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
|
||||
|
|
@ -158,7 +176,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
kwargs["messages"] = compressed["messages"]
|
||||
if compressed_tools:
|
||||
kwargs["tools"] = self._merge_tools(
|
||||
existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")),
|
||||
existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")),
|
||||
compressed_tools=compressed_tools,
|
||||
)
|
||||
call_id = cast(str | None, kwargs.get("litellm_call_id"))
|
||||
|
|
@ -191,14 +209,14 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
) -> tuple[bool, dict]:
|
||||
kwargs: Mapping[str, object],
|
||||
) -> tuple[bool, dict[str, object]]:
|
||||
if not self.enabled:
|
||||
return False, {}
|
||||
if not self._has_retrieval_tool(tools):
|
||||
|
|
@ -216,19 +234,19 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict,
|
||||
tools: Mapping[str, object],
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: Any,
|
||||
messages: list[dict[str, object]],
|
||||
response: object,
|
||||
anthropic_messages_provider_config: object,
|
||||
anthropic_messages_optional_request_params: Mapping[str, object],
|
||||
logging_obj: _AgenticLoopLoggingObj | None,
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> AgenticLoopPlan:
|
||||
self._prune_expired_cache()
|
||||
tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", []))
|
||||
thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", []))
|
||||
tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", []))
|
||||
thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", []))
|
||||
|
||||
call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
|
||||
cache: Final = self._get_cache(call_id=call_id)
|
||||
|
|
@ -271,7 +289,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
full_model_name = model
|
||||
if logging_obj is not None:
|
||||
agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
|
||||
full_model_name = cast(str, agentic_params.get("model", model))
|
||||
full_model_name = agentic_params.get("model", model)
|
||||
|
||||
request_patch: Final = AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
|
|
@ -306,15 +324,15 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
return {}
|
||||
return cache_entry[0]
|
||||
|
||||
def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None:
|
||||
def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None:
|
||||
if logging_obj is not None:
|
||||
logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None)
|
||||
if isinstance(logging_call_id, str) and logging_call_id:
|
||||
return logging_call_id
|
||||
kwargs_call_id: Final = kwargs.get("litellm_call_id")
|
||||
return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None)
|
||||
return kwargs_call_id if isinstance(kwargs_call_id, str) else None
|
||||
|
||||
def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str:
|
||||
def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str:
|
||||
raw_input: Final = tool_call.get("input", {})
|
||||
key = ""
|
||||
if isinstance(raw_input, dict):
|
||||
|
|
@ -325,7 +343,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
return cache[key]
|
||||
return f"[compressed content key '{key}' not found]"
|
||||
|
||||
def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
def _extract_retrieval_tool_calls(
|
||||
self, response: object
|
||||
) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
|
||||
if isinstance(response, dict):
|
||||
content = response.get("content", [])
|
||||
else:
|
||||
|
|
@ -334,8 +354,8 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
if not isinstance(content, list):
|
||||
return [], []
|
||||
|
||||
tool_calls: Final[list[dict[str, Any]]] = []
|
||||
thinking_blocks: Final[list[dict[str, Any]]] = []
|
||||
tool_calls: Final[list[dict[str, object]]] = []
|
||||
thinking_blocks: Final[list[dict[str, object]]] = []
|
||||
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
|
|
@ -382,13 +402,13 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
return tool_calls, thinking_blocks
|
||||
|
||||
def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
internal_keys: Final = {"litellm_logging_obj"}
|
||||
return {
|
||||
k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys
|
||||
}
|
||||
|
||||
def _has_retrieval_tool(self, tools: Any) -> bool:
|
||||
def _has_retrieval_tool(self, tools: object) -> bool:
|
||||
if not isinstance(tools, list):
|
||||
return False
|
||||
for tool in tools:
|
||||
|
|
@ -404,9 +424,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
def _merge_tools(
|
||||
self,
|
||||
existing_tools: list[dict[str, Any]] | None,
|
||||
compressed_tools: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
existing_tools: Sequence[Mapping[str, object]] | None,
|
||||
compressed_tools: Sequence[Mapping[str, object]],
|
||||
) -> list[Mapping[str, object]]:
|
||||
merged: Final = list(existing_tools or [])
|
||||
if self._has_retrieval_tool(merged):
|
||||
return merged
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
# On success, logs events to Promptlayer
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from collections.abc import AsyncGenerator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -123,11 +123,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
return []
|
||||
|
||||
callbacks: Final = AllCallbacks()
|
||||
callback_info: Final = getattr(callbacks, lookup_name, None)
|
||||
callback_info: Final[object] = getattr(callbacks, lookup_name, None)
|
||||
if callback_info is None:
|
||||
return []
|
||||
|
||||
params: Final = getattr(callback_info, "litellm_callback_params", None)
|
||||
params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None)
|
||||
if not params:
|
||||
return []
|
||||
|
||||
|
|
@ -851,7 +851,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
- Converting to string and then truncating the logged content catches this
|
||||
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
|
||||
"""
|
||||
field_value: Final = standard_logging_object.get(field_name)
|
||||
field_value: Final[object] = standard_logging_object.get(field_name)
|
||||
if field_value:
|
||||
str_value: Final = str(field_value)
|
||||
if len(str_value) > max_length:
|
||||
|
|
@ -1005,8 +1005,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
• Keep untyped or text content.
|
||||
• Recursively redact inline base64 blobs in *any* string field, at any depth.
|
||||
"""
|
||||
raw_messages: Final[Any] = payload.get("messages", [])
|
||||
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
raw_messages: Final[object] = payload.get("messages", [])
|
||||
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
|
||||
|
||||
if messages:
|
||||
|
|
@ -1037,8 +1037,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
• Keep untyped or text content.
|
||||
• Recursively redact inline base64 blobs in *any* string field, at any depth.
|
||||
"""
|
||||
raw_messages: Final[Any] = payload.get("messages", [])
|
||||
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
raw_messages: Final[object] = payload.get("messages", [])
|
||||
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
|
||||
|
||||
if messages:
|
||||
|
|
@ -1059,7 +1059,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
value: Any,
|
||||
depth: int = 0,
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""Recursively redact inline base64 from any nested structure with a max recursion depth limit."""
|
||||
if depth > max_depth:
|
||||
verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth)
|
||||
|
|
@ -1090,16 +1090,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
def _process_messages(
|
||||
self,
|
||||
messages: list[Any],
|
||||
messages: list[object],
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> list[dict[str, Any]]:
|
||||
filtered_messages: Final[list[dict[str, Any]]] = []
|
||||
) -> list[dict[str, object]]:
|
||||
filtered_messages: Final[list[dict[str, object]]] = []
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
contents: Any = msg.get("content")
|
||||
contents: object = msg.get("content")
|
||||
if isinstance(contents, list):
|
||||
cleaned: list[Any] = []
|
||||
cleaned: list[object] = []
|
||||
for c in contents:
|
||||
if self._should_keep_content(content=c):
|
||||
cleaned.append(self._redact_base64(value=c, max_depth=max_depth))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -17,12 +17,12 @@ def build_trace_payload(
|
|||
end_time: datetime,
|
||||
input_data: Any,
|
||||
output_data: Any,
|
||||
metadata: dict[str, Any],
|
||||
metadata: dict[str, object],
|
||||
tags: list[str],
|
||||
thread_id: str | None,
|
||||
) -> types.TracePayload:
|
||||
"""Build a complete trace payload."""
|
||||
trace_name: Final = response_obj.get("object", "unknown type")
|
||||
trace_name: Final[str] = response_obj.get("object", "unknown type")
|
||||
|
||||
return types.TracePayload(
|
||||
project_name=project_name,
|
||||
|
|
@ -47,7 +47,7 @@ def build_span_payload(
|
|||
end_time: datetime,
|
||||
input_data: Any,
|
||||
output_data: Any,
|
||||
metadata: dict[str, Any],
|
||||
metadata: dict[str, object],
|
||||
tags: list[str],
|
||||
usage: dict[str, int],
|
||||
provider: str | None = None,
|
||||
|
|
@ -56,9 +56,9 @@ def build_span_payload(
|
|||
"""Build a complete span payload."""
|
||||
span_id: Final = utils.create_uuid7()
|
||||
|
||||
model: Final = response_obj.get("model", "unknown-model")
|
||||
obj_type: Final = response_obj.get("object", "unknown-object")
|
||||
created: Final = response_obj.get("created", 0)
|
||||
model: Final[str] = response_obj.get("model", "unknown-model")
|
||||
obj_type: Final[str] = response_obj.get("object", "unknown-object")
|
||||
created: Final[int] = response_obj.get("created", 0)
|
||||
span_name: Final = f"{model}_{obj_type}_{created}"
|
||||
|
||||
_logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Provider / exporter factory + the Baggage span processor."""
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from opentelemetry import _logs, baggage, metrics
|
||||
from opentelemetry._events import EventLogger
|
||||
|
|
@ -135,14 +135,36 @@ def parse_headers(raw: str | None) -> dict[str, str]:
|
|||
return dict(parse_env_headers(raw, liberal=True))
|
||||
|
||||
|
||||
_IN_MEMORY_KINDS: Final = ("in_memory", "inmemory", "memory")
|
||||
_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", "http/json")
|
||||
_OTLP_GRPC_KINDS: Final = ("otlp_grpc", "grpc")
|
||||
|
||||
|
||||
def exporter_transport(kind: str) -> Literal["http", "grpc", "headerless"]:
|
||||
"""How an exporter of this ``kind`` carries credentials, per ``_exporter_from_spec``.
|
||||
|
||||
``http``/``grpc`` exporters (and any registered factory, which builds an
|
||||
OTLP exporter) stamp ``spec.headers``; ``console``, ``in_memory``, and any
|
||||
unrecognized kind (which falls back to a header-ignoring console exporter)
|
||||
are ``headerless``. Routability decisions must read this rather than a
|
||||
denylist, so a typo'd or unavailable kind is not mistaken for OTLP.
|
||||
"""
|
||||
resolved: Final = kind.lower()
|
||||
if resolved in _OTLP_HTTP_KINDS or resolved in _EXPORTER_FACTORIES:
|
||||
return "http"
|
||||
if resolved in _OTLP_GRPC_KINDS:
|
||||
return "grpc"
|
||||
return "headerless"
|
||||
|
||||
|
||||
def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
|
||||
kind: Final = (spec.kind or "console").lower()
|
||||
factory: Final = _EXPORTER_FACTORIES.get(kind)
|
||||
if factory is not None:
|
||||
return factory(spec)
|
||||
if kind in ("in_memory", "inmemory", "memory"):
|
||||
if kind in _IN_MEMORY_KINDS:
|
||||
return InMemorySpanExporter()
|
||||
if kind in ("otlp_http", "http", "http/protobuf", "http/json"):
|
||||
if kind in _OTLP_HTTP_KINDS:
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
|
||||
OTLPSpanExporter as HTTPExporter,
|
||||
)
|
||||
|
|
@ -151,7 +173,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
|
|||
endpoint=_otlp_traces_endpoint(spec.endpoint),
|
||||
headers=parse_headers(spec.headers),
|
||||
)
|
||||
if kind in ("otlp_grpc", "grpc"):
|
||||
if kind in _OTLP_GRPC_KINDS:
|
||||
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
|
||||
OTLPSpanExporter as GRPCExporter,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ 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,
|
||||
exporter_transport,
|
||||
get_tracer,
|
||||
)
|
||||
from litellm.integrations.otel.presets import (
|
||||
|
|
@ -121,13 +122,27 @@ def _encoded_header_string(headers: Mapping[str, str]) -> str:
|
|||
class TenantRoute:
|
||||
"""The tracer to create a span on, plus whether it must root its own trace.
|
||||
|
||||
``detached`` is True when project routing engaged. Phoenix assigns a whole
|
||||
``detached`` is True when the routed span exports to a DIFFERENT backend
|
||||
than the request's root span, which always exports through the default
|
||||
tracer. A detached span roots a fresh trace with a link back to the request
|
||||
trace for correlation, so the destination account is not left holding a
|
||||
child whose parent it never received. It is driven by whether routing
|
||||
headers were actually applied to an owned exporter, not merely requested:
|
||||
a credential or project route whose callback owns no exporter those headers
|
||||
can reach exports through the default backend unchanged, so it stays
|
||||
parented like an unrouted span.
|
||||
|
||||
Credential routing (a team/key's own vendor account) is one detaching case:
|
||||
the root, auth, and db spans stay on the operator's default backend while
|
||||
the LLM-call span exports to the tenant's account, so parenting it into the
|
||||
request trace makes the tenant account show a fragmented span with a missing
|
||||
parent. Project routing (Phoenix) is the other: Phoenix assigns a whole
|
||||
trace to one project by whichever of its spans arrives first, so a
|
||||
project-routed span parented into the request trace gets dragged into the
|
||||
project of the default-exported request spans and the header does nothing.
|
||||
The span must therefore start a fresh trace (with a link back to the
|
||||
request trace for correlation) — which is also how the v1 Phoenix logger
|
||||
behaved, exporting each request under its own Phoenix-local parent span.
|
||||
Both mirror the v1 loggers, which exported each request under its own
|
||||
backend-local root. Service-name routing does NOT detach: it relabels
|
||||
``service.name`` on the SAME operator backend, where the parent is present.
|
||||
"""
|
||||
|
||||
tracer: Tracer
|
||||
|
|
@ -161,11 +176,20 @@ class TenantTracerCache:
|
|||
self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state
|
||||
# Oldest-first so an overflow of draining providers sheds the stalest.
|
||||
self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers
|
||||
self._project_routable = any(
|
||||
spec.owner == callback_name and spec.kind.lower() not in (*_NON_OTLP_KINDS, *_GRPC_KINDS)
|
||||
for spec in config.exporters
|
||||
# An owned exporter is routable only when its kind actually resolves to a
|
||||
# header-carrying OTLP exporter. A denylist would accept a typo'd or
|
||||
# unavailable kind, which ``_exporter_from_spec`` falls back to a
|
||||
# header-ignoring console exporter: detaching such a span would strand it
|
||||
# on the operator's console, never reaching the tenant backend. Project
|
||||
# headers are HTTP-only; credentials ride gRPC metadata too (Arize's
|
||||
# default exporter is gRPC), so they accept either OTLP transport.
|
||||
owned_transports: Final = tuple(
|
||||
exporter_transport(spec.kind) for spec in config.exporters if spec.owner == callback_name
|
||||
)
|
||||
self._project_routable = "http" in owned_transports
|
||||
self._credential_routable = "http" in owned_transports or "grpc" in owned_transports
|
||||
self._warned_project_unroutable = False
|
||||
self._warned_credential_unroutable = False
|
||||
|
||||
def release(self, provider: TracerProvider | None) -> None:
|
||||
"""Drop one open-span count; shut a retired provider down once drained.
|
||||
|
|
@ -207,7 +231,7 @@ class TenantTracerCache:
|
|||
concurrent overflow eviction can't shut it down between selection and
|
||||
the caller's span start. The caller must ``release`` it exactly once.
|
||||
"""
|
||||
credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
|
||||
credential_headers: Final = self._credential_headers(dynamic_params)
|
||||
project_headers: Final = self._project_headers(auth_metadata)
|
||||
service_name: Final = tenant_service_name(auth_metadata)
|
||||
if not credential_headers and not project_headers and service_name is None:
|
||||
|
|
@ -231,7 +255,7 @@ class TenantTracerCache:
|
|||
_shutdown_provider(evicted)
|
||||
return TenantRoute(
|
||||
tracer=get_tracer(provider, self._tracer_name),
|
||||
detached=bool(project_headers),
|
||||
detached=bool(project_headers) or bool(credential_headers),
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
|
|
@ -275,6 +299,26 @@ class TenantTracerCache:
|
|||
self._open_span_counts.pop(overflowed, None)
|
||||
return overflowed
|
||||
|
||||
def _credential_headers(self, dynamic_params: StandardCallbackDynamicParams | None) -> Mapping[str, str]:
|
||||
"""The per-request dynamic OTLP credentials, if this cache can apply them.
|
||||
|
||||
A callback owning only a console/in_memory exporter has nowhere to stamp
|
||||
them, so the span would export to the operator's default backend
|
||||
unchanged; routing there and detaching would orphan it on the very
|
||||
backend that holds its parent. Warn once and keep the default tracer.
|
||||
"""
|
||||
requested: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
|
||||
if not requested or self._credential_routable:
|
||||
return requested
|
||||
if not self._warned_credential_unroutable:
|
||||
self._warned_credential_unroutable = True
|
||||
verbose_logger.warning(
|
||||
"OTel V2: %s request carries dynamic credentials, but the callback owns no "
|
||||
"OTLP exporter to stamp them onto; spans export to the default backend.",
|
||||
self._callback_name,
|
||||
)
|
||||
return _NO_HEADERS
|
||||
|
||||
def _project_headers(self, auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]:
|
||||
"""The per-request project-routing headers, if this cache can apply them.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -8,11 +8,12 @@ import uuid
|
|||
from collections import Counter
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Never, ReadOnly
|
||||
from typing_extensions import Never, ReadOnly, Required
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
Function,
|
||||
|
|
@ -52,17 +54,102 @@ _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
|
|||
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
|
||||
|
||||
|
||||
class _ServiceToolCall(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
class _ModerationToolCall(TypedDict, total=False):
|
||||
id: ReadOnly[Required[str]]
|
||||
|
||||
|
||||
class _ServiceMessage(TypedDict, total=False):
|
||||
class _ModerationMessage(TypedDict, total=False):
|
||||
content: ReadOnly[str | None]
|
||||
tool_calls: ReadOnly[Sequence[_ModerationToolCall] | None]
|
||||
|
||||
|
||||
class _ModerationChoice(TypedDict, total=False):
|
||||
message: ReadOnly[_ModerationMessage | None]
|
||||
|
||||
|
||||
class _ModerationResponse(TypedDict, total=False):
|
||||
choices: ReadOnly[Sequence[_ModerationChoice]]
|
||||
|
||||
|
||||
class _LogEventKwargs(TypedDict, total=False):
|
||||
standard_logging_object: ReadOnly[Required[StandardLoggingPayload]]
|
||||
litellm_call_id: ReadOnly[str]
|
||||
|
||||
|
||||
class _HasCallId(Protocol):
|
||||
def get(self, key: Literal["litellm_call_id"], /) -> str | None: ...
|
||||
|
||||
|
||||
class _HasModelAttr(Protocol):
|
||||
model: str | None
|
||||
|
||||
|
||||
class _ResponseSource(Protocol):
|
||||
def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ...
|
||||
|
||||
|
||||
class _ModelSource(Protocol):
|
||||
def get(self, key: Literal["model"], default: str, /) -> str: ...
|
||||
|
||||
|
||||
class _FallbackSource(Protocol):
|
||||
@overload
|
||||
def get(self, key: Literal["start_time"], /) -> datetime | None: ...
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
|
||||
|
||||
class _RequestContextSource(Protocol):
|
||||
@overload
|
||||
def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ...
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
def __contains__(self, key: object, /) -> bool: ...
|
||||
def __getitem__(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
class _ToolCallLike(Protocol):
|
||||
id: str | None
|
||||
type: str | None
|
||||
function: Function
|
||||
|
||||
|
||||
class _ModerationSourceToolCall(TypedDict, total=False):
|
||||
function: ReadOnly[Mapping[str, object] | None]
|
||||
|
||||
|
||||
class _ModerationSourceMessage(TypedDict, total=False):
|
||||
role: ReadOnly[str]
|
||||
function_call: ReadOnly[Mapping[str, object] | None]
|
||||
tool_calls: ReadOnly[Sequence[_ModerationSourceToolCall | None] | None]
|
||||
|
||||
|
||||
class _FlattenedModerationMessage(TypedDict):
|
||||
role: ReadOnly[str | None]
|
||||
content: ReadOnly[str]
|
||||
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
|
||||
|
||||
|
||||
class _ServiceChoice(TypedDict, total=False):
|
||||
message: ReadOnly[_ServiceMessage]
|
||||
class _CorrelatablePayload(TypedDict):
|
||||
id: str # writable-ok: _apply_correlation_id overwrites the provider id on a deep-copied payload
|
||||
|
||||
|
||||
class _SystemPromptCarrier(TypedDict, total=False):
|
||||
messages: object # writable-ok: _prepend_system_prompt rebinds messages on the copied payload by design
|
||||
|
||||
|
||||
class _BlockFailurePayload(TypedDict, total=False):
|
||||
id: object # writable-ok: correlation id is pinned after copying the base payload
|
||||
model: ReadOnly[object]
|
||||
model_group: ReadOnly[object]
|
||||
model_id: ReadOnly[str]
|
||||
model_parameters: ReadOnly[object]
|
||||
startTime: ReadOnly[float | None]
|
||||
endTime: ReadOnly[float | None]
|
||||
completionStartTime: ReadOnly[float | None]
|
||||
messages: object # writable-ok: passed to _prepend_system_prompt, which rebinds messages
|
||||
metadata: ReadOnly[StandardLoggingUserAPIKeyMetadata]
|
||||
response: str # writable-ok: block failure text replaces the copied response
|
||||
status: ReadOnly[str]
|
||||
|
||||
|
||||
class _MalformedToolBlockingResponseError(Exception):
|
||||
|
|
@ -385,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@staticmethod
|
||||
def _stash_block_context(
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
request_data: dict,
|
||||
request_data: dict[str, object],
|
||||
) -> None:
|
||||
"""Stash signals so the deferred success-event skips this request and
|
||||
``async_post_call_failure_hook`` can build the failure payload.
|
||||
|
|
@ -414,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
request_data["_rubrik_logging_obj"] = logging_obj
|
||||
|
||||
@staticmethod
|
||||
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
def _normalize_tool_calls(
|
||||
tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike],
|
||||
) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
|
||||
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall:
|
||||
def _normalize_tool_call(
|
||||
tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike,
|
||||
) -> ChatCompletionMessageToolCall:
|
||||
if isinstance(tc, ChatCompletionMessageToolCall):
|
||||
return tc
|
||||
if isinstance(tc, dict):
|
||||
|
|
@ -460,12 +551,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
``content`` is sent so the webhook can moderate the response text;
|
||||
``None`` when the assistant produced no text (tool-call-only response).
|
||||
"""
|
||||
message: Final[dict[str, object]] = {
|
||||
message: Final[Mapping[str, object]] = {
|
||||
"role": "assistant",
|
||||
"content": content or None,
|
||||
**(
|
||||
{"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)}
|
||||
if tool_calls
|
||||
else _EMPTY_MAPPING
|
||||
),
|
||||
}
|
||||
if tool_calls:
|
||||
message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)
|
||||
return {
|
||||
"id": request_id or f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
|
|
@ -481,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
|
||||
def _flatten_messages_for_moderation(
|
||||
messages: Sequence[AllMessageValues | None] | None,
|
||||
) -> tuple[_FlattenedModerationMessage, ...]:
|
||||
"""Collapse each message's content to a plain string for the webhook.
|
||||
|
||||
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
|
||||
|
|
@ -502,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]:
|
||||
def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]:
|
||||
"""Every attacker-controlled text segment of a message: its content plus
|
||||
the arguments of any tool call or deprecated function call."""
|
||||
fc: Final = message.get("function_call")
|
||||
|
|
@ -530,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
``/v1/messages`` requests too. Optional fields are sent only when
|
||||
present so the payload stays clean.
|
||||
"""
|
||||
payload: Final[dict[str, object]] = {
|
||||
"model": inputs.get("model") or request_data.get("model") or "",
|
||||
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
|
||||
}
|
||||
tools: Final = inputs.get("tools")
|
||||
if tools is not None:
|
||||
payload["tools"] = tools
|
||||
user: Final = request_data.get("user")
|
||||
if user:
|
||||
payload["user"] = user
|
||||
# Fall back to litellm_call_id, the stable cross-provider join key the
|
||||
# response/tool path uses (see _correlation_id). LiteLLM does not
|
||||
# populate request_data["correlation_key"]; it carries litellm_call_id.
|
||||
|
|
@ -547,14 +635,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
# when correlation_key is empty, so without this the block fires but no
|
||||
# log is ever written. An explicit correlation_key still wins.
|
||||
correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id")
|
||||
if correlation_key:
|
||||
payload["correlation_key"] = correlation_key
|
||||
return payload
|
||||
return {
|
||||
"model": inputs.get("model") or request_data.get("model") or "",
|
||||
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
|
||||
**({"tools": tools} if tools is not None else _EMPTY_MAPPING),
|
||||
**({"user": user} if user else _EMPTY_MAPPING),
|
||||
**({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_request_data(
|
||||
call_details: Mapping[str, Any],
|
||||
request_data: Mapping[str, object] | None,
|
||||
call_details: _RequestContextSource,
|
||||
request_data: _RequestContextSource | None,
|
||||
) -> Mapping[str, object]:
|
||||
"""Extract original request data from model_call_details for the
|
||||
response moderation service envelope.
|
||||
|
|
@ -590,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
|
||||
def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object:
|
||||
"""Allowlist only routing fields (``url``, ``method``) when forwarding
|
||||
``proxy_server_request`` to an external webhook, dropping inbound
|
||||
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
|
||||
|
|
@ -600,18 +692,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
|
||||
def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str:
|
||||
"""Get the model name for the ModifyResponseException."""
|
||||
response: Final = request_data.get("response")
|
||||
if response and hasattr(response, "model"):
|
||||
response_model: Final[str | None] = getattr(response, "model", None)
|
||||
return response_model or "unknown"
|
||||
return response.model or "unknown"
|
||||
return call_details.get("model", "unknown")
|
||||
|
||||
# -- Logging hooks ---------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
|
||||
def _correlation_id(
|
||||
call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None
|
||||
) -> str | None:
|
||||
"""The id that joins a blocked request's two S3 logs by filename: the
|
||||
moderation (``_blocking``) log and the failure (response) log.
|
||||
|
||||
|
|
@ -625,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
|
||||
|
||||
@classmethod
|
||||
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
|
||||
def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None:
|
||||
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
|
||||
shares its S3 filename id with the moderation (``_blocking``) and
|
||||
failure logs for the same request -- for every provider.
|
||||
|
|
@ -645,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
payload["id"] = correlated
|
||||
|
||||
@staticmethod
|
||||
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
|
||||
def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None:
|
||||
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
|
||||
|
||||
Builds a NEW messages list rather than mutating ``payload["messages"]``
|
||||
|
|
@ -673,9 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _prepare_log_payload(
|
||||
self, kwargs: Mapping[str, object], event_type: str
|
||||
) -> StandardLoggingPayload | None:
|
||||
async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None:
|
||||
"""Shared logic for success logging (sampled)."""
|
||||
if random.random() > self.sampling_rate:
|
||||
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
|
||||
|
|
@ -684,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
# Deep-copy so mutations don't affect other callbacks sharing this object
|
||||
standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"])
|
||||
|
||||
self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
|
||||
self._apply_correlation_id(standard_logging_payload, kwargs)
|
||||
self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
|
||||
|
||||
return standard_logging_payload
|
||||
|
||||
async def _append_and_maybe_flush(self, payload) -> None:
|
||||
async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None:
|
||||
self._ensure_periodic_flush_task()
|
||||
self.log_queue.append(payload)
|
||||
self._enforce_max_queue_size()
|
||||
|
|
@ -714,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
self._dropped_since_warning = 0
|
||||
self._last_drop_warning_time = now
|
||||
|
||||
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
|
||||
async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str):
|
||||
try:
|
||||
payload: Final = await self._prepare_log_payload(kwargs, event_type)
|
||||
if payload is None:
|
||||
|
|
@ -835,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
logging_obj: "LiteLLMLoggingObj",
|
||||
exception: "ModifyResponseException",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> StandardLoggingPayload:
|
||||
) -> _BlockFailurePayload:
|
||||
"""Build a failure-style payload using the exception text as response.
|
||||
|
||||
Blocked-tool events are security-relevant and **bypass sampling**:
|
||||
|
|
@ -877,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
call_details: Final = logging_obj.model_call_details
|
||||
exception_text: Final = f"{type(exception).__name__}: {exception.message}"
|
||||
|
||||
base: Final = call_details.get("standard_logging_object")
|
||||
base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object")
|
||||
if base is not None:
|
||||
payload: dict[str, object] = safe_deep_copy(base)
|
||||
payload: _BlockFailurePayload = self._copy_block_payload_base(base)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"Rubrik: standard_logging_object not yet on model_call_details "
|
||||
|
|
@ -901,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload:
|
||||
return safe_deep_copy(base)
|
||||
|
||||
@staticmethod
|
||||
def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata:
|
||||
"""Identify the caller whose request was blocked.
|
||||
|
|
@ -923,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@classmethod
|
||||
def _build_fallback_payload(
|
||||
cls,
|
||||
call_details: Mapping[str, Any],
|
||||
call_details: _FallbackSource,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> dict[str, object]:
|
||||
) -> _BlockFailurePayload:
|
||||
# Convert datetime to a Unix float so json.dumps can serialize it.
|
||||
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
|
||||
_raw_start: Final = call_details.get("start_time")
|
||||
|
|
@ -959,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
response: Final = await self.async_httpx_client.post(
|
||||
url=self.logging_endpoint,
|
||||
json=data,
|
||||
headers=self._headers,
|
||||
headers=dict(self._headers),
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
|
|
@ -1013,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
# -- Webhook services ------------------------------------------------------
|
||||
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse:
|
||||
"""POST ``payload`` to a Rubrik webhook and return its dict response.
|
||||
|
||||
Raises:
|
||||
|
|
@ -1023,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
verbose_logger.debug("Sending request to %s: %s", service_name, endpoint)
|
||||
http_response: Final = await self.moderation_client.post(
|
||||
endpoint,
|
||||
json=payload,
|
||||
headers=self._headers,
|
||||
json=dict(payload),
|
||||
headers=dict(self._headers),
|
||||
)
|
||||
http_response.raise_for_status()
|
||||
result: Final[object] = http_response.json()
|
||||
result: Final[_ModerationResponse | None] = http_response.json()
|
||||
if not isinstance(result, dict):
|
||||
raise TypeError(
|
||||
f"{service_name} returned non-dict JSON "
|
||||
|
|
@ -1040,7 +1135,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
self,
|
||||
response_data: Mapping[str, object],
|
||||
request_data: Mapping[str, object],
|
||||
) -> Mapping[str, Any]:
|
||||
) -> _ModerationResponse:
|
||||
"""Post the ``{request, response}`` envelope to the after_completion
|
||||
webhook and return its (possibly rewritten) response.
|
||||
|
||||
|
|
@ -1056,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
"Response moderation service",
|
||||
)
|
||||
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse:
|
||||
"""Post a bare OpenAI request to the before_prompt webhook.
|
||||
|
||||
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
|
||||
|
|
@ -1064,14 +1159,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service")
|
||||
|
||||
@staticmethod
|
||||
def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None:
|
||||
def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None:
|
||||
"""Return the refusal text when the prompt was blocked, else None.
|
||||
|
||||
The before_prompt webhook returns ``{}`` (passthrough) or a synthetic
|
||||
chat.completion whose ``choices[0].message.content`` is the refusal
|
||||
explanation.
|
||||
"""
|
||||
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
|
||||
choices: Final = service_response.get("choices")
|
||||
if not choices:
|
||||
return None
|
||||
message: Final = choices[0].get("message") or _EMPTY_MAPPING
|
||||
|
|
@ -1080,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
@staticmethod
|
||||
def _extract_response_block(
|
||||
service_response: Mapping[str, Any],
|
||||
service_response: _ModerationResponse,
|
||||
all_tool_calls: Sequence[ChatCompletionMessageToolCall],
|
||||
sent_content: str,
|
||||
) -> BlockedResponseResult | None:
|
||||
|
|
@ -1103,7 +1198,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
Expects service_response in OpenAI chat completion format:
|
||||
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
|
||||
"""
|
||||
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
|
||||
choices: Final = service_response.get("choices") or ()
|
||||
if not choices:
|
||||
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import asyncio
|
|||
import math
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
|
|
@ -46,7 +46,13 @@ from litellm.types.integrations.websearch_interception import (
|
|||
AnthropicServerToolUseBlock,
|
||||
WebSearchInterceptionConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.anthropic import AnthropicThinkingParam
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAudioParam,
|
||||
ChatCompletionPredictionContentParam,
|
||||
OpenAIWebSearchOptions,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
AgenticLoopParams,
|
||||
CallTypes,
|
||||
|
|
@ -56,6 +62,8 @@ from litellm.types.utils import (
|
|||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from aiohttp import ClientSession
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
|
|
@ -77,6 +85,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
|
|||
# ``web_search_tool_result`` blocks to inject into the final response.
|
||||
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
|
||||
|
||||
_RESPONSE_CONTENT_FIELD: Final = "content"
|
||||
|
||||
_ResponseT: Final = TypeVar("_ResponseT")
|
||||
|
||||
|
||||
class _PlanMetadataView(TypedDict):
|
||||
websearch_native_blocks: Sequence[Mapping[str, object]] | None
|
||||
|
|
@ -90,23 +102,98 @@ class _WebSearchSettingsView(TypedDict):
|
|||
websearch_interception_params: WebSearchInterceptionConfig
|
||||
|
||||
|
||||
class _SearchToolLitellmParams(TypedDict, total=False):
|
||||
search_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _SearchToolConfig(TypedDict, total=False):
|
||||
search_tool_name: str
|
||||
litellm_params: Mapping[str, object] | None
|
||||
litellm_params: ReadOnly[_SearchToolLitellmParams | None]
|
||||
|
||||
|
||||
class _DeploymentKwargsView(TypedDict):
|
||||
"""Typed reads of the untyped request kwargs seen by the deployment hook."""
|
||||
|
||||
class _LitellmParamsProviderView(TypedDict, total=False):
|
||||
custom_llm_provider: ReadOnly[str]
|
||||
litellm_params: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _DeploymentCallKwargsView(TypedDict):
|
||||
custom_llm_provider: ReadOnly[str]
|
||||
litellm_params: ReadOnly[_LitellmParamsProviderView]
|
||||
model: ReadOnly[str]
|
||||
|
||||
|
||||
class _UserAuthView(TypedDict):
|
||||
"""Typed read of the optional team attached to the caller's auth object."""
|
||||
class _AcreateNamedParams(TypedDict, total=False):
|
||||
metadata: ReadOnly[Never]
|
||||
stop_sequences: ReadOnly[Never]
|
||||
stream: ReadOnly[bool | None]
|
||||
system: ReadOnly[str | None]
|
||||
temperature: ReadOnly[float | None]
|
||||
thinking: ReadOnly[Never]
|
||||
tool_choice: ReadOnly[Never]
|
||||
tools: ReadOnly[Never]
|
||||
top_k: ReadOnly[int | None]
|
||||
top_p: ReadOnly[float | None]
|
||||
container: ReadOnly[Never]
|
||||
|
||||
team_id: ReadOnly[str | None]
|
||||
|
||||
class _AsearchNamedParams(TypedDict, total=False):
|
||||
max_results: ReadOnly[int | None]
|
||||
search_domain_filter: ReadOnly[Never]
|
||||
max_tokens_per_page: ReadOnly[int | None]
|
||||
country: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
api_base: ReadOnly[str | None]
|
||||
timeout: ReadOnly[float | None]
|
||||
extra_headers: ReadOnly[Never]
|
||||
|
||||
|
||||
class _AcompletionNamedParams(TypedDict, total=False):
|
||||
functions: ReadOnly[Never]
|
||||
function_call: ReadOnly[str | None]
|
||||
timeout: ReadOnly[float | None]
|
||||
temperature: ReadOnly[float | None]
|
||||
top_p: ReadOnly[float | None]
|
||||
n: ReadOnly[int | None]
|
||||
stream: ReadOnly[bool | None]
|
||||
stream_options: ReadOnly[Never]
|
||||
stop: ReadOnly[Never]
|
||||
max_tokens: ReadOnly[int | None]
|
||||
max_completion_tokens: ReadOnly[int | None]
|
||||
modalities: ReadOnly[Never]
|
||||
prediction: ReadOnly[ChatCompletionPredictionContentParam | None]
|
||||
audio: ReadOnly[ChatCompletionAudioParam | None]
|
||||
presence_penalty: ReadOnly[float | None]
|
||||
frequency_penalty: ReadOnly[float | None]
|
||||
logit_bias: ReadOnly[Never]
|
||||
user: ReadOnly[str | None]
|
||||
response_format: ReadOnly[Never]
|
||||
seed: ReadOnly[int | None]
|
||||
tools: ReadOnly[Never]
|
||||
tool_choice: ReadOnly[Never]
|
||||
parallel_tool_calls: ReadOnly[bool | None]
|
||||
logprobs: ReadOnly[bool | None]
|
||||
top_logprobs: ReadOnly[int | None]
|
||||
deployment_id: ReadOnly[str | None]
|
||||
reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None]
|
||||
verbosity: ReadOnly[Literal["low", "medium", "high"] | None]
|
||||
safety_identifier: ReadOnly[str | None]
|
||||
service_tier: ReadOnly[str | None]
|
||||
store: ReadOnly[bool | None]
|
||||
prompt_cache_key: ReadOnly[str | None]
|
||||
base_url: ReadOnly[str | None]
|
||||
api_version: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
model_list: ReadOnly[Never]
|
||||
extra_headers: ReadOnly[Never]
|
||||
thinking: ReadOnly[AnthropicThinkingParam | None]
|
||||
web_search_options: ReadOnly[OpenAIWebSearchOptions | None]
|
||||
include_server_side_tool_invocations: ReadOnly[bool | None]
|
||||
shared_session: ReadOnly["ClientSession | None"]
|
||||
enable_json_schema_validation: ReadOnly[bool | None]
|
||||
|
||||
|
||||
_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {}
|
||||
_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {}
|
||||
_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {}
|
||||
|
||||
|
||||
class WebSearchInterceptionLogger(CustomLogger):
|
||||
|
|
@ -308,17 +395,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"""
|
||||
# Check if this is for an enabled provider
|
||||
# Try top-level kwargs first, then nested litellm_params, then derive from model name
|
||||
kwargs_view: Final[_DeploymentKwargsView] = {
|
||||
call_kwargs_view: Final[_DeploymentCallKwargsView] = {
|
||||
"custom_llm_provider": kwargs.get("custom_llm_provider", ""),
|
||||
"litellm_params": kwargs.get("litellm_params", {}),
|
||||
"model": kwargs.get("model", ""),
|
||||
}
|
||||
custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get(
|
||||
custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get(
|
||||
"custom_llm_provider", ""
|
||||
)
|
||||
if not custom_llm_provider:
|
||||
try:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"])
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"])
|
||||
except Exception:
|
||||
custom_llm_provider = ""
|
||||
if custom_llm_provider not in self.enabled_providers:
|
||||
|
|
@ -948,17 +1035,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
|
||||
def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT:
|
||||
"""Prepend native blocks to response content, dict or object form."""
|
||||
if not native_blocks:
|
||||
return response
|
||||
if isinstance(response, dict):
|
||||
existing = response.get("content") or []
|
||||
response["content"] = list(native_blocks) + list(existing)
|
||||
existing = response.get(_RESPONSE_CONTENT_FIELD) or []
|
||||
response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing)
|
||||
return response
|
||||
existing = getattr(response, "content", None) or []
|
||||
existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or []
|
||||
try:
|
||||
response.content = list(native_blocks) + list(existing)
|
||||
setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing))
|
||||
except (AttributeError, TypeError):
|
||||
# Object refused write — fall through and leave the response
|
||||
# untouched rather than crash the request.
|
||||
|
|
@ -1214,10 +1301,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
messages: list[dict],
|
||||
tool_calls: list[dict],
|
||||
thinking_blocks: list[dict],
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
anthropic_messages_optional_request_params: Mapping[str, object],
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> "AnthropicMessagesResponse | AsyncIterator[object]":
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
request_patch, structured_results = await self._build_anthropic_request_patch(
|
||||
|
|
@ -1225,9 +1312,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
thinking_blocks=thinking_blocks,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params),
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
kwargs=dict[str, object](kwargs),
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
raise ValueError("WebSearchInterception: missing follow-up messages")
|
||||
|
|
@ -1242,12 +1329,14 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if max_tokens is None:
|
||||
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
|
||||
|
||||
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
|
||||
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
|
||||
max_tokens=max_tokens,
|
||||
messages=request_patch.messages,
|
||||
model=request_patch.model or model,
|
||||
**_NO_ACREATE_NAMED,
|
||||
**optional_params,
|
||||
**request_patch.kwargs,
|
||||
**patch_kwargs,
|
||||
)
|
||||
|
||||
# Legacy path: the new path goes through the typed plan + core
|
||||
|
|
@ -1389,12 +1478,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
|
||||
search_provider: str | None = None
|
||||
search_litellm_params: dict[str, Any] = {}
|
||||
search_litellm_params: Mapping[str, object] = {}
|
||||
search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
|
||||
if search_tool is not None:
|
||||
await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
|
||||
search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
|
||||
search_provider = search_litellm_params.get("search_provider")
|
||||
tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {}
|
||||
search_litellm_params = dict[str, object](tool_params)
|
||||
search_provider = tool_params.get("search_provider")
|
||||
|
||||
# Fallback to perplexity if no router or no search tools configured
|
||||
if not search_provider:
|
||||
|
|
@ -1422,12 +1512,15 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if key != "search_provider" and value is not None
|
||||
}
|
||||
result: Final = (
|
||||
await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
|
||||
await litellm.asearch(
|
||||
query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
|
||||
)
|
||||
if search_metadata is None
|
||||
else await litellm.asearch(
|
||||
query=query,
|
||||
search_provider=search_provider,
|
||||
litellm_metadata=search_metadata,
|
||||
**_NO_ASEARCH_NAMED,
|
||||
**search_kwargs,
|
||||
)
|
||||
)
|
||||
|
|
@ -1467,8 +1560,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
valid_token=user_api_key_auth,
|
||||
)
|
||||
|
||||
auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)}
|
||||
team_id: Final = auth_view["team_id"]
|
||||
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
|
||||
if team_id:
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
|
|
@ -1583,10 +1675,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
model: str,
|
||||
messages: list[dict],
|
||||
tool_calls: list[dict],
|
||||
optional_params: dict,
|
||||
optional_params: Mapping[str, object],
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
response_format: str = "openai",
|
||||
) -> "ModelResponse | CustomStreamWrapper":
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
|
|
@ -1594,8 +1686,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
optional_params=dict[str, object](optional_params),
|
||||
kwargs=dict[str, object](kwargs),
|
||||
response_format=response_format,
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
|
|
@ -1603,11 +1695,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
params: Final = dict(optional_params)
|
||||
params.update(request_patch.optional_params)
|
||||
params.pop("tool_choice", None)
|
||||
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
|
||||
return await litellm.acompletion(
|
||||
model=request_patch.model or model,
|
||||
messages=request_patch.messages,
|
||||
**_NO_ACOMPLETION_NAMED,
|
||||
**params,
|
||||
**request_patch.kwargs,
|
||||
**patch_kwargs,
|
||||
)
|
||||
|
||||
async def _build_chat_completion_request_patch(
|
||||
|
|
|
|||
|
|
@ -1,19 +1,36 @@
|
|||
"""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, chain
|
||||
from itertools import accumulate, groupby
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
CUE_MAX_TOKENS: Final = 15
|
||||
CUE_MAX_DURATION_MS: Final = 5000
|
||||
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:
|
||||
|
|
@ -31,69 +48,138 @@ class SubtitleCue:
|
|||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CueAccumulator:
|
||||
texts: tuple[str, ...] = ()
|
||||
start_ms: int | None = None
|
||||
end_ms: int | None = None
|
||||
speaker: str | int | None = None
|
||||
class _Word:
|
||||
text: str
|
||||
start_ms: int | None
|
||||
end_ms: int | None
|
||||
speaker: str | int | None
|
||||
|
||||
|
||||
def _completed_cue(accumulator: _CueAccumulator) -> tuple[SubtitleCue, ...]:
|
||||
if not accumulator.texts or accumulator.start_ms is None:
|
||||
return ()
|
||||
text: Final = "".join(accumulator.texts).strip()
|
||||
if not text:
|
||||
return ()
|
||||
end_ms: Final = accumulator.end_ms if accumulator.end_ms is not None else accumulator.start_ms
|
||||
return (SubtitleCue(start_ms=accumulator.start_ms, end_ms=end_ms, text=text),)
|
||||
def _is_cjk(ch: str) -> bool:
|
||||
cp: Final = ord(ch)
|
||||
return any(lo <= cp <= hi for lo, hi in _CJK_RANGES)
|
||||
|
||||
|
||||
def _cue_break_reached(accumulator: _CueAccumulator, token: SubtitleToken) -> bool:
|
||||
if len(accumulator.texts) >= CUE_MAX_TOKENS:
|
||||
return True
|
||||
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 (
|
||||
accumulator.start_ms is not None
|
||||
and token.start_ms is not None
|
||||
and token.start_ms - accumulator.start_ms >= CUE_MAX_DURATION_MS
|
||||
first.isspace()
|
||||
or prev_last.isspace()
|
||||
or token.speaker != prev.speaker
|
||||
or _is_cjk_word_boundary(prev_last, first)
|
||||
)
|
||||
|
||||
|
||||
_AbsorbStep = tuple[tuple[SubtitleCue, ...], _CueAccumulator]
|
||||
|
||||
|
||||
def _absorb_token(accumulator: _CueAccumulator, token: SubtitleToken) -> _AbsorbStep:
|
||||
if token.start_ms is None and accumulator.start_ms is None:
|
||||
return (), accumulator
|
||||
if token.speaker is not None and token.speaker != accumulator.speaker:
|
||||
return _completed_cue(accumulator), _CueAccumulator(
|
||||
texts=(token.text,),
|
||||
start_ms=token.start_ms,
|
||||
end_ms=token.end_ms,
|
||||
speaker=token.speaker,
|
||||
)
|
||||
if _cue_break_reached(accumulator, token):
|
||||
return _completed_cue(accumulator), _CueAccumulator(
|
||||
texts=(token.text,),
|
||||
start_ms=token.start_ms,
|
||||
end_ms=token.end_ms,
|
||||
speaker=accumulator.speaker,
|
||||
)
|
||||
return (), _CueAccumulator(
|
||||
texts=(*accumulator.texts, token.text),
|
||||
start_ms=accumulator.start_ms if accumulator.start_ms is not None else token.start_ms,
|
||||
end_ms=token.end_ms if token.end_ms is not None else accumulator.end_ms,
|
||||
speaker=accumulator.speaker,
|
||||
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 _absorb_step(carry: _AbsorbStep, token: SubtitleToken) -> _AbsorbStep:
|
||||
return _absorb_token(carry[1], token)
|
||||
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, ...]:
|
||||
steps: Final = tuple(accumulate(tokens, _absorb_step, initial=((), _CueAccumulator())))
|
||||
completed: Final = chain.from_iterable(emitted for emitted, _ in steps)
|
||||
return (*completed, *_completed_cue(steps[-1][1]))
|
||||
"""
|
||||
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:
|
||||
|
|
|
|||
55
litellm/litellm_core_utils/aws_partition.py
Normal file
55
litellm/litellm_core_utils/aws_partition.py
Normal 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
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -9,8 +10,11 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
|
|||
DEFAULT_MAX_AGENTIC_LOOPS,
|
||||
validated_max_agentic_loops,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
HEADROOM_CONVERTED_STREAM_KEY,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
|
|
@ -50,6 +54,12 @@ def _post_hook_overridden(callback: CustomLogger) -> bool:
|
|||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool:
|
||||
return bool(
|
||||
kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY)
|
||||
)
|
||||
|
||||
|
||||
def _coerce_int(value: object, default: int) -> int:
|
||||
return int(value) if isinstance(value, (int, str)) else default
|
||||
|
||||
|
|
@ -87,16 +97,24 @@ def _check_agentic_loop_safety(
|
|||
return fingerprint
|
||||
|
||||
|
||||
def _wrap_response_as_fake_stream(response: object) -> object:
|
||||
if getattr(response, "object", None) == "chat.completion.chunk":
|
||||
def _wrap_response_as_fake_stream(
|
||||
response: object,
|
||||
*,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: object,
|
||||
) -> object:
|
||||
if isinstance(response, CustomStreamWrapper):
|
||||
return response
|
||||
if not hasattr(response, "choices"):
|
||||
if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject):
|
||||
return response
|
||||
from litellm.llms.base_llm.base_model_iterator import (
|
||||
convert_model_response_to_streaming,
|
||||
)
|
||||
|
||||
return convert_model_response_to_streaming(cast(ModelResponse, response))
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=MockResponseIterator(model_response=response),
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
|
||||
def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
|
||||
|
|
@ -177,8 +195,13 @@ async def _execute_chat_completion_agentic_plan(
|
|||
model,
|
||||
str(e),
|
||||
)
|
||||
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth:
|
||||
return _wrap_response_as_fake_stream(response_followup)
|
||||
if _converted_stream_requested(kwargs) and not depth:
|
||||
return _wrap_response_as_fake_stream(
|
||||
response_followup,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return response_followup
|
||||
finally:
|
||||
try:
|
||||
|
|
@ -302,9 +325,14 @@ async def maybe_run_chat_completion_agentic_loop(
|
|||
str(e),
|
||||
)
|
||||
|
||||
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"):
|
||||
if _converted_stream_requested(kwargs) and not depth:
|
||||
return cast(
|
||||
"ModelResponse | CustomStreamWrapper",
|
||||
_wrap_response_as_fake_stream(response),
|
||||
_wrap_response_as_fake_stream(
|
||||
response,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
),
|
||||
)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -25,6 +25,13 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, Inter
|
|||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups"
|
||||
"""Where auth records the model access groups that authorized the request, for the spend writer.
|
||||
|
||||
The ``user_api_key`` prefix is load-bearing, not cosmetic: when a request carries both
|
||||
``metadata`` and ``litellm_metadata``, ``get_litellm_metadata_from_kwargs`` returns the latter and
|
||||
copies a key across only when ``user_api_key`` appears in its name."""
|
||||
|
||||
_USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth"
|
||||
|
||||
FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset(
|
||||
|
|
|
|||
|
|
@ -64,7 +64,10 @@ 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.internal_call_metadata import (
|
||||
MODEL_ACCESS_GROUP_METADATA_KEY,
|
||||
is_unbilled_non_inference_call,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
|
||||
cost_breakdown_with_guardrail,
|
||||
guardrail_information_cost,
|
||||
|
|
@ -544,6 +547,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
# Init Caching related details
|
||||
self.caching_details: CachingDetails | None = None
|
||||
# Timing for results that cannot carry ``_hidden_params`` (plain-dict /v1/messages
|
||||
# responses and the bridge stream wrappers); see ``update_response_metadata``.
|
||||
self.response_timing_metrics: Mapping[str, float] = {} # mutable-ok: kept deep-copyable
|
||||
|
||||
# Passthrough endpoint guardrails config for field targeting
|
||||
self.passthrough_guardrails_config: dict[str, Any] | None = None
|
||||
|
|
@ -563,6 +569,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._defer_async_logging: bool = False
|
||||
self._enqueue_deferred_logging: Callable[[], None] | None = None
|
||||
|
||||
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
|
||||
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
|
||||
self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable
|
||||
|
||||
def process_dynamic_callbacks(self):
|
||||
"""
|
||||
Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks
|
||||
|
|
@ -2872,6 +2882,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
batch_cost: Final = kwargs.get("batch_cost", None)
|
||||
batch_usage = kwargs.get("batch_usage", None)
|
||||
batch_models = kwargs.get("batch_models", None)
|
||||
batch_successful_requests: Final = kwargs.get("batch_successful_requests", None)
|
||||
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
|
||||
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
|
||||
|
||||
should_compute_batch_data: Final = (
|
||||
|
|
@ -2880,14 +2892,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if has_explicit_batch_data:
|
||||
result._hidden_params["response_cost"] = batch_cost
|
||||
result._hidden_params["batch_models"] = batch_models
|
||||
result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above
|
||||
result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
|
||||
result.usage = batch_usage
|
||||
|
||||
elif should_compute_batch_data:
|
||||
(
|
||||
response_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
) = await _handle_completed_batch(
|
||||
batch_result: Final = await _handle_completed_batch(
|
||||
batch=result,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
model_name=self.get_deployment_model_for_cost(),
|
||||
|
|
@ -2895,9 +2905,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model_info=self.get_router_deployment_model_info(),
|
||||
)
|
||||
|
||||
result._hidden_params["response_cost"] = response_cost
|
||||
result._hidden_params["batch_models"] = batch_models
|
||||
result.usage = batch_usage
|
||||
result._hidden_params["response_cost"] = batch_result.cost
|
||||
result._hidden_params["batch_models"] = batch_result.models
|
||||
result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
|
||||
result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
|
||||
result.usage = batch_result.usage
|
||||
|
||||
start_time, end_time, result = self._success_handler_helper_fn(
|
||||
start_time=start_time,
|
||||
|
|
@ -5049,6 +5061,42 @@ def is_valid_sha256_hash(value: str) -> bool:
|
|||
return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value))
|
||||
|
||||
|
||||
def coerce_model_access_groups(value: object) -> tuple[str, ...]:
|
||||
"""Model access group names out of untrusted request metadata, deduped and order preserving."""
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return ()
|
||||
return tuple(dict.fromkeys(group for group in value if isinstance(group, str) and group))
|
||||
|
||||
|
||||
def _model_access_groups_on_auth_object(user_api_key_auth: object) -> object:
|
||||
if isinstance(user_api_key_auth, Mapping):
|
||||
return user_api_key_auth.get("matched_model_access_groups")
|
||||
return getattr(user_api_key_auth, "matched_model_access_groups", None)
|
||||
|
||||
|
||||
def _model_access_groups_from_metadata(metadata: Mapping[str, object]) -> tuple[str, ...]:
|
||||
stamped: Final = coerce_model_access_groups(metadata.get(MODEL_ACCESS_GROUP_METADATA_KEY))
|
||||
if stamped:
|
||||
return stamped
|
||||
return coerce_model_access_groups(_model_access_groups_on_auth_object(metadata.get("user_api_key_auth")))
|
||||
|
||||
|
||||
def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, object]) -> tuple[str, ...]:
|
||||
"""Access groups the auth layer stamped onto this request, from whichever metadata field carries them.
|
||||
|
||||
Detached internal sub-calls only inherit the identity keys, so the auth object is the
|
||||
fallback there, exactly as _get_budget_reservation_from_metadata does for reservations.
|
||||
"""
|
||||
for metadata_variable_name in ("metadata", "litellm_metadata"):
|
||||
metadata = litellm_params.get(metadata_variable_name)
|
||||
if not isinstance(metadata, Mapping):
|
||||
continue
|
||||
model_access_groups = _model_access_groups_from_metadata(metadata)
|
||||
if model_access_groups:
|
||||
return model_access_groups
|
||||
return ()
|
||||
|
||||
|
||||
class StandardLoggingPayloadSetup:
|
||||
@staticmethod
|
||||
def cleanup_timestamps(
|
||||
|
|
@ -5422,6 +5470,8 @@ class StandardLoggingPayloadSetup:
|
|||
additional_headers=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
@ -5812,6 +5862,8 @@ def _extract_response_obj_and_hidden_params(
|
|||
response_cost=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
@ -5896,6 +5948,7 @@ def get_standard_logging_object_payload(
|
|||
request_tags: Final = StandardLoggingPayloadSetup._get_request_tags(
|
||||
litellm_params=litellm_params, proxy_server_request=proxy_server_request
|
||||
)
|
||||
request_model_access_groups: Final = request_model_access_groups_from_litellm_params(litellm_params)
|
||||
|
||||
# cleanup timestamps
|
||||
(
|
||||
|
|
@ -5959,6 +6012,13 @@ def get_standard_logging_object_payload(
|
|||
clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params)
|
||||
if clean_hidden_params["response_cost"] is None and raw_response_cost is not None:
|
||||
clean_hidden_params["response_cost"] = llm_response_cost
|
||||
if clean_hidden_params["litellm_overhead_time_ms"] is None and status == "success":
|
||||
# /v1/messages dict results and the bridge stream wrappers keep it on the logging object;
|
||||
# failure payloads stay None like every response type that carries its own _hidden_params
|
||||
timing_metrics: Final = (
|
||||
getattr(logging_obj, "response_timing_metrics", None) or {} # mutable-ok: empty fallback
|
||||
)
|
||||
clean_hidden_params["litellm_overhead_time_ms"] = timing_metrics.get("litellm_overhead_time_ms")
|
||||
|
||||
model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information(
|
||||
base_model=base_model,
|
||||
|
|
@ -6058,6 +6118,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,
|
||||
request_model_access_groups=request_model_access_groups,
|
||||
end_user=end_user_id,
|
||||
api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "",
|
||||
model_group=_model_group,
|
||||
|
|
@ -6228,6 +6289,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
|
|||
additional_headers=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
@ -6269,6 +6332,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
|
|||
cache_key=None,
|
||||
saved_cache_cost=saved_cache_cost,
|
||||
request_tags=[],
|
||||
request_model_access_groups=(),
|
||||
end_user=None,
|
||||
requester_ip_address="127.0.0.1",
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -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, ...]], ...]:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
import datetime
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import LITELLM_DETAILED_TIMING
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
|
|
@ -13,6 +16,39 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
def response_timing_metrics(
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
include_overhead: bool = True,
|
||||
) -> Mapping[str, float]:
|
||||
"""``_response_ms`` for the whole call, plus ``litellm_overhead_time_ms`` when it can be derived.
|
||||
|
||||
On a cache hit the overhead is the total minus the cache read; otherwise it is the total minus
|
||||
the provider call (``llm_api_duration_ms``). It is omitted when neither duration was recorded,
|
||||
and when ``include_overhead`` is False because the two durations cover different windows.
|
||||
"""
|
||||
total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000
|
||||
if not include_overhead:
|
||||
return {"_response_ms": total_response_time_ms} # mutable-ok: read-only timing result
|
||||
caching_details: Final = logging_obj.caching_details
|
||||
cache_duration_ms: Final = (
|
||||
caching_details.get("cache_duration_ms")
|
||||
if caching_details is not None and caching_details.get("cache_hit") is True
|
||||
else None
|
||||
)
|
||||
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
|
||||
if cache_duration_ms is not None:
|
||||
overhead_ms: float | None = total_response_time_ms - cache_duration_ms
|
||||
elif llm_api_duration_ms is not None:
|
||||
overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4)
|
||||
else:
|
||||
overhead_ms = None
|
||||
if overhead_ms is None:
|
||||
return {"_response_ms": total_response_time_ms}
|
||||
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
|
||||
|
||||
|
||||
class ResponseMetadata:
|
||||
"""
|
||||
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
|
||||
|
|
@ -25,11 +61,7 @@ class ResponseMetadata:
|
|||
@property
|
||||
def supports_response_time(self) -> bool:
|
||||
"""Check if response type supports timing metrics"""
|
||||
return (
|
||||
isinstance(self.result, ModelResponse)
|
||||
or isinstance(self.result, EmbeddingResponse)
|
||||
or isinstance(self.result, TranscriptionResponse)
|
||||
)
|
||||
return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse))
|
||||
|
||||
def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None:
|
||||
"""Set hidden parameters on the response"""
|
||||
|
|
@ -45,14 +77,14 @@ class ResponseMetadata:
|
|||
result=self.result, litellm_model_name=model, router_model_id=model_id
|
||||
),
|
||||
"additional_headers": process_response_headers(
|
||||
self._get_value_from_hidden_params("additional_headers") or {},
|
||||
self._get_additional_headers_from_hidden_params() or {},
|
||||
preserve_litellm_internal_headers=True,
|
||||
),
|
||||
"litellm_model_name": model,
|
||||
}
|
||||
self._update_hidden_params(new_params)
|
||||
|
||||
def _update_hidden_params(self, new_params: dict) -> None:
|
||||
def _update_hidden_params(self, new_params: Mapping[str, object]) -> None:
|
||||
"""
|
||||
Update hidden params - handles when self._hidden_params is a dict or HiddenParams object
|
||||
"""
|
||||
|
|
@ -64,51 +96,38 @@ class ResponseMetadata:
|
|||
for key, value in new_params.items():
|
||||
setattr(self._hidden_params, key, value)
|
||||
|
||||
def _get_value_from_hidden_params(self, key: str) -> Any | None:
|
||||
"""Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
|
||||
def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None:
|
||||
"""Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
|
||||
if isinstance(self._hidden_params, dict):
|
||||
return self._hidden_params.get(key, None)
|
||||
return self._hidden_params.get("additional_headers", None)
|
||||
elif isinstance(self._hidden_params, HiddenParams):
|
||||
return getattr(self._hidden_params, key, None)
|
||||
return getattr(self._hidden_params, "additional_headers", None)
|
||||
|
||||
def set_timing_metrics(
|
||||
self,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
include_overhead: bool = True,
|
||||
) -> None:
|
||||
"""Set response timing metrics"""
|
||||
total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000
|
||||
timing_metrics: Final = response_timing_metrics(start_time, end_time, logging_obj, include_overhead)
|
||||
total_response_time_ms: Final = timing_metrics["_response_ms"]
|
||||
|
||||
# Set total response time if supported
|
||||
if self.supports_response_time:
|
||||
self.result._response_ms = total_response_time_ms
|
||||
|
||||
#########################################################
|
||||
# 1. Add _response_ms total duration
|
||||
# 1. Add _response_ms total duration and the LiteLLM overhead within it
|
||||
# (total minus the cache read on a cache hit, else total minus the provider call)
|
||||
#########################################################
|
||||
self._update_hidden_params(
|
||||
{
|
||||
"_response_ms": total_response_time_ms,
|
||||
}
|
||||
)
|
||||
self._update_hidden_params(timing_metrics)
|
||||
|
||||
#########################################################
|
||||
# 2. Add LiteLLM overhead duration
|
||||
# 2. Add callback processing duration
|
||||
#########################################################
|
||||
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
|
||||
if llm_api_duration_ms is not None:
|
||||
overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4)
|
||||
self._update_hidden_params(
|
||||
{
|
||||
"litellm_overhead_time_ms": overhead_ms,
|
||||
}
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# 3. Add callback processing duration
|
||||
#########################################################
|
||||
callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None)
|
||||
callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None)
|
||||
if callback_duration_ms is not None:
|
||||
self._update_hidden_params(
|
||||
{
|
||||
|
|
@ -117,36 +136,21 @@ class ResponseMetadata:
|
|||
)
|
||||
|
||||
#########################################################
|
||||
# 4. Add duration for reading from cache
|
||||
# In this case overhead from litellm is the difference between the cache read duration and the total response time
|
||||
#########################################################
|
||||
if (
|
||||
logging_obj.caching_details is not None
|
||||
and logging_obj.caching_details.get("cache_hit") is True
|
||||
and (cache_duration_ms := logging_obj.caching_details.get("cache_duration_ms")) is not None
|
||||
):
|
||||
overhead_ms = total_response_time_ms - cache_duration_ms
|
||||
self._update_hidden_params(
|
||||
{
|
||||
"litellm_overhead_time_ms": overhead_ms,
|
||||
}
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# 5. Detailed per-phase timing (opt-in via env var)
|
||||
# 3. Detailed per-phase timing (opt-in via env var)
|
||||
#########################################################
|
||||
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
|
||||
if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None:
|
||||
detailed: Final[dict] = {
|
||||
detailed: Final[dict[str, float]] = {
|
||||
"timing_llm_api_ms": round(llm_api_duration_ms, 4),
|
||||
}
|
||||
|
||||
# message copy time from Logging.__init__()
|
||||
msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None)
|
||||
msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None)
|
||||
if msg_copy_ms is not None:
|
||||
detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4)
|
||||
|
||||
# pre-processing = time from request start to LLM API call start
|
||||
api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time")
|
||||
api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time")
|
||||
if api_call_start is not None and start_time is not None:
|
||||
pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000
|
||||
detailed["timing_pre_processing_ms"] = round(pre_ms, 4)
|
||||
|
|
@ -170,6 +174,7 @@ def update_response_metadata(
|
|||
kwargs: dict,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
include_overhead: bool = True,
|
||||
) -> None:
|
||||
"""
|
||||
Updates response metadata including hidden params and timing metrics
|
||||
|
|
@ -177,11 +182,22 @@ def update_response_metadata(
|
|||
- response._hidden_params
|
||||
- response._hidden_params["litellm_overhead_time_ms"]
|
||||
- response.response_time_ms
|
||||
A result that cannot hold ``_hidden_params`` gets its timing on ``logging_obj`` instead.
|
||||
Callers whose ``end_time`` covers more than the recorded provider call (a stream read to
|
||||
completion) pass ``include_overhead=False``, since the overhead cannot be derived there.
|
||||
"""
|
||||
if result is None or not hasattr(result, "_hidden_params"):
|
||||
if result is None:
|
||||
return
|
||||
if not hasattr(result, "_hidden_params"):
|
||||
# /v1/messages returns a plain dict and the Anthropic / Responses bridge stream wrappers
|
||||
# cannot hold ``_hidden_params``: keep only the timing on the logging object (no cost
|
||||
# recompute) so the proxy headers and the standard logging payload can still read it.
|
||||
logging_obj.set_response_timing_metrics(
|
||||
response_timing_metrics(start_time, end_time, logging_obj, include_overhead)
|
||||
)
|
||||
return
|
||||
|
||||
metadata: Final = ResponseMetadata(result)
|
||||
metadata.set_hidden_params(logging_obj, model, kwargs)
|
||||
metadata.set_timing_metrics(start_time, end_time, logging_obj)
|
||||
metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead)
|
||||
metadata.apply()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from collections.abc import Iterable, Mapping, Sequence
|
|||
from itertools import groupby
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, cast
|
||||
|
||||
from openai.types.chat.chat_completion_custom_tool_param import (
|
||||
|
|
@ -1089,6 +1090,162 @@ def sanitize_input_schema_for_anthropic(input_schema: dict) -> "AnthropicInputSc
|
|||
return AnthropicInputSchema(**filtered)
|
||||
|
||||
|
||||
_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("allOf", "anyOf", "oneOf")
|
||||
_OPENAI_REJECTED_TOP_LEVEL_SCHEMA_KEYS: Final = ("enum", "const", "not")
|
||||
_LOCAL_SCHEMA_REF_PREFIXES: Final = (("#/$defs/", "$defs"), ("#/definitions/", "definitions"))
|
||||
_MAX_SCHEMA_FLATTEN_DEPTH: Final = 32
|
||||
_EMPTY_SCHEMA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _schema_properties(schema: Mapping[str, object]) -> Mapping[str, object]:
|
||||
properties: Final = schema.get("properties")
|
||||
return properties if isinstance(properties, dict) else _EMPTY_SCHEMA
|
||||
|
||||
|
||||
def _schema_branches(schema: Mapping[str, object], combinator: str) -> tuple[object, ...]:
|
||||
branches: Final = schema.get(combinator)
|
||||
return tuple(branches) if isinstance(branches, list) else ()
|
||||
|
||||
|
||||
def _schema_required_names(schema: Mapping[str, object]) -> frozenset[str]:
|
||||
required: Final = schema.get("required")
|
||||
if not isinstance(required, list):
|
||||
return frozenset()
|
||||
return frozenset(name for name in required if isinstance(name, str))
|
||||
|
||||
|
||||
def _combinator_required_names(combinator: str, branches: tuple[Mapping[str, object], ...]) -> frozenset[str]:
|
||||
branch_names: Final = tuple(_schema_required_names(branch) for branch in branches)
|
||||
if not branch_names:
|
||||
return frozenset()
|
||||
if combinator == "allOf":
|
||||
return branch_names[0].union(*branch_names[1:])
|
||||
return branch_names[0].intersection(*branch_names[1:])
|
||||
|
||||
|
||||
def _resolve_local_schema_ref(root: Mapping[str, object], ref: str) -> Mapping[str, object] | None:
|
||||
matched: Final = next(
|
||||
((prefix, container) for prefix, container in _LOCAL_SCHEMA_REF_PREFIXES if ref.startswith(prefix)),
|
||||
None,
|
||||
)
|
||||
if matched is None:
|
||||
return None
|
||||
prefix, container = matched
|
||||
definitions: Final = root.get(container)
|
||||
if not isinstance(definitions, dict):
|
||||
return None
|
||||
target: Final = definitions.get(ref[len(prefix) :])
|
||||
return target if isinstance(target, dict) else None
|
||||
|
||||
|
||||
def _mergeable_branch(
|
||||
root: Mapping[str, object],
|
||||
branch: object,
|
||||
seen_refs: frozenset[str],
|
||||
depth: int,
|
||||
expanded_refs: dict[str, Mapping[str, object] | None], # mutable-ok: per-call memo bounding repeated $ref work
|
||||
) -> Mapping[str, object] | None:
|
||||
if not isinstance(branch, dict) or depth > _MAX_SCHEMA_FLATTEN_DEPTH:
|
||||
return None
|
||||
ref: Final = branch.get("$ref")
|
||||
if not isinstance(ref, str):
|
||||
flattened: Final = _flatten_schema_against_root(branch, root, seen_refs, depth, expanded_refs)
|
||||
if any(combinator in flattened for combinator in _TOP_LEVEL_SCHEMA_COMBINATORS):
|
||||
return None
|
||||
return flattened
|
||||
if ref in expanded_refs:
|
||||
return expanded_refs[ref]
|
||||
if ref in seen_refs:
|
||||
return None
|
||||
target: Final = _resolve_local_schema_ref(root, ref)
|
||||
expanded: Final = (
|
||||
None
|
||||
if target is None
|
||||
else _mergeable_branch(root, target, seen_refs | frozenset((ref,)), depth + 1, expanded_refs)
|
||||
)
|
||||
expanded_refs[ref] = expanded
|
||||
return expanded
|
||||
|
||||
|
||||
def _is_object_schema(schema: Mapping[str, object]) -> bool:
|
||||
return schema.get("type") == "object" or ("type" not in schema and "properties" in schema)
|
||||
|
||||
|
||||
def _flatten_schema_against_root(
|
||||
schema: Mapping[str, object],
|
||||
root: Mapping[str, object],
|
||||
seen_refs: frozenset[str],
|
||||
depth: int,
|
||||
expanded_refs: dict[str, Mapping[str, object] | None], # mutable-ok: per-call memo bounding repeated $ref work
|
||||
) -> Mapping[str, object]:
|
||||
raw_branch_groups: Final = tuple(
|
||||
(
|
||||
combinator,
|
||||
tuple(
|
||||
_mergeable_branch(root, branch, seen_refs, depth + 1, expanded_refs)
|
||||
for branch in _schema_branches(schema, combinator)
|
||||
),
|
||||
)
|
||||
for combinator in _TOP_LEVEL_SCHEMA_COMBINATORS
|
||||
if isinstance(schema.get(combinator), list)
|
||||
)
|
||||
dropped: Final = (
|
||||
*(combinator for combinator, _ in raw_branch_groups),
|
||||
*(key for key in _OPENAI_REJECTED_TOP_LEVEL_SCHEMA_KEYS if key in schema),
|
||||
)
|
||||
if not dropped:
|
||||
return schema
|
||||
|
||||
if any(branch is None for _, group in raw_branch_groups for branch in group):
|
||||
return schema
|
||||
branch_groups: Final = tuple(
|
||||
(combinator, tuple(branch for branch in group if branch is not None)) for combinator, group in raw_branch_groups
|
||||
)
|
||||
branches: Final = tuple(branch for _, group in branch_groups for branch in group)
|
||||
is_object_schema: Final = _is_object_schema(schema) or (
|
||||
"type" not in schema and branches != () and all(_is_object_schema(branch) for branch in branches)
|
||||
)
|
||||
if not is_object_schema:
|
||||
return schema
|
||||
|
||||
merged_properties: Final = { # mutable-ok: tool parameters are JSON dicts
|
||||
name: value for source in (*reversed(branches), schema) for name, value in _schema_properties(source).items()
|
||||
}
|
||||
required_names: Final = _schema_required_names(schema).union(
|
||||
*(_combinator_required_names(combinator, group) for combinator, group in branch_groups)
|
||||
)
|
||||
kept: Final = MappingProxyType({key: value for key, value in schema.items() if key not in dropped})
|
||||
required_update: Final = MappingProxyType({"required": sorted(required_names)}) if required_names else _EMPTY_SCHEMA
|
||||
return { # mutable-ok: tool parameters are JSON dicts
|
||||
**kept,
|
||||
"type": "object",
|
||||
"properties": merged_properties,
|
||||
**required_update,
|
||||
}
|
||||
|
||||
|
||||
def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Merge top-level ``allOf``/``anyOf``/``oneOf`` branches into an object tool schema.
|
||||
|
||||
OpenAI's function-calling validator rejects tool ``parameters`` carrying
|
||||
'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level (nested uses
|
||||
are accepted), while lenient backends such as the ChatGPT backend Codex
|
||||
talks to natively accept them, so an MCP tool declaring a top-level union
|
||||
400s through LiteLLM. Branch properties merge without clobbering (the
|
||||
top-level schema wins, then earlier branches); ``required`` becomes the
|
||||
top-level list plus the intersection of the branch lists for anyOf/oneOf
|
||||
or their union for allOf. Branches that are local ``$ref``s
|
||||
(``#/$defs/...`` or ``#/definitions/...``) are resolved first, each ref
|
||||
at most once per call, and branches that are themselves combinators are
|
||||
flattened recursively up to a fixed depth; a branch that cannot be fully
|
||||
merged (a boolean schema, an external or cyclic ``$ref``, a non-object
|
||||
union, or nesting past the depth cap) leaves the whole schema untouched so
|
||||
OpenAI's own validation still applies. Non-object schemas pass through
|
||||
unchanged and the input is never mutated.
|
||||
"""
|
||||
return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo
|
||||
|
||||
|
||||
def _get_image_mime_type_from_url(url: str) -> str | None:
|
||||
"""
|
||||
Get mime type for common image URLs
|
||||
|
|
|
|||
|
|
@ -93,7 +93,7 @@ def print_verbose(print_statement: object):
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkParsed:
|
||||
response_obj: dict[str, Any]
|
||||
response_obj: dict[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -1288,7 +1288,7 @@ class CustomStreamWrapper:
|
|||
for key, value in anthropic_response_obj["provider_specific_fields"].items():
|
||||
setattr(model_response, key, value)
|
||||
|
||||
response_obj = cast(dict[str, Any], anthropic_response_obj)
|
||||
response_obj = cast(dict[str, object], anthropic_response_obj)
|
||||
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
|
||||
response_obj = self.handle_replicate_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
|
|
@ -1444,7 +1444,7 @@ class CustomStreamWrapper:
|
|||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"chunk is not a string: {chunk}")
|
||||
response_obj = cast(
|
||||
dict[str, Any],
|
||||
dict[str, object],
|
||||
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
|
||||
)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
|
|
@ -2551,7 +2551,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
|
|||
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
latest_usage_chunk = None
|
||||
latest_usage_chunk: Usage | Mapping[str, int] | None = None
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
|
||||
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
|
||||
cache_creation_token_details: CacheCreationTokenDetails | None = None
|
||||
|
|
|
|||
|
|
@ -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,16 +407,17 @@ 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
|
||||
return str(httpx.URL(request_url).join(location))
|
||||
|
||||
|
||||
def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
|
||||
def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
|
||||
"""
|
||||
Fetch a user-supplied URL with SSRF protection on every redirect hop.
|
||||
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -412,18 +460,21 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
|
|||
raise SSRFError("Too many redirects")
|
||||
|
||||
|
||||
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
|
||||
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
|
||||
"""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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
||||
import httpx
|
||||
from httpx import Headers, Response
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
|
||||
|
|
@ -12,6 +14,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
|
||||
|
|
@ -19,6 +23,29 @@ else:
|
|||
LoggingClass = Any
|
||||
|
||||
|
||||
class AnthropicBatchRequestCounts(TypedDict, total=False):
|
||||
"""The ``request_counts`` object of an Anthropic Message Batch."""
|
||||
|
||||
processing: ReadOnly[int]
|
||||
succeeded: ReadOnly[int]
|
||||
errored: ReadOnly[int]
|
||||
canceled: ReadOnly[int]
|
||||
expired: ReadOnly[int]
|
||||
|
||||
|
||||
class AnthropicMessageBatch(TypedDict, total=False):
|
||||
"""The fields of an Anthropic Message Batch that map onto an OpenAI Batch."""
|
||||
|
||||
id: ReadOnly[str]
|
||||
processing_status: ReadOnly[str]
|
||||
created_at: ReadOnly[str | None]
|
||||
ended_at: ReadOnly[str | None]
|
||||
expires_at: ReadOnly[str | None]
|
||||
cancel_initiated_at: ReadOnly[str | None]
|
||||
archived_at: ReadOnly[str | None]
|
||||
request_counts: ReadOnly[AnthropicBatchRequestCounts]
|
||||
|
||||
|
||||
class AnthropicBatchesConfig(BaseBatchesConfig):
|
||||
def __init__(self):
|
||||
from ..chat.transformation import AnthropicConfig
|
||||
|
|
@ -83,7 +110,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
create_batch_data: CreateBatchRequest,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> bytes | str | dict[str, Any]:
|
||||
) -> bytes | str | dict[str, object]:
|
||||
"""
|
||||
Transform the batch creation request to Anthropic format.
|
||||
|
||||
|
|
@ -133,7 +160,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
batch_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> bytes | str | dict[str, Any]:
|
||||
) -> bytes | str | dict[str, object]:
|
||||
"""
|
||||
Transform batch retrieval request for Anthropic.
|
||||
|
||||
|
|
@ -152,7 +179,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
) -> LiteLLMBatch:
|
||||
"""Transform Anthropic MessageBatch retrieval response to LiteLLM format."""
|
||||
try:
|
||||
response_data: Final = raw_response.json()
|
||||
response_data: Final[AnthropicMessageBatch] = raw_response.json()
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to parse Anthropic batch response: {e}")
|
||||
|
||||
|
|
@ -161,18 +188,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
processing_status: Final = response_data.get("processing_status", "in_progress")
|
||||
|
||||
# Map Anthropic processing_status to OpenAI status
|
||||
status_mapping: dict[
|
||||
str,
|
||||
Literal[
|
||||
"validating",
|
||||
"failed",
|
||||
"in_progress",
|
||||
"finalizing",
|
||||
"completed",
|
||||
"expired",
|
||||
"cancelling",
|
||||
"cancelled",
|
||||
],
|
||||
status_mapping: Final[
|
||||
Mapping[
|
||||
str,
|
||||
Literal[
|
||||
"validating",
|
||||
"failed",
|
||||
"in_progress",
|
||||
"finalizing",
|
||||
"completed",
|
||||
"expired",
|
||||
"cancelling",
|
||||
"cancelled",
|
||||
],
|
||||
]
|
||||
] = {
|
||||
"in_progress": "in_progress",
|
||||
"canceling": "cancelling",
|
||||
|
|
@ -261,7 +290,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:
|
||||
|
|
@ -279,7 +308,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
|
|||
if not line:
|
||||
continue
|
||||
try:
|
||||
response_json = json.loads(line)
|
||||
response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line)
|
||||
# Update model_response with the parsed JSON
|
||||
completion_response = response_json["result"]["message"]
|
||||
transformed_response = self.anthropic_chat_config.transform_parsed_response(
|
||||
|
|
|
|||
|
|
@ -16,9 +16,9 @@ import json
|
|||
from collections.abc import Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
|
||||
|
||||
from typing_extensions import assert_never
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
|
@ -58,6 +58,8 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
|
|
@ -98,6 +100,48 @@ InputWriteBackTarget = (
|
|||
)
|
||||
|
||||
|
||||
class _SSEDelta(TypedDict, total=False):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
stop_reason: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _SSEEventData(TypedDict, total=False):
|
||||
delta: ReadOnly[_SSEDelta]
|
||||
|
||||
|
||||
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return value
|
||||
|
||||
|
||||
def _content_block_at(blocks: Sequence[object], index: int) -> object:
|
||||
return blocks[index]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _ModelDumpBlock(Protocol):
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _TextAttrBlock(Protocol):
|
||||
text: str
|
||||
|
||||
|
||||
class _WritableMessage(Protocol):
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
|
||||
@overload
|
||||
def get(self, key: str, default: object, /) -> object: ...
|
||||
|
||||
def __setitem__(self, key: str, value: object, /) -> None: ...
|
||||
|
||||
|
||||
def _as_writable(value: _WritableMessage) -> _WritableMessage:
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScannedText:
|
||||
text: str
|
||||
|
|
@ -126,7 +170,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
@staticmethod
|
||||
def _build_streaming_usage_response(
|
||||
responses_so_far: list[object],
|
||||
responses_so_far: Sequence[object],
|
||||
request_data: dict | None,
|
||||
) -> ModelResponse | None:
|
||||
chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes)))
|
||||
|
|
@ -144,7 +188,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: list[object] | None = None,
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> list[bytes]:
|
||||
"""
|
||||
Build an Anthropic SSE sequence delivering the guardrail block message
|
||||
|
|
@ -162,9 +206,22 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
would make Anthropic clients reject the stream.
|
||||
"""
|
||||
if stream_started:
|
||||
return self._block_continuation_chunks(exc, responses_so_far or [])
|
||||
return list(self._block_continuation_chunks(exc, responses_so_far or []))
|
||||
return self._standalone_block_chunks(exc)
|
||||
|
||||
def build_stream_error_items(
|
||||
self,
|
||||
exc: "HTTPException",
|
||||
responses_so_far: Sequence[Any] | None = None,
|
||||
) -> Sequence[Any] | None:
|
||||
from litellm.proxy.common_request_processing import (
|
||||
serialize_http_exception_detail,
|
||||
)
|
||||
from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames
|
||||
|
||||
message, _ = serialize_http_exception_detail(exc.detail)
|
||||
return tuple(anthropic_sse_error_frames(message))
|
||||
|
||||
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
|
||||
import uuid
|
||||
|
||||
|
|
@ -187,7 +244,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
|
||||
|
||||
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]:
|
||||
def _block_continuation_chunks(
|
||||
self, exc: "ModifyResponseException", responses_so_far: Sequence[object]
|
||||
) -> Sequence[bytes]:
|
||||
"""Continue an already-started message: close the open content block,
|
||||
append the block message as a new text block, then end the message --
|
||||
without a second message_start."""
|
||||
|
|
@ -199,7 +258,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
def _sse(event_type: str, payload: dict) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None))["output_tokens"]
|
||||
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None)).get("output_tokens", 0)
|
||||
open_index, max_index = self._content_block_state(responses_so_far)
|
||||
new_index: Final = (max_index + 1) if max_index is not None else 0
|
||||
chunks: list[bytes] = []
|
||||
|
|
@ -237,7 +296,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
@staticmethod
|
||||
def _content_block_state(
|
||||
responses_so_far: list[object],
|
||||
responses_so_far: Sequence[object],
|
||||
) -> tuple[int | None, int | None]:
|
||||
"""From the SSE chunks already sent to the client, return (open
|
||||
content-block index or None, highest content-block index seen or None).
|
||||
|
|
@ -263,7 +322,20 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
return open_index, max_index
|
||||
|
||||
@staticmethod
|
||||
def _iter_sse_events(item: object) -> list[dict[str, object]]:
|
||||
def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]:
|
||||
line: Final = raw_line.strip()
|
||||
if not line.startswith("data:"):
|
||||
return ()
|
||||
try:
|
||||
parsed: Final[object] = json.loads(line[len("data:") :].strip())
|
||||
except json.JSONDecodeError:
|
||||
return ()
|
||||
if not isinstance(parsed, dict):
|
||||
return ()
|
||||
return (_as_str_mapping(parsed),)
|
||||
|
||||
@staticmethod
|
||||
def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]:
|
||||
"""Yield the event-data dicts in one stream chunk.
|
||||
|
||||
Handles both formats this stream can carry (see
|
||||
|
|
@ -271,24 +343,15 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
several events separated by a blank line -- and an already-parsed event
|
||||
``dict``."""
|
||||
if isinstance(item, dict):
|
||||
return [item]
|
||||
return (_as_str_mapping(item),)
|
||||
if not isinstance(item, (bytes, bytearray)):
|
||||
return []
|
||||
events: Final[list[dict[str, object]]] = []
|
||||
for block in item.decode("utf-8", errors="replace").split("\n\n"):
|
||||
for line in block.split("\n"):
|
||||
line = line.strip()
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
try:
|
||||
parsed: str | int | float | bool | None | Sequence[object] | Mapping[str, object] = json.loads(
|
||||
line[len("data:") :].strip()
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(parsed, dict):
|
||||
events.append(parsed)
|
||||
return events
|
||||
return ()
|
||||
return tuple(
|
||||
event
|
||||
for block in item.decode("utf-8", errors="replace").split("\n\n")
|
||||
for line in block.split("\n")
|
||||
for event in AnthropicMessagesHandler._parse_sse_data_line(line)
|
||||
)
|
||||
|
||||
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
|
||||
"""Translate Anthropic request to OpenAI chat completion format."""
|
||||
|
|
@ -321,7 +384,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
data: dict,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> Any:
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Process input messages by applying guardrails to text content.
|
||||
"""
|
||||
|
|
@ -481,7 +544,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
@staticmethod
|
||||
def _openai_system_message_to_anthropic(
|
||||
message: dict[str, object],
|
||||
message: Mapping[str, object],
|
||||
) -> dict[str, object] | None: # mutable-ok: API message payload
|
||||
"""Convert an OpenAI system message to the client's Anthropic-shaped entry."""
|
||||
content: Final = message.get("content")
|
||||
|
|
@ -561,7 +624,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
@staticmethod
|
||||
def _defer_systems_inside_tool_exchanges(
|
||||
structured_messages: list, # mutable-ok: API message payload
|
||||
structured_messages: Sequence[Mapping[str, object]],
|
||||
) -> list:
|
||||
"""Hold a system row until the tool exchange around it completes so the call/result pair converts together."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
|
||||
|
|
@ -755,7 +818,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if scan_only_tool_results:
|
||||
return EMPTY_EXTRACTED_INPUT
|
||||
|
||||
text_str: Final = content_item.get("text", None)
|
||||
text_str: Final[str | None] = content_item.get("text")
|
||||
return ExtractedInput(
|
||||
scanned=(
|
||||
() if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),)
|
||||
|
|
@ -805,7 +868,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
async def _apply_guardrail_responses_to_input(
|
||||
self,
|
||||
messages: list[dict[str, object]],
|
||||
messages: Sequence[_WritableMessage],
|
||||
responses: list[str],
|
||||
scanned: tuple[ScannedText, ...],
|
||||
) -> None:
|
||||
|
|
@ -931,7 +994,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
) -> list[Any]:
|
||||
) -> Sequence[object]:
|
||||
"""
|
||||
Process output streaming response by applying guardrails to text content.
|
||||
|
||||
|
|
@ -1027,7 +1090,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
return request_data
|
||||
|
||||
@staticmethod
|
||||
def _get_response_content(response: object) -> list[Any]:
|
||||
def _get_response_content(response: object) -> Sequence[object]:
|
||||
"""Extract content list from a dict or object response."""
|
||||
if isinstance(response, dict):
|
||||
return response.get("content", []) or []
|
||||
|
|
@ -1037,7 +1100,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
def _extract_from_content_blocks(
|
||||
self,
|
||||
response_content: list[Any],
|
||||
response_content: Sequence[object],
|
||||
texts_to_check: list[str],
|
||||
images_to_check: list[str],
|
||||
task_mappings: list[tuple[int, int | None]],
|
||||
|
|
@ -1045,21 +1108,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
) -> None:
|
||||
"""Extract text, images, and tool calls from content blocks."""
|
||||
for content_idx, content_block in enumerate(response_content):
|
||||
block_dict: dict[str, object] = {}
|
||||
if isinstance(content_block, dict):
|
||||
block_type = content_block.get("type")
|
||||
block_dict = cast(dict[str, object], content_block)
|
||||
elif hasattr(content_block, "type"):
|
||||
block_type = getattr(content_block, "type", None)
|
||||
if hasattr(content_block, "model_dump"):
|
||||
block_dict = content_block.model_dump()
|
||||
else:
|
||||
block_dict = {
|
||||
"type": block_type,
|
||||
"text": getattr(content_block, "text", None),
|
||||
}
|
||||
else:
|
||||
fields = self._output_block_fields(content_block)
|
||||
if fields is None:
|
||||
continue
|
||||
block_type, block_dict = fields
|
||||
|
||||
if block_type in ["text", "tool_use"]:
|
||||
self._extract_output_text_and_images(
|
||||
|
|
@ -1071,6 +1123,21 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
tool_calls_to_check=tool_calls_to_check,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None":
|
||||
if isinstance(content_block, dict):
|
||||
block_dict: Final = _as_str_mapping(content_block)
|
||||
return block_dict.get("type"), block_dict
|
||||
if not hasattr(content_block, "type"):
|
||||
return None
|
||||
block_type: Final = getattr(content_block, "type", None)
|
||||
if isinstance(content_block, _ModelDumpBlock):
|
||||
return block_type, content_block.model_dump()
|
||||
return block_type, {
|
||||
"type": block_type,
|
||||
"text": getattr(content_block, "text", None),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _build_guardrail_inputs(
|
||||
texts_to_check: list[str],
|
||||
|
|
@ -1093,7 +1160,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
inputs["model"] = response_model
|
||||
return inputs
|
||||
|
||||
def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str:
|
||||
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
|
||||
"""
|
||||
Parse streaming responses and extract accumulated text content.
|
||||
|
||||
|
|
@ -1164,7 +1231,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
# Only process content_block_delta events
|
||||
if event_type == "content_block_delta" and data_line:
|
||||
try:
|
||||
data = json.loads(data_line)
|
||||
data: _SSEEventData = json.loads(data_line)
|
||||
delta = data.get("delta", {})
|
||||
if delta.get("type") == "text_delta":
|
||||
text += delta.get("text", "")
|
||||
|
|
@ -1176,7 +1243,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
return text
|
||||
|
||||
def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool:
|
||||
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
|
||||
"""
|
||||
Check if streaming response has ended by looking for non-null stop_reason.
|
||||
|
||||
|
|
@ -1227,7 +1294,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
# Check for message_delta event with stop_reason
|
||||
if event_type == "message_delta" and data_line:
|
||||
try:
|
||||
data = json.loads(data_line)
|
||||
data: _SSEEventData = json.loads(data_line)
|
||||
delta = data.get("delta", {})
|
||||
stop_reason = delta.get("stop_reason")
|
||||
if stop_reason is not None:
|
||||
|
|
@ -1271,7 +1338,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
def _extract_output_text_and_images(
|
||||
self,
|
||||
content_block: dict[str, object],
|
||||
content_block: Mapping[str, object],
|
||||
content_idx: int,
|
||||
texts_to_check: list[str],
|
||||
images_to_check: list[str],
|
||||
|
|
@ -1294,7 +1361,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
task_mappings.append((content_idx, None))
|
||||
|
||||
# Extract tool calls
|
||||
elif content_type == "tool_use":
|
||||
elif content_type == "tool_use" and isinstance(content_block, dict):
|
||||
tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format(
|
||||
anthropic_tool_content=content_block,
|
||||
index=content_idx,
|
||||
|
|
@ -1319,7 +1386,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
content_idx = cast(int, mapping[0])
|
||||
|
||||
# Handle both dict and object responses
|
||||
response_content: list[Any] = []
|
||||
response_content: Sequence[object] = []
|
||||
if isinstance(response, dict):
|
||||
response_content = response.get("content", []) or []
|
||||
elif hasattr(response, "content"):
|
||||
|
|
@ -1335,14 +1402,15 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if content_idx >= len(response_content):
|
||||
continue
|
||||
|
||||
content_block = response_content[content_idx]
|
||||
content_block = _content_block_at(response_content, content_idx)
|
||||
|
||||
# Verify it's a text block and update the text field
|
||||
# Handle both dict and Pydantic object content blocks
|
||||
if isinstance(content_block, dict):
|
||||
if content_block.get("type") == "text":
|
||||
cast(dict[str, object], content_block)["text"] = guardrail_response
|
||||
block = _as_writable(content_block)
|
||||
if block.get("type") == "text":
|
||||
block["text"] = guardrail_response
|
||||
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
|
||||
# Update Pydantic object's text attribute
|
||||
if hasattr(content_block, "text"):
|
||||
if isinstance(content_block, _TextAttrBlock):
|
||||
content_block.text = guardrail_response
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1266,7 +1268,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _cap_thinking_budget_to_max_tokens(
|
||||
def cap_thinking_budget_to_max_tokens(
|
||||
thinking: AnthropicThinkingParam, max_tokens: int | None
|
||||
) -> AnthropicThinkingParam | None:
|
||||
"""Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic
|
||||
|
|
@ -1528,7 +1530,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
llm_provider=self._resolved_provider,
|
||||
)
|
||||
capped_thinking = (
|
||||
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
AnthropicConfig.cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
if legacy_thinking is not None
|
||||
else None
|
||||
)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -1099,6 +1099,25 @@ def is_empty_thinking_block(block: object) -> bool:
|
|||
return not isinstance(thinking, str) or not thinking.strip()
|
||||
|
||||
|
||||
def is_empty_unsigned_thinking_block(block: object) -> bool:
|
||||
"""
|
||||
True for an empty ``{"type": "thinking"}`` block carrying no signature.
|
||||
|
||||
The emit-side predicate: response paths drop a thinking block only when it
|
||||
holds nothing the client could need. A signature-only block is a real
|
||||
provider response (Bedrock Converse under adaptive thinking emits a
|
||||
reasoning block with empty text and only a signature) and the client needs
|
||||
the signature to replay reasoning across tool-use turns, so it must be
|
||||
emitted. Request paths keep using :func:`is_empty_thinking_block`:
|
||||
Anthropic rejects empty thinking blocks in request history regardless of
|
||||
signature, and the inbound strip self-heals a replayed signature-only
|
||||
block.
|
||||
"""
|
||||
if not isinstance(block, dict) or not is_empty_thinking_block(block):
|
||||
return False
|
||||
return not block.get("signature")
|
||||
|
||||
|
||||
def normalize_anthropic_tool_use_id(raw_id: str) -> str:
|
||||
"""
|
||||
Normalize a tool_use / tool_result id for Anthropic's ``^[a-zA-Z0-9_-]+$``
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -1029,7 +1029,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
|
||||
@staticmethod
|
||||
def _is_blank_delta(chunk: "ModelResponseStream") -> bool:
|
||||
from litellm.llms.anthropic.common_utils import is_empty_thinking_block
|
||||
from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block
|
||||
|
||||
choice: Final = chunk.choices[0]
|
||||
if choice.finish_reason is not None:
|
||||
|
|
@ -1041,11 +1041,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
return False
|
||||
if getattr(delta, "reasoning_content", None):
|
||||
return False
|
||||
# thinking_blocks whose entries are all empty (even if signed) must not
|
||||
# thinking_blocks whose entries are all empty AND unsigned must not
|
||||
# open a block: the emitted {"type": "thinking", "thinking": ""} gets
|
||||
# replayed as history and Anthropic rejects it (LIT-6357).
|
||||
# replayed as history and Anthropic rejects it (LIT-6357). A signed
|
||||
# entry opens the block so the client receives the replay signature.
|
||||
thinking_blocks: Final = getattr(delta, "thinking_blocks", None)
|
||||
if thinking_blocks and any(isinstance(b, dict) and not is_empty_thinking_block(b) for b in thinking_blocks):
|
||||
if thinking_blocks and any(
|
||||
isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -90,7 +90,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import (
|
|||
reasoning_effort_from_thinking_budget,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
is_empty_thinking_block,
|
||||
is_empty_unsigned_thinking_block,
|
||||
normalize_anthropic_tool_use_id,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management import (
|
||||
|
|
@ -1267,7 +1267,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks:
|
||||
for thinking_block in choice.message.thinking_blocks:
|
||||
if thinking_block.get("type") == "thinking":
|
||||
if is_empty_thinking_block(thinking_block):
|
||||
if is_empty_unsigned_thinking_block(thinking_block):
|
||||
continue
|
||||
thinking_value = thinking_block.get("thinking", "")
|
||||
signature_value = thinking_block.get("signature", "")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.anthropic import AppliedEdit
|
||||
|
|
@ -11,7 +11,13 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE
|
|||
from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112
|
||||
from .result import PolyfillResult
|
||||
|
||||
EditorFn = Callable[..., Any]
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router import Router
|
||||
|
||||
EditorResult: TypeAlias = "PolyfillResult | tuple[list[dict[str, object]], AppliedEdit | None]"
|
||||
|
||||
EditorFn: TypeAlias = "Callable[..., EditorResult | Awaitable[EditorResult]]"
|
||||
|
||||
_EDITOR_REGISTRY: Final[dict[str, EditorFn]] = {
|
||||
CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919,
|
||||
|
|
@ -19,23 +25,31 @@ _EDITOR_REGISTRY: Final[dict[str, EditorFn]] = {
|
|||
}
|
||||
|
||||
|
||||
def _normalize_spec(
|
||||
spec: dict[str, Any] | list[dict[str, Any]] | None,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Accept Anthropic-native dict form or OpenAI list form; return edits list."""
|
||||
if isinstance(spec, list):
|
||||
# Local import to avoid an import cycle at module load.
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec)
|
||||
|
||||
edits: Final = spec.get("edits") if isinstance(spec, dict) else None
|
||||
def _edits_from(normalized: dict[str, object] | None) -> list[dict[str, object]] | None:
|
||||
edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None
|
||||
if not edits or not isinstance(edits, list):
|
||||
return None
|
||||
return [edit for edit in edits if isinstance(edit, dict)]
|
||||
|
||||
|
||||
def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
|
||||
def _normalize_spec(
|
||||
spec: dict[str, object] | list[dict[str, object]] | None,
|
||||
) -> list[dict[str, object]] | None:
|
||||
"""Accept Anthropic-native dict form or OpenAI list form; return edits list."""
|
||||
if isinstance(spec, list):
|
||||
# Local import to avoid an import cycle at module load.
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
return _edits_from(AnthropicConfig.map_openai_context_management_to_anthropic(spec))
|
||||
|
||||
return _edits_from(spec)
|
||||
|
||||
|
||||
def _wrap_editor_return(
|
||||
raw: EditorResult,
|
||||
*,
|
||||
fallback_system: str | list[dict[str, object]] | None,
|
||||
) -> PolyfillResult:
|
||||
"""Coerce an editor's native return shape into a ``PolyfillResult``.
|
||||
|
||||
v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple
|
||||
|
|
@ -46,7 +60,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
|
|||
return raw
|
||||
# Legacy 2-tuple return — sync editors don't mutate ``system``, so
|
||||
# carry the caller's value forward.
|
||||
messages, applied = cast(tuple[list[dict[str, Any]], Any], raw)
|
||||
messages, applied = raw
|
||||
return PolyfillResult(
|
||||
messages=messages,
|
||||
system=fallback_system,
|
||||
|
|
@ -57,13 +71,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult:
|
|||
async def apply_context_management(
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None,
|
||||
system: Any,
|
||||
context_management_spec: dict[str, Any] | list[dict[str, Any]] | None,
|
||||
litellm_metadata: dict[str, Any] | None = None,
|
||||
llm_router: Any = None,
|
||||
user_api_key_auth: Any = None,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
system: str | list[dict[str, object]] | None,
|
||||
context_management_spec: dict[str, object] | list[dict[str, object]] | None,
|
||||
litellm_metadata: dict[str, object] | None = None,
|
||||
llm_router: "Router | None" = None,
|
||||
user_api_key_auth: "UserAPIKeyAuth | None" = None,
|
||||
) -> PolyfillResult:
|
||||
"""Run edits in order; return a single ``PolyfillResult``.
|
||||
|
||||
|
|
@ -92,22 +106,30 @@ async def apply_context_management(
|
|||
)
|
||||
continue
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": current_messages,
|
||||
"tools": tools,
|
||||
"system": current_system,
|
||||
"edit_spec": edit_spec,
|
||||
}
|
||||
# Only async editors accept these — passing them to sync v0 editors
|
||||
# would break their signature.
|
||||
if inspect.iscoroutinefunction(editor):
|
||||
kwargs["litellm_metadata"] = litellm_metadata
|
||||
kwargs["llm_router"] = llm_router
|
||||
kwargs["user_api_key_auth"] = user_api_key_auth
|
||||
raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs)
|
||||
else:
|
||||
raw_result = editor(**kwargs)
|
||||
editor_is_async = inspect.iscoroutinefunction(editor)
|
||||
editor_return = (
|
||||
editor(
|
||||
model=model,
|
||||
messages=current_messages,
|
||||
tools=tools,
|
||||
system=current_system,
|
||||
edit_spec=edit_spec,
|
||||
litellm_metadata=litellm_metadata,
|
||||
llm_router=llm_router,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
if editor_is_async
|
||||
else editor(
|
||||
model=model,
|
||||
messages=current_messages,
|
||||
tools=tools,
|
||||
system=current_system,
|
||||
edit_spec=edit_spec,
|
||||
)
|
||||
)
|
||||
raw_result = editor_return if isinstance(editor_return, (PolyfillResult, tuple)) else await editor_return
|
||||
|
||||
result = _wrap_editor_return(raw_result, fallback_system=current_system)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.anthropic import AppliedEdit
|
||||
|
|
@ -14,7 +16,18 @@ from ..constants import (
|
|||
from ..placeholders import build_cleared_tool_result_content
|
||||
|
||||
|
||||
def _count_tool_uses(messages: list[dict[str, Any]]) -> int:
|
||||
class ClearToolUsesEditSpec(TypedDict, total=False):
|
||||
"""The ``clear_tool_uses_20250919`` entry of a ``context_management`` spec."""
|
||||
|
||||
type: ReadOnly[str]
|
||||
trigger: ReadOnly[dict[str, object]]
|
||||
keep: ReadOnly[dict[str, object]]
|
||||
clear_at_least: ReadOnly[object]
|
||||
exclude_tools: ReadOnly[object]
|
||||
clear_tool_inputs: ReadOnly[object]
|
||||
|
||||
|
||||
def _count_tool_uses(messages: list[dict[str, object]]) -> int:
|
||||
"""Return the number of tool_use content blocks across all messages.
|
||||
|
||||
Only counts blocks with a string ``id`` to stay consistent with
|
||||
|
|
@ -32,7 +45,7 @@ def _count_tool_uses(messages: list[dict[str, Any]]) -> int:
|
|||
return count
|
||||
|
||||
|
||||
def _collect_tool_use_ids_in_order(messages: list[dict[str, Any]]) -> list[str]:
|
||||
def _collect_tool_use_ids_in_order(messages: list[dict[str, object]]) -> list[str]:
|
||||
"""Return tool_use ids in the chronological order they appear in messages."""
|
||||
ids: Final[list[str]] = []
|
||||
for msg in messages:
|
||||
|
|
@ -47,10 +60,10 @@ def _collect_tool_use_ids_in_order(messages: list[dict[str, Any]]) -> list[str]:
|
|||
|
||||
|
||||
def _trigger_met(
|
||||
trigger: dict[str, Any],
|
||||
trigger: dict[str, object],
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
) -> tuple[bool, int | None]:
|
||||
"""Return (trigger_met, input_tokens if counted for reuse)."""
|
||||
trigger_type: Final = trigger.get("type", "input_tokens")
|
||||
|
|
@ -73,7 +86,7 @@ def _trigger_met(
|
|||
return current_tokens > threshold, current_tokens
|
||||
|
||||
|
||||
def _resolve_keep_count(keep: dict[str, Any]) -> int:
|
||||
def _resolve_keep_count(keep: dict[str, object]) -> int:
|
||||
keep_type: Final = keep.get("type", "tool_uses")
|
||||
if keep_type != "tool_uses":
|
||||
return DEFAULT_KEEP_TOOL_USES
|
||||
|
|
@ -84,7 +97,7 @@ def _resolve_keep_count(keep: dict[str, Any]) -> int:
|
|||
|
||||
|
||||
def _last_completed_tool_use_id(
|
||||
messages: list[dict[str, Any]],
|
||||
messages: list[dict[str, object]],
|
||||
) -> str | None:
|
||||
"""Latest completed tool_result id; never cleared."""
|
||||
last_id: str | None = None
|
||||
|
|
@ -99,17 +112,19 @@ def _last_completed_tool_use_id(
|
|||
return last_id
|
||||
|
||||
|
||||
def _clear_tool_results(messages: list[dict[str, Any]], ids_to_clear: set) -> tuple[list[dict[str, Any]], int]:
|
||||
def _clear_tool_results(
|
||||
messages: list[dict[str, object]], ids_to_clear: set[str]
|
||||
) -> tuple[list[dict[str, object]], int]:
|
||||
"""Clear matching tool_result content; return (messages, cleared_count)."""
|
||||
cleared = 0
|
||||
new_messages: Final[list[dict[str, Any]]] = []
|
||||
new_messages: Final[list[dict[str, object]]] = []
|
||||
for msg in messages:
|
||||
content = msg.get("content")
|
||||
if not isinstance(content, list):
|
||||
new_messages.append(msg)
|
||||
continue
|
||||
|
||||
new_blocks: list[Any] = []
|
||||
new_blocks: list[object] = []
|
||||
mutated = False
|
||||
for block in content:
|
||||
if (
|
||||
|
|
@ -138,11 +153,11 @@ def _clear_tool_results(messages: list[dict[str, Any]], ids_to_clear: set) -> tu
|
|||
def apply_clear_tool_uses_20250919(
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None,
|
||||
system: Any,
|
||||
edit_spec: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], AppliedEdit | None]:
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
system: str | list[dict[str, object]] | None,
|
||||
edit_spec: ClearToolUsesEditSpec,
|
||||
) -> tuple[list[dict[str, object]], AppliedEdit | None]:
|
||||
"""Apply clear_tool_uses; return (messages, AppliedEdit or None)."""
|
||||
ignored_knobs = [knob for knob in ("clear_at_least", "exclude_tools", "clear_tool_inputs") if knob in edit_spec]
|
||||
for ignored_knob in ignored_knobs:
|
||||
|
|
@ -153,11 +168,11 @@ def apply_clear_tool_uses_20250919(
|
|||
CLEAR_TOOL_USES_EDIT_TYPE,
|
||||
)
|
||||
|
||||
trigger: Final = edit_spec.get("trigger") or {
|
||||
trigger: Final[dict[str, object]] = edit_spec.get("trigger") or {
|
||||
"type": "input_tokens",
|
||||
"value": DEFAULT_INPUT_TOKENS_TRIGGER,
|
||||
}
|
||||
keep: Final = edit_spec.get("keep") or {
|
||||
keep: Final[dict[str, object]] = edit_spec.get("keep") or {
|
||||
"type": "tool_uses",
|
||||
"value": DEFAULT_KEEP_TOOL_USES,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,10 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
|
|||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -29,6 +29,7 @@ from litellm.types.llms.anthropic import (
|
|||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse
|
||||
from litellm.router import Router
|
||||
from litellm.types.llms.anthropic import (
|
||||
AllAnthropicPassThroughMessageValues,
|
||||
|
|
@ -84,6 +85,77 @@ _PROPAGATED_METADATA_KEYS: Final = (
|
|||
|
||||
_SUMMARY_TAG_RE: Final = re.compile(r"<summary>(.*?)</summary>", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object])
|
||||
|
||||
|
||||
def _as_object(value: object) -> object:
|
||||
return value
|
||||
|
||||
|
||||
def _is_tool_result_block(block: object) -> bool:
|
||||
return isinstance(block, dict) and block.get("type") in ("tool_result",)
|
||||
|
||||
|
||||
class _SummaryCallKwargs(TypedDict):
|
||||
model: ReadOnly[str]
|
||||
max_tokens: ReadOnly[int]
|
||||
timeout: ReadOnly[float]
|
||||
litellm_metadata: ReadOnly[Mapping[str, object]]
|
||||
user: ReadOnly[NotRequired[str]]
|
||||
allowed_model_region: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class _SummaryOptionalKwargs(TypedDict, total=False):
|
||||
user: ReadOnly[str]
|
||||
allowed_model_region: ReadOnly[str]
|
||||
|
||||
|
||||
class _SummaryAcompletion(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
**kwargs: Unpack[_SummaryCallKwargs], # kwargs-ok: forwarded verbatim to acompletion, which owns them
|
||||
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
|
||||
|
||||
|
||||
class _CreateRateLimitDescriptors(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
data: Mapping[str, str],
|
||||
rpm_limit_type: object,
|
||||
tpm_limit_type: object,
|
||||
model_has_failures: bool,
|
||||
) -> "Sequence[RateLimitDescriptor]": ...
|
||||
|
||||
|
||||
class _AddModelRateLimitDescriptor(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
requested_model: str,
|
||||
descriptors: "Sequence[RateLimitDescriptor]",
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class _CreateOrgRateLimitDescriptors(Protocol):
|
||||
def __call__(
|
||||
self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None
|
||||
) -> "Sequence[RateLimitDescriptor]": ...
|
||||
|
||||
|
||||
class _ShouldRateLimit(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
descriptors: "Sequence[RateLimitDescriptor]",
|
||||
parent_otel_span: object,
|
||||
read_only: bool,
|
||||
) -> "Awaitable[RateLimitResponse]": ...
|
||||
|
||||
|
||||
def _read_summary_model_setting() -> str | None:
|
||||
"""Look up the configured summarization model from proxy general_settings."""
|
||||
|
|
@ -159,11 +231,11 @@ async def _check_summary_model_access(
|
|||
return True
|
||||
|
||||
key_models: Final = list(getattr(user_api_key_auth, "models", None) or [])
|
||||
team_id: Final = getattr(user_api_key_auth, "team_id", None)
|
||||
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
|
||||
team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None)
|
||||
team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or [])
|
||||
user_id: Final = getattr(user_api_key_auth, "user_id", None)
|
||||
project_id: Final = getattr(user_api_key_auth, "project_id", None)
|
||||
user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None)
|
||||
project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None)
|
||||
|
||||
checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = (
|
||||
("key", key_models),
|
||||
|
|
@ -372,7 +444,7 @@ async def _check_summary_model_budget(
|
|||
return False
|
||||
|
||||
end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None)
|
||||
end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None)
|
||||
end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None)
|
||||
if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:
|
||||
try:
|
||||
await model_max_budget_limiter.is_end_user_within_model_budget(
|
||||
|
|
@ -424,40 +496,57 @@ async def _check_summary_model_rate_limit(
|
|||
except Exception:
|
||||
return True
|
||||
|
||||
limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
|
||||
limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
|
||||
should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
|
||||
create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
|
||||
limiter, "_create_rate_limit_descriptors", None
|
||||
)
|
||||
add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
|
||||
limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None
|
||||
)
|
||||
add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
|
||||
limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None
|
||||
)
|
||||
create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr(
|
||||
limiter, "create_organization_rate_limit_descriptor", None
|
||||
)
|
||||
if (
|
||||
limiter is None
|
||||
or not hasattr(limiter, "should_rate_limit")
|
||||
or not hasattr(limiter, "_create_rate_limit_descriptors")
|
||||
or should_rate_limit_check is None
|
||||
or create_descriptors is None
|
||||
or add_team_descriptor is None
|
||||
or add_project_descriptor is None
|
||||
or create_org_descriptors is None
|
||||
):
|
||||
return True
|
||||
|
||||
try:
|
||||
metadata: Final = getattr(user_api_key_auth, "metadata", None) or {}
|
||||
metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {}
|
||||
data: Final = {"model": summary_model}
|
||||
descriptors: Final = limiter._create_rate_limit_descriptors(
|
||||
base_descriptors: Final = create_descriptors(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
data=data,
|
||||
rpm_limit_type=metadata.get("rpm_limit_type"),
|
||||
tpm_limit_type=metadata.get("tpm_limit_type"),
|
||||
model_has_failures=False,
|
||||
)
|
||||
limiter._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
add_team_descriptor(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
requested_model=summary_model,
|
||||
descriptors=descriptors,
|
||||
descriptors=base_descriptors,
|
||||
)
|
||||
limiter._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
add_project_descriptor(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
requested_model=summary_model,
|
||||
descriptors=descriptors,
|
||||
descriptors=base_descriptors,
|
||||
)
|
||||
descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model))
|
||||
descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model))
|
||||
if not descriptors:
|
||||
return True
|
||||
response: Final = await limiter.should_rate_limit(
|
||||
parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None)
|
||||
response: Final[RateLimitResponse] = await should_rate_limit_check(
|
||||
descriptors=descriptors,
|
||||
parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None),
|
||||
parent_otel_span=parent_otel_span,
|
||||
read_only=True,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -471,7 +560,7 @@ async def _check_summary_model_rate_limit(
|
|||
|
||||
|
||||
def _find_latest_compaction_index(
|
||||
messages: list[dict[str, object]],
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
) -> tuple[int | None, int | None]:
|
||||
"""Return (message_index, block_index) of the most recent compaction block.
|
||||
|
||||
|
|
@ -490,8 +579,8 @@ def _find_latest_compaction_index(
|
|||
|
||||
|
||||
def _slice_around_compaction_block(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> tuple[list[dict[str, object]], dict[str, object] | None]:
|
||||
messages: Sequence[_MsgT],
|
||||
) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]:
|
||||
"""Apply Anthropic's "drop everything before the compaction block" rule.
|
||||
|
||||
Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)``
|
||||
|
|
@ -506,19 +595,21 @@ def _slice_around_compaction_block(
|
|||
|
||||
original_msg: Final = messages[msg_idx]
|
||||
original_content: Final = original_msg["content"]
|
||||
compaction_block: Final = cast(dict[str, object], original_content[blk_idx])
|
||||
if not isinstance(original_content, list):
|
||||
return messages, None
|
||||
original_blocks: Final = cast("Sequence[dict[str, object]]", original_content)
|
||||
compaction_block: Final = original_blocks[blk_idx]
|
||||
|
||||
# Per Anthropic's contract everything before the compaction block is
|
||||
# dropped, including earlier blocks within the same assistant message.
|
||||
sliced_content: Final = list(original_content[blk_idx:])
|
||||
sliced_content: Final = list(original_blocks[blk_idx:])
|
||||
|
||||
sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}]
|
||||
sliced_messages.extend(messages[msg_idx + 1 :])
|
||||
sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]]
|
||||
return sliced_messages, compaction_block
|
||||
|
||||
|
||||
def _strip_compaction_blocks(
|
||||
messages: list[dict[str, object]],
|
||||
messages: Sequence[dict[str, object]],
|
||||
) -> list[dict[str, object]]:
|
||||
"""Drop any ``compaction`` content blocks from messages.
|
||||
|
||||
|
|
@ -625,7 +716,7 @@ def _propagate_metadata(
|
|||
|
||||
def _count_effective_tokens(
|
||||
model: str,
|
||||
effective_messages: list[dict[str, object]],
|
||||
effective_messages: Sequence[dict[str, object]],
|
||||
compaction_block: CompactionBlock | None,
|
||||
tools: list[dict[str, object]] | None,
|
||||
system: str | list[dict[str, object]] | None = None,
|
||||
|
|
@ -704,17 +795,18 @@ def _system_to_text(
|
|||
return ""
|
||||
if isinstance(system, str):
|
||||
return system
|
||||
parts: Final[list[str]] = []
|
||||
for block in system:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text = block.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
parts.append(text)
|
||||
return "\n".join(parts)
|
||||
return "\n".join(
|
||||
text
|
||||
for block in system
|
||||
if isinstance(block, dict)
|
||||
and block.get("type") == "text"
|
||||
and isinstance(text := block.get("text"), str)
|
||||
and text
|
||||
)
|
||||
|
||||
|
||||
def _select_last_user_question(
|
||||
messages: list[dict[str, object]],
|
||||
messages: Sequence[dict[str, object]],
|
||||
) -> list[dict[str, object]]:
|
||||
"""Pick the most recent ``user`` turn that is a real question.
|
||||
|
||||
|
|
@ -729,16 +821,18 @@ def _select_last_user_question(
|
|||
turns, or contained no user turns at all). The downstream call always
|
||||
needs a non-empty user message.
|
||||
"""
|
||||
blocks: Sequence[object]
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") != "user":
|
||||
continue
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")]
|
||||
blocks = [*map(_as_object, content)]
|
||||
filtered = [blk for blk in blocks if not _is_tool_result_block(blk)]
|
||||
if not filtered:
|
||||
# Purely tool_result — skip and look for an earlier turn.
|
||||
continue
|
||||
if len(filtered) < len(content):
|
||||
if len(filtered) < len(blocks):
|
||||
return [{**msg, "content": filtered}]
|
||||
return [msg]
|
||||
return [
|
||||
|
|
@ -761,7 +855,7 @@ def _extract_summary_text(raw: str | None) -> str | None:
|
|||
|
||||
def _system_to_openai_message(
|
||||
system: str | list[dict[str, Any]] | None,
|
||||
) -> dict[str, object] | None:
|
||||
) -> Mapping[str, object] | None:
|
||||
"""Translate Anthropic-shaped ``system`` to an OpenAI system message.
|
||||
|
||||
Accepts a bare string or a list of Anthropic content blocks; returns
|
||||
|
|
@ -772,17 +866,19 @@ def _system_to_openai_message(
|
|||
if isinstance(system, str):
|
||||
return {"role": "system", "content": system} if system else None
|
||||
if isinstance(system, list):
|
||||
parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"]
|
||||
parts: Final[tuple[str, ...]] = tuple(
|
||||
block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"
|
||||
)
|
||||
joined: Final = "\n\n".join(part for part in parts if part)
|
||||
return {"role": "system", "content": joined} if joined else None
|
||||
return None
|
||||
|
||||
|
||||
def _build_summary_messages(
|
||||
effective_messages: list[dict[str, object]],
|
||||
effective_messages: Sequence[dict[str, object]],
|
||||
prompt: str,
|
||||
system: str | list[dict[str, object]] | None = None,
|
||||
) -> list[dict[str, object]]:
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
"""Build the OpenAI-shape message list for the summary call.
|
||||
|
||||
The caller's ``system`` prompt is prepended (the default summarization
|
||||
|
|
@ -810,7 +906,7 @@ def _build_summary_messages(
|
|||
)
|
||||
openai_messages = stripped
|
||||
|
||||
summary_messages: Final[list[dict[str, object]]] = []
|
||||
summary_messages: Final[list[Mapping[str, object]]] = []
|
||||
system_message: Final = _system_to_openai_message(system)
|
||||
if system_message is not None:
|
||||
summary_messages.append(system_message)
|
||||
|
|
@ -845,35 +941,17 @@ def _append_text_to_content(content: object, extra_text: str) -> object:
|
|||
if isinstance(content, str):
|
||||
return f"{content}\n\n{extra_text}"
|
||||
if isinstance(content, list):
|
||||
appended: Final[list[object]] = [*content, {"type": "text", "text": extra_text}]
|
||||
appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}]
|
||||
return appended
|
||||
return [content, {"type": "text", "text": extra_text}]
|
||||
|
||||
|
||||
class _SummaryCallUserKwarg(TypedDict, total=False):
|
||||
user: ReadOnly[object]
|
||||
|
||||
|
||||
class _SummaryCallRegionKwarg(TypedDict, total=False):
|
||||
allowed_model_region: ReadOnly[str]
|
||||
|
||||
|
||||
class _SummaryCallKwargs(TypedDict):
|
||||
model: ReadOnly[str]
|
||||
messages: ReadOnly[list[dict[str, object]]]
|
||||
max_tokens: ReadOnly[int]
|
||||
timeout: ReadOnly[float]
|
||||
litellm_metadata: ReadOnly[Mapping[str, object]]
|
||||
user: NotRequired[ReadOnly[object]]
|
||||
allowed_model_region: NotRequired[ReadOnly[str]]
|
||||
|
||||
|
||||
async def _call_summary_model(
|
||||
*,
|
||||
summary_model: str,
|
||||
summary_messages: list[dict[str, object]],
|
||||
summary_messages: Sequence[Mapping[str, object]],
|
||||
metadata: Mapping[str, object],
|
||||
llm_router: Any,
|
||||
llm_router: object,
|
||||
allowed_model_region: str | None = None,
|
||||
max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS,
|
||||
) -> Union["ModelResponse", "CustomStreamWrapper"]:
|
||||
|
|
@ -909,28 +987,37 @@ async def _call_summary_model(
|
|||
# than from ``litellm_metadata``, so without it the summary tokens would not
|
||||
# debit the caller's end-user counters.
|
||||
end_user_id: Final = metadata.get("user_api_key_end_user_id")
|
||||
user_kwargs: Final = (
|
||||
_SummaryOptionalKwargs(user=end_user_id)
|
||||
if isinstance(end_user_id, str) and end_user_id
|
||||
else _SummaryOptionalKwargs()
|
||||
)
|
||||
region_kwargs: Final = (
|
||||
_SummaryOptionalKwargs(allowed_model_region=allowed_model_region)
|
||||
if allowed_model_region is not None
|
||||
else _SummaryOptionalKwargs()
|
||||
)
|
||||
call_kwargs: Final[_SummaryCallKwargs] = {
|
||||
"model": summary_model,
|
||||
"messages": summary_messages,
|
||||
"max_tokens": max_tokens,
|
||||
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
|
||||
"litellm_metadata": metadata,
|
||||
**(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()),
|
||||
**(
|
||||
_SummaryCallRegionKwarg(allowed_model_region=allowed_model_region)
|
||||
if allowed_model_region is not None
|
||||
else _SummaryCallRegionKwarg()
|
||||
),
|
||||
**user_kwargs,
|
||||
**region_kwargs,
|
||||
}
|
||||
if llm_router is not None and hasattr(llm_router, "acompletion"):
|
||||
return await llm_router.acompletion(**call_kwargs)
|
||||
return await litellm.acompletion(**call_kwargs)
|
||||
router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None)
|
||||
if llm_router is not None and router_acompletion is not None:
|
||||
return await router_acompletion(messages=summary_messages, **call_kwargs)
|
||||
return await litellm.acompletion(messages=[*summary_messages], **call_kwargs)
|
||||
|
||||
|
||||
def _extract_response_text(response: Any) -> str | None:
|
||||
def _extract_response_text(response: object) -> str | None:
|
||||
try:
|
||||
choice: Final = response.choices[0]
|
||||
message: Final = choice.message
|
||||
choices: Final[Sequence[object] | None] = getattr(response, "choices", None)
|
||||
if choices is None:
|
||||
return None
|
||||
choice: Final = choices[0]
|
||||
message: Final = getattr(choice, "message", None)
|
||||
content: Final = getattr(message, "content", None)
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
|
|
@ -946,7 +1033,7 @@ def _extract_response_text(response: Any) -> str | None:
|
|||
|
||||
|
||||
def _extract_usage(response: object) -> tuple[int, int]:
|
||||
usage: Final = getattr(response, "usage", None)
|
||||
usage: Final[object] = getattr(response, "usage", None)
|
||||
if usage is None:
|
||||
return 0, 0
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -18,11 +18,14 @@ import asyncio
|
|||
import contextlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0
|
||||
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = (
|
||||
b"event: error\n"
|
||||
|
|
@ -181,7 +184,7 @@ class AgenticAnthropicStreamingIterator:
|
|||
messages: list[dict],
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
hold_back: bool = False,
|
||||
|
|
|
|||
|
|
@ -571,7 +571,34 @@ def anthropic_messages_handler(
|
|||
anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig()
|
||||
if anthropic_messages_provider_config is None:
|
||||
# Route to Responses API for OpenAI / Azure, chat/completions for everything else.
|
||||
_shared_kwargs: Final = dict(
|
||||
if _should_route_to_responses_api(custom_llm_provider, original_model, model):
|
||||
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
|
||||
max_tokens=max_tokens,
|
||||
messages=messages,
|
||||
model=original_model,
|
||||
metadata=metadata,
|
||||
stop_sequences=stop_sequences,
|
||||
stream=stream,
|
||||
system=system,
|
||||
temperature=temperature,
|
||||
thinking=thinking,
|
||||
tool_choice=tool_choice,
|
||||
tools=tools,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
_is_async=is_async,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
client=client,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# The in-gateway context_management polyfill runs inside
|
||||
# ``async_anthropic_messages_handler`` so it can ``await`` the
|
||||
# summarization model for ``compact_20260112``. ``context_management``
|
||||
# is passed through as a regular kwarg.
|
||||
return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
|
||||
max_tokens=max_tokens,
|
||||
messages=messages,
|
||||
model=original_model,
|
||||
|
|
@ -592,16 +619,6 @@ def anthropic_messages_handler(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
**kwargs,
|
||||
)
|
||||
if _should_route_to_responses_api(custom_llm_provider, original_model, model):
|
||||
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(**_shared_kwargs)
|
||||
|
||||
# The in-gateway context_management polyfill runs inside
|
||||
# ``async_anthropic_messages_handler`` so it can ``await`` the
|
||||
# summarization model for ``compact_20260112``. ``context_management``
|
||||
# is passed through as a regular kwarg.
|
||||
return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
|
||||
**_shared_kwargs,
|
||||
)
|
||||
|
||||
if custom_llm_provider is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -40,6 +40,11 @@ DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING: Final = (
|
|||
"minimum thinking budget."
|
||||
)
|
||||
|
||||
DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = (
|
||||
"Dropping `thinking` mapped from reasoning_effort=%s for model=%s: max_tokens=%s "
|
||||
"is too small to fit the minimum thinking budget."
|
||||
)
|
||||
|
||||
|
||||
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
||||
@property
|
||||
|
|
@ -335,11 +340,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
return headers, api_base
|
||||
|
||||
@staticmethod
|
||||
def _translate_reasoning_effort_to_anthropic(model: str, optional_params: dict, custom_llm_provider: str) -> None:
|
||||
def _translate_reasoning_effort_to_anthropic(
|
||||
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str
|
||||
) -> None:
|
||||
"""Map OpenAI-style ``reasoning_effort`` to native Anthropic params.
|
||||
|
||||
Caller-supplied ``thinking`` / ``output_config`` win over the alias.
|
||||
``effort='none'`` clears both. Invalid efforts raise a 400.
|
||||
``effort='none'`` clears both. Invalid efforts raise a 400. A mapped
|
||||
thinking budget is capped below ``max_tokens`` and dropped when even
|
||||
the minimum budget cannot fit.
|
||||
"""
|
||||
from litellm.exceptions import BadRequestError as _BadRequestError
|
||||
from litellm.llms.anthropic.chat.transformation import (
|
||||
|
|
@ -365,7 +374,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
optional_params.pop("output_config", None)
|
||||
return
|
||||
|
||||
optional_params.setdefault("thinking", mapped_thinking)
|
||||
fitted_thinking: Final = AnthropicConfig.cap_thinking_budget_to_max_tokens(mapped_thinking, max_tokens)
|
||||
if fitted_thinking is None:
|
||||
verbose_logger.warning(DROP_UNFITTING_REASONING_EFFORT_WARNING, reasoning_effort, model, max_tokens)
|
||||
return
|
||||
|
||||
optional_params.setdefault("thinking", fitted_thinking)
|
||||
if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
|
||||
mapped_effort: Final = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort)
|
||||
if mapped_effort is None:
|
||||
|
|
@ -510,7 +524,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
except _BadRequestError as e:
|
||||
raise AnthropicError(message=str(e.message), status_code=400)
|
||||
capped_thinking: Final = (
|
||||
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
AnthropicConfig.cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
if legacy_thinking is not None
|
||||
else None
|
||||
)
|
||||
|
|
@ -582,6 +596,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
self._translate_reasoning_effort_to_anthropic(
|
||||
model=model,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
max_tokens=max_tokens,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,10 +5,11 @@ Used when the target model is an OpenAI or Azure model.
|
|||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Coroutine, Mapping
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, TypeAlias
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.anthropic import (
|
||||
AllAnthropicMessageValues,
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicMessagesRequest,
|
||||
AnthropicOutputConfig,
|
||||
|
|
@ -23,6 +24,8 @@ from ..utils import local_model_name
|
|||
from .streaming_iterator import AnthropicResponsesStreamWrapper
|
||||
from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
|
||||
|
||||
AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dict[str, object]]
|
||||
|
||||
_ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter()
|
||||
|
||||
|
||||
|
|
@ -34,22 +37,22 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str,
|
|||
def _build_responses_kwargs(
|
||||
*,
|
||||
max_tokens: int,
|
||||
messages: list[dict],
|
||||
messages: AnthropicRequestMessages,
|
||||
model: str,
|
||||
context_management: dict | None = None,
|
||||
metadata: dict | None = None,
|
||||
context_management: dict[str, object] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
output_config: AnthropicOutputConfig | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
stream: bool | None = False,
|
||||
system: str | None = None,
|
||||
temperature: float | None = None,
|
||||
thinking: dict | None = None,
|
||||
tool_choice: dict | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict] | None = None,
|
||||
thinking: dict[str, object] | None = None,
|
||||
tool_choice: dict[str, object] | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
|
||||
top_k: int | None = None,
|
||||
top_p: float | None = None,
|
||||
output_format: AnthropicOutputSchema | None = None,
|
||||
extra_kwargs: dict[str, Any] | None = None,
|
||||
extra_kwargs: Mapping[str, object] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses().
|
||||
|
|
@ -83,30 +86,32 @@ def _build_responses_kwargs(
|
|||
|
||||
anthropic_request: Final = AnthropicMessagesRequest(**request_data)
|
||||
responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request)
|
||||
forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs)
|
||||
|
||||
# Normalize reasoning effort based on model capabilities
|
||||
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
|
||||
reasoning: Final = responses_kwargs.get("reasoning")
|
||||
if isinstance(reasoning, dict) and "effort" in reasoning:
|
||||
from litellm.llms.anthropic.experimental_pass_through.utils import (
|
||||
normalize_reasoning_effort_value,
|
||||
)
|
||||
if isinstance(reasoning, dict):
|
||||
effort: Final[object] = reasoning.get("effort")
|
||||
if isinstance(effort, str):
|
||||
from litellm.llms.anthropic.experimental_pass_through.utils import (
|
||||
normalize_reasoning_effort_value,
|
||||
)
|
||||
|
||||
effort: Final = reasoning["effort"]
|
||||
normalized: Final = normalize_reasoning_effort_value(
|
||||
effort,
|
||||
model=model,
|
||||
custom_llm_provider=(extra_kwargs or {}).get("custom_llm_provider"),
|
||||
)
|
||||
if normalized != effort:
|
||||
responses_kwargs["reasoning"] = {**reasoning, "effort": normalized}
|
||||
provider_hint: Final = forwarded_kwargs.get("custom_llm_provider")
|
||||
normalized: Final = normalize_reasoning_effort_value(
|
||||
effort,
|
||||
model=model,
|
||||
custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None,
|
||||
)
|
||||
if normalized != effort:
|
||||
responses_kwargs["reasoning"] = {**reasoning, "effort": normalized}
|
||||
|
||||
if stream:
|
||||
responses_kwargs["stream"] = True
|
||||
|
||||
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
|
||||
excluded: Final = {"anthropic_messages"}
|
||||
forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs)
|
||||
for key, value in forwarded_kwargs.items():
|
||||
if key == "litellm_logging_obj" and value is not None:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -140,22 +145,22 @@ class LiteLLMMessagesToResponsesAPIHandler:
|
|||
@staticmethod
|
||||
async def async_anthropic_messages_handler(
|
||||
max_tokens: int,
|
||||
messages: list[dict],
|
||||
messages: AnthropicRequestMessages,
|
||||
model: str,
|
||||
context_management: dict | None = None,
|
||||
metadata: dict | None = None,
|
||||
context_management: dict[str, object] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
output_config: AnthropicOutputConfig | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
stream: bool | None = False,
|
||||
system: str | None = None,
|
||||
temperature: float | None = None,
|
||||
thinking: dict | None = None,
|
||||
tool_choice: dict | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict] | None = None,
|
||||
thinking: dict[str, object] | None = None,
|
||||
tool_choice: dict[str, object] | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
|
||||
top_k: int | None = None,
|
||||
top_p: float | None = None,
|
||||
output_format: AnthropicOutputSchema | None = None,
|
||||
**kwargs,
|
||||
**kwargs: object,
|
||||
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
|
||||
responses_kwargs: Final = _build_responses_kwargs(
|
||||
max_tokens=max_tokens,
|
||||
|
|
@ -193,23 +198,23 @@ class LiteLLMMessagesToResponsesAPIHandler:
|
|||
@staticmethod
|
||||
def anthropic_messages_handler(
|
||||
max_tokens: int,
|
||||
messages: list[dict],
|
||||
messages: AnthropicRequestMessages,
|
||||
model: str,
|
||||
context_management: dict | None = None,
|
||||
metadata: dict | None = None,
|
||||
context_management: dict[str, object] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
output_config: AnthropicOutputConfig | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
stream: bool | None = False,
|
||||
system: str | None = None,
|
||||
temperature: float | None = None,
|
||||
thinking: dict | None = None,
|
||||
tool_choice: dict | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict] | None = None,
|
||||
thinking: dict[str, object] | None = None,
|
||||
tool_choice: dict[str, object] | None = None,
|
||||
tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None,
|
||||
top_k: int | None = None,
|
||||
top_p: float | None = None,
|
||||
output_format: AnthropicOutputSchema | None = None,
|
||||
_is_async: bool = False,
|
||||
**kwargs,
|
||||
**kwargs: object,
|
||||
) -> (
|
||||
AnthropicMessagesResponse
|
||||
| AsyncIterator[bytes]
|
||||
|
|
|
|||
|
|
@ -2,9 +2,10 @@
|
|||
Anthropic Skills API configuration and transformations
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
|
|
@ -22,6 +23,8 @@ from litellm.types.llms.anthropic_skills import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_RAW_JSON_PAYLOAD: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
||||
"""Anthropic-specific Skills API configuration"""
|
||||
|
|
@ -104,10 +107,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> Skill:
|
||||
"""Transform Anthropic response to Skill object"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json())
|
||||
verbose_logger.debug("Transforming create skill response: %s", response_json)
|
||||
|
||||
return Skill(**response_json)
|
||||
return Skill.model_validate(response_json)
|
||||
|
||||
def transform_list_skills_request(
|
||||
self,
|
||||
|
|
@ -122,13 +125,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
url: Final = self.get_complete_url(api_base=api_base, endpoint="skills")
|
||||
|
||||
# Build query parameters
|
||||
query_params: Final[dict[str, Any]] = {}
|
||||
if "limit" in list_params and list_params["limit"]:
|
||||
query_params["limit"] = list_params["limit"]
|
||||
if "page" in list_params and list_params["page"]:
|
||||
query_params["page"] = list_params["page"]
|
||||
if "source" in list_params and list_params["source"]:
|
||||
query_params["source"] = list_params["source"]
|
||||
limit: Final = list_params.get("limit")
|
||||
page: Final = list_params.get("page")
|
||||
source: Final = list_params.get("source")
|
||||
query_params: Final[dict[str, int | str]] = {
|
||||
key: value for key, value in (("limit", limit), ("page", page), ("source", source)) if value
|
||||
}
|
||||
|
||||
verbose_logger.debug(
|
||||
"List skills request made to Anthropic Skills endpoint with params: %s",
|
||||
|
|
@ -143,10 +145,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> ListSkillsResponse:
|
||||
"""Transform Anthropic response to ListSkillsResponse"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json())
|
||||
verbose_logger.debug("Transforming list skills response: %s", response_json)
|
||||
|
||||
return ListSkillsResponse(**response_json)
|
||||
return ListSkillsResponse.model_validate(response_json)
|
||||
|
||||
def transform_get_skill_request(
|
||||
self,
|
||||
|
|
@ -168,10 +170,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> Skill:
|
||||
"""Transform Anthropic response to Skill object"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json())
|
||||
verbose_logger.debug("Transforming get skill response: %s", response_json)
|
||||
|
||||
return Skill(**response_json)
|
||||
return Skill.model_validate(response_json)
|
||||
|
||||
def transform_delete_skill_request(
|
||||
self,
|
||||
|
|
@ -193,7 +195,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> DeleteSkillResponse:
|
||||
"""Transform Anthropic response to DeleteSkillResponse"""
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json())
|
||||
verbose_logger.debug("Transforming delete skill response: %s", response_json)
|
||||
|
||||
return DeleteSkillResponse(**response_json)
|
||||
return DeleteSkillResponse.model_validate(response_json)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final, Union
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.llms.base_llm.text_to_speech.transformation import (
|
||||
BaseTextToSpeechConfig,
|
||||
TextToSpeechRequestData,
|
||||
|
|
@ -238,7 +239,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM):
|
|||
return api_base.rstrip("/") + "/v1/speech"
|
||||
|
||||
aws_region_name: Final = litellm_params.get("aws_region_name", self.DEFAULT_REGION)
|
||||
return f"https://polly.{aws_region_name}.amazonaws.com/v1/speech"
|
||||
return f"https://polly.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/v1/speech"
|
||||
|
||||
def is_ssml_input(self, input: str) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from collections.abc import Coroutine
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from openai import AsyncAzureOpenAI, AzureOpenAI
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -16,6 +16,9 @@ from litellm.utils import (
|
|||
from .azure import AzureChatCompletion
|
||||
from .common_utils import AzureOpenAIError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
class AzureAudioTranscription(AzureChatCompletion):
|
||||
def audio_transcriptions(
|
||||
|
|
@ -23,7 +26,7 @@ class AzureAudioTranscription(AzureChatCompletion):
|
|||
model: str,
|
||||
audio_file: FileTypes,
|
||||
optional_params: dict,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
model_response: TranscriptionResponse,
|
||||
timeout: float,
|
||||
max_retries: int,
|
||||
|
|
@ -112,7 +115,7 @@ class AzureAudioTranscription(AzureChatCompletion):
|
|||
data: dict,
|
||||
model_response: TranscriptionResponse,
|
||||
timeout: float,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
api_version: str | None = None,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
|
|
|
|||
|
|
@ -846,6 +846,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
api_key: str,
|
||||
data: dict,
|
||||
headers: dict,
|
||||
deployment_name: str | None = None,
|
||||
) -> httpx.Response:
|
||||
"""
|
||||
Implemented for azure dall-e-2 image gen calls
|
||||
|
|
@ -957,7 +958,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
content=json.dumps(result).encode("utf-8"),
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
|
||||
return await async_handler.post(
|
||||
url=api_base,
|
||||
json=request_json,
|
||||
|
|
@ -973,6 +974,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
api_key: str,
|
||||
data: dict,
|
||||
headers: dict,
|
||||
deployment_name: str | None = None,
|
||||
) -> httpx.Response:
|
||||
"""
|
||||
Implemented for azure dall-e-2 image gen calls
|
||||
|
|
@ -1073,7 +1075,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
content=json.dumps(result).encode("utf-8"),
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
|
||||
return sync_handler.post(
|
||||
url=api_base,
|
||||
json=request_json,
|
||||
|
|
@ -1091,9 +1093,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
AzureFoundryMAIImageGenerationConfig,
|
||||
)
|
||||
|
||||
api_base: str = azure_client_params.get("azure_endpoint", "") # "https://example-endpoint.openai.azure.com"
|
||||
if api_base.endswith("/"):
|
||||
api_base = api_base.rstrip("/")
|
||||
# deployment-scoped endpoints are moved to "base_url" by select_azure_base_url_or_endpoint
|
||||
api_base: str = (azure_client_params.get("azure_endpoint") or azure_client_params.get("base_url") or "").rstrip(
|
||||
"/"
|
||||
)
|
||||
api_version: Final[str] = azure_client_params.get("api_version", "")
|
||||
if model is None:
|
||||
model = ""
|
||||
|
|
@ -1113,6 +1116,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
api_version=api_version,
|
||||
)
|
||||
|
||||
v1_url: Final = BaseAzureLLM.get_azure_v1_image_url(
|
||||
api_base=api_base,
|
||||
api_version=api_version,
|
||||
route="/openai/images/generations",
|
||||
)
|
||||
if v1_url is not None:
|
||||
return v1_url
|
||||
|
||||
if "/openai/deployments/" in api_base:
|
||||
base_url_with_deployment = api_base
|
||||
else:
|
||||
|
|
@ -1167,6 +1178,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
api_key=api_key,
|
||||
data=data,
|
||||
headers=headers,
|
||||
deployment_name=model,
|
||||
)
|
||||
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
|
||||
|
|
@ -1302,6 +1314,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
api_key=api_key or "",
|
||||
data=data,
|
||||
headers=headers,
|
||||
deployment_name=model,
|
||||
)
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
|
||||
if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig):
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ from ...base_llm.chat.transformation import BaseConfig
|
|||
from ..common_utils import AzureOpenAIError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
LoggingClass = LiteLLMLoggingObj
|
||||
|
|
@ -271,7 +273,7 @@ class AzureOpenAIConfig(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:
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal, NamedTuple, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -789,6 +790,32 @@ class BaseAzureLLM(BaseOpenAILLM):
|
|||
|
||||
return str(final_url)
|
||||
|
||||
@staticmethod
|
||||
def get_azure_v1_image_url(api_base: str, api_version: str | None, route: str) -> str | None:
|
||||
"""
|
||||
Azure's v1 surface serves images at ``/openai/v1/images/{generations,edits}`` and routes by
|
||||
``model`` in the request body, so any deployment path and stale ``api-version`` in
|
||||
``api_base`` have to be dropped.
|
||||
|
||||
Returns None when ``api_version`` is a dated one, which still uses the deployment route.
|
||||
"""
|
||||
if not BaseAzureLLM._is_azure_v1_api_version(api_version):
|
||||
return None
|
||||
|
||||
base_url: Final = httpx.URL(api_base)
|
||||
openai_path_start: Final = base_url.path.find("/openai")
|
||||
resource_base: Final = str(
|
||||
base_url.copy_with(
|
||||
path=base_url.path if openai_path_start == -1 else base_url.path[:openai_path_start],
|
||||
params=httpx.QueryParams(tuple((k, v) for k, v in base_url.params.multi_items() if k != "api-version")),
|
||||
)
|
||||
)
|
||||
return BaseAzureLLM._get_base_azure_url(
|
||||
api_base=resource_base,
|
||||
litellm_params=MappingProxyType({"api_version": api_version}),
|
||||
route=route,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_azure_v1_api_version(api_version: str | None) -> bool:
|
||||
if api_version is None:
|
||||
|
|
|
|||
|
|
@ -193,7 +193,7 @@ class AzureTextCompletion(BaseAzureLLM):
|
|||
data: dict,
|
||||
timeout: Any,
|
||||
model_response: ModelResponse,
|
||||
logging_obj: Any,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
max_retries: int,
|
||||
azure_ad_token: str | None = None,
|
||||
client=None, # this is the AsyncAzureOpenAI
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
verbose_logger.debug("create_file_data=%s", create_file_data)
|
||||
response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data))
|
||||
verbose_logger.debug("create_file_response=%s", response)
|
||||
return OpenAIFileObject(**response.model_dump())
|
||||
return OpenAIFileObject.model_validate(response.model_dump())
|
||||
|
||||
def create_file(
|
||||
self,
|
||||
|
|
@ -60,8 +60,8 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
timeout: float | httpx.Timeout,
|
||||
max_retries: int | None,
|
||||
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
) -> OpenAIFileObject | Coroutine[object, object, OpenAIFileObject]:
|
||||
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
|
||||
litellm_params=litellm_params or {},
|
||||
api_key=api_key,
|
||||
|
|
@ -84,7 +84,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create(
|
||||
**self._prepare_create_file_data(create_file_data)
|
||||
)
|
||||
return OpenAIFileObject(**response.model_dump())
|
||||
return OpenAIFileObject.model_validate(response.model_dump())
|
||||
|
||||
async def afile_content(
|
||||
self,
|
||||
|
|
@ -104,8 +104,8 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
max_retries: int | None,
|
||||
api_version: str | None = None,
|
||||
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]:
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]:
|
||||
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
|
||||
litellm_params=litellm_params or {},
|
||||
api_key=api_key,
|
||||
|
|
@ -150,7 +150,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
max_retries: int | None,
|
||||
api_version: str | None = None,
|
||||
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
):
|
||||
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
|
||||
litellm_params=litellm_params or {},
|
||||
|
|
@ -200,7 +200,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
organization: str | None = None,
|
||||
api_version: str | None = None,
|
||||
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
):
|
||||
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
|
||||
litellm_params=litellm_params or {},
|
||||
|
|
@ -252,7 +252,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
|
|||
purpose: str | None = None,
|
||||
api_version: str | None = None,
|
||||
client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
):
|
||||
openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client(
|
||||
litellm_params=litellm_params or {},
|
||||
|
|
|
|||
|
|
@ -93,8 +93,6 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
|
|||
raise ValueError(
|
||||
f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`"
|
||||
)
|
||||
original_url: Final = httpx.URL(api_base)
|
||||
|
||||
# Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default.
|
||||
# Mirrors the fallback chain used by the Azure chat path in common_utils.py,
|
||||
# so callers that set a global / env api_version don't get an unversioned URL.
|
||||
|
|
@ -105,6 +103,16 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
|
|||
or litellm.AZURE_DEFAULT_API_VERSION
|
||||
)
|
||||
|
||||
v1_url: Final = BaseAzureLLM.get_azure_v1_image_url(
|
||||
api_base=api_base,
|
||||
api_version=api_version,
|
||||
route="/openai/images/edits",
|
||||
)
|
||||
if v1_url is not None:
|
||||
return v1_url
|
||||
|
||||
original_url: Final = httpx.URL(api_base)
|
||||
|
||||
# Create a new dictionary with existing params
|
||||
query_params: Final = dict(original_url.params)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
"""HTTP helpers for Azure OpenAI image generation (REST, not SDK)."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> dict:
|
||||
|
||||
def azure_deployment_image_generation_json_body(api_base: str, data: dict, deployment_name: str | None = None) -> dict:
|
||||
"""
|
||||
Build the JSON body for Azure OpenAI image generation POSTs.
|
||||
|
||||
|
|
@ -9,9 +11,20 @@ def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> di
|
|||
deployment in the URL only; sending ``model`` in the body (especially the deployment
|
||||
name) breaks some models (e.g. gpt-image-2). See LiteLLM #26316.
|
||||
|
||||
For the v1 surface (``.../openai/v1/images/...``), Azure routes by the deployment
|
||||
name in the body ``model`` field, so the deployment name must replace any base
|
||||
model name there or Azure answers 404 DeploymentNotFound.
|
||||
|
||||
Provider-style URLs (e.g. ``/providers/...`` for FLUX on Azure AI) keep all keys
|
||||
so non–OpenAI-deployment payloads still work.
|
||||
"""
|
||||
if "images/generations" in api_base and "/openai/deployments/" in api_base:
|
||||
return {k: v for k, v in data.items() if k != "model"}
|
||||
return data
|
||||
drop_model: Final = "images/generations" in api_base and "/openai/deployments/" in api_base
|
||||
v1_route: Final = "/openai/v1/images/" in api_base and bool(deployment_name)
|
||||
if not drop_model and not v1_route:
|
||||
return data
|
||||
entries: Final = (
|
||||
tuple((k, v) for k, v in data.items() if k != "model")
|
||||
if drop_model
|
||||
else (*data.items(), ("model", deployment_name))
|
||||
)
|
||||
return {k: v for k, v in entries}
|
||||
|
|
|
|||
|
|
@ -34,6 +34,8 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.utils import ModelResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
|
|
@ -295,7 +297,7 @@ class AzureAIAgentsConfig(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:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ The Model Router is a special Azure AI deployment that automatically routes requ
|
|||
to the best available model. It has specific cost tracking requirements.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from httpx import Response
|
||||
|
||||
|
|
@ -14,6 +14,9 @@ from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
|
||||
class AzureModelRouterConfig(AzureAIStudioConfig):
|
||||
"""
|
||||
|
|
@ -56,7 +59,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import copy
|
||||
import enum
|
||||
import re
|
||||
from typing import Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
|
@ -25,6 +25,9 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
from litellm.types.utils import ModelResponse, ProviderField
|
||||
from litellm.utils import _add_path_to_api_base, supports_tool_choice
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
|
||||
class AzureFoundryErrorStrings(str, enum.Enum):
|
||||
SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'"
|
||||
|
|
@ -258,7 +261,7 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from litellm.types.utils import ImageResponse
|
|||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
|
|
@ -199,7 +200,7 @@ class AzureFoundryMAIImageGenerationConfig(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:
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import asyncio
|
|||
import re
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
|
@ -41,6 +41,9 @@ from litellm.llms.base_llm.ocr.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR: Final = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"
|
||||
|
||||
|
||||
|
|
@ -676,7 +679,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
|
|||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
"""
|
||||
|
|
@ -751,7 +754,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
|
|||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import httpx
|
|||
import litellm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponse, TextCompletionResponse
|
||||
|
||||
|
|
@ -19,7 +20,7 @@ class BaseLLM:
|
|||
response: httpx.Response,
|
||||
model_response: "ModelResponse",
|
||||
stream: bool,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
optional_params: dict,
|
||||
api_key: str,
|
||||
data: dict | str,
|
||||
|
|
@ -38,7 +39,7 @@ class BaseLLM:
|
|||
response: httpx.Response,
|
||||
model_response: "TextCompletionResponse",
|
||||
stream: bool,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
optional_params: dict,
|
||||
api_key: str,
|
||||
data: dict | str,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.utils import FileTypes, ModelResponse, TranscriptionResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -110,7 +112,7 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -4,9 +4,10 @@ Bridge for transforming API requests to another API requests
|
|||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm import LiteLLMLoggingObj, ModelResponse
|
||||
|
|
@ -38,7 +39,7 @@ class CompletionTransformationBridge(ABC):
|
|||
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":
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ from litellm.types.llms.openai import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
|
@ -342,7 +344,7 @@ class BaseConfig(ABC):
|
|||
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":
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser
|
|||
from litellm.types.utils import ModelResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -66,7 +68,7 @@ class BaseTextCompletionConfig(BaseConfig, ABC):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
|
|||
from litellm.types.utils import EmbeddingResponse, ModelResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -78,7 +80,7 @@ class BaseEmbeddingConfig(BaseConfig, ABC):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders, ModelResponse
|
|||
from ..chat.transformation import BaseConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
from litellm.router import Router as _Router
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
|
|
@ -207,7 +209,7 @@ class BaseFilesConfig(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:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
|
|
@ -73,6 +76,31 @@ class BaseTranslation(ABC):
|
|||
|
||||
return transformed
|
||||
|
||||
@staticmethod
|
||||
def merge_user_api_key_metadata_into_request(
|
||||
request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place
|
||||
user_api_key_dict: Optional["UserAPIKeyAuth"],
|
||||
) -> None:
|
||||
"""
|
||||
Add the prefixed ``user_api_key_*`` metadata to the request's resolved
|
||||
metadata bucket without overwriting existing keys.
|
||||
|
||||
Writes must go through ``get_or_create_metadata_bucket``: creating a
|
||||
``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat
|
||||
completions) flips the bucket for every later metadata write, and spend
|
||||
logging never sees those writes (e.g. guardrail_information).
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_or_create_metadata_bucket,
|
||||
)
|
||||
|
||||
user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
|
||||
if not user_metadata:
|
||||
return
|
||||
_, metadata_bucket = get_or_create_metadata_bucket(request_data)
|
||||
for key, value in user_metadata.items():
|
||||
metadata_bucket.setdefault(key, value)
|
||||
|
||||
@abstractmethod
|
||||
async def process_input_messages(
|
||||
self,
|
||||
|
|
@ -147,6 +175,26 @@ class BaseTranslation(ABC):
|
|||
"""
|
||||
return None
|
||||
|
||||
def build_stream_error_items(
|
||||
self,
|
||||
exc: "HTTPException",
|
||||
responses_so_far: Sequence[Any] | None = None,
|
||||
) -> Sequence[Any] | None:
|
||||
"""
|
||||
Build the stream items that surface a guardrail HTTPException (a block
|
||||
with the default exception-on-block config, or a failed scan) after the
|
||||
response has already started streaming, in this endpoint's wire format.
|
||||
|
||||
Called only once chunks have been sent: the HTTP status is gone, so the
|
||||
failure must travel as an in-stream error frame. ``responses_so_far``
|
||||
holds the chunks the client has already received, for formats whose
|
||||
error frame continues the stream (e.g. sequence numbers).
|
||||
|
||||
Returns None when the format has no in-stream error frame; the caller
|
||||
then re-raises ``exc``. Override in endpoint subclasses.
|
||||
"""
|
||||
return None
|
||||
|
||||
def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None:
|
||||
"""
|
||||
Convert request data to OpenAI-spec structured messages.
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.utils import ImageResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -91,7 +93,7 @@ class BaseImageGenerationConfig(ABC):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -80,7 +82,7 @@ class BaseImageVariationConfig(BaseConfig, ABC):
|
|||
image: FileTypes,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
encoding: "tiktoken.Encoding | None",
|
||||
api_key: str | None = None,
|
||||
) -> ImageResponse:
|
||||
pass
|
||||
|
|
@ -96,7 +98,7 @@ class BaseImageVariationConfig(BaseConfig, ABC):
|
|||
image: FileTypes,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
encoding: "tiktoken.Encoding | None",
|
||||
api_key: str | None = None,
|
||||
) -> ImageResponse:
|
||||
pass
|
||||
|
|
@ -123,7 +125,7 @@ class BaseImageVariationConfig(BaseConfig, ABC):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -12,6 +12,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 UserAPIKeyAuth
|
||||
from litellm.types.utils import SpecialEnums
|
||||
|
|
@ -157,7 +158,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
|
|||
"resource_object": resource_object,
|
||||
"model_mappings": model_mappings,
|
||||
"flat_model_resource_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,
|
||||
}
|
||||
|
|
@ -179,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
|
|||
"unified_resource_id": unified_resource_id,
|
||||
"model_mappings": json.dumps(model_mappings),
|
||||
"flat_model_resource_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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources.
|
|||
|
||||
Returns a Prisma filter and an ownership check that scope managed resources
|
||||
to the caller's identity: proxy admins see everything, user-keyed callers
|
||||
see records they created, and service-account keys (no user_id) fall back
|
||||
to the resource's owning team. Callers with no admin role and no
|
||||
identifying ids are denied so an empty user_id can never select an
|
||||
unscoped query.
|
||||
see records they created, service-account keys (no user_id) fall back to
|
||||
the resource's owning team, and keys with neither a user_id nor a team_id
|
||||
fall back to their own hashed token so they can still reach the resources
|
||||
they created. Callers with no admin role and no identifying ids at all
|
||||
are denied so an empty user_id can never select an unscoped query.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
|
|
@ -19,6 +20,32 @@ from litellm.proxy._types import (
|
|||
)
|
||||
|
||||
|
||||
def resolve_resource_owner_id(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> str | None:
|
||||
"""Return the identity to stamp on (and match against) a managed
|
||||
resource's ``created_by``.
|
||||
|
||||
A key with neither a user_id nor a team_id would otherwise stamp
|
||||
``created_by=None`` and be locked out of its own resources, so it owns
|
||||
them under its hashed token instead, using the ``key:`` scope prefix
|
||||
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
|
||||
means the caller has no usable identity of its own and must fall back
|
||||
to team scoping, or be denied.
|
||||
"""
|
||||
if user_api_key_dict.user_id is not None:
|
||||
return user_api_key_dict.user_id
|
||||
|
||||
if user_api_key_dict.team_id is not None:
|
||||
return None
|
||||
|
||||
token: Final = user_api_key_dict.token or user_api_key_dict.api_key
|
||||
if token:
|
||||
return f"key:{token}"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
|
||||
"""Build the OpenAI-style paginated list response shape used by managed
|
||||
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
|
||||
|
|
@ -39,7 +66,8 @@ def build_owner_filter(
|
|||
to records the caller is allowed to see.
|
||||
|
||||
- ``{}`` means no scoping (proxy admins).
|
||||
- ``{"created_by": <user_id>}`` for user-keyed callers.
|
||||
- ``{"created_by": <owner_id>}`` for user-keyed callers, and for keys
|
||||
with no user_id and no team_id (owner id is their hashed token).
|
||||
- ``{"team_id": <team_id>}`` for service-account callers
|
||||
that have a team but no user_id.
|
||||
- ``{"OR": [...]}`` when the caller has both — listing must include
|
||||
|
|
@ -62,12 +90,13 @@ def build_owner_filter(
|
|||
]
|
||||
}
|
||||
|
||||
if user_id is not None:
|
||||
return {"created_by": user_id}
|
||||
|
||||
if team_id is not None:
|
||||
return {"team_id": team_id}
|
||||
|
||||
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
|
||||
if owner_id is not None:
|
||||
return {"created_by": owner_id}
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -86,8 +115,8 @@ def can_access_resource(
|
|||
if _user_has_admin_view(user_api_key_dict):
|
||||
return True
|
||||
|
||||
user_id: Final = user_api_key_dict.user_id
|
||||
if user_id is not None and created_by is not None and created_by == user_id:
|
||||
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
|
||||
if owner_id is not None and created_by is not None and created_by == owner_id:
|
||||
return True
|
||||
|
||||
team_id: Final = user_api_key_dict.team_id
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.constants import (
|
|||
BEDROCK_MAX_POLICY_SIZE,
|
||||
STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.secret_managers.main import get_secret, get_secret_str
|
||||
|
||||
|
|
@ -348,7 +349,7 @@ class BaseAWSLLM:
|
|||
def _get_aws_region_from_model_arn(self, model: str | None) -> str | None:
|
||||
try:
|
||||
# First check if the string contains the expected prefix
|
||||
if not isinstance(model, str) or "arn:aws:bedrock" not in model:
|
||||
if not isinstance(model, str) or not contains_bedrock_arn(model):
|
||||
return None
|
||||
|
||||
# Split the ARN and check if we have enough parts
|
||||
|
|
@ -625,24 +626,29 @@ class BaseAWSLLM:
|
|||
return match.group(1) if match else None
|
||||
|
||||
@staticmethod
|
||||
def _resolve_sts_region(aws_sts_endpoint: str | None = None) -> str | None:
|
||||
"""STS signing region: parsed from aws_sts_endpoint else AWS_REGION / AWS_DEFAULT_REGION."""
|
||||
def _resolve_sts_region(
|
||||
aws_sts_endpoint: str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
) -> str | None:
|
||||
"""STS signing region: parsed from aws_sts_endpoint, else AWS_REGION / AWS_DEFAULT_REGION, else the configured aws_region_name."""
|
||||
return (
|
||||
BaseAWSLLM._parse_sts_region_from_endpoint(aws_sts_endpoint)
|
||||
or os.getenv("AWS_REGION")
|
||||
or os.getenv("AWS_DEFAULT_REGION")
|
||||
or aws_region_name
|
||||
)
|
||||
|
||||
def _build_sts_client_kwargs(
|
||||
self,
|
||||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
) -> dict:
|
||||
"""STS client kwargs with aligned endpoint_url and region_name (SigV4)."""
|
||||
kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)}
|
||||
if aws_sts_endpoint is not None:
|
||||
kwargs["endpoint_url"] = aws_sts_endpoint
|
||||
sts_region: Final = self._resolve_sts_region(aws_sts_endpoint)
|
||||
sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name)
|
||||
if sts_region is not None:
|
||||
kwargs["region_name"] = sts_region
|
||||
return kwargs
|
||||
|
|
@ -837,6 +843,7 @@ class BaseAWSLLM:
|
|||
sts_client_kwargs: Final = self._build_sts_client_kwargs(
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
|
||||
with tracer.trace("boto3.client(sts)"):
|
||||
|
|
@ -948,6 +955,7 @@ class BaseAWSLLM:
|
|||
aws_external_id: str | None = None,
|
||||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
) -> dict:
|
||||
"""Handle cross-account role assumption for IRSA."""
|
||||
import boto3
|
||||
|
|
@ -961,6 +969,7 @@ class BaseAWSLLM:
|
|||
irsa_sts_kwargs: Final = self._build_sts_client_kwargs(
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
|
||||
# Create an STS client without credentials
|
||||
|
|
@ -1017,6 +1026,7 @@ class BaseAWSLLM:
|
|||
aws_external_id: str | None = None,
|
||||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
) -> dict:
|
||||
"""Handle same-account role assumption for IRSA."""
|
||||
import boto3
|
||||
|
|
@ -1024,6 +1034,7 @@ class BaseAWSLLM:
|
|||
irsa_sts_kwargs: Final = self._build_sts_client_kwargs(
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
|
||||
verbose_logger.debug("Same account role assumption, using automatic IRSA")
|
||||
|
|
@ -1153,6 +1164,7 @@ class BaseAWSLLM:
|
|||
aws_external_id,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
else:
|
||||
sts_response = self._handle_irsa_same_account(
|
||||
|
|
@ -1161,6 +1173,7 @@ class BaseAWSLLM:
|
|||
aws_external_id,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
|
||||
return self._extract_credentials_and_ttl(sts_response)
|
||||
|
|
@ -1182,6 +1195,7 @@ class BaseAWSLLM:
|
|||
sts_client_kwargs: Final = self._build_sts_client_kwargs(
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
if aws_access_key_id is None and aws_secret_access_key is None:
|
||||
with tracer.trace("boto3.client(sts)"):
|
||||
|
|
@ -1363,14 +1377,15 @@ class BaseAWSLLM:
|
|||
"""
|
||||
Select the default endpoint url based on the endpoint type
|
||||
|
||||
Default endpoint url is https://bedrock-runtime.{aws_region_name}.amazonaws.com
|
||||
Default endpoint url is https://bedrock-runtime.{aws_region_name}.{partition dns suffix}
|
||||
"""
|
||||
dns_suffix: Final = get_aws_dns_suffix(aws_region_name)
|
||||
if endpoint_type == "agent":
|
||||
return f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
|
||||
return f"https://bedrock-agent-runtime.{aws_region_name}.{dns_suffix}"
|
||||
elif endpoint_type == "agentcore":
|
||||
return f"https://bedrock-agentcore.{aws_region_name}.amazonaws.com"
|
||||
return f"https://bedrock-agentcore.{aws_region_name}.{dns_suffix}"
|
||||
else:
|
||||
return f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
|
||||
return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}"
|
||||
|
||||
def _get_boto_credentials_from_optional_params(
|
||||
self, optional_params: dict, model: str | None = None
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from openai.types.batch import BatchRequestCounts
|
||||
from openai.types.batch import Metadata as OpenAIBatchMetadata
|
||||
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -68,6 +70,19 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N
|
|||
return f"{output_prefix}{job_id}/{input_basename}.out"
|
||||
|
||||
|
||||
def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequestCounts | None:
|
||||
total_records: Final = response.get("totalRecordCount")
|
||||
success_records: Final = response.get("successRecordCount")
|
||||
if not isinstance(total_records, int) or not isinstance(success_records, int):
|
||||
return None
|
||||
error_records: Final = response.get("errorRecordCount")
|
||||
return BatchRequestCounts(
|
||||
total=total_records,
|
||||
completed=success_records,
|
||||
failed=error_records if isinstance(error_records, int) else 0,
|
||||
)
|
||||
|
||||
|
||||
def _to_epoch(value: Any) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
|
|
@ -271,11 +286,11 @@ class BedrockBatchesHandler:
|
|||
``aws_external_id``). Unknown keys are ignored.
|
||||
|
||||
Returns:
|
||||
``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that
|
||||
``request_counts`` is always ``(0, 0, 0)`` because
|
||||
``GetModelInvocationJob`` does not surface per-record counts;
|
||||
callers that need accurate counts should parse
|
||||
``manifest.json.out`` from the output S3 prefix.
|
||||
``LiteLLMBatch`` shaped like an OpenAI Batch resource.
|
||||
``request_counts`` maps ``GetModelInvocationJob``'s
|
||||
``totalRecordCount`` / ``successRecordCount`` / ``errorRecordCount``
|
||||
when the provider reports them, and is ``None`` when it does not
|
||||
(older botocore, or a status that omits counts).
|
||||
"""
|
||||
try:
|
||||
import boto3
|
||||
|
|
@ -323,7 +338,9 @@ class BedrockBatchesHandler:
|
|||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {"jobIdentifier": batch_id},
|
||||
"api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"),
|
||||
"api_base": (
|
||||
f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}"
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -386,7 +403,7 @@ class BedrockBatchesHandler:
|
|||
failed_at=completed_at if openai_status == "failed" else None,
|
||||
cancelled_at=completed_at if openai_status == "cancelled" else None,
|
||||
expired_at=completed_at if openai_status == "expired" else None,
|
||||
request_counts=BatchRequestCounts(total=0, completed=0, failed=0),
|
||||
request_counts=_record_counts_from_response(response),
|
||||
metadata=openai_batch_metadata,
|
||||
completion_window="24h",
|
||||
endpoint="/v1/chat/completions",
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Final, Literal, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
||||
from httpx import Headers, Response
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
)
|
||||
|
|
@ -34,6 +35,9 @@ from ..common_utils import (
|
|||
resolve_s3_encryption_key_id,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
# Bedrock batch input files are uploaded as
|
||||
# s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see
|
||||
# BedrockFilesTransformation._get_s3_object_name). A uuid4 is always 36 hex/dash
|
||||
|
|
@ -138,8 +142,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
aws_region_name: Final = self._get_aws_region_name(request_params, model)
|
||||
|
||||
# Bedrock model invocation job endpoint
|
||||
# Format: https://bedrock.{region}.amazonaws.com/model-invocation-job
|
||||
bedrock_endpoint: Final = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job"
|
||||
# Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job
|
||||
bedrock_endpoint: Final = (
|
||||
f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job"
|
||||
)
|
||||
|
||||
return bedrock_endpoint
|
||||
|
||||
|
|
@ -238,8 +244,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
# For Bedrock, we need to return a pre-signed request with AWS auth headers
|
||||
# Use common utility for AWS signing
|
||||
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
|
||||
aws_region_name: Final = self._get_aws_region_name(request_params, model)
|
||||
endpoint_url: Final = (
|
||||
f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job"
|
||||
f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job"
|
||||
)
|
||||
signed_headers, signed_data = self.common_utils.sign_aws_request(
|
||||
service_name="bedrock",
|
||||
|
|
@ -261,7 +268,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
self,
|
||||
model: str | None,
|
||||
raw_response: Response,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: dict,
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
|
|
@ -371,7 +378,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
"""
|
||||
# For Bedrock, batch_id should be the full job ARN
|
||||
# The GetModelInvocationJob API expects the full ARN as the identifier
|
||||
if not batch_id.startswith("arn:aws:bedrock:"):
|
||||
if not is_bedrock_arn(batch_id):
|
||||
raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}")
|
||||
|
||||
# Extract the job identifier from the ARN - use the full ARN path part
|
||||
|
|
@ -390,7 +397,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
import urllib.parse as _ul
|
||||
|
||||
encoded_arn: Final = _ul.quote(batch_id, safe="")
|
||||
endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}"
|
||||
endpoint_url: Final = (
|
||||
f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}"
|
||||
)
|
||||
|
||||
# Use common utility for AWS signing
|
||||
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
|
||||
|
|
@ -527,7 +536,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
self,
|
||||
model: str | None,
|
||||
raw_response: Response,
|
||||
logging_obj: Any,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: dict,
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue