mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
merge: preserve callback imports while updating main
This commit is contained in:
commit
ae79d48cff
1140 changed files with 43834 additions and 12532 deletions
|
|
@ -1291,7 +1291,7 @@ jobs:
|
|||
name: Install MCP SDK1 peer
|
||||
command: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
|
||||
- when:
|
||||
condition:
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
*) has_backend=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
|
||||
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json)
|
||||
has_cost_map=true ;;
|
||||
tests/unit/proxy/*) : ;;
|
||||
*) outside_cost_map_set=true ;;
|
||||
|
|
|
|||
2
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
2
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -4,8 +4,6 @@
|
|||
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
|
||||
"## Affected release" with nothing under it must not appear in the final description -->
|
||||
|
||||
## Description
|
||||
|
||||
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
|
||||
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
|
||||
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
|
|
|
|||
8
.github/workflows/test-unit.yml
vendored
8
.github/workflows/test-unit.yml
vendored
|
|
@ -139,6 +139,7 @@ jobs:
|
|||
test-path: |-
|
||||
tests/unit/batch_completion
|
||||
tests/unit/decisions
|
||||
tests/unit/extraction
|
||||
tests/unit/litellm_core_utils
|
||||
tests/unit/systemone
|
||||
python-version: "3.12"
|
||||
|
|
@ -433,6 +434,13 @@ jobs:
|
|||
|
||||
- shard: rust-bridge-harness
|
||||
test-path: >-
|
||||
tests/test_litellm_rust/test_dispatch.py
|
||||
tests/test_litellm_rust/test_public_entries.py
|
||||
tests/test_litellm_rust/chat_completions/test_dispatch.py
|
||||
tests/test_litellm_rust/embeddings/test_dispatch.py
|
||||
tests/test_litellm_rust/messages/test_dispatch.py
|
||||
tests/test_litellm_rust/ocr/test_dispatch.py
|
||||
tests/test_litellm_rust/responses/test_dispatch.py
|
||||
tests/unit/compression
|
||||
tests/unit/harness
|
||||
tests/unit/rust_bridge
|
||||
|
|
|
|||
|
|
@ -4,8 +4,7 @@ Every pull request whose diff against its merge base touches one of the two cost
|
|||
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
|
||||
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
|
||||
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
|
||||
get the file checks and additionally may only touch those two files or the legacy schema
|
||||
path and may only add or update models.
|
||||
get the file checks and additionally may only touch those two files and may only add or update models.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -21,7 +20,6 @@ from typing import Final
|
|||
|
||||
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
|
||||
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
|
||||
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
|
||||
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
|
||||
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
|
||||
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
|
||||
|
|
@ -110,7 +108,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
|
|||
*(
|
||||
f"bot PRs may only change the cost map files, not {path}"
|
||||
for path in changed_files
|
||||
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
|
||||
if path not in GUARDED_PATHS
|
||||
),
|
||||
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
|
||||
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
|
||||
|
|
|
|||
41
cookbook/litellm_proxy_server/batch_api/bedrock/README.md
Normal file
41
cookbook/litellm_proxy_server/batch_api/bedrock/README.md
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
# Bedrock batch endpoints
|
||||
|
||||
Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: bedrock-batch
|
||||
litellm_params:
|
||||
model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0
|
||||
aws_region_name: us-east-1
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
|
||||
aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN
|
||||
s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME
|
||||
aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT
|
||||
```
|
||||
|
||||
Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region
|
||||
|
||||
`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately
|
||||
|
||||
The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above
|
||||
|
||||
Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read
|
||||
|
||||
```bash
|
||||
curl -sS "$LITELLM_API_BASE/v1/batches" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}"
|
||||
```
|
||||
|
||||
Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint
|
||||
|
||||
```bash
|
||||
curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
|
||||
|
||||
curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
|
||||
```
|
||||
|
|
@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
|
|||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
ManagedObjectRow,
|
||||
row_metadata,
|
||||
)
|
||||
|
||||
import litellm
|
||||
|
|
@ -90,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
|
|||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _row_metadata(row: object) -> dict[str, object]:
|
||||
metadata: Final = getattr(row, "metadata", None)
|
||||
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -852,20 +848,26 @@ class CheckBatchCost:
|
|||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
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,
|
||||
batch_prompt_cost=batch_result.prompt_cost,
|
||||
batch_completion_cost=batch_result.completion_cost,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job, original_status, original_file_object)
|
||||
raise
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
|
||||
|
||||
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
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,
|
||||
batch_prompt_cost=batch_result.prompt_cost,
|
||||
batch_completion_cost=batch_result.completion_cost,
|
||||
)
|
||||
finally:
|
||||
if accounting.retryable_failure:
|
||||
await self._release_job_claim(job, original_status, original_file_object)
|
||||
|
||||
if accounting.retryable_failure:
|
||||
return None
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
if prom_logger and response.completed_at and response.created_at:
|
||||
|
|
@ -1057,7 +1059,9 @@ class CheckBatchCost:
|
|||
if (has_key_row and key_row is None) or (team_id and team_row is None):
|
||||
return True
|
||||
settings: Final = get_dynamic_logging_metadata(
|
||||
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
|
||||
UserAPIKeyAuth(
|
||||
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
|
||||
),
|
||||
proxy_config,
|
||||
)
|
||||
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
|
||||
|
|
@ -1186,7 +1190,7 @@ class CheckBatchCost:
|
|||
continue
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to track cost for batch {batch_id} "
|
||||
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
|
||||
f"(job {job.id}): {tracking_err}"
|
||||
)
|
||||
self._record_error(prom_logger, "cost_tracking_error")
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che
|
|||
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
|
||||
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
|
||||
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
|
||||
The spend event reaches the creating key's or team's own logging callbacks the way that creator's
|
||||
requests do. Its message content is redacted when the creator's key or team logging turns message
|
||||
logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before
|
||||
the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or
|
||||
whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only
|
||||
and with its message content redacted.
|
||||
A read that finds no row, because the create could not store one, stores it under the reading key
|
||||
and bills it the same way, so a response the create left untracked is still billed once. A streamed
|
||||
background create is free on its own; its stream stores the row under the creating key and bills
|
||||
|
|
@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast
|
|||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
ManagedObjectRow,
|
||||
ResponseCreator,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
|
|
@ -54,11 +61,13 @@ from litellm.types.utils import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.router import Router
|
||||
|
||||
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
|
||||
_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
class _RoutedModelInfo(BaseModel):
|
||||
|
|
@ -73,6 +82,8 @@ class _RoutedDeployment(BaseModel):
|
|||
|
||||
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
|
||||
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
|
||||
_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"})
|
||||
_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -89,6 +100,7 @@ class BilledBackgroundResponse:
|
|||
custom_llm_provider: str | None
|
||||
litellm_params: Mapping[str, object]
|
||||
metadata: Mapping[str, object]
|
||||
creator_logging: Mapping[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -107,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
|
|||
return unpriced
|
||||
|
||||
|
||||
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
def _string_callbacks(callbacks: object) -> tuple[str, ...]:
|
||||
if not isinstance(callbacks, list):
|
||||
return ()
|
||||
return tuple(
|
||||
callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str)
|
||||
)
|
||||
|
||||
logging_obj: Final = LiteLLMLogging(
|
||||
|
||||
def _running_proxy_config() -> "ProxyConfig":
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
return proxy_config
|
||||
|
||||
|
||||
def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields
|
||||
|
||||
if not creator.settings_read:
|
||||
return _REJECTED_CREATOR_LOGGING
|
||||
try:
|
||||
return key_and_team_logging_fields(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
metadata=dict(creator.key_metadata),
|
||||
team_id=creator.team_id,
|
||||
team_metadata=dict(creator.team_metadata),
|
||||
),
|
||||
proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's "
|
||||
f"callbacks get the billed event: {e}"
|
||||
)
|
||||
return _REJECTED_CREATOR_LOGGING
|
||||
|
||||
|
||||
def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging":
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.utils import split_dynamic_success_callbacks
|
||||
|
||||
sync_callbacks, async_callbacks = split_dynamic_success_callbacks(
|
||||
_string_callbacks(creator_logging.get("success_callback"))
|
||||
)
|
||||
return LiteLLMLogging(
|
||||
model=billed.model,
|
||||
messages=[{"role": "user", "content": "<retrieve_response>"}],
|
||||
stream=False,
|
||||
|
|
@ -119,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
|
|||
start_time=datetime.now(),
|
||||
litellm_call_id=str(uuid.uuid4()),
|
||||
function_id=str(uuid.uuid4()),
|
||||
dynamic_success_callbacks=list(sync_callbacks),
|
||||
dynamic_async_success_callbacks=list(async_callbacks) or None,
|
||||
kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS},
|
||||
)
|
||||
|
||||
|
||||
def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging":
|
||||
try:
|
||||
return _billed_logging_obj(billed, billed.creator_logging)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's "
|
||||
f"callbacks get the billed event: {e}"
|
||||
)
|
||||
return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING)
|
||||
|
||||
|
||||
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
logging_obj: Final = _creator_or_proxy_logging_obj(billed)
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
|
||||
optional_params={},
|
||||
|
|
@ -187,6 +259,7 @@ class CheckResponsesCost:
|
|||
prisma_client: "PrismaClient",
|
||||
llm_router: "Router",
|
||||
log_billed_response: BilledResponseLogger = log_billed_background_response,
|
||||
proxy_config: "ProxyConfig | None" = None,
|
||||
):
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -196,6 +269,7 @@ class CheckResponsesCost:
|
|||
self.llm_router: Router = llm_router
|
||||
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
|
||||
self._log_billed_response: Final = log_billed_response
|
||||
self._proxy_config: Final = proxy_config
|
||||
|
||||
def _resolve_deployment(self, response_id: str) -> Deployment | None:
|
||||
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
|
||||
|
|
@ -307,11 +381,12 @@ class CheckResponsesCost:
|
|||
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
|
||||
)
|
||||
|
||||
async def _billed_response(
|
||||
def _billed_response(
|
||||
self,
|
||||
job: ManagedObjectRow,
|
||||
target: _PollTarget,
|
||||
response: ResponsesAPIResponse,
|
||||
creator: ResponseCreator,
|
||||
) -> BilledBackgroundResponse:
|
||||
deployment: Final = target.deployment
|
||||
return BilledBackgroundResponse(
|
||||
|
|
@ -323,13 +398,14 @@ class CheckResponsesCost:
|
|||
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
|
||||
metadata=MappingProxyType(
|
||||
{
|
||||
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
|
||||
**creator.metadata,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
**target.model_metadata,
|
||||
**(_deployment_metadata(deployment) if deployment is not None else {}),
|
||||
**self._requested_model_group(target, job.team_id),
|
||||
}
|
||||
),
|
||||
creator_logging=creator_logging_fields(creator, self._proxy_config),
|
||||
)
|
||||
|
||||
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
|
||||
|
|
@ -354,7 +430,8 @@ class CheckResponsesCost:
|
|||
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
|
||||
return None
|
||||
try:
|
||||
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
|
||||
creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id)
|
||||
charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator))
|
||||
except BaseException:
|
||||
await self._release_billing_claim(job)
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -10,12 +10,16 @@ from __future__ import annotations
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import CLI_SESSION_KEY_PREFIX
|
||||
from litellm.repositories.table_repositories import (
|
||||
DeletedVerificationTokenRepository,
|
||||
DeprecatedVerificationTokenRepository,
|
||||
)
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
|
|
@ -70,6 +74,7 @@ class ManagedObjectRow(Protocol):
|
|||
|
||||
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
|
||||
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
|
||||
_MAX_ROTATION_HOPS: Final = 16
|
||||
|
||||
|
||||
def _string_tags(request_tags: object) -> list[str] | None:
|
||||
|
|
@ -96,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL
|
|||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _rotated_token_table(
|
||||
prisma_client: PrismaClient,
|
||||
) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]:
|
||||
return DeprecatedVerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]:
|
||||
return DeletedVerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]:
|
||||
return TeamRepository(prisma_client).deleted_table
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CreatorRows:
|
||||
key: prisma_models.LiteLLM_VerificationToken | None
|
||||
|
|
@ -103,6 +122,15 @@ class _CreatorRows:
|
|||
user: prisma_models.LiteLLM_UserTable | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponseCreator:
|
||||
metadata: Mapping[str, object]
|
||||
key_metadata: Mapping[str, object]
|
||||
team_id: str | None
|
||||
team_metadata: Mapping[str, object]
|
||||
settings_read: bool = True
|
||||
|
||||
|
||||
class ManagedObjectCreatorAttribution:
|
||||
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
|
||||
job's spend row is attributed like the creating request: key hash, team, organization,
|
||||
|
|
@ -117,7 +145,12 @@ class ManagedObjectCreatorAttribution:
|
|||
A background response also carries its end user, the creating key's project and its alias,
|
||||
and the key, team, and user budgets and per-model budgets, the way the read that billed it
|
||||
before did, and never names a key-scoped creator as its user. A key that sets no budget of
|
||||
its own takes its budget tier's, the same rule auth applies.
|
||||
its own takes its budget tier's, the same rule auth applies. Its creator's key and team
|
||||
metadata come along too, so the billed event reaches the key's and team's logging callbacks
|
||||
like that read did. A rotated key's metadata is read from the live key its rotations lead
|
||||
to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the
|
||||
creator's settings unread and the attribution to the rows that still load, so the
|
||||
response still bills and still counts toward their budgets.
|
||||
"""
|
||||
|
||||
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
|
||||
|
|
@ -128,31 +161,70 @@ class ManagedObjectCreatorAttribution:
|
|||
rows: Final = await self._rows(job, object_id)
|
||||
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
|
||||
|
||||
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
|
||||
rows: Final = await self._rows(job, object_id)
|
||||
end_user_id: Final = getattr(job, "end_user_id", None)
|
||||
return {
|
||||
"user_api_key_user_id": _creator_user_id(job.created_by),
|
||||
**_shared_metadata(job, rows),
|
||||
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
|
||||
"user_api_key_project_id": getattr(rows.key, "project_id", None),
|
||||
"user_api_key_project_alias": getattr(
|
||||
getattr(rows.key, "litellm_project_table", None), "project_alias", None
|
||||
async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
|
||||
try:
|
||||
return await self._read_response_creator(job, object_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their "
|
||||
f"settings: {e}"
|
||||
)
|
||||
return ResponseCreator(
|
||||
metadata=_response_creator_metadata(job, await self._rows(job, object_id)),
|
||||
key_metadata={},
|
||||
team_id=getattr(job, "team_id", None),
|
||||
team_metadata={},
|
||||
settings_read=False,
|
||||
)
|
||||
|
||||
async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
team_id: Final = getattr(job, "team_id", None)
|
||||
rows: Final = _CreatorRows(
|
||||
key=await self._live_key_row(api_key, job.created_by),
|
||||
team=await self._live_team_row(team_id),
|
||||
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
|
||||
)
|
||||
return ResponseCreator(
|
||||
metadata=_response_creator_metadata(job, rows),
|
||||
key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key),
|
||||
team_id=team_id,
|
||||
team_metadata=(
|
||||
row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id)
|
||||
),
|
||||
**_key_budget_fields(rows.key),
|
||||
**_budget_fields(
|
||||
rows.team,
|
||||
"user_api_key_team_spend",
|
||||
"user_api_key_team_max_budget",
|
||||
"user_api_key_team_model_max_budget",
|
||||
),
|
||||
**_budget_fields(
|
||||
rows.user,
|
||||
"user_api_key_user_spend",
|
||||
"user_api_key_user_max_budget",
|
||||
"user_api_key_user_model_max_budget",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]:
|
||||
if not _names_a_key_row(api_key):
|
||||
return {}
|
||||
current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS)
|
||||
current_key: Final = (
|
||||
await _token_table(self._prisma_client).find_unique(where={"token": current_token})
|
||||
if current_token != api_key
|
||||
else None
|
||||
)
|
||||
if current_key is not None:
|
||||
return row_metadata(current_key)
|
||||
return row_metadata(
|
||||
await _deleted_token_table(self._prisma_client).find_first(
|
||||
where={"token": current_token}, order={"deleted_at": "desc"}
|
||||
)
|
||||
)
|
||||
|
||||
async def _token_rotated_to(self, token: str, hops_left: int) -> str:
|
||||
rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token})
|
||||
if rotation is None or hops_left == 0:
|
||||
return token
|
||||
return await self._token_rotated_to(rotation.active_token_id, hops_left - 1)
|
||||
|
||||
async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]:
|
||||
if not team_id:
|
||||
return {}
|
||||
return row_metadata(
|
||||
await _deleted_team_table(self._prisma_client).find_first(
|
||||
where={"team_id": team_id}, order={"deleted_at": "desc"}
|
||||
)
|
||||
)
|
||||
|
||||
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
|
||||
return _CreatorRows(
|
||||
|
|
@ -175,25 +247,57 @@ class ManagedObjectCreatorAttribution:
|
|||
async def _key_row(
|
||||
self, object_id: str, api_key: str | None, created_by: str | None
|
||||
) -> prisma_models.LiteLLM_VerificationToken | None:
|
||||
if not api_key or _is_cli_session_key(api_key, created_by):
|
||||
return None
|
||||
try:
|
||||
return await _token_table(self._prisma_client).find_unique(
|
||||
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
|
||||
)
|
||||
return await self._live_key_row(api_key, created_by)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
|
||||
if not team_id:
|
||||
async def _live_key_row(
|
||||
self, api_key: str | None, created_by: str | None
|
||||
) -> prisma_models.LiteLLM_VerificationToken | None:
|
||||
if not api_key or _is_cli_session_key(api_key, created_by):
|
||||
return None
|
||||
return await _token_table(self._prisma_client).find_unique(
|
||||
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
|
||||
)
|
||||
|
||||
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
|
||||
try:
|
||||
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
|
||||
return await self._live_team_row(team_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
|
||||
if not team_id:
|
||||
return None
|
||||
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
|
||||
|
||||
|
||||
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
|
||||
end_user_id: Final = getattr(job, "end_user_id", None)
|
||||
return {
|
||||
"user_api_key_user_id": _creator_user_id(job.created_by),
|
||||
**_shared_metadata(job, rows),
|
||||
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
|
||||
"user_api_key_project_id": getattr(rows.key, "project_id", None),
|
||||
"user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None),
|
||||
**_key_budget_fields(rows.key),
|
||||
**_budget_fields(
|
||||
rows.team,
|
||||
"user_api_key_team_spend",
|
||||
"user_api_key_team_max_budget",
|
||||
"user_api_key_team_model_max_budget",
|
||||
),
|
||||
**_budget_fields(
|
||||
rows.user,
|
||||
"user_api_key_user_spend",
|
||||
"user_api_key_user_max_budget",
|
||||
"user_api_key_user_model_max_budget",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
|
|
@ -221,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
|
|||
}
|
||||
|
||||
|
||||
def row_metadata(row: object) -> Mapping[str, object]:
|
||||
metadata: Final = getattr(row, "metadata", None)
|
||||
if not isinstance(metadata, dict):
|
||||
return {}
|
||||
return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column
|
||||
|
||||
|
||||
def _names_a_key_row(api_key: object) -> TypeGuard[str]:
|
||||
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
|
||||
|
||||
return isinstance(api_key, str) and is_valid_sha256_hash(api_key)
|
||||
|
||||
|
||||
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
|
||||
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.78"
|
||||
version = "0.1.79"
|
||||
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.78"
|
||||
version = "0.1.79"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Path allowlist for the gateway component.
|
||||
|
||||
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
|
||||
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
|
||||
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
|
||||
responses, vector stores, passthrough providers, realtime websockets, MCP
|
||||
tool-call endpoints, and operational endpoints (/health, /metrics, and the
|
||||
/debug/memory/summary read of the serving worker's RSS).
|
||||
|
|
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/decisions",
|
||||
"/v1/systemone",
|
||||
"/systemone",
|
||||
"/v1/extraction",
|
||||
"/extraction",
|
||||
"/v1/ocr",
|
||||
"/ocr",
|
||||
"/v1/rag/",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.110"
|
||||
version = "0.4.111"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.110"
|
||||
version = "0.4.111"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ mod selection;
|
|||
|
||||
pub(crate) use native::NativeCacheHandle;
|
||||
pub(crate) use python::{CacheCall, PythonCache};
|
||||
pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native};
|
||||
pub(crate) use selection::{Cached, Selection, configure, configured_native};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use pyo3::{
|
||||
|
|
|
|||
|
|
@ -279,7 +279,7 @@ pub(in crate::cache) fn configured(
|
|||
)> {
|
||||
let handle = native_handle(configured)?.ok_or_else(|| {
|
||||
pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"the configured cache changed to a Python cache after native admission",
|
||||
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
|
||||
)
|
||||
})?;
|
||||
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
|
||||
|
|
|
|||
|
|
@ -76,21 +76,6 @@ fn selected_cache<'py>(
|
|||
Ok(Some(configured))
|
||||
}
|
||||
|
||||
pub(crate) fn admit_native(
|
||||
py: Python<'_>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
call_type: &str,
|
||||
) -> PyResult<()> {
|
||||
if let Some(configured) = selected_cache(py, kwargs, call_type)?
|
||||
&& native::v2::native_handle(&configured)?.is_none()
|
||||
{
|
||||
return Err(crate::errors::RustBridgeDeclined::new_err(
|
||||
"the configured cache requires Python inference",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn configured_native(
|
||||
py: Python<'_>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
|
|
|
|||
|
|
@ -177,7 +177,6 @@ where
|
|||
} else {
|
||||
R::SYNC_CALL_TYPE
|
||||
};
|
||||
crate::cache::admit_native(py, &call.kwargs, call_type)?;
|
||||
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
|
||||
super::run_public_call(
|
||||
py,
|
||||
|
|
|
|||
|
|
@ -8,59 +8,19 @@ use litellm_auth::AuthServices;
|
|||
use litellm_cache_response::{CachePolicy, ScopedCache};
|
||||
use litellm_callbacks_legacy_python::LoggingOperation;
|
||||
use litellm_host::{call::HostedMachine, protocol::Protocol};
|
||||
use litellm_host_python::present;
|
||||
use litellm_inference_responses::{ResponsesRoute, route::Responses};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
pub(crate) use websocket::ResponsesWebSocketConnection;
|
||||
|
||||
use super::{
|
||||
NativeCall,
|
||||
inference::{InferenceHost, InferenceRoute, run_inference},
|
||||
};
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
|
||||
|
||||
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
|
||||
let resolved = call.resolved()?;
|
||||
if let Some(reason) = py
|
||||
.import(ROUTE_HOST_MODULE)?
|
||||
.getattr("decline_reason")?
|
||||
.call1((&resolved,))?
|
||||
.extract::<Option<String>>()?
|
||||
{
|
||||
return Err(RustBridgeDeclined::new_err(reason));
|
||||
}
|
||||
let argument = |name: &str| present(&call.kwargs, &call.base, name);
|
||||
let model = argument("model")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
|
||||
.extract::<String>()?;
|
||||
let provider = argument("custom_llm_provider")?
|
||||
.map(|value| value.extract::<String>())
|
||||
.transpose()?;
|
||||
if provider
|
||||
.as_deref()
|
||||
.is_some_and(|provider| provider != "openai")
|
||||
|| model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(&model)
|
||||
.contains('/')
|
||||
{
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"native HTTP responses provider",
|
||||
));
|
||||
}
|
||||
if argument("stream")?
|
||||
.map(|value| litellm_host_python::from_py::<Value>(&value))
|
||||
.transpose()?
|
||||
.is_some_and(|value| value == Value::Bool(true))
|
||||
{
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"native Python responses streaming",
|
||||
));
|
||||
}
|
||||
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
|
||||
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost::new(host))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1505,6 +1505,7 @@ from .rag.main import *
|
|||
from .sandbox.main import *
|
||||
from .decisions.main import *
|
||||
from .systemone.main import *
|
||||
from .extraction.main import *
|
||||
from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop
|
||||
from .search.main import *
|
||||
from .realtime_api.main import (
|
||||
|
|
@ -1711,6 +1712,9 @@ if TYPE_CHECKING:
|
|||
from .llms.cloudflare.decisions.transformation import (
|
||||
CloudflareDecisionsConfig as CloudflareDecisionsConfig,
|
||||
)
|
||||
from .llms.cloudflare.extraction.transformation import (
|
||||
CloudflareExtractionConfig as CloudflareExtractionConfig,
|
||||
)
|
||||
from .llms.strands_decider.decisions.transformation import (
|
||||
StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -166,6 +166,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"OpenRouterDecisionsConfig",
|
||||
"FireworksAIDecisionsConfig",
|
||||
"CloudflareDecisionsConfig",
|
||||
"CloudflareExtractionConfig",
|
||||
"StrandsDeciderDecisionsConfig",
|
||||
"DatabricksDecisionsConfig",
|
||||
"HostedVLLMDecisionsConfig",
|
||||
|
|
@ -735,6 +736,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
|
||||
"FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),
|
||||
"CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"),
|
||||
"CloudflareExtractionConfig": (".llms.cloudflare.extraction.transformation", "CloudflareExtractionConfig"),
|
||||
"StrandsDeciderDecisionsConfig": (
|
||||
".llms.strands_decider.decisions.transformation",
|
||||
"StrandsDeciderDecisionsConfig",
|
||||
|
|
|
|||
|
|
@ -80,7 +80,7 @@ class A2ACompletionBridgeHandler:
|
|||
@staticmethod
|
||||
def _build_completion_params(
|
||||
params: dict[str, Any],
|
||||
litellm_params: Mapping[str, Any],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_base: str | None,
|
||||
agent_extra_headers: Mapping[str, str] | None,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -95,7 +95,10 @@ class InMemoryCache(BaseCache):
|
|||
"""
|
||||
Check if a specific key is expired
|
||||
"""
|
||||
return key in self.ttl_dict and self._clock() > self.ttl_dict[key]
|
||||
try:
|
||||
return self._clock() > self.ttl_dict[key]
|
||||
except KeyError:
|
||||
return False
|
||||
|
||||
def _cache_keys(self) -> tuple[str, ...]:
|
||||
return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state
|
||||
|
|
@ -227,7 +230,10 @@ class InMemoryCache(BaseCache):
|
|||
if key in self.cache_dict:
|
||||
if self.evict_element_if_expired(key):
|
||||
return None
|
||||
original_cached_response: Final = self.cache_dict[key]
|
||||
try:
|
||||
original_cached_response: Final = self.cache_dict[key]
|
||||
except KeyError:
|
||||
return None
|
||||
try:
|
||||
cached_response = json.loads(original_cached_response)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -1,15 +1,19 @@
|
|||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
|
||||
|
||||
from litellm import main
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge import detectors
|
||||
from litellm.rust_bridge.catalog import Feature, Route, RouteContext
|
||||
from litellm.rust_bridge.chat_completions import detectors as route_detectors
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import (
|
||||
NATIVE_ACOMPLETION,
|
||||
NATIVE_COMPLETION,
|
||||
)
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
|
|
@ -44,6 +48,7 @@ def _python_acompletion() -> PythonAcompletion:
|
|||
|
||||
_PYTHON_COMPLETION: Final = _python_completion()
|
||||
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
|
||||
_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION)
|
||||
_PYTHON_ACOMPLETION: Final = _python_acompletion()
|
||||
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
|
||||
|
||||
|
|
@ -61,11 +66,32 @@ def _public_request(
|
|||
return native_call(legacy, args, kwargs)
|
||||
|
||||
|
||||
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
|
||||
{
|
||||
**detectors.COMMON,
|
||||
Feature.MOCK_RESPONSE: route_detectors.chat_mock_response,
|
||||
Feature.AGENTIC_LOOP_HOOK: route_detectors.chat_agentic_loop_hook,
|
||||
Feature.RESPONSE_CACHE: detectors.response_cache,
|
||||
Feature.STREAMING: detectors.streaming,
|
||||
Feature.DROP_PARAMS: detectors.drop_params,
|
||||
Feature.MODIFY_PARAMS: detectors.modify_params,
|
||||
Feature.EXTRA_BODY: detectors.extra_body,
|
||||
Feature.RETRIES_FALLBACKS: route_detectors.chat_retries_fallbacks,
|
||||
Feature.PROMPT_MANAGEMENT: detectors.prompt_management,
|
||||
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
|
||||
Feature.MODEL_PARAM_GATING: route_detectors.chat_model_param_gating,
|
||||
Feature.MODEL_LIST: route_detectors.model_list,
|
||||
Feature.MODEL_ALIAS_MAP: route_detectors.model_alias_map,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _context(request: NativeCall) -> RouteContext:
|
||||
return RouteContext(
|
||||
Route.CHAT_COMPLETIONS,
|
||||
provider=optional_str(request.resolved.get("custom_llm_provider")),
|
||||
model=str(request.resolved["model"]),
|
||||
features=detectors.detected(_FEATURES, request),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -73,7 +99,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.CHAT_COMPLETIONS,
|
||||
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("acompletion") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -88,9 +113,11 @@ def completion(
|
|||
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
|
||||
) -> ChatResult | Coroutine[object, object, ChatResult]:
|
||||
python: Final = _PYTHON_COMPLETION
|
||||
if _FLAG.requested(kwargs):
|
||||
return acompletion(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_COMPLETION,
|
||||
native=native_call_hook,
|
||||
|
|
@ -101,7 +128,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
|
|||
python: Final = _PYTHON_ACOMPLETION
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_ACOMPLETION,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
|
|||
MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
|
||||
GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
|
||||
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
|
||||
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
|
||||
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
|
||||
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
|
||||
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
|
||||
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
|
||||
|
|
@ -154,6 +156,9 @@ DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLIN
|
|||
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
|
||||
DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)
|
||||
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
|
||||
HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range(
|
||||
"HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400
|
||||
)
|
||||
|
||||
# Maximum wall-clock seconds a streaming response is allowed to run.
|
||||
# Streams exceeding this duration are terminated with a Timeout error.
|
||||
|
|
@ -256,6 +261,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
|
|||
|
||||
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
|
||||
MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
|
||||
MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range(
|
||||
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize
|
||||
)
|
||||
MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
|
||||
MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3)
|
||||
MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float(
|
||||
|
|
@ -1649,6 +1657,7 @@ BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
|
|||
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
||||
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
|
||||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
|
||||
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
|
||||
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
|
||||
|
|
@ -1812,6 +1821,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED
|
|||
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
|
||||
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
|
||||
LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token"
|
||||
CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1."
|
||||
CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session"
|
||||
CLI_SSO_SESSION_TTL_SECONDS: Final = 600
|
||||
CLI_SESSION_KEY_PREFIX: Final = "cli-session"
|
||||
|
|
@ -2320,6 +2330,7 @@ HARNESS_STDERR_TAIL_LINES: Final = 40
|
|||
HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024
|
||||
HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024
|
||||
HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0
|
||||
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS: Final = 30.0
|
||||
HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
|
||||
{
|
||||
".git",
|
||||
|
|
|
|||
|
|
@ -112,6 +112,7 @@ from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_toke
|
|||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage
|
||||
from litellm.types.extraction import ExtractionResponse, ExtractionUsage
|
||||
from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel
|
||||
from litellm.types.llms.openai import (
|
||||
HttpxBinaryResponseContent,
|
||||
|
|
@ -906,11 +907,15 @@ def _bedrock_pricing_info(
|
|||
return model_info_with_pricing_overrides(profile, registration[1])
|
||||
|
||||
|
||||
def _declares_rates(model_info: ModelInfo | None) -> bool:
|
||||
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
|
||||
def cost_map_key_prices_anything(key: str) -> bool:
|
||||
entry: Final = _raw_cost_map_entry(key)
|
||||
return entry is not None and _cost_map_entry_prices_anything(entry)
|
||||
|
||||
|
||||
def _declares_rates(model_info: ModelInfo | None) -> bool:
|
||||
return model_info is not None and cost_map_key_prices_anything(model_info["key"])
|
||||
|
||||
|
||||
def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | None) -> str | None:
|
||||
"""The provider to price an explicit base_model under, when the deployment's own provider has
|
||||
no rates for it (e.g. chatgpt/ subscription rows) but the base model's own row does"""
|
||||
|
|
@ -1167,6 +1172,8 @@ def get_usage_object(
|
|||
|
||||
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
|
||||
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
|
||||
if isinstance(completion_response, ExtractionResponse):
|
||||
return _extraction_usage(completion_response.usage)
|
||||
if isinstance(completion_response, HttpxBinaryResponseContent):
|
||||
return completion_response.usage
|
||||
if usage_obj is None:
|
||||
|
|
@ -1223,6 +1230,14 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage:
|
|||
)
|
||||
|
||||
|
||||
def _extraction_usage(usage: ExtractionUsage) -> Usage:
|
||||
return Usage(
|
||||
prompt_tokens=usage.input_tokens,
|
||||
completion_tokens=usage.output_tokens,
|
||||
total_tokens=usage.input_tokens + usage.output_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _is_known_usage_objects(usage_obj):
|
||||
"""Returns True if the usage obj is a known Usage type"""
|
||||
return (
|
||||
|
|
@ -1230,6 +1245,7 @@ def _is_known_usage_objects(usage_obj):
|
|||
or isinstance(usage_obj, ResponseAPIUsage)
|
||||
or isinstance(usage_obj, DecisionsUsage)
|
||||
or isinstance(usage_obj, OpenAIDecisionUsage)
|
||||
or isinstance(usage_obj, ExtractionUsage)
|
||||
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
|
||||
)
|
||||
|
||||
|
|
@ -1259,6 +1275,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj
|
|||
return "send_message"
|
||||
elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
|
||||
return "adecisions"
|
||||
elif isinstance(completion_response, ExtractionResponse):
|
||||
return "aextraction"
|
||||
|
||||
return call_type
|
||||
|
||||
|
|
@ -1507,6 +1525,7 @@ def completion_cost(
|
|||
litellm_model_name: str | None = None,
|
||||
router_model_id: str | None = None,
|
||||
litellm_logging_obj: LitellmLoggingObject | None = None,
|
||||
pricing_provider: str | None = None,
|
||||
### SERVICE TIER ###
|
||||
service_tier: str | None = None, # for OpenAI service tier pricing
|
||||
### DATA RESIDENCY ###
|
||||
|
|
@ -1527,6 +1546,8 @@ def completion_cost(
|
|||
"""
|
||||
try:
|
||||
call_type = infer_call_type(call_type, completion_response) or "completion"
|
||||
if pricing_provider is not None:
|
||||
custom_llm_provider = pricing_provider # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups
|
||||
|
||||
if call_type == CallTypes.aresponses_websocket.value and isinstance(
|
||||
completion_response, LiteLLMRealtimeStreamLoggingObject
|
||||
|
|
@ -1552,6 +1573,7 @@ def completion_cost(
|
|||
optional_params=optional_params,
|
||||
custom_pricing=custom_pricing,
|
||||
base_model=base_model,
|
||||
pricing_provider=pricing_provider,
|
||||
standard_built_in_tools_params=standard_built_in_tools_params,
|
||||
litellm_model_name=litellm_model_name,
|
||||
router_model_id=router_model_id,
|
||||
|
|
@ -1645,14 +1667,16 @@ def completion_cost(
|
|||
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
|
||||
)
|
||||
|
||||
for idx, (model, pricing_provider) in enumerate(potential_model_names):
|
||||
for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names):
|
||||
try:
|
||||
if verbose_logger.isEnabledFor(logging.DEBUG):
|
||||
verbose_logger.debug("selected model name for cost calculation: %s", model)
|
||||
|
||||
if completion_response is not None and (
|
||||
isinstance(completion_response, BaseModel) or isinstance(completion_response, dict)
|
||||
): # tts returns a custom class
|
||||
isinstance(completion_response, (BaseModel, dict))
|
||||
or isinstance(completion_response, HttpxBinaryResponseContent)
|
||||
and completion_response.usage is not None
|
||||
):
|
||||
if isinstance(completion_response, dict):
|
||||
usage_obj: dict | Usage | None = completion_response.get("usage", {})
|
||||
else:
|
||||
|
|
@ -1666,6 +1690,8 @@ def completion_cost(
|
|||
)
|
||||
if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
|
||||
_usage = _decisions_usage(usage_obj).model_dump()
|
||||
elif isinstance(usage_obj, ExtractionUsage):
|
||||
_usage = _extraction_usage(usage_obj).model_dump()
|
||||
elif usage_obj is None:
|
||||
_usage = {}
|
||||
elif isinstance(usage_obj, BaseModel):
|
||||
|
|
@ -1712,7 +1738,9 @@ def completion_cost(
|
|||
|
||||
hidden_params = getattr(completion_response, "_hidden_params", None)
|
||||
if hidden_params is not None:
|
||||
custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None)
|
||||
custom_llm_provider = pricing_provider or hidden_params.get(
|
||||
"custom_llm_provider", custom_llm_provider or None
|
||||
)
|
||||
region_name = hidden_params.get("region_name", region_name)
|
||||
|
||||
else:
|
||||
|
|
@ -1998,7 +2026,7 @@ def completion_cost(
|
|||
prompt_tokens=prompt_tokens or 0,
|
||||
completion_tokens=completion_tokens or 0,
|
||||
custom_llm_provider=_pricing_provider_for_cost_candidate(
|
||||
candidate_provider=pricing_provider,
|
||||
candidate_provider=candidate_pricing_provider,
|
||||
request_provider=request_pricing_provider,
|
||||
current_provider=custom_llm_provider,
|
||||
),
|
||||
|
|
@ -2102,7 +2130,7 @@ def completion_cost(
|
|||
_token_type_breakdown = get_token_type_cost_breakdown(
|
||||
model=model,
|
||||
custom_llm_provider=_pricing_provider_for_cost_candidate(
|
||||
candidate_provider=pricing_provider,
|
||||
candidate_provider=candidate_pricing_provider,
|
||||
request_provider=request_pricing_provider,
|
||||
current_provider=custom_llm_provider,
|
||||
),
|
||||
|
|
@ -2189,7 +2217,8 @@ def response_cost_calculator(
|
|||
| Response
|
||||
| SearchResponse
|
||||
| DecisionsResponse
|
||||
| OpenAIDecisionResponse,
|
||||
| OpenAIDecisionResponse
|
||||
| ExtractionResponse,
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
call_type: Literal[
|
||||
|
|
@ -2215,6 +2244,8 @@ def response_cost_calculator(
|
|||
"adecisions",
|
||||
"systemone",
|
||||
"asystemone",
|
||||
"extraction",
|
||||
"aextraction",
|
||||
],
|
||||
optional_params: dict,
|
||||
cache_hit: bool | None = None,
|
||||
|
|
@ -2225,6 +2256,7 @@ def response_cost_calculator(
|
|||
litellm_model_name: str | None = None,
|
||||
router_model_id: str | None = None,
|
||||
litellm_logging_obj: LitellmLoggingObject | None = None,
|
||||
pricing_provider: str | None = None,
|
||||
### SERVICE TIER ###
|
||||
service_tier: str | None = None, # for OpenAI service tier pricing
|
||||
### DATA RESIDENCY ###
|
||||
|
|
@ -2265,6 +2297,7 @@ def response_cost_calculator(
|
|||
litellm_model_name=litellm_model_name,
|
||||
router_model_id=router_model_id,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
pricing_provider=pricing_provider,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from pydantic_core import to_jsonable_python
|
|||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
from litellm.cost_calculator import cost_map_key_prices_anything
|
||||
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
|||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.decisions.transformation import (
|
||||
BaseDecisionsConfig,
|
||||
DecisionsResponseValidationError,
|
||||
image_media_type,
|
||||
image_param,
|
||||
input_images,
|
||||
|
|
@ -687,6 +689,13 @@ def _request_body(call: DecisionsCall) -> Mapping[str, object]:
|
|||
return body
|
||||
|
||||
|
||||
def _pricing_model(response: DecisionsIRResponse, call: DecisionsCall) -> str:
|
||||
returned: Final = get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
|
||||
if returned is not None and cost_map_key_prices_anything(returned):
|
||||
return returned
|
||||
return f"{call.custom_llm_provider}/{call.model}"
|
||||
|
||||
|
||||
def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
|
||||
formatted: Final = _formatted_response(response, call)
|
||||
provider_cost: Final = call.provider_config.provider_reported_cost(response)
|
||||
|
|
@ -694,10 +703,7 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci
|
|||
{
|
||||
"model": f"{call.custom_llm_provider}/{call.model}",
|
||||
"custom_llm_provider": call.custom_llm_provider,
|
||||
"provider_response_model": (
|
||||
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
|
||||
or f"{call.custom_llm_provider}/{call.model}"
|
||||
),
|
||||
"provider_response_model": _pricing_model(response, call),
|
||||
**({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}),
|
||||
}
|
||||
)
|
||||
|
|
@ -727,18 +733,21 @@ def _formatted_response(
|
|||
|
||||
|
||||
def _map_upstream_exception(error: Exception, call: DecisionsCall) -> Exception:
|
||||
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
|
||||
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
|
||||
return litellm.APIConnectionError(
|
||||
message=f"{provider_label} - {error.message}",
|
||||
llm_provider=call.custom_llm_provider,
|
||||
model=f"{call.custom_llm_provider}/{call.model}",
|
||||
)
|
||||
return litellm.exception_type(
|
||||
model=f"{call.custom_llm_provider}/{call.model}",
|
||||
custom_llm_provider=call.custom_llm_provider,
|
||||
original_exception=error,
|
||||
)
|
||||
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
|
||||
model: Final = f"{call.custom_llm_provider}/{call.model}"
|
||||
match error:
|
||||
case DecisionsResponseValidationError():
|
||||
return litellm.InternalServerError(
|
||||
message=f"{provider_label} - {error.message}", llm_provider=call.custom_llm_provider, model=model
|
||||
)
|
||||
case BaseLLMException(status_code_is_synthesized=True):
|
||||
return litellm.APIConnectionError(
|
||||
message=f"{provider_label} - {error.message}", llm_provider=call.custom_llm_provider, model=model
|
||||
)
|
||||
case _:
|
||||
return litellm.exception_type(
|
||||
model=model, custom_llm_provider=call.custom_llm_provider, original_exception=error
|
||||
)
|
||||
|
||||
|
||||
async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
|
||||
|
|
|
|||
|
|
@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import (
|
|||
NATIVE_AEMBEDDING,
|
||||
NATIVE_EMBEDDING,
|
||||
)
|
||||
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
native_call_hook,
|
||||
optional_str,
|
||||
signature,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
__all__ = ("aembedding", "embedding")
|
||||
|
|
@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi
|
|||
PythonAembedding, main.aembedding
|
||||
)
|
||||
_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
|
||||
_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE)
|
||||
|
||||
|
||||
def _public_request(
|
||||
|
|
@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.EMBEDDINGS,
|
||||
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("aembedding") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -64,9 +72,11 @@ def embedding(
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: preserve the public embedding call shape
|
||||
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
|
||||
if _FLAG.requested(kwargs):
|
||||
return aembedding(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=_PYTHON_EMBEDDING,
|
||||
binding=NATIVE_EMBEDDING,
|
||||
native=native_call_hook,
|
||||
|
|
@ -76,7 +86,7 @@ def embedding(
|
|||
async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=_PYTHON_AEMBEDDING,
|
||||
binding=NATIVE_AEMBEDDING,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -122,6 +122,7 @@ class SpeechToCompletionBridgeHandler:
|
|||
if isinstance(result, ModelResponse):
|
||||
return self.transformation_handler.transform_response(
|
||||
model_response=result,
|
||||
model=model,
|
||||
response_format=requested_response_format if isinstance(requested_response_format, str) else None,
|
||||
)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final, cast
|
|||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -38,6 +39,11 @@ class ChatAudioParam(TypedDict):
|
|||
format: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class ChatSpeechConfigAudioParam(TypedDict):
|
||||
speech_config: ReadOnly[Mapping[str, object]]
|
||||
format: ReadOnly[str]
|
||||
|
||||
|
||||
class SpeechToCompletionBridgeTransformationHandler:
|
||||
def _validate_response_format(
|
||||
self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object]
|
||||
|
|
@ -49,6 +55,20 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
return
|
||||
from litellm.exceptions import BadRequestError
|
||||
|
||||
if custom_llm_provider.startswith("vertex_ai"):
|
||||
from litellm.llms.vertex_ai.text_to_speech.transformation import VertexAITextToSpeechConfig
|
||||
|
||||
vertex_supported: Final = ", ".join(
|
||||
sorted({*GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS, *VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS})
|
||||
)
|
||||
raise BadRequestError(
|
||||
message=(
|
||||
f"Vertex AI Gemini TTS does not support response_format='{response_format}'."
|
||||
f" Supported response formats: {vertex_supported}."
|
||||
),
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
)
|
||||
supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS))
|
||||
raise BadRequestError(
|
||||
message=(
|
||||
|
|
@ -75,8 +95,27 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
return response_format if isinstance(response_format, str) else None
|
||||
|
||||
def _chat_audio_param(
|
||||
self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object]
|
||||
) -> ChatAudioParam | None:
|
||||
self,
|
||||
model: str,
|
||||
voice: str | Mapping[str, object] | None,
|
||||
optional_params: Mapping[str, object],
|
||||
drop_params: bool,
|
||||
custom_llm_provider: str,
|
||||
) -> ChatAudioParam | ChatSpeechConfigAudioParam | None:
|
||||
import litellm
|
||||
|
||||
if isinstance(voice, Mapping) and self._is_gemini_tts_model(litellm.model_alias_map.get(model, model)):
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
normalize_gemini_speech_config,
|
||||
reject_unknown_gemini_tts_voice_keys,
|
||||
)
|
||||
|
||||
reject_unknown_gemini_tts_voice_keys(model, voice, drop_params, custom_llm_provider)
|
||||
speech_audio: Final[ChatSpeechConfigAudioParam] = {
|
||||
"speech_config": normalize_gemini_speech_config(voice, drops_unsupported_params(drop_params)),
|
||||
"format": GEMINI_TTS_CHAT_AUDIO_FORMAT,
|
||||
}
|
||||
return speech_audio
|
||||
if not isinstance(voice, str):
|
||||
return None
|
||||
audio_format: Final = self._chat_audio_format(model, optional_params)
|
||||
|
|
@ -90,9 +129,9 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
self,
|
||||
model: str,
|
||||
input: str,
|
||||
voice: str | dict | None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
voice: str | Mapping[str, object] | None,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
headers: dict,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
custom_llm_provider: str,
|
||||
|
|
@ -104,7 +143,13 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
"messages": [user_message],
|
||||
"modalities": ["audio"],
|
||||
**self._chat_completion_params(optional_params),
|
||||
"audio": self._chat_audio_param(model, voice, optional_params),
|
||||
"audio": self._chat_audio_param(
|
||||
model,
|
||||
voice,
|
||||
optional_params,
|
||||
drops_unsupported_params(litellm_params.get("drop_params")),
|
||||
custom_llm_provider,
|
||||
),
|
||||
**litellm_params,
|
||||
"headers": headers,
|
||||
"litellm_logging_obj": litellm_logging_obj,
|
||||
|
|
@ -153,8 +198,9 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
return wav_header + pcm_data
|
||||
|
||||
def _is_gemini_tts_model(self, model: str) -> bool:
|
||||
"""Check if the model is a Gemini TTS model that returns PCM16 data."""
|
||||
return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower())
|
||||
from litellm.utils import is_gemini_tts_model
|
||||
|
||||
return is_gemini_tts_model(model)
|
||||
|
||||
def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]:
|
||||
if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT:
|
||||
|
|
@ -162,7 +208,7 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
return self._convert_pcm16_to_wav(decoded_audio), "audio/wav"
|
||||
|
||||
def transform_response(
|
||||
self, model_response: "ModelResponse", response_format: str | None
|
||||
self, model_response: "ModelResponse", model: str, response_format: str | None
|
||||
) -> "HttpxBinaryResponseContent":
|
||||
import base64
|
||||
|
||||
|
|
@ -176,10 +222,9 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
raise ValueError("No audio part found in the response")
|
||||
decoded_audio: Final = base64.b64decode(audio_part.data)
|
||||
|
||||
model: Final = getattr(model_response, "model", "")
|
||||
content, content_type = (
|
||||
self._gemini_tts_response_body(decoded_audio, response_format)
|
||||
if self._is_gemini_tts_model(model)
|
||||
if self._is_gemini_tts_model(model) or self._is_gemini_tts_model(model_response.model or "")
|
||||
else (decoded_audio, "audio/mpeg")
|
||||
)
|
||||
response: Final = httpx.Response(
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_MAX_SSE_EVENT_SIZE_BYTES,
|
||||
MCP_NPM_CACHE_DIR,
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
|
|
@ -292,6 +293,7 @@ class _MCPResponseStream(httpx2.AsyncByteStream):
|
|||
class _MCPHTTPClient(httpx2.AsyncClient):
|
||||
cleanup_scope: anyio.CancelScope | None = None
|
||||
cleanup_errors: tuple[Exception, ...] = ()
|
||||
post_failure: asyncio.Future[Exception] | None = None
|
||||
|
||||
def _record_cleanup_error(self, error: Exception) -> None:
|
||||
if self.cleanup_scope is not None and self.cleanup_scope.shield:
|
||||
|
|
@ -319,7 +321,12 @@ class _MCPHTTPClient(httpx2.AsyncClient):
|
|||
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
|
||||
if request.method == "POST" and response.is_error and response.status_code != 404:
|
||||
await response.aclose()
|
||||
response.raise_for_status()
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx2.HTTPStatusError as error:
|
||||
if self.post_failure is not None and not self.post_failure.done():
|
||||
self.post_failure.set_result(error)
|
||||
raise
|
||||
if stream:
|
||||
response.stream = _MCPResponseStream(
|
||||
cast(httpx2.AsyncByteStream, response.stream), self._record_cleanup_error
|
||||
|
|
@ -641,6 +648,7 @@ class MCPClient:
|
|||
transport_ctx: Final = streamable_http_client(
|
||||
url=self.server_url,
|
||||
http_client=http_client,
|
||||
max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES,
|
||||
)
|
||||
return transport_ctx, http_client
|
||||
|
||||
|
|
@ -757,6 +765,8 @@ class MCPClient:
|
|||
read_stream: Final = transport[0]
|
||||
write_stream: Final = transport[1]
|
||||
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
|
||||
if isinstance(http_client, _MCPHTTPClient):
|
||||
http_client.post_failure = stream_error # rebind-ok: operation-scoped stream error
|
||||
_notify_stream_error(on_stream_error, stream_error)
|
||||
|
||||
session_kwargs: Final = (
|
||||
|
|
@ -823,6 +833,7 @@ class MCPClient:
|
|||
if isinstance(http_client, _MCPHTTPClient):
|
||||
http_client.cleanup_errors = ()
|
||||
http_client.cleanup_scope = None
|
||||
http_client.post_failure = None # rebind-ok: clear operation-scoped stream error
|
||||
await anyio.lowlevel.checkpoint_if_cancelled()
|
||||
if cleanup_scope.cancel_called:
|
||||
raise (
|
||||
|
|
|
|||
3
litellm/extraction/__init__.py
Normal file
3
litellm/extraction/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from litellm.extraction.main import aextraction, extraction
|
||||
|
||||
__all__ = ["aextraction", "extraction"]
|
||||
351
litellm/extraction/main.py
Normal file
351
litellm/extraction/main.py
Normal file
|
|
@ -0,0 +1,351 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.decisions.call import log_caller_request
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.extraction import (
|
||||
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
|
||||
ExtractionAnswer,
|
||||
ExtractionBase64Image,
|
||||
ExtractionField,
|
||||
ExtractionOptions,
|
||||
ExtractionRequestBody,
|
||||
ExtractionResponse,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders, is_litellm_owned_kwarg
|
||||
from litellm.utils import (
|
||||
ProviderConfigManager,
|
||||
client, # pyright: ignore[reportUnknownVariableType] # client is an untyped decorator
|
||||
get_model_cost_key,
|
||||
)
|
||||
|
||||
ExtractionFields: TypeAlias = Mapping[str, ExtractionField | Mapping[str, object]]
|
||||
ExtractionImages: TypeAlias = Sequence[str | ExtractionBase64Image | Mapping[str, object]]
|
||||
ExtractionOptionsParam: TypeAlias = ExtractionOptions | Mapping[str, object]
|
||||
|
||||
_REQUEST_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody)
|
||||
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_HANDLER: Final = BaseLLMHTTPHandler()
|
||||
_LOCAL_KWARGS: Final = frozenset(
|
||||
(
|
||||
*GenericLiteLLMParams.model_fields,
|
||||
"extra_headers",
|
||||
"base_url",
|
||||
"default_headers",
|
||||
"deployment_id",
|
||||
"additional_drop_params",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
class _ExtractionCall:
|
||||
model: str
|
||||
requested_model: str
|
||||
custom_llm_provider: str
|
||||
provider_config: BaseExtractionConfig
|
||||
body: Mapping[str, object]
|
||||
api_base: str
|
||||
api_key: str | None
|
||||
logging_obj: LiteLLMLoggingObj | None
|
||||
headers: Mapping[str, str]
|
||||
timeout: float | httpx.Timeout | None
|
||||
litellm_params: Mapping[str, object]
|
||||
client: HTTPHandler | AsyncHTTPHandler | None
|
||||
|
||||
|
||||
def supported_extraction_providers() -> tuple[str, ...]:
|
||||
return tuple(
|
||||
provider.value
|
||||
for provider in LlmProviders
|
||||
if ProviderConfigManager.get_provider_extraction_config(model="", provider=provider) is not None
|
||||
)
|
||||
|
||||
|
||||
def _provider_config(model: str, custom_llm_provider: str) -> BaseExtractionConfig:
|
||||
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
|
||||
provider_config: Final = (
|
||||
None if provider is None else ProviderConfigManager.get_provider_extraction_config(model, provider)
|
||||
)
|
||||
if provider_config is None:
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
f"Unknown Extraction provider '{custom_llm_provider}'. "
|
||||
f"Supported providers: {', '.join(supported_extraction_providers())}"
|
||||
),
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
)
|
||||
return provider_config
|
||||
|
||||
|
||||
def _request_fields(
|
||||
*,
|
||||
state: str | None,
|
||||
fields: ExtractionFields | None,
|
||||
images: ExtractionImages | None,
|
||||
reference_date: str | None,
|
||||
options: ExtractionOptionsParam | None,
|
||||
) -> Mapping[str, object]:
|
||||
sent: Final[Mapping[str, object | None]] = MappingProxyType(
|
||||
{"state": state, "fields": fields, "images": images, "reference_date": reference_date, "options": options}
|
||||
)
|
||||
return MappingProxyType({name: value for name, value in sent.items() if value is not None})
|
||||
|
||||
|
||||
def _is_unsupported(name: str, value: object, policy: DropPolicy) -> bool:
|
||||
return (
|
||||
value is not None
|
||||
and value is not False
|
||||
and name not in _LOCAL_KWARGS
|
||||
and not is_litellm_owned_kwarg(name)
|
||||
and not policy.drops(name)
|
||||
)
|
||||
|
||||
|
||||
def _unsupported_params(kwargs: Mapping[str, object]) -> tuple[str, ...]:
|
||||
policy: Final = drop_policy(kwargs)
|
||||
return tuple(sorted(name for name, value in kwargs.items() if _is_unsupported(name, value, policy)))
|
||||
|
||||
|
||||
def _prepare_call(
|
||||
*,
|
||||
model: str,
|
||||
request_fields: Mapping[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
custom_llm_provider: str | None,
|
||||
extra_headers: Mapping[str, str] | None,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> _ExtractionCall:
|
||||
log_caller_request(kwargs, request_fields)
|
||||
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
provider_config: Final = _provider_config(model, provider)
|
||||
unsupported: Final = _unsupported_params(kwargs)
|
||||
if unsupported:
|
||||
raise litellm.UnsupportedParamsError(
|
||||
message=(
|
||||
f"{provider} does not support parameters: {list(unsupported)}, for model={model}. "
|
||||
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
|
||||
"`litellm_settings:\n drop_params: true`\n"
|
||||
),
|
||||
model=model,
|
||||
llm_provider=provider,
|
||||
)
|
||||
try:
|
||||
request: Final = _REQUEST_ADAPTER.validate_python(request_fields)
|
||||
except ValidationError as error:
|
||||
raise litellm.BadRequestError(
|
||||
message=f"Invalid Extraction request: {error}", model=model, llm_provider=provider
|
||||
) from error
|
||||
canonical_model: Final = provider_config.canonical_model(upstream_model)
|
||||
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
|
||||
if resolved_api_base is None:
|
||||
raise litellm.BadRequestError(
|
||||
message=provider_config.missing_api_base_message(provider), model=model, llm_provider=provider
|
||||
)
|
||||
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
|
||||
if resolved_api_key is None:
|
||||
raise litellm.BadRequestError(
|
||||
message=f"Missing API key for Extraction provider '{provider}'", model=model, llm_provider=provider
|
||||
)
|
||||
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
http_client: Final = kwargs.get("client")
|
||||
if isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging takes untyped dicts
|
||||
kwargs=dict(kwargs),
|
||||
model=canonical_model,
|
||||
litellm_params={
|
||||
"litellm_call_id": kwargs.get("litellm_call_id"),
|
||||
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
|
||||
},
|
||||
custom_llm_provider=provider,
|
||||
)
|
||||
return _ExtractionCall(
|
||||
model=canonical_model,
|
||||
requested_model=model,
|
||||
custom_llm_provider=provider,
|
||||
provider_config=provider_config,
|
||||
body=provider_config.transform_extraction_request(canonical_model, request),
|
||||
api_base=resolved_api_base,
|
||||
api_key=resolved_api_key,
|
||||
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
|
||||
headers=MappingProxyType(dict(extra_headers or {})),
|
||||
timeout=timeout,
|
||||
litellm_params=kwargs,
|
||||
client=http_client if isinstance(http_client, HTTPHandler | AsyncHTTPHandler) else None,
|
||||
)
|
||||
|
||||
|
||||
def _state_rewritten_by_guardrail(litellm_params: Mapping[str, object]) -> bool:
|
||||
metadata: Final = litellm_params.get(get_metadata_variable_name_from_kwargs(litellm_params))
|
||||
return (
|
||||
isinstance(metadata, Mapping)
|
||||
and _METADATA_ADAPTER.validate_python(metadata).get(EXTRACTION_STATE_REWRITTEN_METADATA_KEY) is True
|
||||
)
|
||||
|
||||
|
||||
def _without_spans(answers: Mapping[str, ExtractionAnswer]) -> Mapping[str, ExtractionAnswer]:
|
||||
return {name: answer.model_copy(update={"span": None, "spans": None}) for name, answer in answers.items()}
|
||||
|
||||
|
||||
def _formatted_response(response: ExtractionResponse, call: _ExtractionCall) -> ExtractionResponse:
|
||||
spans_index_hidden_text: Final = response.state_text is None and _state_rewritten_by_guardrail(call.litellm_params)
|
||||
formatted: Final = response.model_copy(
|
||||
update={
|
||||
"model": call.requested_model,
|
||||
**({"answers": _without_spans(response.answers)} if spans_index_hidden_text else {}),
|
||||
}
|
||||
)
|
||||
formatted.set_hidden_params(
|
||||
{
|
||||
"model": f"{call.custom_llm_provider}/{call.model}",
|
||||
"custom_llm_provider": call.custom_llm_provider,
|
||||
"provider_response_model": (
|
||||
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
|
||||
or f"{call.custom_llm_provider}/{call.model}"
|
||||
),
|
||||
}
|
||||
)
|
||||
return formatted
|
||||
|
||||
|
||||
def _upstream_exception(error: Exception, call: _ExtractionCall) -> Exception:
|
||||
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
|
||||
return litellm.APIConnectionError(
|
||||
message=f"{call.custom_llm_provider.capitalize()}Exception - {error.message}",
|
||||
llm_provider=call.custom_llm_provider,
|
||||
model=f"{call.custom_llm_provider}/{call.model}",
|
||||
)
|
||||
return litellm.exception_type( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
|
||||
model=f"{call.custom_llm_provider}/{call.model}",
|
||||
custom_llm_provider=call.custom_llm_provider,
|
||||
original_exception=error,
|
||||
)
|
||||
|
||||
|
||||
async def _asend(call: _ExtractionCall) -> ExtractionResponse:
|
||||
try:
|
||||
response: Final = await _HANDLER.aextraction(
|
||||
model=call.model,
|
||||
custom_llm_provider=call.custom_llm_provider,
|
||||
logging_obj=call.logging_obj,
|
||||
provider_config=call.provider_config,
|
||||
body=call.body,
|
||||
api_base=call.api_base,
|
||||
api_key=call.api_key,
|
||||
headers=call.headers,
|
||||
timeout=call.timeout,
|
||||
litellm_params=call.litellm_params,
|
||||
client=call.client if isinstance(call.client, AsyncHTTPHandler) else None,
|
||||
)
|
||||
except Exception as error:
|
||||
raise _upstream_exception(error, call) from error
|
||||
return _formatted_response(response, call)
|
||||
|
||||
|
||||
def _send(call: _ExtractionCall) -> ExtractionResponse:
|
||||
try:
|
||||
response: Final = _HANDLER.extraction(
|
||||
model=call.model,
|
||||
custom_llm_provider=call.custom_llm_provider,
|
||||
logging_obj=call.logging_obj,
|
||||
provider_config=call.provider_config,
|
||||
body=call.body,
|
||||
api_base=call.api_base,
|
||||
api_key=call.api_key,
|
||||
headers=call.headers,
|
||||
timeout=call.timeout,
|
||||
litellm_params=call.litellm_params,
|
||||
client=call.client if isinstance(call.client, HTTPHandler) else None,
|
||||
)
|
||||
except Exception as error:
|
||||
raise _upstream_exception(error, call) from error
|
||||
return _formatted_response(response, call)
|
||||
|
||||
|
||||
@client
|
||||
async def aextraction(
|
||||
model: str,
|
||||
fields: ExtractionFields,
|
||||
state: str | None = None,
|
||||
images: ExtractionImages | None = None,
|
||||
reference_date: str | None = None,
|
||||
options: ExtractionOptionsParam | None = None,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: Mapping[str, str] | None = None,
|
||||
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
|
||||
) -> ExtractionResponse:
|
||||
call: Final = await asyncio.to_thread(
|
||||
partial(
|
||||
_prepare_call,
|
||||
model=model,
|
||||
request_fields=_request_fields(
|
||||
state=state, fields=fields, images=images, reference_date=reference_date, options=options
|
||||
),
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
)
|
||||
return await _asend(call)
|
||||
|
||||
|
||||
@client
|
||||
def extraction(
|
||||
model: str,
|
||||
fields: ExtractionFields,
|
||||
state: str | None = None,
|
||||
images: ExtractionImages | None = None,
|
||||
reference_date: str | None = None,
|
||||
options: ExtractionOptionsParam | None = None,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: Mapping[str, str] | None = None,
|
||||
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
|
||||
) -> ExtractionResponse:
|
||||
call: Final = _prepare_call(
|
||||
model=model,
|
||||
request_fields=_request_fields(
|
||||
state=state, fields=fields, images=images, reference_date=reference_date, options=options
|
||||
),
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
return _send(call)
|
||||
|
||||
|
||||
__all__ = ["aextraction", "extraction"]
|
||||
|
|
@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget
|
|||
from litellm.harness.errors import HarnessError, HarnessInstallFailed
|
||||
from litellm.harness.types import Harness, Usage
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -690,18 +691,21 @@ class ModelEndpoint:
|
|||
async def _invoke_sdk(
|
||||
self,
|
||||
route: str,
|
||||
kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs
|
||||
kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs
|
||||
) -> object:
|
||||
if route == ROUTE_MESSAGES:
|
||||
return await litellm.anthropic.messages.acreate(**kwargs)
|
||||
if route == ROUTE_CHAT:
|
||||
if kwargs.get("stream"):
|
||||
stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field
|
||||
kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion
|
||||
"include_usage": True,
|
||||
**stream_options,
|
||||
if not kwargs.get("stream"):
|
||||
return await litellm.acompletion(**kwargs)
|
||||
stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field
|
||||
return await litellm.acompletion(
|
||||
**{
|
||||
**kwargs,
|
||||
"stream_options": {"include_usage": True, **stream_options},
|
||||
STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options,
|
||||
}
|
||||
return await litellm.acompletion(**kwargs)
|
||||
)
|
||||
return await litellm.aresponses(**kwargs)
|
||||
|
||||
async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response:
|
||||
|
|
|
|||
|
|
@ -12,17 +12,27 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import os
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES
|
||||
from litellm.constants import (
|
||||
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS,
|
||||
HARNESS_STDERR_TAIL_LINES,
|
||||
HARNESS_STREAM_READ_CHUNK_BYTES,
|
||||
)
|
||||
from litellm.harness.context import SessionContext
|
||||
from litellm.harness.errors import HarnessInstallFailed, SandboxError
|
||||
from litellm.harness.errors import HarnessError, HarnessInstallFailed, SandboxError
|
||||
from litellm.harness.handlers.base import BaseHarnessHandler
|
||||
from litellm.harness.sandbox.base import Process, Sandbox
|
||||
from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox
|
||||
from litellm.harness.types import Event
|
||||
from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup
|
||||
from litellm.llms.base_llm.harness.transformation import (
|
||||
BaseCLIHarnessConfig,
|
||||
HarnessSessionCopy,
|
||||
HarnessSessionSetup,
|
||||
HarnessTurnError,
|
||||
)
|
||||
from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files
|
||||
|
||||
# Link <private_dir>/<dir> to a LiteLLM-owned cache dir so a later session can resume.
|
||||
|
|
@ -81,6 +91,52 @@ async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> N
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResumeFailed:
|
||||
reason: str
|
||||
|
||||
|
||||
async def bounded_resume_run(
|
||||
sandbox: Sandbox, cmd: Sequence[str], env: Mapping[str, str]
|
||||
) -> CompletedRun | ResumeFailed:
|
||||
try:
|
||||
run: Final = await sandbox.run(cmd, env=env, timeout=HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS)
|
||||
except SandboxError as exc:
|
||||
return ResumeFailed(str(exc))
|
||||
return run if run.exit_code == 0 else ResumeFailed(run.stderr.strip() or f"exit code {run.exit_code}")
|
||||
|
||||
|
||||
async def import_session_copy(sandbox: Sandbox, copy: HarnessSessionCopy, env: Mapping[str, str]) -> str | ResumeFailed:
|
||||
try:
|
||||
await sandbox.write(copy.path, copy.data)
|
||||
except SandboxError as exc:
|
||||
return ResumeFailed(str(exc))
|
||||
imported: Final = await bounded_resume_run(sandbox, copy.import_command, env)
|
||||
return imported if isinstance(imported, ResumeFailed) else copy.session_id
|
||||
|
||||
|
||||
async def resumable_session_id(
|
||||
config: BaseCLIHarnessConfig[object, object],
|
||||
sandbox: Sandbox,
|
||||
setup_env: Mapping[str, str],
|
||||
private_dir: str,
|
||||
native_session_id: str,
|
||||
) -> str | ResumeFailed:
|
||||
export_cmd: Final = config.transform_resume_export_command(native_session_id)
|
||||
if export_cmd is None:
|
||||
return native_session_id
|
||||
exported: Final = await bounded_resume_run(sandbox, export_cmd, setup_env)
|
||||
if isinstance(exported, ResumeFailed):
|
||||
return exported
|
||||
try:
|
||||
copy: Final = config.transform_resume_copy(exported.stdout, private_dir, sandbox.workdir)
|
||||
except HarnessError as exc:
|
||||
return ResumeFailed(str(exc))
|
||||
if copy is None:
|
||||
return native_session_id
|
||||
return await import_session_copy(sandbox, copy, setup_env)
|
||||
|
||||
|
||||
async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None:
|
||||
for skill in skills:
|
||||
name = os.path.basename(os.path.realpath(os.fspath(skill)))
|
||||
|
|
@ -89,13 +145,14 @@ async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str)
|
|||
|
||||
|
||||
class CLIHarnessHandler(BaseHarnessHandler):
|
||||
config: BaseCLIHarnessConfig
|
||||
config: BaseCLIHarnessConfig[object, object]
|
||||
|
||||
def __init__(self, config: BaseCLIHarnessConfig) -> None:
|
||||
super().__init__(config)
|
||||
self._private_dir: str | None = None
|
||||
self._setup: HarnessSessionSetup | None = None
|
||||
self._native_id: str | None = None
|
||||
self._pending_resume: str | None = None
|
||||
self._proc: Process | None = None
|
||||
|
||||
async def start(self, ctx: SessionContext) -> None:
|
||||
|
|
@ -119,6 +176,9 @@ class CLIHarnessHandler(BaseHarnessHandler):
|
|||
async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
|
||||
if self._setup is None or self._private_dir is None:
|
||||
raise RuntimeError("CLIHarnessHandler.turn() called before start()")
|
||||
if self._pending_resume is not None:
|
||||
self._native_id = await self._prepared_session_id(ctx, self._setup, self._private_dir, self._pending_resume)
|
||||
self._pending_resume = None
|
||||
request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id)
|
||||
argv: Final = list(request.argv) # mutable-ok: Sandbox.exec takes list[str]
|
||||
proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd)
|
||||
|
|
@ -158,4 +218,16 @@ class CLIHarnessHandler(BaseHarnessHandler):
|
|||
return self._native_id
|
||||
|
||||
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
|
||||
self._native_id = native_session_id
|
||||
if native_session_id != self._native_id:
|
||||
self._native_id = native_session_id
|
||||
self._pending_resume = native_session_id
|
||||
|
||||
async def _prepared_session_id(
|
||||
self, ctx: SessionContext, setup: HarnessSessionSetup, private_dir: str, saved_id: str
|
||||
) -> str:
|
||||
target: Final = await resumable_session_id(self.config, ctx.sandbox, setup.env, private_dir, saved_id)
|
||||
match target:
|
||||
case ResumeFailed(reason=reason):
|
||||
raise HarnessTurnError(f"could not prepare session {saved_id} for resume: {reason}")
|
||||
case session_id:
|
||||
return session_id
|
||||
|
|
|
|||
|
|
@ -183,7 +183,7 @@ class DeepAgentsHandler(BaseHarnessHandler):
|
|||
break
|
||||
resume: dict[str, object] = {} # mutable-ok: Command(resume=) payload, filled per answered approval
|
||||
for interrupt in state.interrupts:
|
||||
decisions: list[dict[str, Any]] = [] # mutable-ok: HITL decisions collected across awaited approvals
|
||||
decisions: list[dict[str, str]] = [] # mutable-ok: HITL decisions collected across awaited approvals
|
||||
for request in approval_requests(getattr(interrupt, "value", None)):
|
||||
approval = Approval(
|
||||
tool=normalized_tool_name(str(request.get("name") or "")),
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing import Final
|
|||
|
||||
from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS
|
||||
from litellm.harness.errors import SandboxError
|
||||
from litellm.harness.sandbox.base import CompletedRun
|
||||
from litellm.harness.sandbox.base import CompletedRun, Process
|
||||
from litellm.harness.sandbox.snapshot import snapshot_local
|
||||
|
||||
_SECRET_PREFIXES: Final = (
|
||||
|
|
@ -103,7 +103,7 @@ class SubprocessHandle:
|
|||
await self._proc.wait()
|
||||
|
||||
|
||||
async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
|
||||
async def _read_all(handle: Process) -> tuple[bytes, bytes, int]:
|
||||
if handle.stdin is not None:
|
||||
handle.stdin.close()
|
||||
stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
|
||||
|
|
@ -111,13 +111,16 @@ async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
|
|||
return stdout, stderr, exit_code
|
||||
|
||||
|
||||
async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
|
||||
"""Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout."""
|
||||
async def collect_output(handle: Process, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
|
||||
"""Close stdin, read stdout/stderr to EOF; kill on cancel, and kill and raise SandboxError on timeout."""
|
||||
try:
|
||||
stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
await handle.kill()
|
||||
raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}")
|
||||
except asyncio.CancelledError:
|
||||
await handle.kill()
|
||||
raise
|
||||
return CompletedRun(
|
||||
stdout=stdout.decode("utf-8", errors="replace"),
|
||||
stderr=stderr.decode("utf-8", errors="replace"),
|
||||
|
|
@ -180,14 +183,15 @@ class LocalSandbox:
|
|||
self._check_open()
|
||||
if not cmd:
|
||||
raise SandboxError("exec() needs a non-empty command")
|
||||
work_dir: Final = self._resolve_cwd(cwd)
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=self._resolve_cwd(cwd),
|
||||
env=self.child_env(env),
|
||||
cwd=work_dir,
|
||||
env=self.child_env(MappingProxyType({**(env or {}), "PWD": work_dir})),
|
||||
start_new_session=True,
|
||||
)
|
||||
except (FileNotFoundError, PermissionError) as exc:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from datetime import datetime
|
|||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import STREAM_SSE_DONE_STRING
|
||||
|
|
@ -25,6 +26,10 @@ from litellm.types.interactions import (
|
|||
)
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
_INTERACTIONS_STREAMING_CHUNK_ADAPTER: Final = TypeAdapter(
|
||||
object, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
class BaseInteractionsAPIStreamingIterator:
|
||||
"""
|
||||
|
|
@ -83,7 +88,7 @@ class BaseInteractionsAPIStreamingIterator:
|
|||
|
||||
try:
|
||||
# Parse the JSON chunk
|
||||
parsed_chunk: Final = json.loads(stripped_chunk)
|
||||
parsed_chunk: Final = _INTERACTIONS_STREAMING_CHUNK_ADAPTER.validate_python(json.loads(stripped_chunk))
|
||||
|
||||
# Format as InteractionsAPIStreamingResponse
|
||||
if isinstance(parsed_chunk, dict):
|
||||
|
|
|
|||
|
|
@ -329,11 +329,14 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad
|
|||
Helper to get litellm metadata for spend tracking
|
||||
|
||||
PATCH for issue where both `litellm_metadata` and `metadata` are present in the kwargs
|
||||
and user_api_key values are in 'metadata'.
|
||||
and user_api_key values are in 'metadata'. The proxy receive stamp rides along for the same
|
||||
reason: a guardrail hook can plant `litellm_metadata` on a chat request after the proxy stamped
|
||||
`metadata`, and the response timing headers anchor on that stamp.
|
||||
"""
|
||||
potential_spend_tracking_metadata_substring: Final = "user_api_key"
|
||||
for key, value in dict(metadata).items():
|
||||
if potential_spend_tracking_metadata_substring in key:
|
||||
carry_receive_stamp = key == "litellm_received_at" and key not in litellm_metadata
|
||||
if potential_spend_tracking_metadata_substring in key or carry_receive_stamp:
|
||||
litellm_metadata[key] = value
|
||||
return litellm_metadata
|
||||
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
|
|||
"aws_external_id",
|
||||
"aws_session_tags",
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"aws_bedrock_control_plane_endpoint",
|
||||
"aws_bedrock_project_id",
|
||||
}
|
||||
)
|
||||
|
|
@ -76,6 +77,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"otpm",
|
||||
"use_xai_oauth",
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"github_copilot_user_session",
|
||||
"fireworks_forward_user_id",
|
||||
"responses_api_version",
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Helper functions for health check calls.
|
|||
|
||||
import base64
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import DocumentType
|
||||
|
|
@ -25,6 +26,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
|
|||
)
|
||||
|
||||
ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16
|
||||
EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
|
||||
{"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})}
|
||||
)
|
||||
|
||||
|
||||
def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None:
|
||||
|
|
@ -215,6 +219,7 @@ class HealthCheckHelpers:
|
|||
"anthropic_messages",
|
||||
"ocr",
|
||||
"decisions",
|
||||
"extraction",
|
||||
],
|
||||
Callable,
|
||||
]:
|
||||
|
|
@ -237,7 +242,11 @@ class HealthCheckHelpers:
|
|||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
get_audio_file_for_health_check,
|
||||
)
|
||||
from litellm.litellm_core_utils.health_check_utils import DECISIONS_CALL_PARAMS, _filter_model_params
|
||||
from litellm.litellm_core_utils.health_check_utils import (
|
||||
DECISIONS_CALL_PARAMS,
|
||||
EXTRACTION_CALL_PARAMS,
|
||||
_filter_model_params,
|
||||
)
|
||||
from litellm.realtime_api.main import realtime_health_check
|
||||
|
||||
return {
|
||||
|
|
@ -322,4 +331,13 @@ class HealthCheckHelpers:
|
|||
}
|
||||
)
|
||||
),
|
||||
"extraction": lambda: litellm.aextraction(
|
||||
**EXTRACTION_CALL_PARAMS.validate_python(
|
||||
{
|
||||
"state": prompt or "health check",
|
||||
"fields": EXTRACTION_HEALTH_CHECK_FIELDS,
|
||||
**_filter_model_params(model_params=model_params),
|
||||
}
|
||||
)
|
||||
),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,8 +8,10 @@ from typing import Final
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.types.decisions import DecisionsCallParams
|
||||
from litellm.types.extraction import ExtractionCallParams
|
||||
|
||||
DECISIONS_CALL_PARAMS: Final[TypeAdapter[DecisionsCallParams]] = TypeAdapter(DecisionsCallParams)
|
||||
EXTRACTION_CALL_PARAMS: Final[TypeAdapter[ExtractionCallParams]] = TypeAdapter(ExtractionCallParams)
|
||||
OPTIONAL_STR: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import subprocess
|
|||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime as dt_object
|
||||
from functools import lru_cache
|
||||
|
|
@ -117,7 +117,8 @@ from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request
|
|||
from litellm.litellm_core_utils.redact_messages import (
|
||||
redact_message_input_output_from_custom_logger,
|
||||
redact_message_input_output_from_logging,
|
||||
redact_streaming_responses_for_custom_logger,
|
||||
redacted_model_call_details_for_custom_logger,
|
||||
redacted_response_obj_for_custom_logger,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
|
|
@ -131,11 +132,13 @@ from litellm.responses.utils import ResponseAPILoggingUtils
|
|||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
from litellm.types.containers.main import ContainerObject
|
||||
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
|
||||
from litellm.types.extraction import ExtractionResponse
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
|
||||
from litellm.types.interactions import (
|
||||
InteractionsAPIResponse,
|
||||
InteractionsAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
Batch,
|
||||
|
|
@ -337,6 +340,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
|
|||
|
||||
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
|
||||
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
_METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
|
||||
("user_api_key_auth", "user_api_key_budget_reservation")
|
||||
|
|
@ -571,6 +575,27 @@ class UncachedPrice:
|
|||
cost_breakdown: CostBreakdown | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledRoundsSpend:
|
||||
usage: Usage
|
||||
cost: float
|
||||
cost_breakdown: CostBreakdown | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledRoundsDispatch:
|
||||
"""A parked success log for a stream bridged through chat completions that billed several rounds, and the spend
|
||||
its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged
|
||||
stream."""
|
||||
|
||||
dispatch: Coroutine[object, object, None]
|
||||
spend: BilledRoundsSpend
|
||||
|
||||
|
||||
def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float:
|
||||
return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0)
|
||||
|
||||
|
||||
def _timestamp_seconds(moment: object) -> float | None:
|
||||
if isinstance(moment, datetime.datetime):
|
||||
return moment.timestamp()
|
||||
|
|
@ -579,6 +604,10 @@ def _timestamp_seconds(moment: object) -> float | None:
|
|||
return None
|
||||
|
||||
|
||||
def same_attempt_marker(first: object, second: object) -> bool:
|
||||
return first is second or first == second
|
||||
|
||||
|
||||
class Logging(LiteLLMLoggingBaseClass):
|
||||
global \
|
||||
supabaseClient, \
|
||||
|
|
@ -592,6 +621,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
prometheusLogger, \
|
||||
slack_app
|
||||
custom_pricing: bool = False
|
||||
replaced_global_callback_keys: frozenset[str] = frozenset()
|
||||
stream_options = None
|
||||
litellm_request_debug: bool = False
|
||||
streamed_anthropic_message_id: str | None = None
|
||||
|
|
@ -599,6 +629,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
classifier_input: Mapping[str, JsonValue] | None = None
|
||||
baseline_cache_context: "BaselineCacheContext | None" = None
|
||||
baseline_observation: "CapturedBaselineObservation | None" = None
|
||||
_on_failure_payload_built: "Callable[[], None] | None" = None
|
||||
|
||||
@property
|
||||
def _defer_async_logging(self) -> bool:
|
||||
|
|
@ -646,11 +677,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
litellm_call_id: str,
|
||||
function_id: str,
|
||||
litellm_trace_id: str | None = None,
|
||||
dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None,
|
||||
dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
|
||||
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
|
||||
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
|
||||
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
|
||||
dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
|
||||
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
|
||||
dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
|
||||
applied_guardrails: list[str] | None = None,
|
||||
kwargs: dict | None = None,
|
||||
log_raw_request_response: bool = False,
|
||||
|
|
@ -695,6 +726,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._pre_call_session_id: str = session_id_var.get()
|
||||
_sid: Final = kwargs.get("litellm_session_id") if kwargs else None
|
||||
self.litellm_session_id: str = str(_sid) if _sid else ""
|
||||
self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True
|
||||
# supports_correlation_logging is False for calls originating from the
|
||||
# sync client entry point (wrapper() in utils.py): a plain OS thread
|
||||
# has no per-call context isolation the way an asyncio Task does, and
|
||||
|
|
@ -723,15 +755,17 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.raw_request_only = raw_request_only
|
||||
|
||||
# Initialize dynamic callbacks
|
||||
self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks
|
||||
self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks
|
||||
self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
|
||||
dynamic_success_callbacks
|
||||
)
|
||||
self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
|
||||
dynamic_async_success_callbacks
|
||||
)
|
||||
self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks
|
||||
self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = (
|
||||
self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
|
||||
dynamic_failure_callbacks
|
||||
)
|
||||
self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
|
||||
dynamic_async_failure_callbacks
|
||||
)
|
||||
|
||||
|
|
@ -794,6 +828,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.defer_async_logging: bool = False
|
||||
self.enqueue_deferred_logging: Callable[[], None] | None = None
|
||||
self._async_success_scheduled: bool = False
|
||||
self._async_success_enqueued: bool = False
|
||||
self._sync_success_scheduled: bool = False
|
||||
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
|
@ -822,6 +857,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"litellm_call_id": litellm_call_id,
|
||||
"start_time": start_time,
|
||||
"first_api_call_start_time": None,
|
||||
"current_attempt_api_call_start_time": None,
|
||||
"litellm_params": litellm_params,
|
||||
}
|
||||
new_call.streaming_chunks = []
|
||||
|
|
@ -847,10 +883,32 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.dynamic_async_failure_callbacks, callback
|
||||
)
|
||||
|
||||
def replace_global_callback(self, callback: CustomLogger) -> None:
|
||||
self.add_dynamic_callback(callback)
|
||||
self.replaced_global_callback_keys = self.replaced_global_callback_keys | {
|
||||
litellm.logging_callback_manager.get_custom_logger_key(callback)
|
||||
}
|
||||
|
||||
def add_dynamic_function_callbacks(
|
||||
self,
|
||||
*,
|
||||
async_success: Callable[..., object],
|
||||
failure: Callable[..., object],
|
||||
async_failure: Callable[..., object],
|
||||
) -> None:
|
||||
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_success_callbacks, async_success
|
||||
)
|
||||
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure)
|
||||
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_failure_callbacks, async_failure
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _with_dynamic_callback(
|
||||
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
|
||||
) -> list[str | Callable | CustomLogger]:
|
||||
callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None,
|
||||
callback: Callable[..., object] | CustomLogger,
|
||||
) -> list[str | Callable[..., object] | CustomLogger]:
|
||||
existing: Final = tuple(callbacks or ())
|
||||
return [*existing, *(() if callback in existing else (callback,))]
|
||||
|
||||
|
|
@ -887,9 +945,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def _process_dynamic_callback_list(
|
||||
self,
|
||||
callback_list: list[str | Callable | CustomLogger] | None,
|
||||
callback_list: list[str | Callable[..., object] | CustomLogger] | None,
|
||||
dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"],
|
||||
) -> list[str | Callable | CustomLogger] | None:
|
||||
) -> list[str | Callable[..., object] | CustomLogger] | None:
|
||||
"""
|
||||
Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks
|
||||
|
||||
|
|
@ -901,7 +959,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if callback_list is None:
|
||||
return None
|
||||
|
||||
processed_list: Final[list[str | Callable | CustomLogger]] = []
|
||||
processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = []
|
||||
for callback in callback_list:
|
||||
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
|
||||
for callback_instance in self._resolve_dynamic_callback_string(callback):
|
||||
|
|
@ -995,6 +1053,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
return model_id
|
||||
return None
|
||||
|
||||
def get_router_pricing_provider(self) -> str | None:
|
||||
"""The deployment's ``model_info.pricing_provider``, the provider its cost-map lookups run under."""
|
||||
litellm_params: Final = cast( # cast-ok: callback metadata is caller-provided
|
||||
Mapping[str, object] | None, getattr(self, "litellm_params", None)
|
||||
)
|
||||
if not litellm_params:
|
||||
return None
|
||||
candidates: Final = (
|
||||
_METADATA_MAPPING.validate_python(
|
||||
_METADATA_MAPPING.validate_python(litellm_params.get(key) or {}).get("model_info") or {}
|
||||
).get("pricing_provider")
|
||||
for key in ("litellm_metadata", "metadata")
|
||||
)
|
||||
return next((provider for provider in candidates if isinstance(provider, str) and provider), None)
|
||||
|
||||
def _routed_deployment_model(self) -> str | None:
|
||||
litellm_params: Final = self._model_call_details_view.get("litellm_params")
|
||||
metadata: Final = (
|
||||
|
|
@ -1535,7 +1608,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# A proxy-gate failure (auth, rate limit, client disconnect) fires pre_call for its callbacks
|
||||
# without any provider handoff, so it must not stamp the handoff instants.
|
||||
if LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL not in self.model_call_details:
|
||||
self.model_call_details["api_call_start_time"] = datetime.datetime.now()
|
||||
api_call_start_time: Final = datetime.datetime.now()
|
||||
self.model_call_details["api_call_start_time"] = api_call_start_time
|
||||
self.model_call_details["current_attempt_api_call_start_time"] = api_call_start_time
|
||||
# Set-once first provider-handoff instant. api_call_start_time
|
||||
# is overwritten on every retry, so it can't measure one-time
|
||||
# preprocessing; pinning the first attempt excludes retry loops
|
||||
|
|
@ -1543,11 +1618,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# litellm_params["metadata"] (caller request metadata, typed
|
||||
# Dict[str, str], echoed downstream; a datetime breaks it).
|
||||
if self.model_call_details.get("first_api_call_start_time") is None:
|
||||
self.model_call_details["first_api_call_start_time"] = self.model_call_details[
|
||||
"api_call_start_time"
|
||||
]
|
||||
self.model_call_details["first_api_call_start_time"] = api_call_start_time
|
||||
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
|
||||
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
|
||||
callbacks: Final = (
|
||||
*self._without_replaced_global_callbacks(litellm.input_callback),
|
||||
*(self.dynamic_input_callbacks or ()),
|
||||
)
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if callback == "supabase" and supabaseClient is not None:
|
||||
|
|
@ -1720,7 +1796,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
|
||||
|
||||
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
|
||||
callbacks: Final = (
|
||||
*self._without_replaced_global_callbacks(litellm.input_callback),
|
||||
*(self.dynamic_input_callbacks or ()),
|
||||
)
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if callback == "sentry" and add_breadcrumb:
|
||||
|
|
@ -2020,6 +2099,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"prompt": prompt,
|
||||
"standard_built_in_tools_params": self.standard_built_in_tools_params,
|
||||
"router_model_id": router_model_id,
|
||||
"pricing_provider": self.get_router_pricing_provider(),
|
||||
"litellm_logging_obj": self,
|
||||
"data_residency": (
|
||||
self.litellm_params.get("data_residency")
|
||||
|
|
@ -2164,7 +2244,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
pricing: Final = pricing_entry_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=result,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
custom_llm_provider=self.get_router_pricing_provider() or custom_llm_provider,
|
||||
custom_pricing=self.custom_pricing_for(result),
|
||||
base_model=get_base_model_from_metadata(model_call_details=self.model_call_details),
|
||||
router_model_id=router_model_id or self.get_router_model_id(),
|
||||
|
|
@ -2190,20 +2270,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
usage=usage,
|
||||
)
|
||||
|
||||
def custom_pricing_for(self, result: object) -> bool:
|
||||
litellm_params: Final = getattr(self, "litellm_params", None)
|
||||
def spilled_over_ptu_for(self, result: object) -> bool:
|
||||
result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({})
|
||||
additional_headers: Final = (
|
||||
result_hidden_params.get("additional_headers")
|
||||
if isinstance(result_hidden_params, dict)
|
||||
else getattr(result_hidden_params, "additional_headers", None)
|
||||
)
|
||||
spilled_over: Final = is_spilled_over_ptu_request(
|
||||
model_info=_deployment_model_info(litellm_params),
|
||||
return is_spilled_over_ptu_request(
|
||||
model_info=_deployment_model_info(getattr(self, "litellm_params", None)),
|
||||
response_headers=self.model_call_details.get("response_headers"),
|
||||
additional_headers=additional_headers,
|
||||
)
|
||||
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
|
||||
|
||||
def custom_pricing_for(self, result: object) -> bool:
|
||||
if self.spilled_over_ptu_for(result):
|
||||
return False
|
||||
return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None))
|
||||
|
||||
def _prompt_for_cost_calculation(self) -> str:
|
||||
"""
|
||||
|
|
@ -2365,6 +2448,31 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.model_call_details["combined_usage_object"] = usage
|
||||
self.model_call_details["response_cost"] = response_cost
|
||||
|
||||
def record_failure_after_earlier_rounds(
|
||||
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
|
||||
) -> None:
|
||||
"""Bill a failed round on the failure log after the rounds of its request before it were billed on the success
|
||||
log.
|
||||
|
||||
The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the
|
||||
earlier rounds already paid one."""
|
||||
from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin
|
||||
|
||||
failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown
|
||||
repeated_fixed_margin: Final = (
|
||||
0.0
|
||||
if failed_breakdown is None
|
||||
else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost)
|
||||
)
|
||||
self.model_call_details["combined_usage_object"] = failed_usage
|
||||
self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin
|
||||
self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin)
|
||||
|
||||
def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None:
|
||||
"""Bill every round of a fully streamed request on the failure log when a post-call hook rejects it."""
|
||||
self.record_partial_usage_for_failure(spend.usage, spend.cost)
|
||||
self.cost_breakdown = spend.cost_breakdown
|
||||
|
||||
def attempt_scoped_copy(self) -> "Logging":
|
||||
"""A copy with its own model_call_details, so logging one attempt of a request through it leaves
|
||||
the state the request's other attempts share untouched."""
|
||||
|
|
@ -2372,12 +2480,37 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
attempt.model_call_details = dict(self.model_call_details)
|
||||
return attempt
|
||||
|
||||
def dispatch_attempt_failure_handlers(
|
||||
self, exception: Exception, traceback_exception: str
|
||||
) -> Coroutine[None, None, None]:
|
||||
shared_payload: Final = self.model_call_details.get("standard_logging_object")
|
||||
attempt_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time")
|
||||
attempt: Final = self.attempt_scoped_copy()
|
||||
|
||||
def publish() -> None:
|
||||
attempt_payload: Final = attempt.model_call_details.get("standard_logging_object")
|
||||
current_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time")
|
||||
if (
|
||||
attempt_payload is not None
|
||||
and self.model_call_details.get("standard_logging_object") is shared_payload
|
||||
and same_attempt_marker(current_marker, attempt_marker)
|
||||
):
|
||||
self.model_call_details["standard_logging_object"] = attempt_payload
|
||||
|
||||
attempt._on_failure_payload_built = publish
|
||||
|
||||
async def dispatch() -> None:
|
||||
await attempt.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
|
||||
|
||||
return dispatch()
|
||||
|
||||
def carry_failed_attempt_usage(self) -> None:
|
||||
"""Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its
|
||||
own usage while the request's failure log still bills every attempt."""
|
||||
usage: Final = self.model_call_details.pop("combined_usage_object", None)
|
||||
response_cost: Final = self.model_call_details.pop("response_cost", None)
|
||||
if isinstance(usage, Usage):
|
||||
self.model_call_details["current_attempt_api_call_start_time"] = None
|
||||
self.model_call_details["failed_attempts_usage"] = (
|
||||
*self.model_call_details.get("failed_attempts_usage", ()),
|
||||
(usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0),
|
||||
|
|
@ -2447,7 +2580,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
await self.async_failure_handler(exception, traceback_exception)
|
||||
finally:
|
||||
if self._should_run_sync_failure_callbacks_for_async_calls():
|
||||
executor.submit(self.failure_handler, exception, traceback_exception)
|
||||
executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True)
|
||||
|
||||
def should_run_logging(
|
||||
self,
|
||||
|
|
@ -2462,6 +2595,16 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
except Exception:
|
||||
return True
|
||||
|
||||
@property
|
||||
def async_success_log_enqueued(self) -> bool:
|
||||
"""True once the success log task was actually created, not merely claimed, so a later failure raised
|
||||
around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second
|
||||
time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False."""
|
||||
return self._async_success_enqueued
|
||||
|
||||
def mark_async_success_log_enqueued(self) -> None:
|
||||
self._async_success_enqueued = True
|
||||
|
||||
def claim_async_success_log(self) -> bool:
|
||||
"""One async success log per request. The innermost @client wrapper always exits first,
|
||||
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
|
||||
|
|
@ -2481,6 +2624,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""
|
||||
return
|
||||
self.model_call_details[f"has_logged_{event_type}"] = True
|
||||
if event_type == "async_failure":
|
||||
self.model_call_details["async_failure_logged_attempt"] = self.model_call_details.get(
|
||||
"current_attempt_api_call_start_time"
|
||||
)
|
||||
return
|
||||
|
||||
def has_run_logging(
|
||||
|
|
@ -2890,6 +3037,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
or isinstance(logging_result, SearchResponse) # Search API
|
||||
or isinstance(logging_result, DecisionsResponse)
|
||||
or isinstance(logging_result, OpenAIDecisionResponse)
|
||||
or isinstance(logging_result, ExtractionResponse)
|
||||
or (
|
||||
isinstance(logging_result, InteractionsAPIResponse)
|
||||
and logging_result.usage is not None
|
||||
|
|
@ -3489,8 +3637,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
): # custom logger class
|
||||
if self.stream and complete_streaming_response is None:
|
||||
callback.log_stream_event(
|
||||
kwargs=call_details,
|
||||
response_obj=result,
|
||||
kwargs=redacted_model_call_details_for_custom_logger(
|
||||
model_call_details=call_details, custom_logger=callback
|
||||
),
|
||||
response_obj=redacted_response_obj_for_custom_logger(result, callback),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
|
@ -3500,8 +3650,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
result = call_details_view["complete_response"]
|
||||
|
||||
callback.log_success_event(
|
||||
kwargs=callback.redact_standard_logging_payload_from_model_call_details(call_details),
|
||||
response_obj=result,
|
||||
kwargs=redacted_model_call_details_for_custom_logger(
|
||||
model_call_details=callback.redact_standard_logging_payload_from_model_call_details(
|
||||
call_details
|
||||
),
|
||||
custom_logger=callback,
|
||||
),
|
||||
response_obj=redacted_response_obj_for_custom_logger(result, callback),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
|
@ -3844,7 +3999,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details=model_call_details
|
||||
)
|
||||
model_call_details = redact_streaming_responses_for_custom_logger(
|
||||
model_call_details = redacted_model_call_details_for_custom_logger(
|
||||
model_call_details=model_call_details, custom_logger=callback
|
||||
)
|
||||
##################################
|
||||
|
|
@ -3855,7 +4010,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
kwargs=model_call_details,
|
||||
response_obj=model_call_details["async_complete_streaming_response"]
|
||||
if self.stream
|
||||
else result,
|
||||
else redacted_response_obj_for_custom_logger(result, callback),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
guardrail_snapshot=self.shadow_eval_request_snapshot,
|
||||
|
|
@ -3871,14 +4026,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
else:
|
||||
await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function
|
||||
kwargs=model_call_details,
|
||||
response_obj=result,
|
||||
response_obj=redacted_response_obj_for_custom_logger(result, callback),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await callback.async_log_success_event(
|
||||
kwargs=model_call_details,
|
||||
response_obj=result,
|
||||
response_obj=redacted_response_obj_for_custom_logger(result, callback),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
|
@ -4054,15 +4209,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
*,
|
||||
for_async_request: bool = False,
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
logging (including any nested calls its callbacks trigger) is fully done.
|
||||
|
||||
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
|
||||
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
|
||||
try:
|
||||
return self._failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
for_async_request=for_async_request,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
|
@ -4073,6 +4234,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
for_async_request: bool = False,
|
||||
) -> None:
|
||||
verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback)
|
||||
if not self.should_run_logging(event_type="sync_failure"): # prevent double logging
|
||||
|
|
@ -4165,6 +4327,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
isinstance(callback, CustomLogger)
|
||||
and is_sync_request
|
||||
and self.call_type != CallTypes.pass_through.value
|
||||
and not (for_async_request and self._async_handler_logs_to(callback))
|
||||
): # custom logger class
|
||||
callback.log_failure_event(
|
||||
start_time=start_time,
|
||||
|
|
@ -4282,7 +4445,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
if self._on_failure_payload_built is not None:
|
||||
self._on_failure_payload_built()
|
||||
callbacks: Final = self.get_combined_callback_list(
|
||||
dynamic_success_callbacks=self.dynamic_async_failure_callbacks,
|
||||
global_callbacks=litellm._async_failure_callback,
|
||||
|
|
@ -4467,14 +4631,26 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
_filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks)
|
||||
return len(_filtered_failure_callbacks) > 0
|
||||
|
||||
def _without_replaced_global_callbacks(
|
||||
self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger]
|
||||
) -> tuple[str | Callable[..., object] | CustomLogger, ...]:
|
||||
if not self.replaced_global_callback_keys:
|
||||
return tuple(global_callbacks)
|
||||
return tuple(
|
||||
litellm.logging_callback_manager.without_custom_logger_keys(
|
||||
global_callbacks, self.replaced_global_callback_keys
|
||||
)
|
||||
)
|
||||
|
||||
def get_combined_callback_list(
|
||||
self,
|
||||
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None,
|
||||
global_callbacks: list[str | Callable[..., object] | CustomLogger],
|
||||
) -> list[str | Callable[..., object] | CustomLogger]:
|
||||
kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks)
|
||||
if dynamic_success_callbacks is None:
|
||||
return list(global_callbacks)
|
||||
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
|
||||
return list(kept_global_callbacks)
|
||||
return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks)))
|
||||
|
||||
def _remove_internal_litellm_callbacks(self, callbacks: list) -> list:
|
||||
"""
|
||||
|
|
@ -4527,6 +4703,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
cb_name: Final = self._get_callback_name(cb)
|
||||
return any(prefix in cb_name for prefix in INTERNAL_PREFIXES)
|
||||
|
||||
def _async_handler_logs_to(self, logger: CustomLogger) -> bool:
|
||||
"""True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``."""
|
||||
if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped
|
||||
return False
|
||||
return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor
|
||||
|
||||
def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list:
|
||||
"""
|
||||
Removes internal custom logger callbacks from the list.
|
||||
|
|
@ -7195,13 +7377,16 @@ def get_standard_logging_object_payload(
|
|||
saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0
|
||||
|
||||
## Get model cost information ##
|
||||
base_model = get_base_model_from_metadata(model_call_details=kwargs)
|
||||
# The router overrides completion_response.model to the model-group alias before
|
||||
# this payload is built, so cost-map lookup via that alias always misses.
|
||||
# Fall back to the actual deployment model set by the router in metadata.
|
||||
if base_model is None:
|
||||
base_model = metadata.get("deployment")
|
||||
custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj)
|
||||
configured_base_model: Final = get_base_model_from_metadata(model_call_details=kwargs)
|
||||
spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj)
|
||||
# The router overrides completion_response.model to the model-group alias before this payload is
|
||||
# built, so cost-map lookup via that alias misses and the deployment model stands in for base_model.
|
||||
# A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates,
|
||||
# which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there.
|
||||
base_model: Final = (
|
||||
metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model
|
||||
)
|
||||
custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj)
|
||||
raw_response_cost: Final = kwargs.get("response_cost")
|
||||
llm_response_cost: Final[float] = raw_response_cost or 0.0
|
||||
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from functools import reduce
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import LITELLM_DETAILED_TIMING
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
|
||||
|
|
@ -101,6 +102,32 @@ def response_timing_metrics(
|
|||
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
|
||||
|
||||
|
||||
_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _HIDDEN_PARAMS_DICT.validate_python(hidden_params)
|
||||
except ValidationError:
|
||||
return {}
|
||||
|
||||
|
||||
def _existing_provider_response_model(result: object) -> str | None:
|
||||
"""The pricing model a response already carries, like the decisions path sets, wins over the response model."""
|
||||
hidden_params: Final[object] = getattr(result, "_hidden_params", None)
|
||||
carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr(
|
||||
hidden_params, "provider_response_model", None
|
||||
)
|
||||
return carried if isinstance(carried, str) and carried else None
|
||||
|
||||
|
||||
def _provider_response_model(result: object) -> str | None:
|
||||
"""The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the
|
||||
client asked for and cost-map lookups for the standard logging payload would otherwise miss."""
|
||||
model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None
|
||||
return model if isinstance(model, str) and model else None
|
||||
|
||||
|
||||
class ResponseMetadata:
|
||||
"""
|
||||
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
|
||||
|
|
@ -147,6 +174,8 @@ class ResponseMetadata:
|
|||
preserve_litellm_internal_headers=True,
|
||||
),
|
||||
"litellm_model_name": model,
|
||||
"provider_response_model": _existing_provider_response_model(self.result)
|
||||
or _provider_response_model(self.result),
|
||||
}
|
||||
self._update_hidden_params(new_params)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
from collections.abc import Callable
|
||||
import threading
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from types import MethodType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -15,6 +17,7 @@ else:
|
|||
_custom_logger_compatible_callbacks_literal = str
|
||||
|
||||
_generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {}
|
||||
_callback_lists_lock: Final = threading.RLock()
|
||||
|
||||
|
||||
class LoggingCallbackManager:
|
||||
|
|
@ -43,16 +46,18 @@ class LoggingCallbackManager:
|
|||
Add a input callback to litellm.input_callback.
|
||||
Auto-routes async callbacks to litellm._async_input_callback.
|
||||
"""
|
||||
if not isinstance(callback, str) and self._is_async_callable(callback):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
|
||||
with _callback_lists_lock:
|
||||
if not isinstance(callback, str) and self._is_async_callable(callback):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
|
||||
|
||||
def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]):
|
||||
"""
|
||||
Add a service callback to litellm.service_callback
|
||||
"""
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
|
||||
with _callback_lists_lock:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
|
||||
|
||||
def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]):
|
||||
"""
|
||||
|
|
@ -60,10 +65,11 @@ class LoggingCallbackManager:
|
|||
|
||||
Ensures no duplicates are added.
|
||||
"""
|
||||
self._safe_add_callback_to_list(
|
||||
callback=callback,
|
||||
parent_list=litellm.callbacks,
|
||||
)
|
||||
with _callback_lists_lock:
|
||||
self._safe_add_callback_to_list(
|
||||
callback=callback,
|
||||
parent_list=litellm.callbacks,
|
||||
)
|
||||
|
||||
def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]):
|
||||
"""
|
||||
|
|
@ -71,37 +77,41 @@ class LoggingCallbackManager:
|
|||
Auto-routes async callbacks to litellm._async_success_callback.
|
||||
Special-cases 'dynamodb' and 'openmeter' as async callbacks.
|
||||
"""
|
||||
if (
|
||||
isinstance(callback, str)
|
||||
and callback in ("dynamodb", "openmeter")
|
||||
or not isinstance(callback, str)
|
||||
and self._is_async_callable(callback)
|
||||
):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
|
||||
with _callback_lists_lock:
|
||||
if (
|
||||
isinstance(callback, str)
|
||||
and callback in ("dynamodb", "openmeter")
|
||||
or not isinstance(callback, str)
|
||||
and self._is_async_callable(callback)
|
||||
):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
|
||||
|
||||
def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]):
|
||||
"""
|
||||
Add a failure callback to `litellm.failure_callback`.
|
||||
Auto-routes async callbacks to litellm._async_failure_callback.
|
||||
"""
|
||||
if not isinstance(callback, str) and self._is_async_callable(callback):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
|
||||
with _callback_lists_lock:
|
||||
if not isinstance(callback, str) and self._is_async_callable(callback):
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
|
||||
else:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
|
||||
|
||||
def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str):
|
||||
"""
|
||||
Add a success callback to litellm._async_success_callback
|
||||
"""
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
|
||||
with _callback_lists_lock:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
|
||||
|
||||
def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str):
|
||||
"""
|
||||
Add a failure callback to litellm._async_failure_callback
|
||||
"""
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
|
||||
with _callback_lists_lock:
|
||||
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
|
||||
|
||||
def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True):
|
||||
"""
|
||||
|
|
@ -110,13 +120,14 @@ class LoggingCallbackManager:
|
|||
if not isinstance(callback_list, list): # Not list -> do nothing
|
||||
return
|
||||
|
||||
if require_self:
|
||||
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
|
||||
else:
|
||||
remove_list = [c for c in callback_list if c == obj]
|
||||
with _callback_lists_lock:
|
||||
if require_self:
|
||||
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
|
||||
else:
|
||||
remove_list = [c for c in callback_list if c == obj]
|
||||
|
||||
for c in remove_list:
|
||||
callback_list.remove(c)
|
||||
for c in remove_list:
|
||||
callback_list.remove(c)
|
||||
|
||||
def remove_callbacks_by_type(self, callback_list, callback_type):
|
||||
"""
|
||||
|
|
@ -134,10 +145,11 @@ class LoggingCallbackManager:
|
|||
if not isinstance(callback_list, list):
|
||||
return
|
||||
|
||||
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
|
||||
with _callback_lists_lock:
|
||||
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
|
||||
|
||||
for c in remove_list:
|
||||
callback_list.remove(c)
|
||||
for c in remove_list:
|
||||
callback_list.remove(c)
|
||||
|
||||
def _add_string_callback_to_list(
|
||||
self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str]
|
||||
|
|
@ -336,6 +348,55 @@ class LoggingCallbackManager:
|
|||
|
||||
return "-".join(key_parts)
|
||||
|
||||
get_custom_logger_key = _get_custom_logger_key
|
||||
callback_lists_lock: Final = _callback_lists_lock
|
||||
|
||||
def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]:
|
||||
dropped_ids: Final = frozenset(map(id, callbacks))
|
||||
owner_ids: Final = frozenset(map(id, method_owners))
|
||||
|
||||
def keep(callback: object) -> bool:
|
||||
if id(callback) in dropped_ids:
|
||||
return False
|
||||
return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids)
|
||||
|
||||
with _callback_lists_lock:
|
||||
vacated_keys: Final = frozenset(
|
||||
self._get_custom_logger_key(callback)
|
||||
for callback in litellm.callbacks
|
||||
if isinstance(callback, CustomLogger) and id(callback) in dropped_ids
|
||||
)
|
||||
litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)]
|
||||
litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)]
|
||||
litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)]
|
||||
litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)]
|
||||
litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)]
|
||||
litellm._async_success_callback = [
|
||||
callback for callback in litellm._async_success_callback if keep(callback)
|
||||
]
|
||||
litellm._async_failure_callback = [
|
||||
callback for callback in litellm._async_failure_callback if keep(callback)
|
||||
]
|
||||
return vacated_keys
|
||||
|
||||
def with_custom_logger_keys(
|
||||
self, candidates: Iterable[CustomLogger], keys: frozenset[str]
|
||||
) -> tuple[CustomLogger, ...]:
|
||||
if not keys:
|
||||
return ()
|
||||
class_names: Final = frozenset(key.split("-", 1)[0] for key in keys)
|
||||
same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names)
|
||||
return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys)
|
||||
|
||||
def without_custom_logger_keys(
|
||||
self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str]
|
||||
) -> Iterator[str | Callable[..., object] | CustomLogger]:
|
||||
return (
|
||||
callback
|
||||
for callback in callbacks
|
||||
if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys)
|
||||
)
|
||||
|
||||
def _reset_all_callbacks(self):
|
||||
"""
|
||||
Reset all callbacks to an empty list
|
||||
|
|
@ -363,19 +424,21 @@ class LoggingCallbackManager:
|
|||
|
||||
_get_all_callbacks = get_all_callbacks
|
||||
|
||||
def remove_callback_from_all_lists(self, obj, require_self=False) -> None:
|
||||
def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None:
|
||||
"""
|
||||
Remove a callback object from every callback list it may have been
|
||||
promoted into, so a re-initialized callback leaves no stale instance behind.
|
||||
"""
|
||||
for callback_list in (
|
||||
litellm.callbacks,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
):
|
||||
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
|
||||
with _callback_lists_lock:
|
||||
for callback_list in (
|
||||
litellm.callbacks,
|
||||
litellm.input_callback,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
):
|
||||
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
|
||||
|
||||
def get_active_additional_logging_utils_from_custom_logger(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -10,10 +10,12 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import inspect
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import replace
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -40,36 +42,79 @@ else:
|
|||
LiteLLMLoggingObject = Any
|
||||
|
||||
|
||||
def _opted_out_of_message_logging(custom_logger: CustomLogger) -> bool:
|
||||
return hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True
|
||||
|
||||
|
||||
def redact_message_input_output_from_custom_logger(
|
||||
litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger
|
||||
):
|
||||
if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True:
|
||||
) -> object:
|
||||
if _opted_out_of_message_logging(custom_logger):
|
||||
return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False)
|
||||
return result
|
||||
|
||||
|
||||
def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
|
||||
def redacted_model_call_details_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
|
||||
"""
|
||||
Returns a copy of model_call_details whose streaming response entries are redacted deepcopies
|
||||
when the custom logger has opted out of message logging. The shared model_call_details is left
|
||||
untouched so other callbacks still receive the unredacted response.
|
||||
Returns a copy of model_call_details with the request messages, instructions and streaming responses
|
||||
redacted when the custom logger has opted out of message logging. The shared model_call_details is left
|
||||
untouched so other callbacks and the spend log still receive the unredacted values.
|
||||
"""
|
||||
if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True):
|
||||
if not _opted_out_of_message_logging(custom_logger):
|
||||
return model_call_details
|
||||
redacted_entries: Final = {
|
||||
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
|
||||
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
|
||||
if model_call_details.get(streaming_key) is not None
|
||||
return {
|
||||
**{key: value for key, value in model_call_details.items() if key != SERVED_OUTPUT_TEXTS_KEY},
|
||||
**_redacted_entries_for_opted_out_logger(model_call_details),
|
||||
}
|
||||
if not redacted_entries:
|
||||
return model_call_details
|
||||
return {**model_call_details, **redacted_entries}
|
||||
|
||||
|
||||
def _redacted_streaming_response_copy(streaming_response):
|
||||
redacted_response: Final = copy.deepcopy(streaming_response)
|
||||
_redact_streaming_response(redacted_response)
|
||||
return redacted_response
|
||||
def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, object]) -> Mapping[str, object]:
|
||||
optional_params: Final = _as_mapping(model_call_details.get("optional_params"))
|
||||
litellm_params: Final = _as_mapping(model_call_details.get("litellm_params"))
|
||||
standard_logging_object: Final = _as_mapping(model_call_details.get("standard_logging_object"))
|
||||
return {
|
||||
"messages": [{"role": "user", "content": REDACTED_BY_LITELLM}],
|
||||
"prompt": "",
|
||||
"input": "",
|
||||
**({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}),
|
||||
**({"optional_params": _with_redacted_instructions(optional_params)} if optional_params is not None else {}),
|
||||
**({"litellm_params": redacted_litellm_params(litellm_params)} if litellm_params is not None else {}),
|
||||
**(
|
||||
{"standard_logging_object": _redact_standard_logging_object(standard_logging_object)}
|
||||
if standard_logging_object is not None
|
||||
else {}
|
||||
),
|
||||
**{
|
||||
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
|
||||
for streaming_key in (
|
||||
"complete_streaming_response",
|
||||
"async_complete_streaming_response",
|
||||
"complete_response",
|
||||
)
|
||||
if model_call_details.get(streaming_key) is not None
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def redacted_response_obj_for_custom_logger(response_obj: object, custom_logger: CustomLogger) -> object:
|
||||
if not _opted_out_of_message_logging(custom_logger) or not _is_redactable_response(response_obj):
|
||||
return response_obj
|
||||
return perform_redaction({}, response_obj)
|
||||
|
||||
|
||||
def _is_redactable_response(response_obj: object) -> bool:
|
||||
return isinstance(
|
||||
response_obj, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse)
|
||||
) or (isinstance(response_obj, dict) and ("choices" in response_obj or "output" in response_obj))
|
||||
|
||||
|
||||
def _redacted_streaming_response_copy(streaming_response: object) -> object:
|
||||
redacted_response: Final = (
|
||||
streaming_response
|
||||
if isinstance(streaming_response, litellm.ResponsesAPIResponse)
|
||||
else copy.deepcopy(streaming_response)
|
||||
)
|
||||
return _redact_streaming_response(redacted_response)
|
||||
|
||||
|
||||
def _redact_streaming_response(streaming_response):
|
||||
|
|
@ -77,10 +122,13 @@ def _redact_streaming_response(streaming_response):
|
|||
for choice in streaming_response.choices:
|
||||
_redact_choice_content(choice)
|
||||
redact_vertex_ai_metadata_from_logged_object(streaming_response)
|
||||
elif isinstance(streaming_response, litellm.ResponsesAPIResponse):
|
||||
return perform_redaction({}, streaming_response)
|
||||
elif hasattr(streaming_response, "output"):
|
||||
_redact_responses_api_output(streaming_response.output)
|
||||
if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None:
|
||||
streaming_response.reasoning = None
|
||||
return streaming_response
|
||||
|
||||
|
||||
def _redact_tool_calls(tool_calls) -> None:
|
||||
|
|
@ -177,31 +225,75 @@ def _redact_responses_api_output(output_items):
|
|||
output_item.input = REDACTED_BY_LITELLM
|
||||
|
||||
|
||||
def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
||||
"""Helper to redact ResponsesAPIResponse output items in dict form."""
|
||||
for output_item in output_items:
|
||||
if not isinstance(output_item, dict):
|
||||
continue
|
||||
_RESPONSE_ITEMS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
if output_item.get("text") is not None:
|
||||
output_item["text"] = redacted_str
|
||||
|
||||
if isinstance(output_item.get("content"), list):
|
||||
for content_item in output_item["content"]:
|
||||
if isinstance(content_item, dict) and content_item.get("text") is not None:
|
||||
content_item["text"] = redacted_str
|
||||
if isinstance(content_item, dict) and content_item.get("refusal") is not None:
|
||||
content_item["refusal"] = redacted_str
|
||||
def _redacted_responses_api_output_part(part: object, fields: tuple[str, ...], redacted_str: str) -> object:
|
||||
if (part_mapping := _as_mapping(part)) is None:
|
||||
return part
|
||||
return {key: redacted_str if key in fields and value is not None else value for key, value in part_mapping.items()}
|
||||
|
||||
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
|
||||
for summary_item in output_item["summary"]:
|
||||
if isinstance(summary_item, dict) and summary_item.get("text") is not None:
|
||||
summary_item["text"] = redacted_str
|
||||
|
||||
if output_item.get("type") == "function_call" and "arguments" in output_item:
|
||||
output_item["arguments"] = redacted_str
|
||||
if output_item.get("type") == "custom_tool_call" and "input" in output_item:
|
||||
output_item["input"] = redacted_str
|
||||
def _redacted_responses_api_output_item(item: object, redacted_str: str) -> object:
|
||||
if (item_mapping := _as_mapping(item)) is None:
|
||||
return item
|
||||
content: Final = item_mapping.get("content")
|
||||
summary: Final = item_mapping.get("summary")
|
||||
return {
|
||||
**item_mapping,
|
||||
**({"text": redacted_str} if item_mapping.get("text") is not None else {}),
|
||||
**(
|
||||
{
|
||||
"content": [
|
||||
_redacted_responses_api_output_part(part, ("text", "refusal"), redacted_str)
|
||||
for part in _RESPONSE_ITEMS.validate_python(content)
|
||||
]
|
||||
}
|
||||
if isinstance(content, list)
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"summary": [
|
||||
_redacted_responses_api_output_part(part, ("text",), redacted_str)
|
||||
for part in _RESPONSE_ITEMS.validate_python(summary)
|
||||
]
|
||||
}
|
||||
if item_mapping.get("type") == "reasoning" and isinstance(summary, list)
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{"arguments": redacted_str}
|
||||
if item_mapping.get("type") == "function_call" and "arguments" in item_mapping
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{"input": redacted_str}
|
||||
if item_mapping.get("type") == "custom_tool_call" and "input" in item_mapping
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _redact_responses_api_output_dict(output_items: Sequence[object], redacted_str: str) -> Sequence[object]:
|
||||
return [_redacted_responses_api_output_item(item, redacted_str) for item in output_items]
|
||||
|
||||
|
||||
def _with_redacted_instructions(params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return params if params.get("instructions") is None else {**params, "instructions": REDACTED_BY_LITELLM}
|
||||
|
||||
|
||||
def _redacted_responses_api_response(
|
||||
response: Mapping[str, object], *, redact_instructions: bool = True
|
||||
) -> Mapping[str, object]:
|
||||
output: Final = response.get("output")
|
||||
redacted: Final = {
|
||||
**response,
|
||||
"output": _redact_responses_api_output_dict(_RESPONSE_ITEMS.validate_python(output), REDACTED_BY_LITELLM)
|
||||
if isinstance(output, list)
|
||||
else output,
|
||||
}
|
||||
return _with_redacted_instructions(redacted) if redact_instructions else redacted
|
||||
|
||||
|
||||
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -215,20 +307,25 @@ def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[
|
|||
return _redact_standard_logging_object(payload)
|
||||
|
||||
|
||||
def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]:
|
||||
def _redact_standard_logging_object(
|
||||
payload: Mapping[str, object], *, redact_instructions: bool = True
|
||||
) -> dict[str, object]:
|
||||
standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload))
|
||||
redacted_str: Final = REDACTED_BY_LITELLM
|
||||
|
||||
if standard_logging_object.get("messages") is not None:
|
||||
standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}]
|
||||
model_parameters: Final = _as_mapping(standard_logging_object.get("model_parameters"))
|
||||
if model_parameters is not None and redact_instructions:
|
||||
standard_logging_object["model_parameters"] = _with_redacted_instructions(model_parameters)
|
||||
|
||||
response: Final = standard_logging_object.get("response")
|
||||
if response is not None:
|
||||
if isinstance(response, dict) and "output" in response:
|
||||
# ResponsesAPIResponse format - redact content in output items
|
||||
if isinstance(response.get("output"), list):
|
||||
_redact_responses_api_output_dict(response["output"], redacted_str)
|
||||
redact_vertex_ai_metadata_from_logged_object(response)
|
||||
standard_logging_object["response"] = _redacted_responses_api_response(
|
||||
response, redact_instructions=redact_instructions
|
||||
)
|
||||
redact_vertex_ai_metadata_from_logged_object(standard_logging_object["response"])
|
||||
elif isinstance(response, dict) and "choices" in response:
|
||||
# ModelResponse dict format - redact content in choices
|
||||
if isinstance(response.get("choices"), list):
|
||||
|
|
@ -320,11 +417,11 @@ def _redact_hidden_provider_specific_fields(params: Mapping[str, object]) -> Non
|
|||
def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body: bool = True) -> dict[str, object]:
|
||||
"""
|
||||
Copy of litellm_params without classifier audit fields or baseline route request parameters.
|
||||
redact_request_body also redacts the nested proxy_server_request body.
|
||||
redact_request_body also redacts the request instructions and the nested proxy_server_request body.
|
||||
"""
|
||||
request: Final = _as_mapping(params.get("proxy_server_request"))
|
||||
return {
|
||||
**params,
|
||||
**(_with_redacted_instructions(params) if redact_request_body else params),
|
||||
**{
|
||||
key: _redacted_baseline_metadata(value)
|
||||
for key, value in params.items()
|
||||
|
|
@ -344,8 +441,9 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
|
||||
redact_streaming_responses=False skips the in-place redaction of the shared streaming
|
||||
response entries; per-callback redaction hands each opted-out callback its own redacted
|
||||
copy via redact_streaming_responses_for_custom_logger instead. It also leaves the shared
|
||||
proxy request body and hidden provider fields as they are, since the spend log reads them.
|
||||
copy via redacted_model_call_details_for_custom_logger instead. It also leaves the shared
|
||||
request instructions, proxy request body and hidden provider fields as they are, since the
|
||||
spend log and the other callbacks read them.
|
||||
"""
|
||||
# Redact model_call_details
|
||||
params: Final = _as_mapping(model_call_details.get("litellm_params"))
|
||||
|
|
@ -358,16 +456,31 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
optional_params: Final = _as_mapping(model_call_details.get("optional_params"))
|
||||
if redact_streaming_responses:
|
||||
model_call_details.update(
|
||||
{
|
||||
**({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}),
|
||||
**({"optional_params": _with_redacted_instructions(optional_params)} if optional_params else {}),
|
||||
}
|
||||
)
|
||||
model_call_details.pop(SERVED_OUTPUT_TEXTS_KEY, None)
|
||||
standard_logging_object: Final = model_call_details.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, Mapping):
|
||||
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
|
||||
model_call_details["standard_logging_object"] = _redact_standard_logging_object(
|
||||
standard_logging_object, redact_instructions=redact_streaming_responses
|
||||
)
|
||||
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
|
||||
|
||||
# Redact streaming response
|
||||
if redact_streaming_responses and model_call_details.get("stream", False) is True:
|
||||
for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"):
|
||||
_redact_streaming_response(model_call_details.get(_streaming_key))
|
||||
model_call_details.update(
|
||||
{
|
||||
streaming_key: _redact_streaming_response(model_call_details[streaming_key])
|
||||
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
|
||||
if model_call_details.get(streaming_key) is not None
|
||||
}
|
||||
)
|
||||
|
||||
# Redact result
|
||||
if result is not None:
|
||||
|
|
@ -381,10 +494,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
# For async objects, return a simple redacted response without deepcopy
|
||||
return {"text": REDACTED_BY_LITELLM}
|
||||
|
||||
if not (
|
||||
isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse))
|
||||
or (isinstance(result, dict) and ("choices" in result or "output" in result))
|
||||
):
|
||||
if not _is_redactable_response(result):
|
||||
return {"text": REDACTED_BY_LITELLM}
|
||||
|
||||
_result: Final = copy.deepcopy(result)
|
||||
|
|
@ -399,14 +509,18 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
_redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM)
|
||||
redact_vertex_ai_metadata_from_logged_object(_result)
|
||||
elif isinstance(_result, dict) and "output" in _result:
|
||||
if isinstance(_result.get("output"), list):
|
||||
_redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM)
|
||||
return _redacted_responses_api_response(_result, redact_instructions=redact_streaming_responses)
|
||||
elif isinstance(_result, litellm.ResponsesAPIResponse):
|
||||
if hasattr(_result, "output"):
|
||||
_redact_responses_api_output(_result.output)
|
||||
# Redact reasoning field in ResponsesAPIResponse
|
||||
if hasattr(_result, "reasoning") and _result.reasoning is not None:
|
||||
_result.reasoning = None
|
||||
return _result.model_copy(
|
||||
update={"instructions": REDACTED_BY_LITELLM}
|
||||
if redact_streaming_responses and _result.instructions is not None
|
||||
else {}
|
||||
)
|
||||
elif isinstance(_result, litellm.EmbeddingResponse):
|
||||
if hasattr(_result, "data") and _result.data is not None:
|
||||
_result.data = []
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import traceback
|
|||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
|
@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base
|
|||
from .llm_response_utils.get_api_base import get_api_base
|
||||
from .rules import Rules
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .litellm_logging import BilledRoundsSpend
|
||||
|
||||
# Constants for special delta attribute names
|
||||
AUDIO_ATTRIBUTE: Final = "audio"
|
||||
IMAGE_ATTRIBUTE: Final = "images"
|
||||
|
|
@ -259,6 +262,8 @@ def _provider_hidden_params(
|
|||
|
||||
|
||||
class CustomStreamWrapper:
|
||||
_earlier_rounds_spend: "BilledRoundsSpend | None" = None
|
||||
|
||||
@property
|
||||
def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation
|
||||
return self._hidden_params
|
||||
|
|
@ -342,6 +347,10 @@ class CustomStreamWrapper:
|
|||
self.messages = getattr(logging_obj, "messages", None)
|
||||
self.sent_stream_usage = False
|
||||
self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
|
||||
self.caller_requested_stream_usage: Final = (
|
||||
self.check_send_stream_usage(getattr(logging_obj, "stream_options", None))
|
||||
and getattr(logging_obj, "stream_usage_injected", False) is not True
|
||||
)
|
||||
self.tool_call = False
|
||||
# keep track of the returned chunks - used for calculating the input/output tokens for stream options
|
||||
self.chunks: list[ModelResponseStream] = []
|
||||
|
|
@ -2070,7 +2079,7 @@ class CustomStreamWrapper:
|
|||
except StopIteration:
|
||||
if self.sent_last_chunk is True:
|
||||
complete_streaming_response: Final = self._build_complete_streaming_response()
|
||||
response = self.model_response_creator()
|
||||
response = self._new_usage_chunk()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
|
||||
|
||||
|
|
@ -2222,27 +2231,10 @@ class CustomStreamWrapper:
|
|||
)
|
||||
|
||||
if _has_usage:
|
||||
# Store a copy ONLY when usage stripping below will mutate
|
||||
# the chunk. For non-usage chunks (vast majority), store
|
||||
# directly to avoid expensive model_copy() per chunk.
|
||||
self.chunks.append(processed_chunk.model_copy())
|
||||
|
||||
# Strip usage from the outgoing chunk so it's not sent twice
|
||||
# (once in the chunk, once in _hidden_params).
|
||||
obj_dict = processed_chunk.model_dump()
|
||||
if "usage" in obj_dict:
|
||||
del obj_dict["usage"]
|
||||
processed_chunk = self.model_response_creator(
|
||||
chunk=obj_dict,
|
||||
hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage
|
||||
Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR)
|
||||
),
|
||||
)
|
||||
is_empty = is_model_response_stream_empty(
|
||||
model_response=cast(ModelResponseStream, processed_chunk)
|
||||
)
|
||||
if is_empty:
|
||||
usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
|
||||
if usage_free_chunk is None:
|
||||
continue
|
||||
processed_chunk = usage_free_chunk
|
||||
else:
|
||||
# No usage data — safe to store directly without copying
|
||||
self.chunks.append(processed_chunk)
|
||||
|
|
@ -2273,25 +2265,25 @@ class CustomStreamWrapper:
|
|||
if processed_chunk is None:
|
||||
continue
|
||||
|
||||
choice = processed_chunk.choices[0]
|
||||
if isinstance(choice, StreamingChoices):
|
||||
self._accumulate_streamed_delta(choice.delta)
|
||||
if processed_chunk.choices:
|
||||
choice = processed_chunk.choices[0]
|
||||
if isinstance(choice, StreamingChoices):
|
||||
self._accumulate_streamed_delta(choice.delta)
|
||||
self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
|
||||
# RETURN RESULT
|
||||
self.chunks.append(processed_chunk)
|
||||
return processed_chunk
|
||||
if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None:
|
||||
self.chunks.append(processed_chunk)
|
||||
return processed_chunk
|
||||
usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
|
||||
if usage_free_chunk is None:
|
||||
continue
|
||||
return usage_free_chunk
|
||||
except (StopAsyncIteration, StopIteration):
|
||||
return await self._finalize_completed_stream(cache_hit=cache_hit)
|
||||
except httpx.TimeoutException as e: # if httpx read timeout error occues
|
||||
traceback_exception = traceback.format_exc()
|
||||
## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT
|
||||
traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}"
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
self._failure_logging = asyncio.create_task(
|
||||
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
|
||||
)
|
||||
self._dispatch_stream_failure(e, traceback_exception)
|
||||
await self.aclose()
|
||||
self._handle_stream_fallback_error(e)
|
||||
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
|
||||
|
|
@ -2306,7 +2298,7 @@ class CustomStreamWrapper:
|
|||
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
|
||||
if self.sent_last_chunk is True:
|
||||
complete_streaming_response: Final = await self._async_build_complete_streaming_response()
|
||||
response: Final = self.model_response_creator()
|
||||
response: Final = self._new_usage_chunk()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
|
||||
|
||||
|
|
@ -2392,29 +2384,36 @@ class CustomStreamWrapper:
|
|||
return processed_chunk
|
||||
|
||||
def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn:
|
||||
traceback_exception: Final = traceback.format_exc()
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
self._failure_logging = asyncio.create_task(
|
||||
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
|
||||
)
|
||||
self._dispatch_stream_failure(e, traceback.format_exc())
|
||||
self._handle_stream_fallback_error(e)
|
||||
|
||||
async def wait_for_failure_logging(self, timeout: float) -> None:
|
||||
def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None:
|
||||
self._record_partial_usage_for_failure()
|
||||
self._failure_logging = asyncio.create_task(
|
||||
self.logging_obj.dispatch_attempt_failure_handlers(e, traceback_exception)
|
||||
)
|
||||
|
||||
@property
|
||||
def has_dispatched_failure_logging(self) -> bool:
|
||||
return self._failure_logging is not None
|
||||
|
||||
async def wait_for_failure_logging(self, timeout: float | None) -> None:
|
||||
"""Let this stream's failure log read its own deployment before a fallback reuses the logging object."""
|
||||
if self._failure_logging is None:
|
||||
return
|
||||
await asyncio.wait((self._failure_logging,), timeout=timeout)
|
||||
|
||||
def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
|
||||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
"""
|
||||
A stream that breaks mid-flight still billed the provider for the chunks
|
||||
already delivered. Recover that partial usage from the chunks seen so
|
||||
far and stash it, with its cost, on the logging object so the failure
|
||||
handler records the real partial spend instead of zero. A request that
|
||||
later recovers via a router fallback overwrites this with the combined
|
||||
success log on the same request id, so this never double counts.
|
||||
handler records the real partial spend instead of zero. The router
|
||||
moves it into terminal failure usage before opening the next attempt,
|
||||
so a recovered attempt logs its own usage.
|
||||
"""
|
||||
if self.logging_obj is None or not self.chunks:
|
||||
return
|
||||
|
|
@ -2433,10 +2432,12 @@ class CustomStreamWrapper:
|
|||
if self.model:
|
||||
partial_response.model = self.model
|
||||
backfill_missing_cache_usage_fields(usage)
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = usage
|
||||
self.logging_obj.model_call_details["response_cost"] = (
|
||||
self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
|
||||
)
|
||||
partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
|
||||
if self._earlier_rounds_spend is None:
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = usage
|
||||
self.logging_obj.model_call_details["response_cost"] = partial_cost
|
||||
return
|
||||
self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost)
|
||||
except Exception as recover_error:
|
||||
verbose_logger.debug(
|
||||
"could not recover partial usage for interrupted stream: %s",
|
||||
|
|
@ -2463,6 +2464,22 @@ class CustomStreamWrapper:
|
|||
return False
|
||||
return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS)
|
||||
|
||||
def _new_usage_chunk(self) -> ModelResponseStream:
|
||||
usage_chunk: Final = self.model_response_creator()
|
||||
if self.caller_requested_stream_usage:
|
||||
usage_chunk.choices = []
|
||||
return usage_chunk
|
||||
|
||||
def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None:
|
||||
self.chunks.append(chunk.model_copy())
|
||||
usage_free_chunk: Final = self.model_response_creator(
|
||||
chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"},
|
||||
hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage
|
||||
Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR)
|
||||
),
|
||||
)
|
||||
return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk
|
||||
|
||||
def _accumulate_streamed_delta(self, delta: object) -> None:
|
||||
# Shared by the sync, async, and non-aiohttp iteration sites so answer
|
||||
# text and the disqualifying-content latch stay in step across all three.
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ Pattern Overview:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, MutableSequence, Sequence
|
||||
from collections.abc import Iterable, Mapping, MutableSequence, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain, repeat
|
||||
|
|
@ -27,8 +27,10 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
||||
OPENAI_MAX_TOOL_NAME_LENGTH,
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
is_provider_native_tool_dict,
|
||||
truncate_tool_name,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import (
|
||||
keep_non_string_tool_result_texts,
|
||||
|
|
@ -202,6 +204,60 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _tool_use_name(block: object) -> str | None:
|
||||
fields: Final = as_str_mapping(block)
|
||||
if fields is None or fields.get("type") != "tool_use":
|
||||
return None
|
||||
return anthropic_tool_name(fields)
|
||||
|
||||
|
||||
def _tool_use_names(message: object) -> tuple[str, ...]:
|
||||
return tuple(name for name in map(_tool_use_name, list_field(message, "content")) if name is not None)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OriginalToolNames:
|
||||
"""Caller tool names over the guardrail's limit, keyed by the shortened name the guardrail sees."""
|
||||
|
||||
declared: Mapping[str, str]
|
||||
used: Mapping[str, str]
|
||||
|
||||
|
||||
def _restorable_names(names: Iterable[str]) -> Mapping[str, str]:
|
||||
listed: Final = tuple(names)
|
||||
taken: Final = frozenset(listed)
|
||||
shortened: Final = tuple(
|
||||
(truncate_tool_name(name), name) for name in listed if len(name) > OPENAI_MAX_TOOL_NAME_LENGTH
|
||||
)
|
||||
return MappingProxyType({alias: name for alias, name in shortened if alias not in taken})
|
||||
|
||||
|
||||
def _original_tool_names(request: Mapping[str, object]) -> _OriginalToolNames:
|
||||
return _OriginalToolNames(
|
||||
declared=_restorable_names(chain.from_iterable(map(anthropic_tool_names, list_field(request, "tools")))),
|
||||
used=_restorable_names(chain.from_iterable(map(_tool_use_names, list_field(request, "messages")))),
|
||||
)
|
||||
|
||||
|
||||
def _restore_tool_use_names(messages: Sequence[object], original_names: Mapping[str, str]) -> None:
|
||||
for message in messages:
|
||||
for block in list_field(message, "content"):
|
||||
if (name := _tool_use_name(block)) is not None and name in original_names and isinstance(block, dict):
|
||||
block["name"] = original_names[name]
|
||||
|
||||
|
||||
def _with_original_name(tool: ChatCompletionToolParam, original_names: Mapping[str, str]) -> ChatCompletionToolParam:
|
||||
if "function" not in tool:
|
||||
return tool
|
||||
function: Final = tool["function"]
|
||||
fields: Final = as_str_mapping(function)
|
||||
name: Final = None if fields is None else fields.get("name")
|
||||
original: Final = original_names.get(name) if isinstance(name, str) else None
|
||||
if original is None:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "name": original}}
|
||||
|
||||
|
||||
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
|
||||
if not isinstance(system, list):
|
||||
return
|
||||
|
|
@ -629,6 +685,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
model: Final = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
original_tool_names: Final = _original_tool_names(
|
||||
cast(Mapping[str, object], data) # cast-ok: Anthropic request body
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
|
|
@ -643,7 +702,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
anthropic_config: Final = AnthropicConfig()
|
||||
anthropic_tools: Final[list[AllAnthropicToolsValues]] = []
|
||||
for tool in guardrailed_tools:
|
||||
converted_tool, mcp_server = anthropic_config.map_tool_helper(tool)
|
||||
converted_tool, mcp_server = anthropic_config.map_tool_helper(
|
||||
_with_original_name(tool, original_tool_names.declared)
|
||||
)
|
||||
if converted_tool is not None:
|
||||
anthropic_tools.append(converted_tool)
|
||||
# Note: MCP servers are handled separately in the main transformation
|
||||
|
|
@ -676,6 +737,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
),
|
||||
hoisted_system_message=hoisted_system_message,
|
||||
preserve_system_messages=has_midturn_system_message,
|
||||
original_tool_names=original_tool_names.used,
|
||||
)
|
||||
else:
|
||||
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
|
||||
|
|
@ -840,6 +902,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
structured_messages: list, # mutable-ok: API message payload
|
||||
hoisted_system_message: object = None,
|
||||
preserve_system_messages: bool = False,
|
||||
*,
|
||||
original_tool_names: Mapping[str, str],
|
||||
) -> None:
|
||||
"""Write a guardrail's structured-message rewrite back without losing corrections."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -903,6 +967,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "thinking":
|
||||
block.pop("cache_control", None)
|
||||
_restore_tool_use_names(
|
||||
cast(Sequence[object], converted), # cast-ok: Anthropic message payload
|
||||
original_tool_names,
|
||||
)
|
||||
data["messages"] = keep_non_string_tool_result_texts(converted, data.get("messages"))
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -712,7 +712,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
optional_params: dict,
|
||||
model: str,
|
||||
param: str,
|
||||
value: Any,
|
||||
value: object,
|
||||
drop_params: bool,
|
||||
output_key: str,
|
||||
) -> None:
|
||||
|
|
@ -1788,6 +1788,21 @@ def strip_thinking_blocks_from_anthropic_messages_request_dict(
|
|||
data.pop("thinking", None)
|
||||
|
||||
|
||||
_DROP_ANTHROPIC_MESSAGE: Final = object()
|
||||
|
||||
|
||||
def _strip_empty_content_blocks_from_message(message: object) -> object:
|
||||
if not isinstance(message, dict) or not isinstance(message.get("content"), list):
|
||||
return message
|
||||
content: Final = message["content"]
|
||||
filtered: Final = [
|
||||
block for block in content if not _is_empty_text_block(block) and not is_empty_unsigned_thinking_block(block)
|
||||
]
|
||||
if len(filtered) == len(content):
|
||||
return message
|
||||
return {**message, "content": filtered} if filtered else _DROP_ANTHROPIC_MESSAGE
|
||||
|
||||
|
||||
def strip_empty_content_blocks_from_anthropic_messages(
|
||||
messages: Sequence[object],
|
||||
) -> list[Any]:
|
||||
|
|
@ -1822,18 +1837,11 @@ def strip_empty_content_blocks_from_anthropic_messages(
|
|||
The caller's list and its content blocks are never mutated; modified
|
||||
messages are returned as shallow copies with a fresh content list.
|
||||
"""
|
||||
out: Final[list[Any]] = []
|
||||
for m in messages:
|
||||
if not isinstance(m, dict) or not isinstance(m.get("content"), list):
|
||||
out.append(m)
|
||||
continue
|
||||
content = m["content"]
|
||||
filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_unsigned_thinking_block(b)]
|
||||
if len(filtered) == len(content):
|
||||
out.append(m)
|
||||
elif filtered:
|
||||
out.append({**m, "content": filtered})
|
||||
return out
|
||||
return [
|
||||
stripped_message
|
||||
for message in messages
|
||||
if (stripped_message := _strip_empty_content_blocks_from_message(message)) is not _DROP_ANTHROPIC_MESSAGE
|
||||
]
|
||||
|
||||
|
||||
def _is_empty_text_block(block: object) -> bool:
|
||||
|
|
@ -1921,18 +1929,15 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: Sequence[object]) -> l
|
|||
and ``:`` — valid on the upstream provider but rejected by Anthropic when
|
||||
the session is switched to a native Anthropic deployment.
|
||||
"""
|
||||
out: Final[list[Any]] = []
|
||||
for m in messages:
|
||||
if not isinstance(m, dict) or not isinstance(m.get("content"), list):
|
||||
out.append(m)
|
||||
continue
|
||||
content = m["content"]
|
||||
new_content = [_sanitize_tool_use_id_content_block(b) for b in content]
|
||||
if new_content == content:
|
||||
out.append(m)
|
||||
else:
|
||||
out.append({**m, "content": new_content})
|
||||
return out
|
||||
return [_sanitize_tool_use_ids_in_message(message) for message in messages]
|
||||
|
||||
|
||||
def _sanitize_tool_use_ids_in_message(message: object) -> object:
|
||||
if not isinstance(message, dict) or not isinstance(message.get("content"), list):
|
||||
return message
|
||||
content: Final = message["content"]
|
||||
new_content: Final = [_sanitize_tool_use_id_content_block(block) for block in content]
|
||||
return message if new_content == content else {**message, "content": new_content}
|
||||
|
||||
|
||||
class _ReplayedSearchQuery(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from typing import (
|
|||
Protocol,
|
||||
cast,
|
||||
get_args,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -65,6 +66,14 @@ class _ChunkStream(Protocol):
|
|||
def __aiter__(self) -> "AsyncIterator[ModelResponseStream]": ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _FailureLoggingWaiter(Protocol):
|
||||
@property
|
||||
def has_dispatched_failure_logging(self) -> bool: ...
|
||||
|
||||
async def wait_for_failure_logging(self, timeout: float | None) -> None: ...
|
||||
|
||||
|
||||
def _optional_attr(obj: object, name: str) -> object:
|
||||
return getattr(obj, name, None)
|
||||
|
||||
|
|
@ -189,6 +198,14 @@ class _CombinedChunkSplitter:
|
|||
self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({})
|
||||
self._live_call_index: int | None = None
|
||||
|
||||
@property
|
||||
def has_dispatched_failure_logging(self) -> bool:
|
||||
return isinstance(self._stream, _FailureLoggingWaiter) and self._stream.has_dispatched_failure_logging is True
|
||||
|
||||
async def wait_for_failure_logging(self, timeout: float | None) -> None:
|
||||
if isinstance(self._stream, _FailureLoggingWaiter):
|
||||
await self._stream.wait_for_failure_logging(timeout=timeout)
|
||||
|
||||
@property
|
||||
def chunks(self) -> "list[ModelResponseStream] | None":
|
||||
return cast( # cast-ok: chunks is a list of ModelResponseStream on the inner stream
|
||||
|
|
@ -522,7 +539,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# Wrap the upstream stream so chunks that carry both content and a
|
||||
# finish_reason (fake-streamed providers) are split into two — see
|
||||
# _CombinedChunkSplitter.
|
||||
super().__init__(_CombinedChunkSplitter(completion_stream))
|
||||
combined_chunk_splitter: Final = _CombinedChunkSplitter(completion_stream)
|
||||
super().__init__(combined_chunk_splitter)
|
||||
self._combined_chunk_splitter: Final = combined_chunk_splitter
|
||||
self.model = model
|
||||
self._message_id: str = f"msg_{uuid.uuid4()}"
|
||||
if litellm_logging_obj is not None:
|
||||
|
|
@ -1244,7 +1263,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
if provider_error is e:
|
||||
raise
|
||||
raise provider_error from e
|
||||
if logging_obj is not None:
|
||||
stream_has_dispatched_failure_logging: Final = self._combined_chunk_splitter.has_dispatched_failure_logging
|
||||
if stream_has_dispatched_failure_logging:
|
||||
await self._combined_chunk_splitter.wait_for_failure_logging(timeout=None)
|
||||
elif logging_obj is not None:
|
||||
try:
|
||||
await logging_obj.dispatch_failure_handlers(
|
||||
exception=provider_error,
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ How it works:
|
|||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -48,7 +48,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
|
|||
|
||||
def can_handle(
|
||||
self,
|
||||
tools: list[dict] | None,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
custom_llm_provider: str | None,
|
||||
) -> bool:
|
||||
if not tools:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
|
|
@ -21,7 +21,7 @@ class MessagesInterceptor(ABC):
|
|||
@abstractmethod
|
||||
def can_handle(
|
||||
self,
|
||||
tools: list[dict] | None,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
custom_llm_provider: str | None,
|
||||
) -> bool:
|
||||
"""Return True if this interceptor should handle the request."""
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
## Translates OpenAI call to Anthropic `/v1/messages` format
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import ConfigDict, field_validator
|
||||
|
|
@ -21,7 +21,10 @@ from litellm.llms.anthropic.pass_through.messages.utils import (
|
|||
)
|
||||
from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
|
||||
from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
|
||||
from litellm.responses.streaming_iterator import stream_error_status_and_message
|
||||
from litellm.responses.streaming_iterator import (
|
||||
partial_responses_incomplete_event_on_disconnect,
|
||||
stream_error_status_and_message,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
|
|
@ -33,6 +36,7 @@ from .transformation import (
|
|||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator
|
||||
from litellm.types.llms.openai import ResponseIncompleteEvent
|
||||
|
||||
|
||||
class _UpstreamFailure(LiteLLMBaseModel):
|
||||
|
|
@ -488,8 +492,38 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]:
|
||||
"""Yield SSE-encoded bytes for each Anthropic event chunk."""
|
||||
def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream":
|
||||
"""SSE-encoded bytes for each Anthropic event chunk."""
|
||||
return AnthropicResponsesSSEByteStream(self)
|
||||
|
||||
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
|
||||
safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator)
|
||||
async for chunk in self:
|
||||
yield anthropic_sse_frame(await safeguard_results.observe(chunk))
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
|
||||
|
||||
|
||||
class AnthropicResponsesSSEByteStream:
|
||||
"""
|
||||
The byte stream the proxy relays for a default-route /v1/messages request.
|
||||
Keeps the Responses API iterator reachable so a client disconnect can bill
|
||||
the partial output the way native /v1/responses streams do.
|
||||
"""
|
||||
|
||||
def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
|
||||
self._wrapper: Final = wrapper
|
||||
self._frames: Final = wrapper._sse_frames() # pyright: ignore[reportPrivateUsage] # the wrapper's own byte stream
|
||||
|
||||
def __aiter__(self) -> "AnthropicResponsesSSEByteStream":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
return await self._frames.__anext__()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._frames.aclose()
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return self._wrapper.partial_response_on_disconnect()
|
||||
|
|
|
|||
|
|
@ -464,33 +464,32 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
|
||||
return input_items
|
||||
|
||||
@staticmethod
|
||||
def _translate_tool_to_responses_api(tool_dict: Mapping[str, object]) -> dict[str, object]:
|
||||
tool_type: Final = tool_dict.get("type", "")
|
||||
tool_name: Final = tool_dict.get("name", "")
|
||||
# web_search tool
|
||||
if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search":
|
||||
return {"type": "web_search_preview"}
|
||||
# Responses turns strict mode on when `strict` is omitted, silently rewriting
|
||||
# `required` to every property. Anthropic tools are non-strict unless asked.
|
||||
func_tool: Final[dict[str, object]] = {
|
||||
"type": "function",
|
||||
"name": tool_name,
|
||||
"strict": bool(tool_dict.get("strict")),
|
||||
}
|
||||
if "description" in tool_dict:
|
||||
func_tool["description"] = tool_dict["description"]
|
||||
if "input_schema" in tool_dict:
|
||||
func_tool["parameters"] = json_copy_or_original(tool_dict["input_schema"])
|
||||
return func_tool
|
||||
|
||||
def translate_tools_to_responses_api(
|
||||
self,
|
||||
tools: list[AllAnthropicToolsValues],
|
||||
) -> list[dict[str, object]]:
|
||||
"""Convert Anthropic tool definitions to Responses API function tools."""
|
||||
result: Final[list[dict[str, object]]] = []
|
||||
for tool in tools:
|
||||
tool_dict = cast(dict[str, Any], tool)
|
||||
tool_type = tool_dict.get("type", "")
|
||||
tool_name = tool_dict.get("name", "")
|
||||
# web_search tool
|
||||
if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search":
|
||||
result.append({"type": "web_search_preview"})
|
||||
continue
|
||||
# Responses turns strict mode on when `strict` is omitted, silently rewriting
|
||||
# `required` to every property. Anthropic tools are non-strict unless asked.
|
||||
func_tool: dict[str, object] = {
|
||||
"type": "function",
|
||||
"name": tool_name,
|
||||
"strict": bool(tool_dict.get("strict")),
|
||||
}
|
||||
if "description" in tool_dict:
|
||||
func_tool["description"] = tool_dict["description"]
|
||||
if "input_schema" in tool_dict:
|
||||
func_tool["parameters"] = json_copy_or_original(tool_dict["input_schema"])
|
||||
result.append(func_tool)
|
||||
return result
|
||||
return [self._translate_tool_to_responses_api(tool) for tool in tools]
|
||||
|
||||
@staticmethod
|
||||
def translate_tool_choice_to_responses_api(
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Final
|
|||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic_core import ErrorDetails
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -394,7 +395,9 @@ _MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, ob
|
|||
_LABELS: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...])
|
||||
|
||||
|
||||
def _normalize_answer(value: object) -> Mapping[str, object]:
|
||||
def _normalize_answer(value: object) -> object:
|
||||
if not isinstance(value, Mapping):
|
||||
return value
|
||||
answer: Final = _MAPPING.validate_python(value)
|
||||
legend: Final = answer.get("legend")
|
||||
if answer.get("type") != "score" or not isinstance(legend, list):
|
||||
|
|
@ -403,11 +406,36 @@ def _normalize_answer(value: object) -> Mapping[str, object]:
|
|||
return {**answer, "legend": {str(index): label for index, label in enumerate(labels)}}
|
||||
|
||||
|
||||
def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse:
|
||||
def _normalize_systemone_response(payload: object) -> object:
|
||||
if not isinstance(payload, Mapping):
|
||||
return payload
|
||||
response: Final = _MAPPING.validate_python(payload)
|
||||
answers: Final = _MAPPING.validate_python(response.get("answers"))
|
||||
normalized: Final = {**response, "answers": {name: _normalize_answer(answer) for name, answer in answers.items()}}
|
||||
return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(normalized), request)
|
||||
answers: Final = response.get("answers")
|
||||
if not isinstance(answers, Mapping):
|
||||
return response
|
||||
return {
|
||||
**response,
|
||||
"answers": {name: _normalize_answer(answer) for name, answer in _MAPPING.validate_python(answers).items()},
|
||||
}
|
||||
|
||||
|
||||
def _validation_detail(detail: ErrorDetails) -> str:
|
||||
location: Final = ".".join(str(part) for part in detail["loc"])
|
||||
return f"{location}: {detail['msg']}" if location else detail["msg"]
|
||||
|
||||
|
||||
def _validation_summary(error: ValidationError) -> str:
|
||||
return "; ".join(_validation_detail(detail) for detail in error.errors(include_url=False, include_input=False))
|
||||
|
||||
|
||||
class DecisionsResponseValidationError(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse:
|
||||
return systemone_response_to_ir(
|
||||
_SYSTEMONE_RESPONSE_ADAPTER.validate_python(_normalize_systemone_response(payload)), request
|
||||
)
|
||||
|
||||
|
||||
def _systemone_answer(answer: DecisionsIRAnswer) -> DecisionAnswer | None:
|
||||
|
|
@ -627,9 +655,10 @@ class BaseDecisionsConfig(ABC):
|
|||
try:
|
||||
return self.parse_response(payload, request)
|
||||
except ValidationError as error:
|
||||
raise BaseLLMException(
|
||||
raise DecisionsResponseValidationError(
|
||||
status_code=500,
|
||||
message=f"Decisions provider '{custom_llm_provider}' returned an unexpected response: {error}",
|
||||
message=f"Decisions provider '{custom_llm_provider}' returned an unexpected response: "
|
||||
f"{error.title} {_validation_summary(error)}",
|
||||
) from error
|
||||
|
||||
def get_error_class(
|
||||
|
|
|
|||
3
litellm/llms/base_llm/extraction/__init__.py
Normal file
3
litellm/llms/base_llm/extraction/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .transformation import BaseExtractionConfig
|
||||
|
||||
__all__ = ["BaseExtractionConfig"]
|
||||
86
litellm/llms/base_llm/extraction/transformation.py
Normal file
86
litellm/llms/base_llm/extraction/transformation.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.extraction import ExtractionRequestBody, ExtractionResponse
|
||||
|
||||
_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
|
||||
_RESPONSE_ADAPTER: Final[TypeAdapter[ExtractionResponse]] = TypeAdapter(ExtractionResponse)
|
||||
_RESERVED_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "content-type"})
|
||||
|
||||
|
||||
class BaseExtractionConfig(ABC):
|
||||
api_key_env: tuple[str, ...] = ()
|
||||
api_base_env: tuple[str, ...] = ()
|
||||
|
||||
def get_default_api_base(self) -> str | None:
|
||||
return None
|
||||
|
||||
def missing_api_base_message(self, custom_llm_provider: str) -> str:
|
||||
return f"api_base is required for Extraction provider '{custom_llm_provider}'"
|
||||
|
||||
def resolve_api_base(self, api_base: str | None) -> str | None:
|
||||
return api_base or self._first_secret(self.api_base_env) or self.get_default_api_base()
|
||||
|
||||
def resolve_api_key(self, api_key: str | None) -> str | None:
|
||||
return api_key or self._first_secret(self.api_key_env)
|
||||
|
||||
@staticmethod
|
||||
def _first_secret(names: tuple[str, ...]) -> str | None:
|
||||
return next((value for value in (get_secret_str(name) for name in names) if value), None)
|
||||
|
||||
def canonical_model(self, model: str) -> str:
|
||||
return model
|
||||
|
||||
def request_model(self, model: str) -> str:
|
||||
return model
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
api_key: str | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> Mapping[str, str]:
|
||||
return {
|
||||
**{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS},
|
||||
**({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}),
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
def get_complete_url(self, api_base: str, model: str) -> str: ...
|
||||
|
||||
def transform_extraction_request(self, model: str, request: ExtractionRequestBody) -> Mapping[str, object]:
|
||||
return {**request.model_dump(mode="json", exclude_unset=True), "model": self.request_model(model)}
|
||||
|
||||
def unwrap_response(self, payload: object) -> object:
|
||||
return payload
|
||||
|
||||
def transform_extraction_response(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
raw_response: httpx.Response,
|
||||
) -> ExtractionResponse:
|
||||
payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content)
|
||||
try:
|
||||
return _RESPONSE_ADAPTER.validate_python(self.unwrap_response(payload))
|
||||
except ValidationError as error:
|
||||
raise BaseLLMException(
|
||||
status_code=500,
|
||||
message=f"Extraction provider '{custom_llm_provider}' returned an unexpected response: {error}",
|
||||
) from error
|
||||
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: Mapping[str, str] | httpx.Headers,
|
||||
) -> BaseLLMException:
|
||||
return BaseLLMException(status_code=status_code, message=error_message, headers=httpx.Headers(headers))
|
||||
|
|
@ -185,7 +185,7 @@ class BaseTranslation(ABC):
|
|||
request_data: dict | None = None,
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
deliver_ended_stream_rewrites: bool = False,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""
|
||||
Process output streaming response with guardrails.
|
||||
|
||||
|
|
|
|||
|
|
@ -68,6 +68,16 @@ class HarnessTurnRequest:
|
|||
cwd: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HarnessSessionCopy:
|
||||
"""A saved session under fresh ids: written to `path` in the sandbox, then loaded by `import_command`."""
|
||||
|
||||
session_id: str
|
||||
path: str
|
||||
data: bytes
|
||||
import_command: Sequence[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HarnessTurnResponse:
|
||||
"""What the runtime produced for one turn, after the process exited."""
|
||||
|
|
@ -128,6 +138,18 @@ class BaseCLIHarnessConfig(BaseHarnessConfig[OptionsT], Generic[OptionsT, Stream
|
|||
) -> HarnessTurnRequest:
|
||||
"""argv / env / stdin for one turn. native_session_id is set after the first turn."""
|
||||
|
||||
def transform_resume_export_command(self, native_session_id: str) -> Sequence[str] | None:
|
||||
"""Command that prints a saved session before resuming one this handler did not run.
|
||||
|
||||
None resumes the saved session as is. A non-zero exit or a timeout fails the turn, and the next turn
|
||||
tries again.
|
||||
"""
|
||||
return None
|
||||
|
||||
def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None:
|
||||
"""The copy to import and continue, from the export command's stdout, or None to continue the original."""
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def create_stream_state(self) -> StreamStateT:
|
||||
"""Fresh per-turn parser state."""
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.utils import PrismaClient as _PrismaClient
|
||||
from litellm.router import Router as _Router
|
||||
|
||||
Span = _Span | Any
|
||||
Span = _Span
|
||||
InternalUsageCache = _InternalUsageCache
|
||||
PrismaClient = _PrismaClient
|
||||
Router = _Router
|
||||
|
|
@ -550,7 +550,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
|
|||
limit: int | None = None,
|
||||
after: str | None = None,
|
||||
additional_filters: Mapping[str, object] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
List resources created by a user.
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,10 @@ from typing import TYPE_CHECKING, Final, Literal
|
|||
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.llms.bedrock.common_utils import (
|
||||
bedrock_control_plane_base_url,
|
||||
configured_bedrock_control_plane_endpoint,
|
||||
)
|
||||
from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
|
@ -156,10 +159,14 @@ class BedrockBatchesHandler:
|
|||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
creds: Final = BedrockBatchesConfig().resolve_credentials(auth_params, region)
|
||||
control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint(
|
||||
kwargs.get("aws_bedrock_control_plane_endpoint")
|
||||
)
|
||||
|
||||
client: Final = boto3.client(
|
||||
"bedrock",
|
||||
region_name=region,
|
||||
endpoint_url=control_plane_endpoint,
|
||||
aws_access_key_id=creds.access_key,
|
||||
aws_secret_access_key=creds.secret_key,
|
||||
aws_session_token=creds.token,
|
||||
|
|
@ -171,6 +178,7 @@ class BedrockBatchesHandler:
|
|||
batch_id=batch_id,
|
||||
aws_region_name=region,
|
||||
logging_obj=logging_obj,
|
||||
aws_bedrock_control_plane_endpoint=control_plane_endpoint,
|
||||
**auth_params.model_dump(),
|
||||
)
|
||||
|
||||
|
|
@ -270,7 +278,7 @@ class BedrockBatchesHandler:
|
|||
batch_id: str,
|
||||
aws_region_name: str | None = None,
|
||||
logging_obj=None,
|
||||
**kwargs,
|
||||
**kwargs: object,
|
||||
) -> "LiteLLMBatch":
|
||||
"""
|
||||
Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch
|
||||
|
|
@ -294,7 +302,10 @@ class BedrockBatchesHandler:
|
|||
``aws_session_token``, ``aws_profile_name``,
|
||||
``aws_role_name``, ``aws_session_name``,
|
||||
``aws_web_identity_token``, ``aws_sts_endpoint``,
|
||||
``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored.
|
||||
``aws_external_id``, ``aws_session_tags``), plus
|
||||
``aws_bedrock_control_plane_endpoint`` to send the call to a VPC
|
||||
endpoint or gateway instead of the regional ``bedrock`` host.
|
||||
Unknown keys are ignored.
|
||||
|
||||
Returns:
|
||||
``LiteLLMBatch`` shaped like an OpenAI Batch resource.
|
||||
|
|
@ -318,10 +329,14 @@ class BedrockBatchesHandler:
|
|||
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
|
||||
|
||||
creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region)
|
||||
control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint(
|
||||
kwargs.get("aws_bedrock_control_plane_endpoint")
|
||||
)
|
||||
|
||||
client: Final = boto3.client(
|
||||
"bedrock",
|
||||
region_name=region,
|
||||
endpoint_url=control_plane_endpoint,
|
||||
aws_access_key_id=creds.access_key,
|
||||
aws_secret_access_key=creds.secret_key,
|
||||
aws_session_token=creds.token,
|
||||
|
|
@ -334,14 +349,13 @@ class BedrockBatchesHandler:
|
|||
# `GetModelInvocationJob` accepts either form, but only the bare id
|
||||
# produces a sensible-looking URL in logs.
|
||||
url_path_id: Final = _extract_job_id_from_arn(batch_id) or batch_id
|
||||
control_plane_base: Final = bedrock_control_plane_base_url(control_plane_endpoint, region)
|
||||
logging_obj.pre_call(
|
||||
input=batch_id,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {"jobIdentifier": batch_id},
|
||||
"api_base": (
|
||||
f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}"
|
||||
),
|
||||
"api_base": f"{control_plane_base}/model-invocation-job/{url_path_id}",
|
||||
},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, 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.aws_partition import is_bedrock_arn
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
)
|
||||
|
|
@ -32,6 +32,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage
|
|||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import (
|
||||
CommonBatchFilesUtils,
|
||||
bedrock_control_plane_base_url,
|
||||
merge_bedrock_aws_request_params,
|
||||
resolve_s3_bucket_owner,
|
||||
resolve_s3_encryption_key_id,
|
||||
|
|
@ -42,7 +43,7 @@ if TYPE_CHECKING:
|
|||
|
||||
# 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
|
||||
# _batch_s3_object_name in litellm/llms/bedrock/files/transformation.py). A uuid4 is always 36 hex/dash
|
||||
# characters, so it can be stripped off the end unambiguously even though the
|
||||
# model name itself may contain dashes.
|
||||
_S3_BATCH_FILE_UUID_SUFFIX_PATTERN: Final = re.compile(
|
||||
|
|
@ -177,11 +178,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
|
||||
aws_region_name: Final = self._get_aws_region_name(request_params, model)
|
||||
|
||||
# Bedrock model invocation job endpoint
|
||||
# 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"
|
||||
control_plane_base: Final = bedrock_control_plane_base_url(
|
||||
request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name
|
||||
)
|
||||
bedrock_endpoint: Final = f"{control_plane_base}/model-invocation-job"
|
||||
|
||||
return bedrock_endpoint
|
||||
|
||||
|
|
@ -279,9 +279,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
# 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.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job"
|
||||
control_plane_base: Final = bedrock_control_plane_base_url(
|
||||
request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name
|
||||
)
|
||||
endpoint_url: Final = f"{control_plane_base}/model-invocation-job"
|
||||
signed_headers, signed_data = self.common_utils.sign_aws_request(
|
||||
service_name="bedrock",
|
||||
data=bedrock_request,
|
||||
|
|
@ -431,12 +432,13 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
import urllib.parse as _ul
|
||||
|
||||
encoded_arn: Final = _ul.quote(batch_id, safe="")
|
||||
endpoint_url: Final = (
|
||||
f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}"
|
||||
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
|
||||
control_plane_base: Final = bedrock_control_plane_base_url(
|
||||
request_params.get("aws_bedrock_control_plane_endpoint"), region
|
||||
)
|
||||
endpoint_url: Final = f"{control_plane_base}/model-invocation-job/{encoded_arn}"
|
||||
|
||||
# Use common utility for AWS signing
|
||||
request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params)
|
||||
signed_headers, _ = self.common_utils.sign_aws_request(
|
||||
service_name="bedrock",
|
||||
data={}, # GET request has no body
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None:
|
|||
return headers.get(name, None)
|
||||
|
||||
|
||||
def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
|
||||
def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
|
||||
try:
|
||||
guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
|
||||
except ValidationError as e:
|
||||
|
|
@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
guardrail_headers: Final = (
|
||||
()
|
||||
if raw_guardrail_config is None
|
||||
else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
|
||||
else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
|
||||
)
|
||||
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
|
||||
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
|
||||
|
|
|
|||
|
|
@ -138,6 +138,25 @@ def merge_bedrock_aws_request_params(
|
|||
return request_params
|
||||
|
||||
|
||||
def configured_bedrock_control_plane_endpoint(endpoint: object) -> str | None:
|
||||
"""An ``aws_bedrock_control_plane_endpoint`` value without a trailing slash, or None when unset."""
|
||||
if isinstance(endpoint, str) and endpoint.strip():
|
||||
return endpoint.strip().rstrip("/")
|
||||
return None
|
||||
|
||||
|
||||
def bedrock_control_plane_base_url(endpoint: object, aws_region_name: str) -> str:
|
||||
"""Base URL for Bedrock control-plane calls such as model invocation jobs.
|
||||
|
||||
A configured ``aws_bedrock_control_plane_endpoint`` (a VPC endpoint or a compatible gateway)
|
||||
wins over the regional ``bedrock`` host. ``api_base`` is deliberately not used here, since on a
|
||||
Bedrock deployment it points at the runtime host, which does not serve these APIs.
|
||||
"""
|
||||
return configured_bedrock_control_plane_endpoint(endpoint) or (
|
||||
f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
|
||||
)
|
||||
|
||||
|
||||
def s3_static_key_pair(params: Mapping[str, object]) -> tuple[str, str] | None:
|
||||
"""The s3_access_key_id / s3_secret_access_key pair when both are set, otherwise None."""
|
||||
s3_access_key_id: Final = params.get("s3_access_key_id")
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import os
|
|||
import posixpath
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
|
||||
from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
|
@ -23,6 +23,7 @@ from typing_extensions import ReadOnly
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.files.utils import FilesAPIUtils
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
|
|
@ -324,7 +325,10 @@ _REJECTED_FILE_ID_REQUEST_URL: Final = "https://litellm.ai"
|
|||
|
||||
|
||||
def _rejected_file_id(reason: ValueError) -> BedrockError:
|
||||
message: Final = str(reason)
|
||||
return _client_error(str(reason))
|
||||
|
||||
|
||||
def _client_error(message: str) -> BedrockError:
|
||||
return BedrockError(
|
||||
status_code=400,
|
||||
message=message,
|
||||
|
|
@ -336,6 +340,92 @@ def _rejected_file_id(reason: ValueError) -> BedrockError:
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BatchInputLine:
|
||||
line_number: int
|
||||
record: _OpenAIBatchRecord
|
||||
|
||||
|
||||
_BATCH_RECORD_CONVERSION_ERRORS: Final = (AttributeError, LookupError, TypeError, ValueError, BadRequestError)
|
||||
|
||||
|
||||
def _batch_line_error(line_number: int, problem: str) -> BedrockError:
|
||||
return _client_error(f"Batch input file line {line_number} {problem}")
|
||||
|
||||
|
||||
def _decode_batch_input(content: str | bytes) -> str:
|
||||
if isinstance(content, str):
|
||||
return content.removeprefix("\ufeff")
|
||||
try:
|
||||
return content.decode("utf-8-sig")
|
||||
except UnicodeDecodeError as e:
|
||||
raise _batch_line_error(content.count(b"\n", 0, e.start) + 1, f"is not valid UTF-8: {e.reason}") from e
|
||||
|
||||
|
||||
def _load_batch_line(line_number: int, line: str) -> _OpenAIBatchRecord:
|
||||
try:
|
||||
return json.loads(line)
|
||||
except json.JSONDecodeError as e:
|
||||
raise _batch_line_error(line_number, f"is not valid JSON: {e.msg} at column {e.colno}") from e
|
||||
except (ValueError, RecursionError) as e:
|
||||
raise _batch_line_error(line_number, f"could not be parsed: {e}") from e
|
||||
|
||||
|
||||
def _json_object(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, dict) else None
|
||||
|
||||
|
||||
def _batch_json_object(line_number: int, line: str) -> _OpenAIBatchRecord:
|
||||
record: Final = _load_batch_line(line_number, line)
|
||||
if _json_object(record) is None:
|
||||
raise _batch_line_error(line_number, "is not a JSON object")
|
||||
return record
|
||||
|
||||
|
||||
def _batch_body_problem(record: Mapping[str, object]) -> str | None:
|
||||
body: Final = _json_object(record.get("body"))
|
||||
if body is None:
|
||||
return "has a body that is missing or not a JSON object"
|
||||
if "model" in body and not isinstance(body["model"], str):
|
||||
return "has a body.model that is not a string"
|
||||
return None
|
||||
|
||||
|
||||
def _batch_input_line(line_number: int, line: str) -> _BatchInputLine:
|
||||
record: Final = _batch_json_object(line_number, line)
|
||||
problem: Final = _batch_body_problem(record)
|
||||
if problem is not None:
|
||||
raise _batch_line_error(line_number, problem)
|
||||
return _BatchInputLine(line_number=line_number, record=record)
|
||||
|
||||
|
||||
def _non_blank_lines(file_content: str) -> Iterator[tuple[int, str]]:
|
||||
numbered_lines: Final = enumerate(file_content.splitlines(), start=1)
|
||||
return ((line_number, line) for line_number, line in numbered_lines if line.strip())
|
||||
|
||||
|
||||
def _parse_batch_jsonl(file_content: str) -> tuple[_BatchInputLine, ...]:
|
||||
return tuple(_batch_input_line(line_number, line) for line_number, line in _non_blank_lines(file_content))
|
||||
|
||||
|
||||
def _batch_file_model(file_content: str) -> str | None:
|
||||
records: Final = tuple(
|
||||
_batch_json_object(line_number, line) for line_number, line in _non_blank_lines(file_content)
|
||||
)
|
||||
if not records:
|
||||
return None
|
||||
body: Final = _json_object(records[0].get("body"))
|
||||
model: Final = body.get("model") if body is not None else None
|
||||
return model if isinstance(model, str) else ""
|
||||
|
||||
|
||||
def _batch_s3_object_name(model: str) -> str:
|
||||
safe_model: Final = sanitize_cloud_object_component(
|
||||
model.removeprefix("bedrock/").replace(":", "-"), fallback="model"
|
||||
)
|
||||
return f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl"
|
||||
|
||||
|
||||
def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]:
|
||||
configured_bucket_names: Final = get_configured_s3_bucket_names(litellm_params)
|
||||
allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params)
|
||||
|
|
@ -545,29 +635,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
# File-like objects need to be read
|
||||
content = file_content.read()
|
||||
|
||||
# Ensure content is string
|
||||
if isinstance(content, bytes):
|
||||
content = content.decode("utf-8")
|
||||
|
||||
return content
|
||||
|
||||
def _get_s3_object_name_from_batch_jsonl(
|
||||
self,
|
||||
openai_jsonl_content: Sequence[_OpenAIBatchRecord],
|
||||
) -> str:
|
||||
"""
|
||||
Gets a unique S3 object name for the Bedrock batch processing job
|
||||
|
||||
named as: litellm-bedrock-files/{model}/{uuid}
|
||||
"""
|
||||
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
|
||||
# Remove bedrock/ prefix if present
|
||||
_model = _model.removeprefix("bedrock/")
|
||||
|
||||
safe_model: Final = sanitize_cloud_object_component(_model.replace(":", "-"), fallback="model")
|
||||
|
||||
object_name: Final = f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl"
|
||||
return object_name
|
||||
return _decode_batch_input(content)
|
||||
|
||||
def get_object_name(self, extracted_file_data: ExtractedFileData, purpose: str) -> str:
|
||||
"""
|
||||
|
|
@ -581,11 +649,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
if purpose == "batch":
|
||||
## 1. If jsonl, check if there's a model name
|
||||
file_content: Final = self._get_content_from_openai_file(extracted_file_data_content)
|
||||
|
||||
# Split into lines and parse each line as JSON
|
||||
openai_jsonl_content: Final = [json.loads(line) for line in file_content.splitlines() if line.strip()]
|
||||
if len(openai_jsonl_content) > 0:
|
||||
return self._get_s3_object_name_from_batch_jsonl(openai_jsonl_content)
|
||||
batch_model: Final = _batch_file_model(file_content)
|
||||
if batch_model is not None:
|
||||
return _batch_s3_object_name(batch_model)
|
||||
|
||||
## 2. If not jsonl, store under a server-generated managed object name
|
||||
filename: Final = extracted_file_data.get("filename")
|
||||
|
|
@ -1080,42 +1146,60 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
}
|
||||
}
|
||||
"""
|
||||
return [
|
||||
self._transform_openai_batch_record(idx, openai_record, target_model)
|
||||
for idx, openai_record in enumerate(openai_jsonl_content)
|
||||
]
|
||||
|
||||
def transform_batch_input_lines(
|
||||
self, batch_input_lines: Sequence[_BatchInputLine], target_model: str = ""
|
||||
) -> tuple[_BedrockBatchRecord, ...]:
|
||||
return tuple(
|
||||
self._transform_batch_input_line(idx, batch_input_line, target_model)
|
||||
for idx, batch_input_line in enumerate(batch_input_lines)
|
||||
)
|
||||
|
||||
def _transform_batch_input_line(
|
||||
self, idx: int, batch_input_line: _BatchInputLine, target_model: str
|
||||
) -> _BedrockBatchRecord:
|
||||
try:
|
||||
return self._transform_openai_batch_record(idx, batch_input_line.record, target_model)
|
||||
except _BATCH_RECORD_CONVERSION_ERRORS as e:
|
||||
raise _batch_line_error(
|
||||
batch_input_line.line_number, f"could not be converted to a Bedrock batch record: {e}"
|
||||
) from e
|
||||
|
||||
def _transform_openai_batch_record(
|
||||
self, idx: int, openai_record: _OpenAIBatchRecord, target_model: str
|
||||
) -> _BedrockBatchRecord:
|
||||
import litellm
|
||||
|
||||
bedrock_jsonl_content: Final = []
|
||||
for idx, _openai_jsonl_content in enumerate(openai_jsonl_content):
|
||||
# Extract the request body from OpenAI format
|
||||
openai_body = _openai_jsonl_content.get("body", {})
|
||||
record_model = openai_body.get("model", "")
|
||||
resolved_model = litellm.model_alias_map.get(record_model, record_model)
|
||||
model_for_transform, provider = self._resolve_batch_record_model_and_provider(
|
||||
record_model=resolved_model, target_model=target_model
|
||||
openai_body: Final = openai_record.get("body", {})
|
||||
record_model: Final = openai_body.get("model", "")
|
||||
resolved_model: Final = litellm.model_alias_map.get(record_model, record_model)
|
||||
model_for_transform, provider = self._resolve_batch_record_model_and_provider(
|
||||
record_model=resolved_model, target_model=target_model
|
||||
)
|
||||
|
||||
# Route to the embedding transformer when the OpenAI batch line
|
||||
# targets /v1/embeddings; every other endpoint shape is normalized
|
||||
# to chat completions first. We branch here (rather than inside
|
||||
# `_map_openai_to_bedrock_params`) so the chat helper keeps its
|
||||
# narrow contract and the embedding helper can evolve independently.
|
||||
record_kind: Final = self._classify_batch_record(openai_record)
|
||||
model_input: Final = (
|
||||
self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body, model=model_for_transform)
|
||||
if record_kind is BedrockBatchRecordKind.EMBEDDING
|
||||
else self._map_openai_to_bedrock_params(
|
||||
openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind),
|
||||
model=model_for_transform,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
# Route to the embedding transformer when the OpenAI batch line
|
||||
# targets /v1/embeddings; every other endpoint shape is normalized
|
||||
# to chat completions first. We branch here (rather than inside
|
||||
# `_map_openai_to_bedrock_params`) so the chat helper keeps its
|
||||
# narrow contract and the embedding helper can evolve independently.
|
||||
record_kind = self._classify_batch_record(_openai_jsonl_content)
|
||||
if record_kind is BedrockBatchRecordKind.EMBEDDING:
|
||||
model_input = self._map_openai_embedding_to_bedrock_params(
|
||||
openai_request_body=openai_body, model=model_for_transform
|
||||
)
|
||||
else:
|
||||
model_input = self._map_openai_to_bedrock_params(
|
||||
openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind),
|
||||
model=model_for_transform,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
# Create Bedrock batch record
|
||||
record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}")
|
||||
bedrock_record = {"recordId": record_id, "modelInput": model_input}
|
||||
|
||||
bedrock_jsonl_content.append(bedrock_record)
|
||||
return bedrock_jsonl_content
|
||||
)
|
||||
return _BedrockBatchRecord(
|
||||
recordId=openai_record.get("custom_id", f"CALL{str(idx).zfill(7)}"),
|
||||
modelInput=model_input,
|
||||
)
|
||||
|
||||
def transform_openai_jsonl_content_to_bedrock_jsonl_content(
|
||||
self,
|
||||
|
|
@ -1151,11 +1235,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
):
|
||||
## Transform JSONL content to Bedrock format
|
||||
original_file_content: Final = self._get_content_from_openai_file(extracted_file_data_content)
|
||||
openai_jsonl_content = [json.loads(line) for line in original_file_content.splitlines() if line.strip()]
|
||||
batch_input_lines: Final = _parse_batch_jsonl(original_file_content)
|
||||
litellm_params_model: Final = litellm_params.get("model")
|
||||
target_model: Final = model or (litellm_params_model if isinstance(litellm_params_model, str) else "")
|
||||
bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(
|
||||
openai_jsonl_content, target_model=target_model
|
||||
bedrock_jsonl_content: Final = self.transform_batch_input_lines(
|
||||
batch_input_lines, target_model=target_model
|
||||
)
|
||||
file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content)
|
||||
elif isinstance(extracted_file_data_content, bytes):
|
||||
|
|
@ -1662,37 +1746,11 @@ class BedrockJsonlFilesTransformation:
|
|||
# Read the content of the file
|
||||
file_content: Final = self._get_content_from_openai_file(openai_file_content)
|
||||
|
||||
# Split into lines and parse each line as JSON
|
||||
openai_jsonl_content: Final = [json.loads(line) for line in file_content.splitlines() if line.strip()]
|
||||
bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content)
|
||||
batch_input_lines: Final = _parse_batch_jsonl(file_content)
|
||||
bedrock_jsonl_content: Final = BedrockFilesConfig().transform_batch_input_lines(batch_input_lines)
|
||||
bedrock_jsonl_string: Final = "\n".join(json.dumps(item) for item in bedrock_jsonl_content)
|
||||
object_name: Final = self._get_s3_object_name(openai_jsonl_content=openai_jsonl_content)
|
||||
return bedrock_jsonl_string, object_name
|
||||
|
||||
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
|
||||
self, openai_jsonl_content: Sequence[_OpenAIBatchRecord]
|
||||
):
|
||||
"""
|
||||
Delegate to the main BedrockFilesConfig transformation method
|
||||
"""
|
||||
config: Final = BedrockFilesConfig()
|
||||
return config.transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content)
|
||||
|
||||
def _get_s3_object_name(
|
||||
self,
|
||||
openai_jsonl_content: Sequence[_OpenAIBatchRecord],
|
||||
) -> str:
|
||||
"""
|
||||
Gets a unique S3 object name for the Bedrock batch processing job
|
||||
|
||||
named as: litellm-bedrock-files-{model}-{uuid}
|
||||
"""
|
||||
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
|
||||
# Remove bedrock/ prefix if present
|
||||
_model = _model.removeprefix("bedrock/")
|
||||
safe_model: Final = sanitize_cloud_object_component(_model.replace(":", "-"), fallback="model")
|
||||
object_name: Final = f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl"
|
||||
return object_name
|
||||
first_model: Final = batch_input_lines[0].record.get("body", {}).get("model", "") if batch_input_lines else ""
|
||||
return bedrock_jsonl_string, _batch_s3_object_name(first_model)
|
||||
|
||||
def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str:
|
||||
"""
|
||||
|
|
@ -1725,11 +1783,7 @@ class BedrockJsonlFilesTransformation:
|
|||
# File-like objects need to be read
|
||||
content = file_content.read()
|
||||
|
||||
# Ensure content is string
|
||||
if isinstance(content, bytes):
|
||||
content = content.decode("utf-8")
|
||||
|
||||
return content
|
||||
return _decode_batch_input(content)
|
||||
|
||||
def transform_s3_bucket_response_to_openai_file_object(
|
||||
self, create_file_data: CreateFileRequest, s3_upload_response: _S3UploadResponse
|
||||
|
|
|
|||
|
|
@ -14,6 +14,9 @@ from litellm.constants import AWS_ERROR_TYPE_HEADER
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
|
||||
from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
bedrock_invoke_guardrail_headers,
|
||||
)
|
||||
from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
|
|
@ -42,6 +45,9 @@ if TYPE_CHECKING:
|
|||
BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"})
|
||||
_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
|
||||
_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"})
|
||||
_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset(
|
||||
{"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"}
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_endpoint_action(endpoint: str) -> str:
|
||||
|
|
@ -123,6 +129,31 @@ def _with_invoke_request_metadata(
|
|||
return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body
|
||||
|
||||
|
||||
def _converse_guardrail_config(action: str, guardrail_config: object) -> object:
|
||||
fields: Final = as_json_object(guardrail_config)
|
||||
if action == "converse-stream" or fields is None:
|
||||
return guardrail_config
|
||||
return {key: value for key, value in fields.items() if key != "streamProcessingMode"}
|
||||
|
||||
|
||||
def _with_deployment_guardrail(
|
||||
action: str,
|
||||
headers: Mapping[str, str],
|
||||
body: Mapping[str, object],
|
||||
guardrail_config: object,
|
||||
) -> tuple[Mapping[str, str], Mapping[str, object]]:
|
||||
if guardrail_config is None:
|
||||
return headers, body
|
||||
if action in _BEDROCK_CONVERSE_ACTIONS:
|
||||
return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)}
|
||||
if action in _BEDROCK_INVOKE_ACTIONS:
|
||||
guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items())
|
||||
return merge_bedrock_invoke_headers(
|
||||
dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS
|
||||
), body
|
||||
return headers, body
|
||||
|
||||
|
||||
def _is_cache_point(entry: object) -> bool:
|
||||
fields: Final = as_json_object(entry)
|
||||
return fields is not None and "cachePoint" in fields
|
||||
|
|
@ -453,11 +484,14 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
|
|||
body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
guarded_headers, guarded_body = _with_deployment_guardrail(
|
||||
action, signed_headers, signed_body, litellm_params.get("guardrailConfig")
|
||||
)
|
||||
return self._sign_request(
|
||||
service_name="bedrock",
|
||||
headers=dict(signed_headers),
|
||||
headers=dict(guarded_headers),
|
||||
optional_params=optional_params,
|
||||
request_data=dict(signed_body),
|
||||
request_data=dict(guarded_body),
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
api_key=self.get_bedrock_bearer_token(optional_params),
|
||||
|
|
|
|||
|
|
@ -256,12 +256,11 @@ def merge_bedrock_invoke_headers(
|
|||
"""
|
||||
Merge the ``X-Amzn-*`` headers the Invoke paths derive from params.
|
||||
|
||||
``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is
|
||||
the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's
|
||||
headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry
|
||||
proxy-authenticated identity into an AWS billing record that the caller must not be able to
|
||||
write. Names are compared case-insensitively so a caller cannot leave a second spelling in
|
||||
the dict and let the transport pick the winner.
|
||||
``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped
|
||||
from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names
|
||||
the caller must not be able to write, such as proxy-authenticated identity headed for an AWS
|
||||
billing record or a deployment's guardrail. Names are compared case-insensitively so a caller
|
||||
cannot leave a second spelling in the dict and let the transport pick the winner.
|
||||
"""
|
||||
if not caller_owned and not proxy_owned and not proxy_owned_names:
|
||||
return headers
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ class ChatGPTConfig(OpenAIConfig):
|
|||
default_headers: Final = get_chatgpt_default_headers(api_key or "", account_id, session_id)
|
||||
return {**default_headers, **validated_headers}
|
||||
|
||||
def post_stream_processing(self, stream: Any) -> Any:
|
||||
def post_stream_processing(self, stream: Any) -> ChatGPTToolCallNormalizer:
|
||||
return ChatGPTToolCallNormalizer(stream)
|
||||
|
||||
def map_openai_params(
|
||||
|
|
|
|||
53
litellm/llms/cloudflare/common_utils.py
Normal file
53
litellm/llms/cloudflare/common_utils.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.secret_managers.main import (
|
||||
get_secret_str,
|
||||
normalize_nonempty_secret_str,
|
||||
)
|
||||
|
||||
CLOUDFLARE_MISSING_API_BASE_MESSAGE: Final = (
|
||||
"Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
|
||||
)
|
||||
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def workers_ai_default_api_base() -> str | None:
|
||||
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
|
||||
if account_id is None:
|
||||
return None
|
||||
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
|
||||
|
||||
|
||||
def workers_ai_canonical_model(model: str) -> str:
|
||||
if model.startswith("@cf/"):
|
||||
return model
|
||||
return f"@cf/cloudflare/{model}"
|
||||
|
||||
|
||||
def workers_ai_request_model(model: str) -> str:
|
||||
return model.rsplit("/", maxsplit=1)[-1]
|
||||
|
||||
|
||||
def workers_ai_run_url(api_base: str, model: str) -> str:
|
||||
normalized_api_base: Final = api_base.rstrip("/")
|
||||
if normalized_api_base.endswith("/ai/v1"):
|
||||
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
|
||||
if normalized_api_base.endswith("/ai/run"):
|
||||
return f"{normalized_api_base}/{model}"
|
||||
return f"{normalized_api_base}/ai/run/{model}"
|
||||
|
||||
|
||||
def unwrap_workers_ai_result(payload: object) -> object:
|
||||
try:
|
||||
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
|
||||
except ValidationError:
|
||||
return payload
|
||||
if "answers" in response_mapping:
|
||||
return response_mapping
|
||||
try:
|
||||
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
|
||||
except ValidationError:
|
||||
return response_mapping
|
||||
|
|
@ -1,16 +1,13 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
|
||||
from litellm.secret_managers.main import (
|
||||
get_secret_str,
|
||||
normalize_nonempty_secret_str,
|
||||
from litellm.llms.cloudflare.common_utils import (
|
||||
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
|
||||
unwrap_workers_ai_result,
|
||||
workers_ai_canonical_model,
|
||||
workers_ai_default_api_base,
|
||||
workers_ai_request_model,
|
||||
workers_ai_run_url,
|
||||
)
|
||||
|
||||
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class CloudflareDecisionsConfig(BaseDecisionsConfig):
|
||||
api_key_env = ("CLOUDFLARE_API_KEY",)
|
||||
|
|
@ -23,38 +20,19 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig):
|
|||
inlines_remote_images = True
|
||||
|
||||
def get_default_api_base(self) -> str | None:
|
||||
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
|
||||
if account_id is None:
|
||||
return None
|
||||
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
|
||||
return workers_ai_default_api_base()
|
||||
|
||||
def missing_api_base_message(self, custom_llm_provider: str) -> str:
|
||||
return "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
|
||||
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
|
||||
|
||||
def canonical_model(self, model: str) -> str:
|
||||
if model.startswith("@cf/"):
|
||||
return model
|
||||
return f"@cf/cloudflare/{model}"
|
||||
return workers_ai_canonical_model(model)
|
||||
|
||||
def request_model(self, model: str) -> str:
|
||||
return model.rsplit("/", maxsplit=1)[-1]
|
||||
return workers_ai_request_model(model)
|
||||
|
||||
def get_complete_url(self, api_base: str, model: str) -> str:
|
||||
normalized_api_base: Final = api_base.rstrip("/")
|
||||
if normalized_api_base.endswith("/ai/v1"):
|
||||
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
|
||||
if normalized_api_base.endswith("/ai/run"):
|
||||
return f"{normalized_api_base}/{model}"
|
||||
return f"{normalized_api_base}/ai/run/{model}"
|
||||
return workers_ai_run_url(api_base, model)
|
||||
|
||||
def unwrap_response(self, payload: object) -> object:
|
||||
try:
|
||||
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
|
||||
except ValidationError:
|
||||
return payload
|
||||
if "answers" in response_mapping:
|
||||
return response_mapping
|
||||
try:
|
||||
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
|
||||
except ValidationError:
|
||||
return response_mapping
|
||||
return unwrap_workers_ai_result(payload)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,11 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
guardrail_translation_mappings: Final = {
|
||||
CallTypes.extraction: ExtractionHandler,
|
||||
CallTypes.aextraction: ExtractionHandler,
|
||||
}
|
||||
|
||||
__all__ = ["ExtractionHandler", "guardrail_translation_mappings"]
|
||||
|
|
@ -0,0 +1,276 @@
|
|||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias, TypeVar
|
||||
|
||||
from pydantic import Field, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.types.extraction import (
|
||||
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
|
||||
ExtractionAnswer,
|
||||
ExtractionBase64Image,
|
||||
ExtractionResponse,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
|
||||
_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter(
|
||||
Mapping[str, Mapping[str, object]]
|
||||
)
|
||||
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_IMAGES_ADAPTER: Final[TypeAdapter[tuple[ExtractionBase64Image | JsonValue, ...]]] = TypeAdapter(
|
||||
tuple[Annotated[ExtractionBase64Image | JsonValue, Field(union_mode="left_to_right")], ...]
|
||||
)
|
||||
|
||||
|
||||
_AnswerKey: TypeAlias = Literal["value", "guess", "mean", "median", "p10", "p90", "quantiles"]
|
||||
_ESTIMATE_KEYS: Final[tuple[_AnswerKey, ...]] = ("value", "guess", "mean", "median", "p10", "p90", "quantiles")
|
||||
_STATISTIC_KEYS: Final[frozenset[_AnswerKey]] = frozenset(("mean", "median", "p10", "p90", "quantiles"))
|
||||
_NULLED_ANSWER: Final[Mapping[str, None]] = MappingProxyType(
|
||||
dict.fromkeys(("value", "guess", "mean", "median", "p10", "p90", "quantiles", "steps", "span", "spans"), None)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AnswerText:
|
||||
name: str
|
||||
key: _AnswerKey
|
||||
item: int | None
|
||||
text: str
|
||||
|
||||
|
||||
_SPEC_TEXT_KEYS: Final = ("instructions", "unit")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SpecText:
|
||||
name: str
|
||||
key: str
|
||||
text: str
|
||||
|
||||
|
||||
def _spec_texts_of(name: str, spec: Mapping[str, object]) -> tuple[_SpecText, ...]:
|
||||
return tuple(
|
||||
_SpecText(name=name, key=key, text=text) for key in _SPEC_TEXT_KEYS if isinstance(text := spec.get(key), str)
|
||||
)
|
||||
|
||||
|
||||
def _spec_texts(fields: Mapping[str, Mapping[str, object]]) -> tuple[_SpecText, ...]:
|
||||
return tuple(chain.from_iterable(_spec_texts_of(name, spec) for name, spec in fields.items()))
|
||||
|
||||
|
||||
_NamedText = TypeVar("_NamedText", _AnswerText, _SpecText)
|
||||
|
||||
|
||||
def _scans_by_name(scanned: Iterable[tuple[_NamedText, str]]) -> Mapping[str, tuple[tuple[_NamedText, str], ...]]:
|
||||
return {name: tuple(group) for name, group in groupby(scanned, key=lambda pair: pair[0].name)}
|
||||
|
||||
|
||||
def _guardrailed_spec(spec: Mapping[str, object], scanned: Sequence[tuple[_SpecText, str]]) -> Mapping[str, object]:
|
||||
return {**spec, **{text.key: guardrailed for text, guardrailed in scanned}}
|
||||
|
||||
|
||||
def _metadata_key(data: Mapping[str, object]) -> Literal["metadata", "litellm_metadata"]:
|
||||
return get_metadata_variable_name_from_kwargs(data)
|
||||
|
||||
|
||||
def _metadata(data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return _METADATA_ADAPTER.validate_python(data.get(_metadata_key(data)) or {})
|
||||
|
||||
|
||||
def _is_text_answer(answer: ExtractionAnswer, spec: Mapping[str, object] | None) -> bool:
|
||||
return answer.type == "text" or (answer.type == "list" and (spec is None or spec.get("items") == "text"))
|
||||
|
||||
|
||||
def _image_urls(images: object) -> tuple[str, ...]:
|
||||
if not isinstance(images, list):
|
||||
return ()
|
||||
return tuple(
|
||||
f"data:{image.content_type};base64,{image.base64}" if isinstance(image, ExtractionBase64Image) else image
|
||||
for image in _IMAGES_ADAPTER.validate_python(images)
|
||||
if isinstance(image, (str, ExtractionBase64Image))
|
||||
)
|
||||
|
||||
|
||||
def _scalar_text(value: JsonValue, whole: bool) -> str | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, float) and whole:
|
||||
return format(value, ".0f")
|
||||
return str(value) if isinstance(value, (str, int, float)) else None
|
||||
|
||||
|
||||
def _item_texts(items: Sequence[JsonValue], whole: bool) -> tuple[tuple[int, str], ...]:
|
||||
return tuple((index, text) for index, item in enumerate(items) if (text := _scalar_text(item, whole)) is not None)
|
||||
|
||||
|
||||
def _value_texts(name: str, key: _AnswerKey, value: JsonValue) -> tuple[_AnswerText, ...]:
|
||||
whole: Final = key in _STATISTIC_KEYS
|
||||
if isinstance(value, list):
|
||||
return tuple(
|
||||
_AnswerText(name=name, key=key, item=index, text=text) for index, text in _item_texts(value, whole)
|
||||
)
|
||||
text: Final = _scalar_text(value, whole)
|
||||
return () if text is None else (_AnswerText(name=name, key=key, item=None, text=text),)
|
||||
|
||||
|
||||
def _step_texts(name: str, step: JsonValue) -> tuple[_AnswerText, ...]:
|
||||
if not isinstance(step, dict):
|
||||
return ()
|
||||
return tuple(chain.from_iterable(_value_texts(name, key, step.get(key)) for key in _ESTIMATE_KEYS))
|
||||
|
||||
|
||||
def _answer_texts(name: str, answer: ExtractionAnswer) -> tuple[_AnswerText, ...]:
|
||||
return (
|
||||
*_value_texts(name, "value", answer.value),
|
||||
*_value_texts(name, "guess", answer.guess),
|
||||
*_value_texts(name, "mean", answer.mean),
|
||||
*_value_texts(name, "median", answer.median),
|
||||
*_value_texts(name, "p10", answer.p10),
|
||||
*_value_texts(name, "p90", answer.p90),
|
||||
*_value_texts(name, "quantiles", list(answer.quantiles or ())),
|
||||
*chain.from_iterable(_step_texts(name, step) for step in answer.steps or ()),
|
||||
)
|
||||
|
||||
|
||||
def _guardrailed_value(value: JsonValue, scanned: Mapping[int | None, str]) -> JsonValue:
|
||||
if isinstance(value, list):
|
||||
return [scanned.get(index, item) for index, item in enumerate(value)]
|
||||
return scanned.get(None, value)
|
||||
|
||||
|
||||
def _scans_of(key: _AnswerKey, scanned: Sequence[tuple[_AnswerText, str]]) -> Mapping[int | None, str]:
|
||||
return {text.item: guardrailed for text, guardrailed in scanned if text.key == key}
|
||||
|
||||
|
||||
def _guardrailed_answer(
|
||||
answer: ExtractionAnswer,
|
||||
scanned: Sequence[tuple[_AnswerText, str]],
|
||||
drop_spans: bool,
|
||||
is_text: bool,
|
||||
) -> ExtractionAnswer:
|
||||
if not is_text and any(text.text != guardrailed for text, guardrailed in scanned):
|
||||
return answer.model_copy(update=_NULLED_ANSWER)
|
||||
value_scans: Final = _scans_of("value", scanned if is_text else ())
|
||||
guess_scans: Final = _scans_of("guess", scanned if is_text else ())
|
||||
update: Final = {
|
||||
**({"value": _guardrailed_value(answer.value, value_scans)} if value_scans else {}),
|
||||
**({"guess": _guardrailed_value(answer.guess, guess_scans)} if guess_scans else {}),
|
||||
**({"span": None, "spans": None} if drop_spans else {}),
|
||||
}
|
||||
return answer.model_copy(update=update) if update else answer
|
||||
|
||||
|
||||
async def _scanned_texts(
|
||||
texts: Sequence[str],
|
||||
model: object,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict
|
||||
input_type: Literal["request", "response"],
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
images: Sequence[str] = (),
|
||||
) -> tuple[str, ...]:
|
||||
unique_texts: Final = tuple(dict.fromkeys(texts))
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None)
|
||||
if images:
|
||||
inputs["images"] = list(images)
|
||||
guardrailed: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type=input_type,
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
returned: Final = guardrailed.get("texts", [])
|
||||
guardrailed_by_text: Final = MappingProxyType(
|
||||
{text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)}
|
||||
)
|
||||
return tuple(guardrailed_by_text[text] for text in texts)
|
||||
|
||||
|
||||
class ExtractionHandler(BaseTranslation):
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> Mapping[str, object]:
|
||||
state: Final = data.get("state")
|
||||
states: Final = (state,) if isinstance(state, str) and state else ()
|
||||
fields: Final = _FIELDS_ADAPTER.validate_python(data.get("fields") or {})
|
||||
spec_texts: Final = _spec_texts(fields)
|
||||
images: Final = _image_urls(data.get("images"))
|
||||
if not states and not spec_texts and not images:
|
||||
return data
|
||||
scanned: Final = await _scanned_texts(
|
||||
(*states, *(text.text for text in spec_texts)),
|
||||
data.get("model"),
|
||||
guardrail_to_apply,
|
||||
data,
|
||||
"request",
|
||||
litellm_logging_obj,
|
||||
images,
|
||||
)
|
||||
spec_scans: Final = _scans_by_name(zip(spec_texts, scanned[len(states) :]))
|
||||
state_rewritten: Final = bool(states) and scanned[0] != states[0]
|
||||
return {
|
||||
**data,
|
||||
**({"state": scanned[0]} if states else {}),
|
||||
**(
|
||||
{_metadata_key(data): {**_metadata(data), EXTRACTION_STATE_REWRITTEN_METADATA_KEY: True}}
|
||||
if state_rewritten
|
||||
else {}
|
||||
),
|
||||
"fields": {name: _guardrailed_spec(spec, spec_scans.get(name, ())) for name, spec in fields.items()},
|
||||
}
|
||||
|
||||
async def process_output_response(
|
||||
self,
|
||||
response: object,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict
|
||||
) -> object:
|
||||
if not isinstance(response, ExtractionResponse):
|
||||
return response
|
||||
texts: Final = tuple(
|
||||
chain.from_iterable(_answer_texts(name, answer) for name, answer in response.answers.items())
|
||||
)
|
||||
state_texts: Final = (response.state_text,) if response.state_text else ()
|
||||
request: Final = request_data or {}
|
||||
scanned: Final = (
|
||||
await _scanned_texts(
|
||||
(*(text.text for text in texts), *state_texts),
|
||||
response.model,
|
||||
guardrail_to_apply,
|
||||
request,
|
||||
"response",
|
||||
litellm_logging_obj,
|
||||
)
|
||||
if texts or state_texts
|
||||
else ()
|
||||
)
|
||||
answer_scans: Final = _scans_by_name(zip(texts, scanned))
|
||||
state_text: Final = next(iter(scanned[len(texts) :]), response.state_text)
|
||||
specs: Final = _FIELDS_ADAPTER.validate_python(request.get("fields") or {})
|
||||
return response.model_copy(
|
||||
update={
|
||||
"answers": {
|
||||
name: _guardrailed_answer(
|
||||
answer,
|
||||
answer_scans.get(name, ()),
|
||||
drop_spans=state_text != response.state_text,
|
||||
is_text=_is_text_answer(answer, specs.get(name)),
|
||||
)
|
||||
for name, answer in response.answers.items()
|
||||
},
|
||||
**({"state_text": state_text} if state_texts else {}),
|
||||
}
|
||||
)
|
||||
32
litellm/llms/cloudflare/extraction/transformation.py
Normal file
32
litellm/llms/cloudflare/extraction/transformation.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
|
||||
from litellm.llms.cloudflare.common_utils import (
|
||||
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
|
||||
unwrap_workers_ai_result,
|
||||
workers_ai_canonical_model,
|
||||
workers_ai_default_api_base,
|
||||
workers_ai_request_model,
|
||||
workers_ai_run_url,
|
||||
)
|
||||
|
||||
|
||||
class CloudflareExtractionConfig(BaseExtractionConfig):
|
||||
api_key_env = ("CLOUDFLARE_API_KEY",)
|
||||
api_base_env = ("CLOUDFLARE_API_BASE",)
|
||||
|
||||
def get_default_api_base(self) -> str | None:
|
||||
return workers_ai_default_api_base()
|
||||
|
||||
def missing_api_base_message(self, custom_llm_provider: str) -> str:
|
||||
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
|
||||
|
||||
def canonical_model(self, model: str) -> str:
|
||||
return workers_ai_canonical_model(model)
|
||||
|
||||
def request_model(self, model: str) -> str:
|
||||
return workers_ai_request_model(model)
|
||||
|
||||
def get_complete_url(self, api_base: str, model: str) -> str:
|
||||
return workers_ai_run_url(api_base, model)
|
||||
|
||||
def unwrap_response(self, payload: object) -> object:
|
||||
return unwrap_workers_ai_result(payload)
|
||||
|
|
@ -76,6 +76,7 @@ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
|
|||
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
|
||||
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
|
||||
from litellm.llms.base_llm.files.transformation import (
|
||||
BaseFilesConfig,
|
||||
BaseFileUploadStream,
|
||||
|
|
@ -128,6 +129,7 @@ from litellm.types.containers.main import (
|
|||
DeleteContainerResult,
|
||||
)
|
||||
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
|
||||
from litellm.types.extraction import ExtractionResponse
|
||||
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
|
|
@ -1700,6 +1702,115 @@ class BaseLLMHTTPHandler:
|
|||
model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request
|
||||
)
|
||||
|
||||
def _prepare_extraction_request(
|
||||
self,
|
||||
model: str,
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
provider_config: BaseExtractionConfig,
|
||||
body: Mapping[str, object],
|
||||
api_base: str,
|
||||
api_key: str | None,
|
||||
headers: Mapping[str, str],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[str, Mapping[str, str]]:
|
||||
outbound_headers: Final = provider_config.validate_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
url: Final = provider_config.get_complete_url(api_base=api_base, model=model)
|
||||
if logging_obj is not None:
|
||||
logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # Logging.pre_call is untyped
|
||||
input=dict(body),
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
additional_args={"api_base": url, "complete_input_dict": dict(body), "headers": outbound_headers},
|
||||
)
|
||||
return url, outbound_headers
|
||||
|
||||
def extraction(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
provider_config: BaseExtractionConfig,
|
||||
body: Mapping[str, object],
|
||||
api_base: str,
|
||||
api_key: str | None,
|
||||
headers: Mapping[str, str],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
client: HTTPHandler | None = None,
|
||||
) -> ExtractionResponse:
|
||||
url, outbound_headers = self._prepare_extraction_request(
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
provider_config=provider_config,
|
||||
body=body,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
headers=headers,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
sync_httpx_client: Final = client if client is not None else get_httpx_client()
|
||||
try:
|
||||
response: Final = sync_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
|
||||
url,
|
||||
json=dict(body),
|
||||
headers=dict(outbound_headers),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
return provider_config.transform_extraction_response(
|
||||
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
|
||||
)
|
||||
|
||||
async def aextraction(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
provider_config: BaseExtractionConfig,
|
||||
body: Mapping[str, object],
|
||||
api_base: str,
|
||||
api_key: str | None,
|
||||
headers: Mapping[str, str],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> ExtractionResponse:
|
||||
url, outbound_headers = self._prepare_extraction_request(
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
provider_config=provider_config,
|
||||
body=body,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
headers=headers,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
async_httpx_client: Final = (
|
||||
client
|
||||
if client is not None
|
||||
else get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider))
|
||||
)
|
||||
try:
|
||||
response: Final = await async_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
|
||||
url,
|
||||
json=dict(body),
|
||||
headers=dict(outbound_headers),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
return provider_config.transform_extraction_response(
|
||||
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
|
||||
)
|
||||
|
||||
def _prepare_audio_transcription_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -6462,6 +6573,7 @@ class BaseLLMHTTPHandler:
|
|||
provider_config: Union[
|
||||
BaseConfig,
|
||||
BaseDecisionsConfig,
|
||||
BaseExtractionConfig,
|
||||
BaseRerankConfig,
|
||||
BaseResponsesAPIConfig,
|
||||
BaseImageEditConfig,
|
||||
|
|
|
|||
|
|
@ -82,9 +82,9 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
return "tts" in model
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
sampling_params: Final = ("temperature", "top_p") if self.supports_sampling_params(model) else ()
|
||||
supported_params: Final = [
|
||||
"temperature",
|
||||
"top_p",
|
||||
*sampling_params,
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"stream",
|
||||
|
|
|
|||
78
litellm/llms/github_copilot/github_host.py
Normal file
78
litellm/llms/github_copilot/github_host.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
|
||||
from .authenticator import (
|
||||
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
DEFAULT_GITHUB_CLIENT_ID,
|
||||
DEFAULT_GITHUB_DEVICE_CODE_URL,
|
||||
)
|
||||
|
||||
_GHE_SUBDOMAIN_PATTERN: Final = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GithubHostEndpoints:
|
||||
host: str
|
||||
device_code_url: str
|
||||
access_token_url: str
|
||||
api_key_url: str
|
||||
user_url: str
|
||||
client_id: str
|
||||
copilot_api_hosts: frozenset[str]
|
||||
copilot_api_host_suffix: str | None
|
||||
default_copilot_api_base: str
|
||||
|
||||
|
||||
def normalized_github_host(raw: object) -> str | None:
|
||||
if raw is None or raw == "":
|
||||
return GITHUB_COPILOT_DEFAULT_HOST
|
||||
if not isinstance(raw, str):
|
||||
return None
|
||||
|
||||
host: Final = raw.lower()
|
||||
if host == GITHUB_COPILOT_DEFAULT_HOST:
|
||||
return host
|
||||
if not host.endswith(".ghe.com"):
|
||||
return None
|
||||
|
||||
subdomain: Final = host.removesuffix(".ghe.com")
|
||||
return host if _GHE_SUBDOMAIN_PATTERN.fullmatch(subdomain) is not None else None
|
||||
|
||||
|
||||
def github_host_endpoints(host: str) -> GithubHostEndpoints:
|
||||
normalized: Final = normalized_github_host(host)
|
||||
if normalized is None:
|
||||
raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
if normalized == GITHUB_COPILOT_DEFAULT_HOST:
|
||||
return GithubHostEndpoints(
|
||||
host=normalized,
|
||||
device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
|
||||
access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
|
||||
api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
|
||||
user_url="https://api.github.com/user",
|
||||
client_id=client_id,
|
||||
copilot_api_hosts=frozenset({"githubcopilot.com"}),
|
||||
copilot_api_host_suffix=".githubcopilot.com",
|
||||
default_copilot_api_base="https://api.githubcopilot.com",
|
||||
)
|
||||
|
||||
subdomain: Final = normalized.removesuffix(".ghe.com")
|
||||
api_host: Final = f"api.{subdomain}.ghe.com"
|
||||
copilot_api_host: Final = f"copilot-api.{subdomain}.ghe.com"
|
||||
return GithubHostEndpoints(
|
||||
host=normalized,
|
||||
device_code_url=f"https://{normalized}/login/device/code",
|
||||
access_token_url=f"https://{normalized}/login/oauth/access_token",
|
||||
api_key_url=f"https://{api_host}/copilot_internal/v2/token",
|
||||
user_url=f"https://{api_host}/user",
|
||||
client_id=client_id,
|
||||
copilot_api_hosts=frozenset({copilot_api_host}),
|
||||
copilot_api_host_suffix=None,
|
||||
default_copilot_api_base=f"https://{copilot_api_host}",
|
||||
)
|
||||
|
|
@ -5,8 +5,7 @@ touched on the per-user path.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import hmac
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -27,6 +26,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import (
|
||||
GITHUB_COPILOT_AUTH_TYPE_KEY,
|
||||
GITHUB_COPILOT_HOST_KEY,
|
||||
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
|
||||
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
|
||||
)
|
||||
|
|
@ -38,14 +38,8 @@ from litellm.exceptions import (
|
|||
ServiceUnavailableError,
|
||||
)
|
||||
|
||||
from .authenticator import (
|
||||
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
DEFAULT_GITHUB_CLIENT_ID,
|
||||
DEFAULT_GITHUB_DEVICE_CODE_URL,
|
||||
github_api_headers,
|
||||
)
|
||||
from .common_utils import DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
from .authenticator import github_api_headers
|
||||
from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host
|
||||
|
||||
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session"
|
||||
_LLM_PROVIDER: Final = "github_copilot"
|
||||
|
|
@ -98,7 +92,7 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess
|
|||
)
|
||||
_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
|
||||
_IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking
|
||||
dict[str, "asyncio.Future[GithubCopilotUserSession]"]
|
||||
dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]]
|
||||
] = {}
|
||||
|
||||
|
||||
|
|
@ -141,7 +135,7 @@ _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
|
|||
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
|
||||
_GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
|
||||
_COPILOT_ENDPOINTS: Final = TypeAdapter(_CopilotEndpointsPayload)
|
||||
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, float])
|
||||
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, str, float])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
|
|
@ -178,53 +172,72 @@ class GithubCopilotDeviceFlowPoll:
|
|||
access_token: str | None = field(default=None, repr=False)
|
||||
|
||||
|
||||
def validated_copilot_api_base(endpoints_api: object) -> str:
|
||||
"""Only a genuine *.githubcopilot.com HTTPS endpoint may receive a Copilot
|
||||
token; anything else falls back to the default base."""
|
||||
def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str:
|
||||
if not isinstance(endpoints_api, str):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
parsed: Final = urlparse(endpoints_api)
|
||||
return endpoints.default_copilot_api_base
|
||||
try:
|
||||
parsed: Final = urlparse(endpoints_api)
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
return endpoints.default_copilot_api_base
|
||||
hostname: Final = (parsed.hostname or "").lower()
|
||||
if (
|
||||
parsed.scheme.lower() != "https"
|
||||
or not (hostname == "githubcopilot.com" or hostname.endswith(".githubcopilot.com"))
|
||||
or not (
|
||||
hostname in endpoints.copilot_api_hosts
|
||||
or (endpoints.copilot_api_host_suffix is not None and hostname.endswith(endpoints.copilot_api_host_suffix))
|
||||
)
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
or port not in (None, 443)
|
||||
):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
try:
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
if port not in (None, 443):
|
||||
return DEFAULT_GITHUB_COPILOT_API_BASE
|
||||
return endpoints.default_copilot_api_base
|
||||
return endpoints_api.rstrip("/")
|
||||
|
||||
|
||||
def _session_cache_key(user_id: str, github_token: str) -> str:
|
||||
return hashlib.sha256(f"{user_id}\x00{github_token}".encode()).hexdigest()
|
||||
def github_copilot_host_endpoints(credential_name: str) -> GithubHostEndpoints:
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
|
||||
credential: Final = CredentialAccessor.find_credential(credential_name)
|
||||
host_raw: Final[object] = (
|
||||
credential.credential_values.get(GITHUB_COPILOT_HOST_KEY) if credential is not None else None
|
||||
)
|
||||
host: Final = normalized_github_host(host_raw)
|
||||
if host is None:
|
||||
raise BadRequestError(
|
||||
message=f"Credential '{credential_name}' has an invalid {GITHUB_COPILOT_HOST_KEY}",
|
||||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
)
|
||||
return github_host_endpoints(host)
|
||||
|
||||
|
||||
def evict_copilot_user_session(user_id: str, github_token: str) -> None:
|
||||
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, github_token))
|
||||
def _session_cache_key(user_id: str, credential_name: str, github_host: str) -> str:
|
||||
return f"{user_id}\x00{credential_name}\x00{github_host}"
|
||||
|
||||
|
||||
def _cached_session(cache_key: str) -> GithubCopilotUserSession | None:
|
||||
def evict_copilot_user_session(user_id: str, credential_name: str, github_host: str) -> None:
|
||||
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, github_host))
|
||||
|
||||
|
||||
def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSession | None:
|
||||
cached: Final[object] = _SESSION_CACHE.get_cache(cache_key)
|
||||
try:
|
||||
token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
|
||||
cached_github_token, token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
|
||||
except ValidationError:
|
||||
return None
|
||||
if not hmac.compare_digest(cached_github_token, github_token):
|
||||
return None
|
||||
if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time():
|
||||
return None
|
||||
return GithubCopilotUserSession(token=token, api_base=api_base)
|
||||
|
||||
|
||||
def _cache_session(cache_key: str, session: GithubCopilotUserSession, expires_at: float) -> None:
|
||||
def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None:
|
||||
ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time()
|
||||
if ttl <= 0:
|
||||
return
|
||||
_SESSION_CACHE.set_cache(cache_key, (session.token, session.api_base, expires_at), ttl=ttl)
|
||||
_SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl)
|
||||
|
||||
|
||||
def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError:
|
||||
|
|
@ -238,7 +251,13 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr
|
|||
)
|
||||
|
||||
|
||||
def _session_from_response(response: httpx.Response, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
def _session_from_response(
|
||||
response: httpx.Response,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
if response.status_code in (401, 403, 404):
|
||||
_SESSION_CACHE.delete_cache(cache_key)
|
||||
raise _reject_connection(credential_name)
|
||||
|
|
@ -263,15 +282,17 @@ def _session_from_response(response: httpx.Response, credential_name: str, cache
|
|||
model="",
|
||||
) from e
|
||||
|
||||
endpoints: Final[object] = payload.get("endpoints")
|
||||
payload_endpoints: Final[object] = payload.get("endpoints")
|
||||
endpoints_api: Final[object] = (
|
||||
_COPILOT_ENDPOINTS.validate_python(endpoints).get("api") if isinstance(endpoints, Mapping) else None
|
||||
_COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
|
||||
if isinstance(payload_endpoints, Mapping)
|
||||
else None
|
||||
)
|
||||
api_base: Final = validated_copilot_api_base(endpoints_api)
|
||||
api_base: Final = validated_copilot_api_base(endpoints_api, endpoints)
|
||||
session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
|
||||
expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
|
||||
if expires_at > 0:
|
||||
_cache_session(cache_key, session, expires_at)
|
||||
_cache_session(cache_key, github_token, session, expires_at)
|
||||
return session
|
||||
|
||||
|
||||
|
|
@ -303,7 +324,12 @@ def _expires_at_seconds(raw: object) -> float:
|
|||
return 0.0
|
||||
|
||||
|
||||
def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
def _do_exchange(
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
import litellm
|
||||
|
||||
try:
|
||||
|
|
@ -311,7 +337,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
|
|||
_SyncGetClient, litellm.module_level_client
|
||||
)
|
||||
response: Final = client.get(
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
endpoints.api_key_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -320,7 +346,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
|
|||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
) from e
|
||||
return _session_from_response(response, credential_name, cache_key)
|
||||
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
def _new_async_github_client() -> _AsyncGitHubClient:
|
||||
|
|
@ -332,11 +358,16 @@ def _new_async_github_client() -> _AsyncGitHubClient:
|
|||
)(_LLM_PROVIDER)
|
||||
|
||||
|
||||
async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
|
||||
async def _do_aexchange(
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
cache_key: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
client: Final = _new_async_github_client()
|
||||
try:
|
||||
response: Final = await client.get(
|
||||
DEFAULT_GITHUB_API_KEY_URL,
|
||||
endpoints.api_key_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -345,43 +376,50 @@ async def _do_aexchange(github_token: str, credential_name: str, cache_key: str)
|
|||
llm_provider=_LLM_PROVIDER,
|
||||
model="",
|
||||
) from e
|
||||
return _session_from_response(response, credential_name, cache_key)
|
||||
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
def exchange_github_token(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
cache_key: Final = _session_cache_key(user_id, github_token)
|
||||
cached: Final = _cached_session(cache_key)
|
||||
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
|
||||
cached: Final = _cached_session(cache_key, github_token)
|
||||
if cached is not None:
|
||||
return cached
|
||||
lock: Final = _EXCHANGE_LOCKS[int(cache_key, 16) % _SYNC_LOCK_COUNT]
|
||||
lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
|
||||
with lock:
|
||||
rechecked: Final = _cached_session(cache_key)
|
||||
rechecked: Final = _cached_session(cache_key, github_token)
|
||||
if rechecked is not None:
|
||||
return rechecked
|
||||
return _do_exchange(github_token, credential_name, cache_key)
|
||||
return _do_exchange(github_token, credential_name, cache_key, endpoints)
|
||||
|
||||
|
||||
async def aexchange_github_token(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
cache_key: Final = _session_cache_key(user_id, github_token)
|
||||
cached: Final = _cached_session(cache_key)
|
||||
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
|
||||
cached: Final = _cached_session(cache_key, github_token)
|
||||
if cached is not None:
|
||||
return cached
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
in_flight: Final = _IN_FLIGHT.get(cache_key)
|
||||
if in_flight is not None and not in_flight.done() and in_flight.get_loop() is loop:
|
||||
return await asyncio.shield(in_flight)
|
||||
in_flight_entry: Final = _IN_FLIGHT.get(cache_key)
|
||||
if in_flight_entry is not None:
|
||||
in_flight_token, in_flight = in_flight_entry
|
||||
if not hmac.compare_digest(in_flight_token, github_token):
|
||||
return await _do_aexchange(github_token, credential_name, cache_key, endpoints)
|
||||
if not in_flight.done() and in_flight.get_loop() is loop:
|
||||
return await asyncio.shield(in_flight)
|
||||
future: Final = loop.create_future()
|
||||
_IN_FLIGHT[cache_key] = future
|
||||
entry: Final = (github_token, future)
|
||||
_IN_FLIGHT[cache_key] = entry
|
||||
try:
|
||||
session: Final = await _do_aexchange(github_token, credential_name, cache_key)
|
||||
session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints)
|
||||
except BaseException as e:
|
||||
if not future.done():
|
||||
future.set_exception(e)
|
||||
|
|
@ -392,19 +430,17 @@ async def aexchange_github_token(
|
|||
future.set_result(session)
|
||||
return session
|
||||
finally:
|
||||
if _IN_FLIGHT.get(cache_key) is future:
|
||||
if _IN_FLIGHT.get(cache_key) is entry:
|
||||
_IN_FLIGHT.pop(cache_key, None)
|
||||
|
||||
|
||||
async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
|
||||
async def astart_device_flow(endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowStart:
|
||||
client: Final = _new_async_github_client()
|
||||
device_code_url: Final = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL)
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
device_code_url,
|
||||
endpoints.device_code_url,
|
||||
headers=github_api_headers(),
|
||||
json={"client_id": client_id, "scope": _DEVICE_FLOW_SCOPE},
|
||||
json={"client_id": endpoints.client_id, "scope": _DEVICE_FLOW_SCOPE},
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise ServiceUnavailableError(
|
||||
|
|
@ -435,16 +471,14 @@ async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
|
|||
)
|
||||
|
||||
|
||||
async def apoll_device_flow(device_code: str) -> GithubCopilotDeviceFlowPoll:
|
||||
async def apoll_device_flow(device_code: str, endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowPoll:
|
||||
client: Final = _new_async_github_client()
|
||||
access_token_url: Final = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL)
|
||||
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
access_token_url,
|
||||
endpoints.access_token_url,
|
||||
headers=github_api_headers(),
|
||||
json={
|
||||
"client_id": client_id,
|
||||
"client_id": endpoints.client_id,
|
||||
"device_code": device_code,
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
|
||||
},
|
||||
|
|
@ -500,11 +534,11 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice
|
|||
)
|
||||
|
||||
|
||||
async def afetch_github_login(github_token: str) -> str:
|
||||
async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str:
|
||||
client: Final = _new_async_github_client()
|
||||
try:
|
||||
response: Final = await client.get(
|
||||
"https://api.github.com/user",
|
||||
endpoints.user_url,
|
||||
headers=github_api_headers(github_token),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -530,11 +564,17 @@ async def afetch_github_login(github_token: str) -> str:
|
|||
return payload["login"]
|
||||
|
||||
|
||||
async def acheck_copilot_seat(user_id: str, github_token: str, credential_name: str) -> GithubCopilotUserSession:
|
||||
async def acheck_copilot_seat(
|
||||
user_id: str,
|
||||
github_token: str,
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> GithubCopilotUserSession:
|
||||
return await aexchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -567,7 +607,11 @@ def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError:
|
|||
)
|
||||
|
||||
|
||||
def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tuple[str, str]:
|
||||
def _caller_connection(
|
||||
kwargs: Mapping[str, object],
|
||||
credential_name: str,
|
||||
endpoints: GithubHostEndpoints,
|
||||
) -> tuple[str, str]:
|
||||
secret_fields_raw: Final[object] = kwargs.get("secret_fields")
|
||||
secret_fields: Final = (
|
||||
cast(Mapping[object, object], secret_fields_raw) # cast-ok: value is Mapping-checked or dict-shaped at runtime
|
||||
|
|
@ -585,8 +629,21 @@ def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tu
|
|||
if isinstance(credentials_raw, Mapping)
|
||||
else None
|
||||
)
|
||||
token: Final[object] = credentials.get(credential_name) if credentials is not None else None
|
||||
if not isinstance(user_id, str) or not user_id or not isinstance(token, str) or not token:
|
||||
connection_raw: Final[object] = credentials.get(credential_name) if credentials is not None else None
|
||||
connection: Final = (
|
||||
cast(Mapping[object, object], connection_raw) # cast-ok: runtime Mapping check validates the stored connection
|
||||
if isinstance(connection_raw, Mapping)
|
||||
else None
|
||||
)
|
||||
token: Final[object] = connection.get("access_token") if connection is not None else None
|
||||
github_host: Final = normalized_github_host(connection.get("github_host") if connection is not None else None)
|
||||
if (
|
||||
not isinstance(user_id, str)
|
||||
or not user_id
|
||||
or not isinstance(token, str)
|
||||
or not token
|
||||
or github_host != endpoints.host
|
||||
):
|
||||
raise _connect_error(credential_name)
|
||||
return user_id, token
|
||||
|
||||
|
|
@ -705,11 +762,13 @@ def attach_github_copilot_user_session(
|
|||
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
|
||||
return
|
||||
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
|
||||
session: Final = exchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
|
||||
_strip_caller_authorization(kwargs)
|
||||
|
|
@ -725,11 +784,13 @@ async def aattach_github_copilot_user_session(
|
|||
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
|
||||
return
|
||||
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name)
|
||||
endpoints: Final = github_copilot_host_endpoints(credential_name)
|
||||
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
|
||||
session: Final = await aexchange_github_token(
|
||||
user_id=user_id,
|
||||
github_token=github_token,
|
||||
credential_name=credential_name,
|
||||
endpoints=endpoints,
|
||||
)
|
||||
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
|
||||
_strip_caller_authorization(kwargs)
|
||||
|
|
|
|||
|
|
@ -16,7 +16,9 @@ The handler is intentionally a thin orchestration layer:
|
|||
3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS)
|
||||
vs self-hosted (any host:port, optional TLS) modes.
|
||||
4. Stream the audio through Riva's ``streaming_response_generator`` and
|
||||
aggregate ``is_final`` results into a single transcript.
|
||||
aggregate ``is_final`` results into a single transcript, or, when
|
||||
``riva_offline`` is set, send the whole buffer through the unary
|
||||
``offline_recognize`` RPC for NIMs that only ship offline engines.
|
||||
5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed
|
||||
on ``_hidden_params`` so cost calculation works.
|
||||
|
||||
|
|
@ -26,13 +28,19 @@ without the optional STT extras installed.
|
|||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Callable, Iterable
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
get_audio_file_name,
|
||||
process_audio_file,
|
||||
)
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.llms.nvidia_riva.audio_transcription.audio_utils import (
|
||||
resample_to_riva_pcm,
|
||||
)
|
||||
|
|
@ -45,6 +53,7 @@ from litellm.llms.nvidia_riva.common_utils import (
|
|||
NvidiaRivaException,
|
||||
grpc_error_to_litellm_exception,
|
||||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.utils import FileTypes, TranscriptionResponse
|
||||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
|
|
@ -67,10 +76,20 @@ class _RivaAuth(Protocol):
|
|||
"""Opaque ``riva.client.Auth`` handle."""
|
||||
|
||||
|
||||
class _RecognizeFuture(Protocol):
|
||||
def result(self, timeout: float | None = None) -> object: ...
|
||||
|
||||
def done(self) -> bool: ...
|
||||
|
||||
def cancel(self) -> bool: ...
|
||||
|
||||
|
||||
class _AsrService(Protocol):
|
||||
@property
|
||||
def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
|
||||
|
||||
def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ...
|
||||
|
||||
|
||||
class _EndpointingConfig(Protocol):
|
||||
"""Opaque ``EndpointingConfig`` protobuf message."""
|
||||
|
|
@ -128,6 +147,38 @@ def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding:
|
|||
return riva_asr_module.AudioEncoding
|
||||
|
||||
|
||||
def _parse_riva_offline(value: object, drop_params: bool) -> bool:
|
||||
match value:
|
||||
case None:
|
||||
return False
|
||||
case bool():
|
||||
return value
|
||||
case str() if (parsed := str_to_bool(value)) is not None:
|
||||
return parsed
|
||||
case _ if drop_params:
|
||||
return False
|
||||
case _:
|
||||
raise NvidiaRivaException(
|
||||
status_code=400,
|
||||
message=f"riva_offline must be true or false, got {value!r}. Set drop_params to ignore it.",
|
||||
)
|
||||
|
||||
|
||||
def _deadline_seconds(timeout: float | str | httpx.Timeout, drop_params: bool) -> float:
|
||||
try:
|
||||
limit: Final = CompletionTimeout.as_limit(timeout)
|
||||
except ValueError:
|
||||
if drop_params:
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
raise NvidiaRivaException(
|
||||
status_code=400,
|
||||
message=f"timeout must be a number of seconds, got {timeout!r}. Set drop_params to ignore it.",
|
||||
) from None
|
||||
if not isinstance(limit, httpx.Timeout):
|
||||
return limit
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS if limit.read is None else limit.read
|
||||
|
||||
|
||||
class NvidiaRivaAudioTranscription:
|
||||
"""Sync + async entry point for Riva ASR."""
|
||||
|
||||
|
|
@ -138,7 +189,7 @@ class NvidiaRivaAudioTranscription:
|
|||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
model_response: TranscriptionResponse,
|
||||
timeout: float,
|
||||
timeout: float | str | httpx.Timeout | None,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
|
|
@ -183,7 +234,7 @@ class NvidiaRivaAudioTranscription:
|
|||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
model_response: TranscriptionResponse,
|
||||
timeout: float,
|
||||
timeout: float | str | httpx.Timeout | None,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
|
|
@ -213,7 +264,7 @@ class NvidiaRivaAudioTranscription:
|
|||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
model_response: TranscriptionResponse,
|
||||
timeout: float,
|
||||
timeout: float | str | httpx.Timeout | None,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
|
|
@ -259,6 +310,10 @@ class NvidiaRivaAudioTranscription:
|
|||
|
||||
response_format: Final = request_payload.get("response_format") or "json"
|
||||
timestamp_granularities: Final = request_payload.get("timestamp_granularities")
|
||||
drop_params: Final = (
|
||||
"drop_params" in litellm_params and litellm_params["drop_params"] is True
|
||||
) or litellm.drop_params is True
|
||||
riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params)
|
||||
|
||||
riva_module, riva_asr_module = _import_riva()
|
||||
auth_obj: Final = self._construct_auth(
|
||||
|
|
@ -272,10 +327,6 @@ class NvidiaRivaAudioTranscription:
|
|||
riva_asr_module=riva_asr_module,
|
||||
recognition_config_dict=recognition_config_dict,
|
||||
)
|
||||
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
|
||||
config=recognition_config, interim_results=False
|
||||
)
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=None,
|
||||
api_key=api_key,
|
||||
|
|
@ -286,24 +337,31 @@ class NvidiaRivaAudioTranscription:
|
|||
"recognition_config": recognition_config_dict,
|
||||
"nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")),
|
||||
"use_ssl": optional_params.get("use_ssl"),
|
||||
"riva_offline": riva_offline,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
|
||||
audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes)
|
||||
stream_kwargs: Final[dict[str, object]] = {
|
||||
"audio_chunks": audio_chunks,
|
||||
"streaming_config": streaming_config,
|
||||
}
|
||||
# Forward the deadline so the stream cannot block forever if the
|
||||
# server stalls. Older riva-client versions do not accept a
|
||||
# ``timeout`` kwarg, so pass it only when supported.
|
||||
if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator):
|
||||
stream_kwargs["timeout"] = float(timeout)
|
||||
stream: Final = asr_service.streaming_response_generator(**stream_kwargs)
|
||||
final_results: Final = self._collect_final_results(stream)
|
||||
responses: Final = (
|
||||
self._offline_responses(
|
||||
asr_service=asr_service,
|
||||
pcm_bytes=resampled.pcm_bytes,
|
||||
recognition_config=recognition_config,
|
||||
timeout=None if timeout is None else _deadline_seconds(timeout, drop_params),
|
||||
)
|
||||
if riva_offline
|
||||
else self._streaming_responses(
|
||||
asr_service=asr_service,
|
||||
riva_asr_module=riva_asr_module,
|
||||
pcm_bytes=resampled.pcm_bytes,
|
||||
recognition_config=recognition_config,
|
||||
timeout=timeout,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
)
|
||||
final_results: Final = self._collect_final_results(responses, require_final=not riva_offline)
|
||||
except NvidiaRivaException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
@ -409,6 +467,54 @@ class NvidiaRivaAudioTranscription:
|
|||
|
||||
return config
|
||||
|
||||
def _streaming_responses(
|
||||
self,
|
||||
asr_service: _AsrService,
|
||||
riva_asr_module: _RivaAsrModule,
|
||||
pcm_bytes: bytes,
|
||||
recognition_config: _RecognitionConfig,
|
||||
timeout: float | str | httpx.Timeout | None,
|
||||
drop_params: bool,
|
||||
) -> Iterable[object]:
|
||||
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
|
||||
config=recognition_config, interim_results=False
|
||||
)
|
||||
return asr_service.streaming_response_generator(
|
||||
audio_chunks=self._iter_audio_chunks(pcm_bytes),
|
||||
streaming_config=streaming_config,
|
||||
**self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params),
|
||||
)
|
||||
|
||||
def _offline_responses(
|
||||
self,
|
||||
asr_service: _AsrService,
|
||||
pcm_bytes: bytes,
|
||||
recognition_config: _RecognitionConfig,
|
||||
timeout: float | None,
|
||||
) -> Iterable[object]:
|
||||
pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True)
|
||||
try:
|
||||
return (pending.result(timeout=timeout),)
|
||||
except Exception as e:
|
||||
if pending.done():
|
||||
raise
|
||||
pending.cancel()
|
||||
raise NvidiaRivaException(
|
||||
status_code=408,
|
||||
message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds",
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def _timeout_kwargs(
|
||||
cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool
|
||||
) -> Mapping[str, float]:
|
||||
# Forward the deadline so the call cannot block forever if the
|
||||
# server stalls. Older riva-client versions do not accept a
|
||||
# ``timeout`` kwarg, so pass it only when supported.
|
||||
if timeout is None or not cls._supports_timeout_kwarg(rpc):
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)})
|
||||
|
||||
@staticmethod
|
||||
def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
|
||||
try:
|
||||
|
|
@ -429,9 +535,9 @@ class NvidiaRivaAudioTranscription:
|
|||
yield chunk
|
||||
|
||||
@staticmethod
|
||||
def _collect_final_results(stream) -> list[dict[str, object]]:
|
||||
def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]:
|
||||
"""
|
||||
Walk the gRPC stream, ignore empty / non-final chunks, and return a
|
||||
Walk the gRPC responses, ignore empty / non-final chunks, and return a
|
||||
list of normalized final-result dicts. Matching the user's note: the
|
||||
``id`` blocks with no ``results`` are streaming heartbeats and must
|
||||
be skipped.
|
||||
|
|
@ -440,7 +546,7 @@ class NvidiaRivaAudioTranscription:
|
|||
for response in stream:
|
||||
results = getattr(response, "results", None) or []
|
||||
for result in results:
|
||||
if not getattr(result, "is_final", False):
|
||||
if require_final and not getattr(result, "is_final", False):
|
||||
continue
|
||||
alternatives = getattr(result, "alternatives", None) or []
|
||||
if not alternatives:
|
||||
|
|
|
|||
|
|
@ -41,6 +41,9 @@ _GRPC_STATUS_CODE_TO_HTTP: Final[dict] = {
|
|||
}
|
||||
|
||||
|
||||
_GRPC_MESSAGE_TOO_LARGE: Final = "message larger than max"
|
||||
|
||||
|
||||
def _extract_grpc_status_name(error: object) -> str | None:
|
||||
"""
|
||||
Best-effort extraction of a gRPC StatusCode name from an arbitrary error.
|
||||
|
|
@ -79,8 +82,16 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException:
|
|||
a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code.
|
||||
"""
|
||||
status_name: Final = _extract_grpc_status_name(error)
|
||||
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
|
||||
|
||||
detail: Final = _extract_grpc_details(error) or str(error)
|
||||
if status_name == "RESOURCE_EXHAUSTED" and _GRPC_MESSAGE_TOO_LARGE in detail:
|
||||
return NvidiaRivaException(
|
||||
status_code=413,
|
||||
message=(
|
||||
f"NVIDIA Riva gRPC error ({status_name}): {detail}. The audio is too large for one offline "
|
||||
"request to this Riva server: send shorter audio, or drop riva_offline on models that support streaming."
|
||||
),
|
||||
)
|
||||
|
||||
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
|
||||
message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}"
|
||||
return NvidiaRivaException(status_code=http_status, message=message)
|
||||
|
|
|
|||
|
|
@ -8,17 +8,31 @@ repo's own opencode.json cannot redirect model calls. The token is never in argv
|
|||
the config references it with `{file:<private_dir>/token}`. XDG dirs point at a persisted
|
||||
LiteLLM-owned root so the user's opencode config and auth are never read, and the session
|
||||
DB outlives a session for resume. Verified against opencode 1.14.41.
|
||||
|
||||
opencode files each session under the folder it ran in and only streams a resumed session's
|
||||
events from that folder, so resuming one filed elsewhere (by an older LiteLLM that left PWD
|
||||
unset, or from another sandbox) hangs. The first turn after a resume runs `opencode export` in
|
||||
the sandbox, and when the session is filed elsewhere it imports a copy under fresh ids, filed
|
||||
under the sandbox folder, and continues the copy. Import overwrites the folder of a session id it
|
||||
already has, so reusing the ids would break the original wherever it was saved. Before 1.15.6,
|
||||
import files a session under the folder named in its JSON, so the copy names the sandbox folder.
|
||||
A failed or timed-out export or import fails that turn instead of hanging, and the next turn
|
||||
tries again. Verified against opencode 1.14.41, 1.15.5, 1.15.6 and 1.19.3.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
import secrets
|
||||
import string
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.harness.errors import CapabilityUnsupported, HarnessError, OptionsMismatch
|
||||
from litellm.harness.options import OpenCodeOptions
|
||||
from litellm.harness.types import (
|
||||
|
|
@ -33,6 +47,7 @@ from litellm.harness.types import (
|
|||
)
|
||||
from litellm.llms.base_llm.harness.transformation import (
|
||||
BaseCLIHarnessConfig,
|
||||
HarnessSessionCopy,
|
||||
HarnessSessionSetup,
|
||||
HarnessTurnError,
|
||||
HarnessTurnRequest,
|
||||
|
|
@ -57,6 +72,10 @@ OPENCODE_PROVIDER_NPM: Final = "@ai-sdk/openai-compatible"
|
|||
# A fixed title skips opencode's extra title-generation model call on the first turn.
|
||||
OPENCODE_SESSION_TITLE: Final = "litellm-harness"
|
||||
TOKEN_FILENAME: Final = "token"
|
||||
RESUME_COPY_FILENAME: Final = "resume-session.json"
|
||||
ID_ORDERED_PREFIX_LENGTH: Final = 16
|
||||
ID_RANDOM_LENGTH: Final = 14
|
||||
ID_RANDOM_ALPHABET: Final = string.ascii_letters + string.digits
|
||||
INSTRUCTIONS_FILENAME: Final = "instructions.md"
|
||||
XDG_DIRNAME: Final = "xdg"
|
||||
XDG_SUBDIRS: Final = ("config", "data", "state", "cache")
|
||||
|
|
@ -162,6 +181,94 @@ class OpenCodeStreamState:
|
|||
step_texts: Sequence[str] = ()
|
||||
|
||||
|
||||
class ExportedPart(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
|
||||
|
||||
class ExportedMessageInfo(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
|
||||
|
||||
class ExportedMessage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
info: ExportedMessageInfo
|
||||
parts: tuple[ExportedPart, ...]
|
||||
|
||||
|
||||
class ExportedSessionInfo(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
directory: str
|
||||
|
||||
|
||||
class ExportedSession(BaseModel):
|
||||
"""The ids and folder of an `opencode export`; every other field is copied as is."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
info: ExportedSessionInfo
|
||||
messages: tuple[ExportedMessage, ...]
|
||||
|
||||
def ids(self) -> Iterator[str]:
|
||||
yield self.info.id
|
||||
for message in self.messages:
|
||||
yield message.info.id
|
||||
yield from (part.id for part in message.parts)
|
||||
|
||||
|
||||
EXPORT_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
def fresh_id(old: str) -> str:
|
||||
"""Keeps the type prefix and time-ordered part of an opencode id, so the copy sorts like the original."""
|
||||
suffix: Final = "".join(secrets.choice(ID_RANDOM_ALPHABET) for _ in range(ID_RANDOM_LENGTH))
|
||||
return f"{old[:ID_ORDERED_PREFIX_LENGTH]}{suffix}"
|
||||
|
||||
|
||||
def with_id_fields_replaced(value: JsonValue, fresh: Mapping[str, str], key: str = "") -> JsonValue:
|
||||
match value:
|
||||
case str() if key == "id" or key.endswith("ID"):
|
||||
return fresh.get(value, value)
|
||||
case list():
|
||||
return [with_id_fields_replaced(item, fresh, key) for item in value]
|
||||
case dict():
|
||||
return {name: with_id_fields_replaced(item, fresh, name) for name, item in value.items()}
|
||||
case _:
|
||||
return value
|
||||
|
||||
|
||||
def filed_under(exported: JsonValue, workdir: str) -> JsonValue:
|
||||
match exported:
|
||||
case {"info": dict() as info}:
|
||||
return {**exported, "info": {**info, "directory": workdir}}
|
||||
case _:
|
||||
return exported
|
||||
|
||||
|
||||
def session_copy(exported: str, path: str, workdir: str) -> HarnessSessionCopy | None:
|
||||
"""A copy of an exported session filed under workdir, or None when the session is already filed there."""
|
||||
try:
|
||||
raw: Final = EXPORT_JSON.validate_json(exported)
|
||||
session: Final = ExportedSession.model_validate(raw)
|
||||
except ValidationError as exc:
|
||||
raise HarnessError(f"opencode export did not print a session: {exc.errors()[0]['msg']}") from exc
|
||||
if session.info.directory == workdir:
|
||||
return None
|
||||
fresh: Final = MappingProxyType({old: fresh_id(old) for old in session.ids()})
|
||||
return HarnessSessionCopy(
|
||||
session_id=fresh[session.info.id],
|
||||
path=path,
|
||||
data=json.dumps(filed_under(with_id_fields_replaced(raw, fresh), workdir)).encode("utf-8"),
|
||||
import_command=(OPENCODE_BINARY, "--pure", "import", path),
|
||||
)
|
||||
|
||||
|
||||
def _as_dict(value: object) -> Mapping[str, Any]:
|
||||
return value if isinstance(value, dict) else MappingProxyType({})
|
||||
|
||||
|
|
@ -362,6 +469,12 @@ class OpenCodeHarnessConfig(BaseCLIHarnessConfig):
|
|||
# The prompt goes on stdin; opencode appends non-TTY stdin to the message.
|
||||
return HarnessTurnRequest(argv=argv, env=setup.env, stdin=turn_prompt(ctx, prompt), cwd=ctx.sandbox.workdir)
|
||||
|
||||
def transform_resume_export_command(self, native_session_id: str) -> Sequence[str]:
|
||||
return (OPENCODE_BINARY, "--pure", "export", native_session_id)
|
||||
|
||||
def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None:
|
||||
return session_copy(exported, f"{private_dir}/{RESUME_COPY_FILENAME}", workdir)
|
||||
|
||||
def create_stream_state(self) -> OpenCodeStreamState:
|
||||
return OpenCodeStreamState()
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import time
|
|||
from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping
|
||||
from contextlib import aclosing
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import IO, Any, Final, TypedDict
|
||||
from urllib.parse import quote, unquote
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.litellm_core_utils.cloud_storage_security import (
|
|||
split_configured_cloud_bucket_name,
|
||||
validate_managed_cloud_file_id,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_data,
|
||||
|
|
@ -1103,21 +1105,26 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
def _map_openai_to_vertex_params(
|
||||
self,
|
||||
openai_request_body: dict[str, Any],
|
||||
deployment_drop_params: object = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
wrapper to call VertexGeminiConfig.map_openai_params
|
||||
"""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
config: Final = VertexGeminiConfig()
|
||||
_model: Final = openai_request_body.get("model", "")
|
||||
raw_model: Final = openai_request_body.get("model", "")
|
||||
model: Final = raw_model if isinstance(raw_model, str) else ""
|
||||
sampling_params: Final = tuple(
|
||||
p for p in ("temperature", "top_p", "top_k") if openai_request_body.get(p) is not None
|
||||
)
|
||||
if sampling_params and not VertexGeminiConfig.supports_sampling_params(model):
|
||||
VertexGeminiConfig.reject_unsupported_sampling_param(
|
||||
model=model, param=", ".join(sampling_params), drop_params=False
|
||||
)
|
||||
vertex_params: Final = config.map_openai_params(
|
||||
model=_model,
|
||||
model=model,
|
||||
non_default_params=openai_request_body,
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
drop_params=drops_unsupported_params(openai_request_body.get("drop_params", deployment_drop_params)),
|
||||
allow_integer_enums=False,
|
||||
)
|
||||
return vertex_params
|
||||
|
|
@ -1153,11 +1160,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
create_file_data=create_file_data,
|
||||
content_type=content_type,
|
||||
):
|
||||
deployment_drop_params: Final[object] = litellm_params.get("drop_params")
|
||||
return {
|
||||
"streaming_media_upload": StreamingMediaUploadConfig(
|
||||
body_stream=_OpenAIToVertexBatchUploadStream(
|
||||
file_data,
|
||||
self._map_openai_to_vertex_params,
|
||||
partial(self._map_openai_to_vertex_params, deployment_drop_params=deployment_drop_params),
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,17 +2,20 @@ import json
|
|||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, chain
|
||||
from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys
|
||||
from typing import Final, Literal, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import Timeout
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int_or_none
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
Candidates,
|
||||
Citation,
|
||||
CitationMetadata,
|
||||
GenerateContentResponseBody,
|
||||
GenerationConfig,
|
||||
GroundingMetadata,
|
||||
GroundingSegment,
|
||||
GroundingSupport,
|
||||
|
|
@ -48,8 +51,14 @@ HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout)
|
|||
HOP_END_FINISH_REASON: Final = "CONTINUATION"
|
||||
CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON})
|
||||
|
||||
DEFAULT_OUTPUT_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS"
|
||||
_OUTPUT_LIMIT_KEYS: Final = frozenset({"max_output_tokens", "maxOutputTokens"})
|
||||
_REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET: Final = frozenset({"generation_config", "continuationToken"})
|
||||
_TEXT_ONLY_MODALITIES: Final = TypeAdapter(list[Literal["TEXT"]])
|
||||
|
||||
_VERTEX_V1_PATH: Final = "/v1/projects/"
|
||||
_VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/"
|
||||
_CANDIDATE_COUNT_KEYS: Final = ("candidate_count", "candidateCount")
|
||||
|
||||
LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...]
|
||||
AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]]
|
||||
|
|
@ -76,15 +85,53 @@ def _int_or_none(value: object) -> int | None:
|
|||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _asks_for_one_candidate(generation_config: Mapping[str, object]) -> bool:
|
||||
return all(
|
||||
_int_or_none(generation_config[key] or 1) == 1 for key in _CANDIDATE_COUNT_KEYS if key in generation_config
|
||||
)
|
||||
|
||||
|
||||
def continuation_cap_of(request_body: RequestBody) -> int | None:
|
||||
generation_config: Final = request_body.get("generationConfig")
|
||||
if generation_config is None or "max_output_tokens" not in generation_config:
|
||||
return None
|
||||
if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1:
|
||||
if not _asks_for_one_candidate(generation_config):
|
||||
return None
|
||||
return _int_or_none(generation_config["max_output_tokens"])
|
||||
|
||||
|
||||
def _asks_for_text_only(generation_config: Mapping[str, object]) -> bool:
|
||||
if "response_modalities" in generation_config:
|
||||
return False
|
||||
try:
|
||||
_ = _TEXT_ONLY_MODALITIES.validate_python(generation_config.get("responseModalities", []))
|
||||
except ValidationError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def with_default_output_budget(
|
||||
request_body: RequestBody, model_max_output_tokens: Callable[[], int | None]
|
||||
) -> RequestBody:
|
||||
generation_config: Final = request_body.get("generationConfig", GenerationConfig())
|
||||
request_json: Final[Mapping[str, object]] = request_body
|
||||
budget: Final = get_env_int_or_none(DEFAULT_OUTPUT_BUDGET_ENV)
|
||||
if budget is None or budget <= 0:
|
||||
return request_body
|
||||
if not _REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET.isdisjoint(request_body):
|
||||
return request_body
|
||||
if not isinstance(request_json.get("generationConfig", {}), dict):
|
||||
return request_body
|
||||
if not _OUTPUT_LIMIT_KEYS.isdisjoint(generation_config):
|
||||
return request_body
|
||||
if not _asks_for_text_only(generation_config):
|
||||
return request_body
|
||||
model_max: Final = model_max_output_tokens()
|
||||
if model_max is None:
|
||||
return request_body
|
||||
return {**request_body, "generationConfig": {**generation_config, "max_output_tokens": min(budget, model_max)}}
|
||||
|
||||
|
||||
def client_continuation_token(candidate: Candidates) -> str | None:
|
||||
if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -7,15 +7,19 @@ Why separate file? Make it easy to see how transformation works
|
|||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_image_mime_type_from_url,
|
||||
)
|
||||
|
|
@ -64,6 +68,7 @@ from ..common_utils import (
|
|||
get_supports_response_schema,
|
||||
get_supports_system_message,
|
||||
)
|
||||
from .long_decode import with_default_output_budget
|
||||
|
||||
# VertexBase is imported only for type checking to avoid introducing a
|
||||
# module-load-time cyclic import to vertex_llm_base. The instance is lazily
|
||||
|
|
@ -868,7 +873,7 @@ def gemini_convert_messages_with_history(
|
|||
"""
|
||||
has_text_in_content = check_text_in_content(user_content)
|
||||
if has_text_in_content is False:
|
||||
verbose_logger.warning(
|
||||
verbose_logger.debug(
|
||||
"No text in user content. Adding a blank text to user content, to ensure Gemini doesn't fail the request. Relevant Issue - https://github.com/BerriAI/litellm/issues/5515"
|
||||
)
|
||||
user_content.append(
|
||||
|
|
@ -1189,7 +1194,18 @@ def transform_request_body(
|
|||
# labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
|
||||
labels: Final = pop_vertex_request_labels(optional_params, litellm_params)
|
||||
|
||||
filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))}
|
||||
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
drops_top_k: Final = "top_k" in optional_params and not VertexGeminiConfig.supports_sampling_params(model)
|
||||
if drops_top_k:
|
||||
request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature
|
||||
VertexGeminiConfig.reject_unsupported_sampling_param(
|
||||
model=model,
|
||||
param="top_k",
|
||||
drop_params=drops_unsupported_params(request_params.get("drop_params")),
|
||||
)
|
||||
config_keys: Final = set(config_fields) - {"top_k"} if drops_top_k else set(config_fields)
|
||||
filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, config_keys)}
|
||||
|
||||
generation_config: Final[GenerationConfig | None] = GenerationConfig(**filtered_params)
|
||||
|
||||
|
|
@ -1292,13 +1308,75 @@ def sync_transform_request_body(
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
)
|
||||
|
||||
return transform_request_body(
|
||||
messages=messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
cached_content=cached_content,
|
||||
optional_params=optional_params,
|
||||
return _with_default_output_budget(
|
||||
transform_request_body(
|
||||
messages=messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
cached_content=cached_content,
|
||||
optional_params=optional_params,
|
||||
),
|
||||
model,
|
||||
custom_llm_provider,
|
||||
litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict
|
||||
)
|
||||
|
||||
|
||||
class _CostMapOutputModalities(TypedDict):
|
||||
supported_output_modalities: NotRequired[ReadOnly[tuple[str, ...]]]
|
||||
|
||||
|
||||
_COST_MAP_OUTPUT_MODALITIES: Final = TypeAdapter(_CostMapOutputModalities)
|
||||
|
||||
|
||||
def _answers_text(cost_map_key: str) -> bool:
|
||||
try:
|
||||
row: Final = _COST_MAP_OUTPUT_MODALITIES.validate_python(litellm.model_cost.get(cost_map_key)) # pyright: ignore[reportUnknownMemberType] # cost map rows are untyped JSON
|
||||
except ValidationError:
|
||||
return False
|
||||
return "text" in row.get("supported_output_modalities", ("text",))
|
||||
|
||||
|
||||
class _OutputLimit(TypedDict):
|
||||
max_output_tokens: NotRequired[ReadOnly[PositiveInt | None]]
|
||||
|
||||
|
||||
_OUTPUT_LIMIT: Final = TypeAdapter(_OutputLimit)
|
||||
|
||||
|
||||
def _output_limit(fields: object) -> int | None:
|
||||
try:
|
||||
return _OUTPUT_LIMIT.validate_python(fields).get("max_output_tokens")
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _catalog_capped_output_limit(cost_map_key: str, override: object) -> int | None:
|
||||
limits: Final = (
|
||||
_output_limit(GetModelCostMap.loaded_model_cost_map().get(cost_map_key)),
|
||||
_output_limit(override),
|
||||
)
|
||||
return min((limit for limit in limits if limit is not None), default=None)
|
||||
|
||||
|
||||
def _text_chat_model_max_output_tokens(
|
||||
model: str, custom_llm_provider: str, deployment_model_info: object
|
||||
) -> int | None:
|
||||
info: Final = litellm.utils.lookup_model_info_for_registration(model, custom_llm_provider)
|
||||
if isinstance(info, litellm.utils.ModelInfoLookupMiss) or info.get("mode") != "chat":
|
||||
return None
|
||||
if not _answers_text(info["key"]):
|
||||
return None
|
||||
return _catalog_capped_output_limit(info["key"], deployment_model_info or info)
|
||||
|
||||
|
||||
def _with_default_output_budget(
|
||||
request_body: RequestBody, model: str, custom_llm_provider: str, deployment_model_info: object
|
||||
) -> RequestBody:
|
||||
return with_default_output_budget(
|
||||
request_body,
|
||||
partial(_text_chat_model_max_output_tokens, model, custom_llm_provider, deployment_model_info),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1368,21 +1446,8 @@ async def async_transform_request_body(
|
|||
messages, should_inline=should_inline_remote_media(custom_llm_provider)
|
||||
)
|
||||
|
||||
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
|
||||
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
|
||||
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
|
||||
# whole sync transformation on a worker thread so it does not block the
|
||||
# async event loop.
|
||||
return await asyncify(transform_request_body)(
|
||||
messages=inlined_messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
cached_content=cached_content,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
|
||||
return transform_request_body(
|
||||
transform: Final = partial(
|
||||
transform_request_body,
|
||||
messages=inlined_messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
@ -1390,6 +1455,21 @@ async def async_transform_request_body(
|
|||
cached_content=cached_content,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
|
||||
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
|
||||
# whole sync transformation on a worker thread so it does not block the
|
||||
# async event loop.
|
||||
request_body: Final[RequestBody] = (
|
||||
await asyncify(transform)()
|
||||
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages)
|
||||
else transform()
|
||||
)
|
||||
return _with_default_output_budget(
|
||||
request_body,
|
||||
model,
|
||||
custom_llm_provider,
|
||||
litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict
|
||||
)
|
||||
|
||||
|
||||
def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage:
|
||||
|
|
|
|||
|
|
@ -3,15 +3,16 @@
|
|||
## Initial implementation - covers gemini + image gen calls
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from functools import lru_cache, partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -27,7 +28,8 @@ from litellm.constants import (
|
|||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO,
|
||||
)
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
encode_tool_call_id_with_signature,
|
||||
|
|
@ -196,6 +198,29 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
|
|||
)
|
||||
|
||||
|
||||
_SAMPLING_DEPRECATION_UNMAPPED_MODELS: Final = "models outside litellm's model map"
|
||||
_sampling_deprecation_lock: Final = threading.Lock()
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _log_gemini_sampling_params_deprecated(model: str) -> None:
|
||||
verbose_logger.warning(
|
||||
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
|
||||
model,
|
||||
)
|
||||
|
||||
|
||||
def _sampling_deprecation_notice_key(model: str) -> str:
|
||||
catalog: Final = GetModelCostMap.loaded_model_cost_map()
|
||||
is_cataloged: Final = any(name in catalog for name in (model, f"vertex_ai/{model}", f"gemini/{model}"))
|
||||
return model if is_cataloged else _SAMPLING_DEPRECATION_UNMAPPED_MODELS
|
||||
|
||||
|
||||
def _warn_gemini_sampling_params_deprecated(model: str) -> None:
|
||||
with _sampling_deprecation_lock:
|
||||
_log_gemini_sampling_params_deprecated(_sampling_deprecation_notice_key(model))
|
||||
|
||||
|
||||
def _cache_creation_tokens(logging_obj: LoggingClass) -> int:
|
||||
recorded: Final[object] = getattr(logging_obj, "vertex_ai_cache_creation_input_tokens", 0)
|
||||
return recorded if isinstance(recorded, int) and recorded > 0 else 0
|
||||
|
|
@ -207,6 +232,274 @@ def _served_model_name(model_version: object) -> str | None:
|
|||
return model_version.split("@", 1)[0]
|
||||
|
||||
|
||||
_GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"multi_speaker_voice_config": "multiSpeakerVoiceConfig",
|
||||
"speaker_voice_configs": "speakerVoiceConfigs",
|
||||
"voice_config": "voiceConfig",
|
||||
"prebuilt_voice_config": "prebuiltVoiceConfig",
|
||||
"voice_name": "voiceName",
|
||||
"replicated_voice_config": "replicatedVoiceConfig",
|
||||
"mime_type": "mimeType",
|
||||
"voice_sample_audio": "voiceSampleAudio",
|
||||
"language_code": "languageCode",
|
||||
"speaker_alias": "speakerAlias",
|
||||
"speaker_id": "speakerId",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_SPEECH_CONFIG_ITEMS: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
def _normalize_gemini_speech_config_item(value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
normalized: Final[dict[str, object]] = {} # mutable-ok: iterative traversal fills nested output without recursion
|
||||
|
||||
pending: Final[ # mutable-ok: the explicit worklist must grow while avoiding recursive traversal
|
||||
list[tuple[dict[str, object] | list[object], Mapping[str, object] | list[object]]]
|
||||
] = [(normalized, value)]
|
||||
while pending:
|
||||
target, source = pending.pop()
|
||||
items = source.items() if isinstance(source, Mapping) else enumerate(source)
|
||||
for key, item in items:
|
||||
if isinstance(source, Mapping):
|
||||
if not isinstance(key, str):
|
||||
continue
|
||||
normalized_key: str | int = _GEMINI_SPEECH_CONFIG_KEY_MAP.get(key, key)
|
||||
else:
|
||||
normalized_key = len(target)
|
||||
|
||||
child: object
|
||||
if (mapping_item := as_str_mapping(item)) is not None:
|
||||
pending.append((child_mapping := dict[str, object](), mapping_item))
|
||||
child = child_mapping
|
||||
elif isinstance(item, list):
|
||||
pending.append((child_items := list[object](), _SPEECH_CONFIG_ITEMS.validate_python(item)))
|
||||
child = child_items
|
||||
else:
|
||||
child = (
|
||||
OPENAI_TO_GEMINI_TTS_VOICES.get(item.lower(), item)
|
||||
if normalized_key == "voiceName" and isinstance(item, str)
|
||||
else item
|
||||
)
|
||||
|
||||
if isinstance(target, dict):
|
||||
assert isinstance(normalized_key, str)
|
||||
target[normalized_key] = child
|
||||
else:
|
||||
target.append(child)
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def _gemini_speaker_voice_config(item: object) -> object:
|
||||
speaker_config: Final = as_str_mapping(item)
|
||||
if speaker_config is None:
|
||||
return item
|
||||
speaker: Final = speaker_config.get("speakerAlias")
|
||||
voice_name: Final = speaker_config.get("speakerId")
|
||||
if not isinstance(speaker, str) or not isinstance(voice_name, str):
|
||||
return item
|
||||
return {
|
||||
**{key: value for key, value in speaker_config.items() if key not in ("speakerAlias", "speakerId")},
|
||||
"speaker": speaker,
|
||||
"voiceConfig": {
|
||||
"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _with_gemini_speaker_voice_configs(config: Mapping[str, object]) -> Mapping[str, object]:
|
||||
multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig"))
|
||||
if multi_speaker_config is None:
|
||||
return config
|
||||
speaker_voice_configs: Final = multi_speaker_config.get("speakerVoiceConfigs")
|
||||
if not isinstance(speaker_voice_configs, list):
|
||||
return config
|
||||
return {
|
||||
**config,
|
||||
"multiSpeakerVoiceConfig": {
|
||||
**multi_speaker_config,
|
||||
"speakerVoiceConfigs": [
|
||||
_gemini_speaker_voice_config(item)
|
||||
for item in _SPEECH_CONFIG_ITEMS.validate_python(speaker_voice_configs)
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
_GEMINI_SPEECH_CONFIG_FIELDS: Final = frozenset({"voiceConfig", "multiSpeakerVoiceConfig", "languageCode"})
|
||||
|
||||
_GEMINI_VOICE_CONFIG_FIELDS: Final = MappingProxyType(
|
||||
{
|
||||
"prebuiltVoiceConfig": frozenset({"voiceName"}),
|
||||
"replicatedVoiceConfig": frozenset({"mimeType", "voiceSampleAudio"}),
|
||||
"voice": None,
|
||||
}
|
||||
)
|
||||
|
||||
_GEMINI_SPEAKER_VOICE_CONFIG_FIELDS: Final = frozenset({"speaker", "voiceConfig"})
|
||||
|
||||
|
||||
def _known_leaf_fields(value: object, fields: frozenset[str] | None) -> object:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if fields is None or mapping is None:
|
||||
return value
|
||||
return {key: item for key, item in mapping.items() if key in fields}
|
||||
|
||||
|
||||
def _known_voice_config_fields(value: object) -> object:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if mapping is None:
|
||||
return value
|
||||
return {
|
||||
key: _known_leaf_fields(item, _GEMINI_VOICE_CONFIG_FIELDS[key])
|
||||
for key, item in mapping.items()
|
||||
if key in _GEMINI_VOICE_CONFIG_FIELDS
|
||||
}
|
||||
|
||||
|
||||
def _known_speaker_voice_config_fields(value: object) -> object:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if mapping is None:
|
||||
return value
|
||||
return {
|
||||
key: _known_voice_config_fields(item) if key == "voiceConfig" else item
|
||||
for key, item in mapping.items()
|
||||
if key in _GEMINI_SPEAKER_VOICE_CONFIG_FIELDS
|
||||
}
|
||||
|
||||
|
||||
def _known_speaker_voice_configs(value: object) -> object:
|
||||
if not isinstance(value, list):
|
||||
return value
|
||||
return [_known_speaker_voice_config_fields(item) for item in _SPEECH_CONFIG_ITEMS.validate_python(value)]
|
||||
|
||||
|
||||
def _known_multi_speaker_voice_config_fields(value: object) -> object:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if mapping is None:
|
||||
return value
|
||||
return {key: _known_speaker_voice_configs(item) for key, item in mapping.items() if key == "speakerVoiceConfigs"}
|
||||
|
||||
|
||||
def _known_speech_config_field(key: str, value: object) -> object:
|
||||
match key:
|
||||
case "voiceConfig":
|
||||
return _known_voice_config_fields(value)
|
||||
case "multiSpeakerVoiceConfig":
|
||||
return _known_multi_speaker_voice_config_fields(value)
|
||||
case _:
|
||||
return value
|
||||
|
||||
|
||||
def _known_gemini_speech_fields(
|
||||
config: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict
|
||||
return {
|
||||
key: _known_speech_config_field(key, item)
|
||||
for key, item in config.items()
|
||||
if key in _GEMINI_SPEECH_CONFIG_FIELDS
|
||||
}
|
||||
|
||||
|
||||
_GEMINI_TTS_VOICE_KEYS: Final = frozenset(
|
||||
{
|
||||
"format",
|
||||
"name",
|
||||
"voice",
|
||||
"voiceName",
|
||||
"voice_name",
|
||||
"languageCode",
|
||||
"language_code",
|
||||
"modelName",
|
||||
"model_name",
|
||||
"speechConfig",
|
||||
"speech_config",
|
||||
"voiceConfig",
|
||||
"voice_config",
|
||||
"multiSpeakerVoiceConfig",
|
||||
"multi_speaker_voice_config",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_NAMED_UNKNOWN_VOICE_KEYS: Final = 10
|
||||
|
||||
_NAMED_VOICE_KEY_CHARS: Final = 100
|
||||
|
||||
|
||||
def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
camel_config: Final = as_str_mapping(value.get("speechConfig"))
|
||||
return camel_config if camel_config is not None else as_str_mapping(value.get("speech_config"))
|
||||
|
||||
|
||||
def named_voice_keys(keys: Sequence[str]) -> str:
|
||||
named: Final = ", ".join(key[:_NAMED_VOICE_KEY_CHARS] for key in keys[:_NAMED_UNKNOWN_VOICE_KEYS])
|
||||
unnamed_count: Final = len(keys) - _NAMED_UNKNOWN_VOICE_KEYS
|
||||
return f"{named}, and {unnamed_count} more" if unnamed_count > 0 else named
|
||||
|
||||
|
||||
def reject_unknown_gemini_tts_voice_keys(
|
||||
model: str, voice: Mapping[str, object], drop_params: bool, llm_provider: str | None
|
||||
) -> None:
|
||||
nested_config: Final = _gemini_speech_config_wrapper(voice) or {}
|
||||
ignored_wrapper: Final = (
|
||||
("speech_config",)
|
||||
if as_str_mapping(voice.get("speechConfig")) is not None and voice.get("speech_config") is not None
|
||||
else ()
|
||||
)
|
||||
unknown_keys: Final = sorted(
|
||||
{*(key for key in (*voice, *nested_config) if key not in _GEMINI_TTS_VOICE_KEYS), *ignored_wrapper}
|
||||
)
|
||||
if not unknown_keys or drops_unsupported_params(drop_params):
|
||||
return
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not support the voice keys {named_voice_keys(unknown_keys)}. Remove them, or set "
|
||||
"drop_params to ignore them."
|
||||
),
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
|
||||
def normalize_gemini_speech_config(
|
||||
value: Mapping[str, object], drop_unknown_fields: bool = False
|
||||
) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict
|
||||
speech_config: Final = {**_with_gemini_speaker_voice_configs(mapped_gemini_speech_config(value))}
|
||||
return _known_gemini_speech_fields(speech_config) if drop_unknown_fields else speech_config
|
||||
|
||||
|
||||
def _merge_gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
nested_config: Final = _gemini_speech_config_wrapper(value)
|
||||
if nested_config is None:
|
||||
return value
|
||||
return {
|
||||
**{key: item for key, item in value.items() if key not in ("speechConfig", "speech_config")},
|
||||
**nested_config,
|
||||
}
|
||||
|
||||
|
||||
def mapped_gemini_speech_config(
|
||||
value: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict
|
||||
normalized: Final = _normalize_gemini_speech_config_item(_merge_gemini_speech_config_wrapper(value))
|
||||
speech_config: Final = {key: item for key, item in normalized.items() if key in _GEMINI_SPEECH_CONFIG_FIELDS}
|
||||
voice_name: Final = next(
|
||||
(normalized[key] for key in ("name", "voiceName", "voice") if isinstance(normalized.get(key), str)),
|
||||
None,
|
||||
)
|
||||
if "voiceConfig" in speech_config or "multiSpeakerVoiceConfig" in speech_config or not isinstance(voice_name, str):
|
||||
return speech_config
|
||||
return {
|
||||
**speech_config,
|
||||
"voiceConfig": {
|
||||
"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class VertexAIBaseConfig:
|
||||
def get_mapped_special_auth_params(self) -> dict:
|
||||
"""
|
||||
|
|
@ -404,11 +697,43 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def supports_sampling_params(model: str) -> bool:
|
||||
"""Whether Google still honors ``temperature`` / ``top_p`` / ``top_k`` for ``model``.
|
||||
|
||||
Driven by ``supports_sampling_params`` in the model map: Gemini 3.6 Flash, 3.5 Flash-Lite
|
||||
and later accept the fields but ignore them, and a later generation will 400 on them.
|
||||
A model the map says nothing about keeps forwarding them. ``temperature`` and ``top_p``
|
||||
leave ``get_supported_openai_params`` when the flag is off, so the shared unsupported-param
|
||||
check raises or drops them and ``allowed_openai_params`` still forwards them verbatim;
|
||||
``top_k`` is a provider-specific kwarg that bypasses that check and is gated while the
|
||||
request body is built."""
|
||||
rows: Final = cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: JSON rows by model id
|
||||
bare: Final = model.split("/")[-1]
|
||||
for candidate in (model, bare, f"vertex_ai/{bare}", f"gemini/{bare}"):
|
||||
if isinstance(flag := rows.get(candidate, {}).get("supports_sampling_params"), bool):
|
||||
return flag
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def reject_unsupported_sampling_param(model: str, param: str, drop_params: bool) -> None:
|
||||
"""Return under ``drop_params`` so the caller drops ``param``, else raise the client-side 400
|
||||
that names the escape hatch."""
|
||||
if litellm.drop_params or drop_params:
|
||||
return
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not support {param}: the model map flags it `supports_sampling_params: false` "
|
||||
"because Google ignores sampling params on this model. "
|
||||
"To drop unsupported params, set `litellm.drop_params = True`."
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
sampling_params: Final = ("temperature", "top_p", "top_k") if self.supports_sampling_params(model) else ()
|
||||
supported_params: Final = [
|
||||
"temperature",
|
||||
"top_p",
|
||||
"top_k",
|
||||
*sampling_params,
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"stream",
|
||||
|
|
@ -1180,7 +1505,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
status_code=400,
|
||||
)
|
||||
|
||||
def _map_audio_params(self, value: dict) -> dict:
|
||||
def _map_audio_params(
|
||||
self, value: Mapping[str, object], drop_unknown_fields: bool = False
|
||||
) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary
|
||||
"""
|
||||
Expected input:
|
||||
{
|
||||
|
|
@ -1199,12 +1526,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
languageCode: "en-US",
|
||||
}
|
||||
"""
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
PrebuiltVoiceConfig,
|
||||
SpeechConfig,
|
||||
VoiceConfig,
|
||||
)
|
||||
|
||||
# Validate audio format - Gemini TTS only supports pcm16
|
||||
audio_format: Final = value.get("format")
|
||||
if audio_format is not None and audio_format != "pcm16":
|
||||
|
|
@ -1214,22 +1535,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
f"Please set audio format to 'pcm16'."
|
||||
)
|
||||
|
||||
# Map OpenAI audio parameter to Gemini speech config
|
||||
speech_config: Final[SpeechConfig] = {}
|
||||
camel_speech_config: Final = value.get("speechConfig")
|
||||
snake_speech_config: Final = value.get("speech_config")
|
||||
speech_config_value: Final = (
|
||||
camel_speech_config if isinstance(camel_speech_config, Mapping) else snake_speech_config
|
||||
)
|
||||
language_code: Final = value.get("language_code", value.get("languageCode"))
|
||||
if isinstance(speech_config_value, Mapping):
|
||||
return normalize_gemini_speech_config(value, drop_unknown_fields)
|
||||
|
||||
if "voice" in value:
|
||||
voice: Final = value["voice"]
|
||||
voice_name: Final = (
|
||||
OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) if isinstance(voice, str) else voice
|
||||
)
|
||||
prebuilt_voice_config: Final[PrebuiltVoiceConfig] = {"voiceName": voice_name}
|
||||
voice_config: Final[VoiceConfig] = {"prebuiltVoiceConfig": prebuilt_voice_config}
|
||||
speech_config["voiceConfig"] = voice_config
|
||||
voice: Final = value.get("voice")
|
||||
speech_config_without_language: Final[ # mutable-ok: provider request serialization requires a concrete dict
|
||||
dict[str, object]
|
||||
] = (
|
||||
{
|
||||
"voiceConfig": {
|
||||
"prebuiltVoiceConfig": {
|
||||
"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice)
|
||||
if isinstance(voice, str)
|
||||
else voice,
|
||||
}
|
||||
}
|
||||
}
|
||||
if "voice" in value
|
||||
else {}
|
||||
)
|
||||
|
||||
if "language_code" in value:
|
||||
speech_config["languageCode"] = value["language_code"]
|
||||
if language_code is not None:
|
||||
return {
|
||||
**speech_config_without_language,
|
||||
"languageCode": language_code,
|
||||
}
|
||||
|
||||
return cast(dict, speech_config)
|
||||
return speech_config_without_language
|
||||
|
||||
def _map_chat_audio_params(
|
||||
self, model: str, value: Mapping[str, object], drop_params: bool
|
||||
) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary
|
||||
if _gemini_speech_config_wrapper(value) is not None:
|
||||
reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider)
|
||||
return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params))
|
||||
|
||||
def map_audio_params(
|
||||
self,
|
||||
|
|
@ -1266,8 +1611,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
allow_integer_enums: bool = True,
|
||||
) -> dict:
|
||||
self._apply_include_server_side_tool_invocations(non_default_params, optional_params)
|
||||
gemini_sampling_params_warned: bool = False
|
||||
supports_sampling_params: Final = self.supports_sampling_params(model)
|
||||
for param, value in non_default_params.items():
|
||||
if param in ("temperature", "top_p", "top_k") and not supports_sampling_params:
|
||||
continue
|
||||
if param == "temperature":
|
||||
if VertexGeminiConfig._is_gemini_3_or_newer(model):
|
||||
if value is not None and value < 1.0:
|
||||
|
|
@ -1275,28 +1622,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
"Warning: Setting temperature < 1.0 for Gemini 3 models (%s) can cause infinite loops, degraded reasoning performance, and failure on complex tasks. Strongly recommended to use temperature = 1.0 (default).",
|
||||
model,
|
||||
)
|
||||
if not gemini_sampling_params_warned:
|
||||
verbose_logger.warning(
|
||||
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
|
||||
model,
|
||||
)
|
||||
gemini_sampling_params_warned = True
|
||||
_warn_gemini_sampling_params_deprecated(model)
|
||||
optional_params["temperature"] = value
|
||||
elif param == "top_p":
|
||||
if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned:
|
||||
verbose_logger.warning(
|
||||
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
|
||||
model,
|
||||
)
|
||||
gemini_sampling_params_warned = True
|
||||
if VertexGeminiConfig._is_gemini_3_or_newer(model):
|
||||
_warn_gemini_sampling_params_deprecated(model)
|
||||
optional_params["top_p"] = value
|
||||
elif param == "top_k":
|
||||
if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned:
|
||||
verbose_logger.warning(
|
||||
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
|
||||
model,
|
||||
)
|
||||
gemini_sampling_params_warned = True
|
||||
if VertexGeminiConfig._is_gemini_3_or_newer(model):
|
||||
_warn_gemini_sampling_params_deprecated(model)
|
||||
optional_params["top_k"] = value
|
||||
elif (
|
||||
param == "stream" and value is True
|
||||
|
|
@ -1305,7 +1639,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
elif param == "n":
|
||||
optional_params["candidate_count"] = value
|
||||
elif param == "audio" and isinstance(value, dict):
|
||||
optional_params["speechConfig"] = self._map_audio_params(value)
|
||||
optional_params["speechConfig"] = self._map_chat_audio_params(model, value, drop_params)
|
||||
elif param == "stop":
|
||||
if isinstance(value, str):
|
||||
optional_params["stop_sequences"] = [value]
|
||||
|
|
|
|||
|
|
@ -6,20 +6,25 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s
|
|||
"""
|
||||
|
||||
import base64
|
||||
from collections.abc import Coroutine, Mapping
|
||||
import math
|
||||
from collections.abc import Coroutine, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, TypeAlias, Union
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
DEFAULT_SPEECH_MEDIA_TYPE,
|
||||
calculate_request_duration,
|
||||
speech_media_type_from_audio_bytes,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.text_to_speech.transformation import (
|
||||
BaseTextToSpeechConfig,
|
||||
TextToSpeechRequestData,
|
||||
|
|
@ -33,6 +38,7 @@ from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
|
|||
from litellm.types.llms.vertex_ai_text_to_speech import (
|
||||
VertexTextToSpeechAudioConfig,
|
||||
VertexTextToSpeechInput,
|
||||
VertexTextToSpeechSpeakerVoiceConfig,
|
||||
VertexTextToSpeechVoice,
|
||||
)
|
||||
|
||||
|
|
@ -49,6 +55,166 @@ _LyriaVoice: TypeAlias = str | dict | None
|
|||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...], config=ConfigDict(hide_input_in_errors=True))
|
||||
_STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True))
|
||||
_GEMINI_CLOUD_TTS_LOCATIONS: Final = MappingProxyType(
|
||||
{
|
||||
**dict.fromkeys(
|
||||
("us", "us-central1", "us-east1", "us-east4", "us-east5", "us-south1", "us-west1", "us-west4"), "us"
|
||||
),
|
||||
**dict.fromkeys(
|
||||
("eu", "europe-central2", "europe-north1", "europe-southwest1", "europe-west1", "europe-west4"), "eu"
|
||||
),
|
||||
"northamerica-northeast1": "northamerica-northeast1",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_DEFAULT_GEMINI_TTS_VOICE: Final = "Kore"
|
||||
|
||||
|
||||
class _GeminiCloudTtsInput(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
text: str | None = None
|
||||
ssml: str | None = None
|
||||
prompt: str | None = None
|
||||
|
||||
|
||||
class _GeminiCloudTtsAudioConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
audioEncoding: str = ""
|
||||
sampleRateHertz: int | None = None
|
||||
|
||||
|
||||
class _GeminiCloudTtsRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
input: _GeminiCloudTtsInput
|
||||
audioConfig: _GeminiCloudTtsAudioConfig
|
||||
|
||||
|
||||
def _fallback_gemini_tts_audio_duration(audio: bytes, encoding: str, sample_rate: int) -> float | None:
|
||||
if encoding == "PCM":
|
||||
return len(audio) / (2 * sample_rate) if sample_rate > 0 else None
|
||||
if encoding in ("ALAW", "MULAW") and not (audio[:4] == b"RIFF" and audio[8:12] == b"WAVE"):
|
||||
return len(audio) / sample_rate if sample_rate > 0 else None
|
||||
|
||||
if encoding in ("LINEAR16", "ALAW", "MULAW") and audio[:4] == b"RIFF" and audio[8:12] == b"WAVE":
|
||||
format_offset: Final = audio.find(b"fmt ", 12)
|
||||
data_offset: Final = audio.find(b"data", 12)
|
||||
if format_offset >= 0 and data_offset >= 0 and format_offset + 20 <= len(audio):
|
||||
byte_rate: Final = int.from_bytes(audio[format_offset + 16 : format_offset + 20], "little")
|
||||
data_size: Final = int.from_bytes(audio[data_offset + 4 : data_offset + 8], "little")
|
||||
if byte_rate > 0 and data_size <= len(audio) - data_offset - 8:
|
||||
return data_size / byte_rate
|
||||
|
||||
if encoding == "MP3" and speech_media_type_from_audio_bytes(audio) == "audio/mpeg":
|
||||
return len(audio) / 4000
|
||||
|
||||
if encoding == "OGG_OPUS" and audio[:4] == b"OggS":
|
||||
opus_header: Final = audio.find(b"OpusHead")
|
||||
last_page: Final = audio.rfind(b"OggS")
|
||||
if opus_header >= 0 and opus_header + 12 <= len(audio) and last_page + 14 <= len(audio):
|
||||
pre_skip: Final = int.from_bytes(audio[opus_header + 10 : opus_header + 12], "little")
|
||||
granule: Final = int.from_bytes(audio[last_page + 6 : last_page + 14], "little")
|
||||
if pre_skip <= granule < 2**64 - 1:
|
||||
return (granule - pre_skip) / 48000
|
||||
|
||||
return None
|
||||
|
||||
|
||||
_SPEAKER_CONFIG_ITEMS: Final = TypeAdapter(list[object])
|
||||
|
||||
_CloudTtsVoiceFields: TypeAlias = "Mapping[str, _CloudTtsVoiceFields]"
|
||||
|
||||
_CLOUD_TTS_VOICE_LEAF: Final[_CloudTtsVoiceFields] = MappingProxyType({})
|
||||
|
||||
_CLOUD_TTS_PREBUILT_VOICE_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType(
|
||||
{"prebuiltVoiceConfig": MappingProxyType({"voiceName": _CLOUD_TTS_VOICE_LEAF})}
|
||||
)
|
||||
|
||||
_CLOUD_TTS_SPEECH_CONFIG_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType(
|
||||
{
|
||||
"languageCode": _CLOUD_TTS_VOICE_LEAF,
|
||||
"voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS,
|
||||
"multiSpeakerVoiceConfig": MappingProxyType(
|
||||
{
|
||||
"speakerVoiceConfigs": MappingProxyType(
|
||||
{
|
||||
"speaker": _CLOUD_TTS_VOICE_LEAF,
|
||||
"speakerAlias": _CLOUD_TTS_VOICE_LEAF,
|
||||
"speakerId": _CLOUD_TTS_VOICE_LEAF,
|
||||
"voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS,
|
||||
}
|
||||
)
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_CLOUD_TTS_LIST_FIELDS: Final = frozenset({"speakerVoiceConfigs"})
|
||||
|
||||
_GEMINI_TTS_VOICE_NAME_KEYS: Final = ("name", "voiceName", "voice_name", "voice")
|
||||
|
||||
|
||||
def _fields_cloud_tts_cannot_send(
|
||||
config: Mapping[str, object], fields: _CloudTtsVoiceFields, path: str
|
||||
) -> Iterator[str]:
|
||||
for key, value in config.items():
|
||||
yield from _field_cloud_tts_cannot_send(f"{path}{key}", value, fields.get(key), key in _CLOUD_TTS_LIST_FIELDS)
|
||||
|
||||
|
||||
def _field_cloud_tts_cannot_send(
|
||||
path: str, value: object, fields: "_CloudTtsVoiceFields | None", is_list_field: bool
|
||||
) -> Iterator[str]:
|
||||
if fields is None:
|
||||
yield path
|
||||
return
|
||||
if value is None:
|
||||
return
|
||||
if not fields:
|
||||
yield from () if isinstance(value, str) else (path,)
|
||||
return
|
||||
items: Final = (
|
||||
_SPEAKER_CONFIG_ITEMS.validate_python(value) if is_list_field and isinstance(value, list) else (value,)
|
||||
)
|
||||
for item in items:
|
||||
mapping = as_str_mapping(item)
|
||||
yield from (path,) if mapping is None else _fields_cloud_tts_cannot_send(mapping, fields, f"{path}.")
|
||||
|
||||
|
||||
def _non_string_voice_names(config: Mapping[str, object], path: str) -> Iterator[str]:
|
||||
return (
|
||||
f"{path}{key}"
|
||||
for key in _GEMINI_TTS_VOICE_NAME_KEYS
|
||||
if config.get(key) is not None and not isinstance(config.get(key), str)
|
||||
)
|
||||
|
||||
|
||||
def _non_mapping_speech_config_wrappers(voice: Mapping[str, object]) -> Iterator[str]:
|
||||
return (
|
||||
key
|
||||
for key in ("speechConfig", "speech_config")
|
||||
if voice.get(key) is not None and as_str_mapping(voice.get(key)) is None
|
||||
)
|
||||
|
||||
|
||||
def _speaker_entries(config: Mapping[str, object]) -> Iterator[Mapping[str, object]]:
|
||||
multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig"))
|
||||
speakers: Final = multi_speaker_config.get("speakerVoiceConfigs") if multi_speaker_config is not None else None
|
||||
items: Final = _SPEAKER_CONFIG_ITEMS.validate_python(speakers) if isinstance(speakers, list) else ()
|
||||
return (entry for item in items if (entry := as_str_mapping(item)) is not None)
|
||||
|
||||
|
||||
def _prebuilt_voices_replaced_by_speaker_id(config: Mapping[str, object]) -> Iterator[str]:
|
||||
return (
|
||||
"multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.prebuiltVoiceConfig"
|
||||
for entry in _speaker_entries(config)
|
||||
if entry.get("speakerId") is not None
|
||||
and (voice_config := as_str_mapping(entry.get("voiceConfig"))) is not None
|
||||
and voice_config.get("prebuiltVoiceConfig") is not None
|
||||
)
|
||||
|
||||
|
||||
class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
||||
|
|
@ -81,14 +247,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
}
|
||||
|
||||
# Response format mappings from OpenAI to Google Cloud audio encoding
|
||||
FORMAT_MAPPINGS = {
|
||||
"mp3": "MP3",
|
||||
"opus": "OGG_OPUS",
|
||||
"aac": "MP3", # Google doesn't have AAC, use MP3
|
||||
"flac": "FLAC",
|
||||
"wav": "LINEAR16",
|
||||
"pcm": "LINEAR16",
|
||||
}
|
||||
FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"mp3": "MP3",
|
||||
"opus": "OGG_OPUS",
|
||||
"aac": "MP3", # Google doesn't have AAC, use MP3
|
||||
"flac": "FLAC",
|
||||
"wav": "LINEAR16",
|
||||
"pcm": "LINEAR16",
|
||||
}
|
||||
)
|
||||
GEMINI_FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType(
|
||||
{"mp3": "MP3", "opus": "OGG_OPUS", "alaw": "ALAW", "mulaw": "MULAW"}
|
||||
)
|
||||
|
||||
def __init__(self) -> None:
|
||||
BaseTextToSpeechConfig.__init__(self)
|
||||
|
|
@ -96,7 +267,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
|
||||
def _map_voice_to_vertex_format(
|
||||
self,
|
||||
voice: str | dict | None,
|
||||
voice: str | dict[str, object] | None,
|
||||
model: str | None = None,
|
||||
drop_params: bool = False,
|
||||
) -> tuple[str | None, dict | None]:
|
||||
"""
|
||||
Map voice to Vertex AI format.
|
||||
|
|
@ -111,6 +284,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
- voice_str: Original string voice (for interface compatibility)
|
||||
- voice_dict: Vertex AI format dict with languageCode and name
|
||||
"""
|
||||
if model is not None and self._is_gemini_tts_model(model):
|
||||
return self._map_gemini_tts_voice_to_vertex_format(model=model, voice=voice, drop_params=drop_params)
|
||||
|
||||
if voice is None:
|
||||
return None, None
|
||||
|
||||
|
|
@ -142,11 +318,225 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
|
||||
return voice_str, voice_dict
|
||||
|
||||
@staticmethod
|
||||
def _is_gemini_tts_model(model: str) -> bool:
|
||||
from litellm.utils import is_gemini_tts_model
|
||||
|
||||
return is_gemini_tts_model(model)
|
||||
|
||||
@staticmethod
|
||||
def _get_str_value(
|
||||
source: Mapping[str, object],
|
||||
*keys: str,
|
||||
) -> str | None:
|
||||
for key in keys:
|
||||
value = source.get(key)
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_dict_value(
|
||||
source: Mapping[str, object],
|
||||
*keys: str,
|
||||
) -> Mapping[str, object] | None:
|
||||
for key in keys:
|
||||
value = source.get(key)
|
||||
if (mapping := as_str_mapping(value)) is not None:
|
||||
return mapping
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _resolve_gemini_tts_voice_name(model: str, voice_name: str, drop_params: bool) -> str:
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES
|
||||
|
||||
gemini_voice_name: Final = OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)
|
||||
if "-" not in gemini_voice_name:
|
||||
return gemini_voice_name
|
||||
if drop_params:
|
||||
return _DEFAULT_GEMINI_TTS_VOICE
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} speaks Gemini voices such as {_DEFAULT_GEMINI_TTS_VOICE}, and {voice_name} is a Cloud "
|
||||
"Text-to-Speech voice. Use a Gemini voice, or set drop_params to use the default Gemini voice."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reject_malformed_gemini_tts_voice(model: str, drop_params: bool) -> None:
|
||||
if drop_params:
|
||||
return
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} needs voice to be a Gemini voice name such as {_DEFAULT_GEMINI_TTS_VOICE} or a speech "
|
||||
"config object. Send one, or set drop_params to use the default Gemini voice."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reject_voice_fields_cloud_tts_cannot_send(model: str, voice: Mapping[str, object], drop_params: bool) -> None:
|
||||
if drop_params:
|
||||
return
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
mapped_gemini_speech_config,
|
||||
named_voice_keys,
|
||||
)
|
||||
|
||||
speech_config: Final = VertexAITextToSpeechConfig._get_dict_value(voice, "speechConfig", "speech_config")
|
||||
mapped_config: Final = mapped_gemini_speech_config(voice)
|
||||
unsendable: Final = sorted(
|
||||
{
|
||||
*_fields_cloud_tts_cannot_send(mapped_config, _CLOUD_TTS_SPEECH_CONFIG_FIELDS, ""),
|
||||
*_prebuilt_voices_replaced_by_speaker_id(mapped_config),
|
||||
*_non_mapping_speech_config_wrappers(voice),
|
||||
*_non_string_voice_names(voice, ""),
|
||||
*_non_string_voice_names(speech_config or {}, "speechConfig."),
|
||||
}
|
||||
)
|
||||
if not unsendable:
|
||||
return
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} cannot send the voice fields {named_voice_keys(unsendable)} as mp3, opus, alaw, or mulaw. "
|
||||
"Request wav or pcm, remove them, or set drop_params to ignore them."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reject_malformed_gemini_tts_speaker(model: str, drop_params: bool) -> None:
|
||||
if drop_params:
|
||||
return
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} needs each entry of speaker_voice_configs to name a speaker and a voice, as speaker with "
|
||||
"voice_config.prebuilt_voice_config.voice_name, or speakerAlias with speakerId."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
def _extract_gemini_tts_speaker_configs(
|
||||
self,
|
||||
model: str,
|
||||
speech_config: Mapping[str, object],
|
||||
drop_params: bool,
|
||||
) -> Sequence[VertexTextToSpeechSpeakerVoiceConfig]:
|
||||
multi_speaker_config: Final = self._get_dict_value(speech_config, "multiSpeakerVoiceConfig")
|
||||
if multi_speaker_config is None:
|
||||
return []
|
||||
raw_speaker_configs: Final = multi_speaker_config.get("speakerVoiceConfigs", [])
|
||||
try:
|
||||
validated_speaker_configs: Final = _SPEAKER_CONFIG_ITEMS.validate_python(raw_speaker_configs, strict=True)
|
||||
except ValidationError:
|
||||
self._reject_malformed_gemini_tts_speaker(model, drop_params)
|
||||
return []
|
||||
speaker_configs: Final[ # mutable-ok: validated payloads are accumulated for serialization
|
||||
list[VertexTextToSpeechSpeakerVoiceConfig]
|
||||
] = []
|
||||
for raw_item in validated_speaker_configs:
|
||||
if (raw_config := as_str_mapping(raw_item)) is None:
|
||||
self._reject_malformed_gemini_tts_speaker(model, drop_params)
|
||||
continue
|
||||
speaker_alias = self._get_str_value(raw_config, "speakerAlias", "speaker")
|
||||
speaker_id = self._get_str_value(raw_config, "speakerId") or self._prebuilt_voice_name(raw_config)
|
||||
if speaker_alias is None or speaker_id is None:
|
||||
self._reject_malformed_gemini_tts_speaker(model, drop_params)
|
||||
continue
|
||||
speaker_configs.append(
|
||||
{
|
||||
"speakerAlias": speaker_alias,
|
||||
"speakerId": self._resolve_gemini_tts_voice_name(model, speaker_id, drop_params),
|
||||
}
|
||||
)
|
||||
return speaker_configs
|
||||
|
||||
def _prebuilt_voice_name(self, config: Mapping[str, object]) -> str | None:
|
||||
voice_config: Final = self._get_dict_value(config, "voiceConfig")
|
||||
prebuilt_voice_config: Final = (
|
||||
self._get_dict_value(voice_config, "prebuiltVoiceConfig") if voice_config is not None else None
|
||||
)
|
||||
if prebuilt_voice_config is None:
|
||||
return None
|
||||
return self._get_str_value(prebuilt_voice_config, "voiceName")
|
||||
|
||||
def _map_gemini_tts_voice_to_vertex_format(
|
||||
self,
|
||||
model: str,
|
||||
voice: object,
|
||||
drop_params: bool = False,
|
||||
) -> tuple[str | None, dict[str, object]]: # mutable-ok: provider request serialization requires a concrete dict
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
mapped_gemini_speech_config,
|
||||
reject_unknown_gemini_tts_voice_keys,
|
||||
)
|
||||
|
||||
default_voice: Final[dict[str, object]] = {
|
||||
"languageCode": self.DEFAULT_LANGUAGE_CODE,
|
||||
"modelName": model,
|
||||
"name": _DEFAULT_GEMINI_TTS_VOICE,
|
||||
}
|
||||
if voice is None:
|
||||
return None, default_voice
|
||||
|
||||
if isinstance(voice, str) and voice:
|
||||
return voice, {
|
||||
"languageCode": self.DEFAULT_LANGUAGE_CODE,
|
||||
"modelName": model,
|
||||
"name": self._resolve_gemini_tts_voice_name(model, voice, drop_params),
|
||||
}
|
||||
|
||||
voice_config: Final = as_str_mapping(voice)
|
||||
if voice_config is None:
|
||||
self._reject_malformed_gemini_tts_voice(model, drop_params)
|
||||
return None, default_voice
|
||||
reject_unknown_gemini_tts_voice_keys(model, voice_config, drop_params, "vertex_ai")
|
||||
|
||||
speech_config: Final = mapped_gemini_speech_config(voice_config)
|
||||
language_code: Final = self._get_str_value(speech_config, "languageCode") or self.DEFAULT_LANGUAGE_CODE
|
||||
model_name: Final = model
|
||||
speaker_configs: Final = self._extract_gemini_tts_speaker_configs(model, speech_config, drop_params)
|
||||
self._reject_voice_fields_cloud_tts_cannot_send(model, voice_config, drop_params)
|
||||
if speaker_configs:
|
||||
return None, {
|
||||
"languageCode": language_code,
|
||||
"modelName": model_name,
|
||||
"multiSpeakerVoiceConfig": {
|
||||
"speakerVoiceConfigs": speaker_configs,
|
||||
},
|
||||
}
|
||||
voice_name: Final = self._prebuilt_voice_name(speech_config)
|
||||
if voice_name is not None:
|
||||
return None, {
|
||||
"languageCode": language_code,
|
||||
"modelName": model_name,
|
||||
"name": self._resolve_gemini_tts_voice_name(model, voice_name, drop_params),
|
||||
}
|
||||
return None, {
|
||||
"languageCode": language_code,
|
||||
"modelName": model_name,
|
||||
"name": _DEFAULT_GEMINI_TTS_VOICE,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dispatch_voice_name(voice: str | Mapping[str, object] | None) -> str | None:
|
||||
if isinstance(voice, str):
|
||||
return voice
|
||||
if not isinstance(voice, Mapping):
|
||||
return None
|
||||
name: Final = voice.get("name")
|
||||
return name if isinstance(name, str) else None
|
||||
|
||||
def dispatch_text_to_speech(
|
||||
self,
|
||||
model: str,
|
||||
input: str,
|
||||
voice: str | dict | None,
|
||||
voice: str | dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
litellm_params_dict: dict[str, object],
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
|
|
@ -175,14 +565,17 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
vertex_project: Final = self.safe_get_vertex_ai_project(litellm_params_dict)
|
||||
vertex_location: Final = self.safe_get_vertex_ai_location(litellm_params_dict)
|
||||
|
||||
# Convert voice to string if it's a dict (extract name)
|
||||
# Actual voice mapping happens in map_openai_params
|
||||
voice_str: str | None = None
|
||||
if isinstance(voice, str):
|
||||
voice_str = voice
|
||||
elif isinstance(voice, dict):
|
||||
# Extract voice name from dict if needed
|
||||
voice_str = voice.get("name") if voice else None
|
||||
mapped_voice, mapped_params = (
|
||||
self.map_openai_params(
|
||||
model=model,
|
||||
voice=voice,
|
||||
optional_params=optional_params,
|
||||
drop_params=drops_unsupported_params(kwargs.get("drop_params")),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
if self._is_gemini_tts_model(model)
|
||||
else (self._dispatch_voice_name(voice), optional_params)
|
||||
)
|
||||
|
||||
# Store credentials in litellm_params for use in transform methods
|
||||
litellm_params_dict.update(
|
||||
|
|
@ -198,9 +591,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
response: Final = base_llm_http_handler.text_to_speech_handler(
|
||||
model=model,
|
||||
input=input,
|
||||
voice=voice_str,
|
||||
text_to_speech_provider_config=self,
|
||||
text_to_speech_optional_params=optional_params,
|
||||
voice=mapped_voice,
|
||||
text_to_speech_provider_config=(
|
||||
VertexAIGeminiCloudTextToSpeechConfig() if self._is_gemini_tts_model(model) else self
|
||||
),
|
||||
text_to_speech_optional_params=mapped_params,
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params=litellm_params_dict,
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -219,13 +614,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
Note: Vertex AI also supports additional parameters like audioConfig
|
||||
which can be passed but are not part of the OpenAI spec
|
||||
"""
|
||||
return ["voice", "response_format", "speed"]
|
||||
return ["voice", "response_format", "speed", *(["instructions"] if self._is_gemini_tts_model(model) else [])]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
voice: str | dict | None = None,
|
||||
voice: str | dict[str, object] | None = None,
|
||||
drop_params: bool = False,
|
||||
kwargs: dict = {},
|
||||
) -> tuple[str | None, dict]:
|
||||
|
|
@ -244,19 +639,28 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
Tuple of (mapped_voice_str, mapped_params)
|
||||
"""
|
||||
mapped_params: Final[dict[str, object]] = {}
|
||||
openai_params: Final = _JSON_OBJECT.validate_python(optional_params)
|
||||
drops_params: Final = drop_params or litellm.drop_params is True
|
||||
|
||||
##########################################################
|
||||
# Map voice using helper
|
||||
##########################################################
|
||||
mapped_voice_str, voice_dict = self._map_voice_to_vertex_format(voice)
|
||||
mapped_voice_str, voice_dict = self._map_voice_to_vertex_format(
|
||||
voice=voice,
|
||||
model=model,
|
||||
drop_params=drops_params,
|
||||
)
|
||||
if voice_dict is not None:
|
||||
mapped_params["vertex_voice_dict"] = voice_dict
|
||||
|
||||
# Map response format
|
||||
if "response_format" in optional_params:
|
||||
format_name: Final = optional_params["response_format"]
|
||||
if format_name in self.FORMAT_MAPPINGS:
|
||||
mapped_params["audioEncoding"] = self.FORMAT_MAPPINGS[format_name]
|
||||
format_mappings: Final = (
|
||||
self.GEMINI_FORMAT_MAPPINGS if self._is_gemini_tts_model(model) else self.FORMAT_MAPPINGS
|
||||
)
|
||||
if format_name in format_mappings:
|
||||
mapped_params["audioEncoding"] = format_mappings[format_name]
|
||||
else:
|
||||
# Try to use it directly as Google Cloud format
|
||||
mapped_params["audioEncoding"] = format_name
|
||||
|
|
@ -265,10 +669,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
mapped_params["audioEncoding"] = self.DEFAULT_AUDIO_ENCODING
|
||||
|
||||
# Map speed (OpenAI: 0.25-4.0, Vertex AI: speakingRate 0.25-4.0)
|
||||
if "speed" in optional_params:
|
||||
speed: Final = optional_params["speed"]
|
||||
if speed is not None:
|
||||
mapped_params["speakingRate"] = str(speed)
|
||||
speed: Final = openai_params.get("speed")
|
||||
if speed is not None:
|
||||
mapped_params["speakingRate"] = str(speed)
|
||||
|
||||
instructions: Final = openai_params.get("instructions")
|
||||
if isinstance(instructions, str) and self._is_gemini_tts_model(model):
|
||||
mapped_params["prompt"] = instructions
|
||||
elif instructions is not None and self._is_gemini_tts_model(model) and not drops_params:
|
||||
raise UnsupportedParamsError(
|
||||
message=f"{model} needs instructions to be a string. Send one, or set drop_params to ignore it.",
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
# Pass through Vertex AI-specific parameters from kwargs
|
||||
if "audioConfig" in kwargs:
|
||||
|
|
@ -311,15 +724,48 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
|
||||
Google Cloud TTS endpoint: https://texttospeech.googleapis.com/v1/text:synthesize
|
||||
"""
|
||||
if self._is_gemini_tts_model(model):
|
||||
return self._get_gemini_tts_url(
|
||||
model=model, api_base=api_base, litellm_params=_JSON_OBJECT.validate_python(litellm_params)
|
||||
)
|
||||
|
||||
if api_base:
|
||||
return api_base
|
||||
|
||||
return self.TTS_API_URL
|
||||
|
||||
def _get_gemini_tts_url(self, model: str, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
|
||||
if api_base:
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which does not use "
|
||||
"this deployment's api_base. Request wav or pcm, or use a deployment without api_base."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
location: Final = litellm_params.get("vertex_location")
|
||||
vertex_location: Final = location if isinstance(location, str) else self.get_default_vertex_location()
|
||||
if vertex_location == "global":
|
||||
return self.TTS_API_URL
|
||||
cloud_tts_location: Final = _GEMINI_CLOUD_TTS_LOCATIONS.get(vertex_location)
|
||||
if cloud_tts_location is None:
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which has no Gemini TTS "
|
||||
f"location for vertex_location {vertex_location}. Use a deployment in global, us, eu, "
|
||||
"northamerica-northeast1, or a US or EU Vertex AI region."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
return f"https://{cloud_tts_location}-texttospeech.googleapis.com/v1/text:synthesize"
|
||||
|
||||
def _validate_vertex_input(
|
||||
self,
|
||||
input_data: VertexTextToSpeechInput,
|
||||
optional_params: dict,
|
||||
detect_ssml: bool,
|
||||
) -> VertexTextToSpeechInput:
|
||||
"""
|
||||
Validate and transform input for Vertex AI TTS
|
||||
|
|
@ -343,7 +789,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
else:
|
||||
# LiteLLM will auto-detect if text is in ssml format
|
||||
# check if "text" is an ssml - in this case we should pass it as ssml instead of text
|
||||
if input_data:
|
||||
if detect_ssml and input_data:
|
||||
_text: Final = input_data.get("text", None) or ""
|
||||
if "<speak>" in _text:
|
||||
input_data["ssml"] = input_data.pop("text")
|
||||
|
|
@ -403,8 +849,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
headers["x-goog-user-project"] = vertex_project
|
||||
|
||||
####### Build the request ################
|
||||
vertex_input = VertexTextToSpeechInput(text=input)
|
||||
vertex_input = self._validate_vertex_input(vertex_input, optional_params)
|
||||
prompt: Final = _JSON_OBJECT.validate_python(optional_params).get("prompt")
|
||||
validated_input: Final = self._validate_vertex_input(
|
||||
VertexTextToSpeechInput(text=input), optional_params, detect_ssml=not self._is_gemini_tts_model(model)
|
||||
)
|
||||
vertex_input: Final[VertexTextToSpeechInput] = (
|
||||
{**validated_input, "prompt": prompt} if isinstance(prompt, str) else validated_input
|
||||
)
|
||||
|
||||
# Build voice configuration
|
||||
# Check for voice dict stored in:
|
||||
|
|
@ -412,8 +863,8 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
# 2. optional_params by map_openai_params
|
||||
voice_dict: Final = litellm_params.get("vertex_voice_dict") or optional_params.get("vertex_voice_dict")
|
||||
if voice_dict is not None and isinstance(voice_dict, dict):
|
||||
vertex_voice = VertexTextToSpeechVoice(**voice_dict)
|
||||
elif voice is not None and isinstance(voice, str):
|
||||
vertex_voice = _JSON_OBJECT.validate_python(voice_dict)
|
||||
elif voice is not None:
|
||||
# Handle string voice (shouldn't normally happen if dispatch was called)
|
||||
parts: Final = voice.split("-")
|
||||
if len(parts) >= 2:
|
||||
|
|
@ -485,8 +936,44 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
|
|||
content=binary_data,
|
||||
)
|
||||
|
||||
# Initialize the HttpxBinaryResponseContent instance
|
||||
return HttpxBinaryResponseContent(response)
|
||||
binary_response: Final = HttpxBinaryResponseContent(response)
|
||||
if self._is_gemini_tts_model(model):
|
||||
from litellm.types.utils import CompletionTokensDetailsWrapper, Usage
|
||||
|
||||
request: Final = _GeminiCloudTtsRequest.model_validate_json(raw_response.request.content)
|
||||
input_data: Final = request.input
|
||||
audio_config: Final = request.audioConfig
|
||||
container_duration: Final = calculate_request_duration(binary_data)
|
||||
sample_rate: Final = audio_config.sampleRateHertz or 24000
|
||||
duration: Final = (
|
||||
container_duration
|
||||
if container_duration is not None
|
||||
else _fallback_gemini_tts_audio_duration(binary_data, audio_config.audioEncoding, sample_rate)
|
||||
)
|
||||
if duration is None:
|
||||
verbose_logger.warning(
|
||||
"Could not measure Gemini TTS audio duration for %s, so its audio tokens are not billed", model
|
||||
)
|
||||
return binary_response
|
||||
input_text: Final = " ".join(
|
||||
value for value in (input_data.text, input_data.ssml, input_data.prompt) if value
|
||||
)
|
||||
prompt_tokens: Final = litellm.token_counter(model=model, text=input_text)
|
||||
audio_tokens: Final = math.ceil(duration * 25)
|
||||
binary_response.usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=audio_tokens,
|
||||
total_tokens=prompt_tokens + audio_tokens,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=audio_tokens),
|
||||
)
|
||||
return binary_response
|
||||
|
||||
|
||||
class VertexAIGeminiCloudTextToSpeechConfig(VertexAITextToSpeechConfig):
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException:
|
||||
return super().get_error_class(
|
||||
error_message=error_message, status_code=400 if status_code == 404 else status_code, headers=headers
|
||||
)
|
||||
|
||||
|
||||
class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
|
||||
|
|
|
|||
|
|
@ -5988,6 +5988,7 @@ def completion(
|
|||
*OPENAI_WIF_KWARGS_KEYS,
|
||||
*OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS,
|
||||
"github_copilot_auth_type",
|
||||
"github_copilot_host",
|
||||
"github_copilot_user_session",
|
||||
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
|
||||
"fireworks_forward_user_id",
|
||||
|
|
@ -8845,8 +8846,11 @@ def speech(
|
|||
|
||||
generic_optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
|
||||
|
||||
# Handle Gemini models separately (they use speech_to_completion_bridge)
|
||||
if "gemini" in model:
|
||||
needs_cloud_text_to_speech: Final = (
|
||||
litellm.utils.is_gemini_tts_model(model)
|
||||
and response_format in VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS
|
||||
)
|
||||
if "gemini" in model and not needs_cloud_text_to_speech:
|
||||
from .endpoints.speech.speech_to_completion_bridge.handler import (
|
||||
speech_to_completion_bridge_handler,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,18 +1,21 @@
|
|||
import inspect
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
|
||||
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.llms.anthropic.pass_through.messages import handler as main
|
||||
from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge import detectors
|
||||
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.messages import detectors as route_detectors
|
||||
from litellm.rust_bridge.messages.entrypoints import (
|
||||
NATIVE_AMESSAGES,
|
||||
NATIVE_MESSAGES,
|
||||
)
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
|
|
@ -46,6 +49,7 @@ def _python_amessages() -> PythonAmessages:
|
|||
|
||||
_PYTHON_MESSAGES: Final = _python_messages()
|
||||
_MESSAGES: Final = signature(_PYTHON_MESSAGES)
|
||||
_FLAG: Final = AsyncFlag("is_async", sync=_MESSAGES, is_set=bool)
|
||||
_PYTHON_AMESSAGES: Final = _python_amessages()
|
||||
_AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
|
||||
|
||||
|
|
@ -71,45 +75,76 @@ def _public_request(
|
|||
return native_call(legacy, args, _native_kwargs(fields, kwargs))
|
||||
|
||||
|
||||
def _resolved_provider(request: NativeCall) -> str | None:
|
||||
try:
|
||||
return get_llm_provider(
|
||||
str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))
|
||||
)[1]
|
||||
except BadRequestError:
|
||||
return optional_str(request.resolved.get("custom_llm_provider"))
|
||||
|
||||
|
||||
def _python_serves(request: NativeCall) -> bool:
|
||||
"""Vertex AI serves only Claude through the Anthropic Messages config, as
|
||||
``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's
|
||||
translation through Chat Completions. A request Python would answer with classifier
|
||||
``safeguard_results`` stays on Python too."""
|
||||
def _resolved(request: NativeCall) -> tuple[str, str | None]:
|
||||
model: Final = str(request.resolved["model"])
|
||||
vertex_without_claude: Final = _resolved_provider(request) == "vertex_ai" and "claude" not in model.lower()
|
||||
return vertex_without_claude or native_route_classifies(model, request.kwargs.get("safeguards"))
|
||||
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
|
||||
try:
|
||||
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
|
||||
except BadRequestError:
|
||||
return model, provider
|
||||
return resolved_model, resolved_provider
|
||||
|
||||
|
||||
def _context(request: NativeCall) -> RouteContext:
|
||||
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
|
||||
requested_model: Final = str(request.resolved["model"])
|
||||
model_info: Final = request.resolved.get("model_info")
|
||||
api_base: Final = request.resolved.get("api_base")
|
||||
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
|
||||
model, provider, model_info, api_base if isinstance(api_base, str) else None
|
||||
)
|
||||
if native_config is not None:
|
||||
return Target.MESSAGES
|
||||
compacting: Final = request.resolved.get("compaction") is not None
|
||||
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
|
||||
return Target.RESPONSES
|
||||
return Target.CHAT_COMPLETIONS
|
||||
|
||||
|
||||
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
|
||||
{
|
||||
**detectors.COMMON,
|
||||
Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier,
|
||||
Feature.RESPONSE_CACHE: detectors.response_cache,
|
||||
Feature.MODIFY_PARAMS: detectors.modify_params,
|
||||
Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity,
|
||||
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
|
||||
Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection,
|
||||
Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry,
|
||||
}
|
||||
)
|
||||
|
||||
_ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
|
||||
{
|
||||
**_FEATURES,
|
||||
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
|
||||
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
|
||||
Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext:
|
||||
requested_model: Final = str(request.resolved["model"])
|
||||
model, provider = _resolved(request)
|
||||
return RouteContext(
|
||||
Route.MESSAGES,
|
||||
provider=_resolved_provider(request),
|
||||
model=str(request.resolved["model"]),
|
||||
provider=provider,
|
||||
model=requested_model,
|
||||
target=_target(request, model, provider),
|
||||
features=detectors.detected(features, request),
|
||||
)
|
||||
|
||||
|
||||
_DISPATCH: Final = PublicDispatch(
|
||||
route=Route.MESSAGES,
|
||||
request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request),
|
||||
context=lambda request: _context(request, _FEATURES),
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
route=Route.MESSAGES,
|
||||
request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
|
||||
context=_context,
|
||||
bypass=_python_serves,
|
||||
context=lambda request: _context(request, _ASYNC_FEATURES),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -118,9 +153,11 @@ def anthropic_messages_handler(
|
|||
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
|
||||
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
|
||||
python: Final = _PYTHON_MESSAGES
|
||||
if _FLAG.requested(kwargs):
|
||||
return anthropic_messages(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_MESSAGES,
|
||||
native=native_call_hook,
|
||||
|
|
@ -131,7 +168,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
|
|||
python: Final = _PYTHON_AMESSAGES
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_AMESSAGES,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -16580,6 +16580,24 @@
|
|||
"source": "https://developers.cloudflare.com/workers-ai/platform/pricing/",
|
||||
"supports_vision": true
|
||||
},
|
||||
"cloudflare/clairic": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cloudflare",
|
||||
"max_input_tokens": 65536,
|
||||
"mode": "extraction",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
|
||||
"supports_vision": true
|
||||
},
|
||||
"cloudflare/@cf/cloudflare/clairic": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "cloudflare",
|
||||
"max_input_tokens": 65536,
|
||||
"mode": "extraction",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
|
||||
"supports_vision": true
|
||||
},
|
||||
"cloudflare/@cf/meta/llama-2-7b-chat-fp16": {
|
||||
"input_cost_per_token": 1.923e-06,
|
||||
"litellm_provider": "cloudflare",
|
||||
|
|
@ -21903,6 +21921,7 @@
|
|||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -21927,6 +21946,7 @@
|
|||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -21958,6 +21978,7 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -22021,6 +22042,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -29417,6 +29439,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -30139,6 +30162,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -30200,6 +30224,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -30268,6 +30293,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -30317,6 +30343,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": false,
|
||||
"supports_url_context": false,
|
||||
|
|
@ -31743,6 +31770,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -31927,6 +31955,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -31989,6 +32018,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -32051,6 +32081,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -32420,6 +32451,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -32481,6 +32513,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -32549,6 +32582,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -32598,6 +32632,7 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": false,
|
||||
"supports_url_context": false,
|
||||
|
|
@ -33192,6 +33227,7 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"github_copilot/gemini-3.8-flash": {
|
||||
|
|
@ -33210,6 +33246,7 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_vision": true
|
||||
},
|
||||
"github_copilot/gpt-3.5-turbo": {
|
||||
|
|
@ -52605,6 +52642,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
|
|
@ -64157,13 +64195,17 @@
|
|||
},
|
||||
"fireworks_ai/accounts/fireworks/models/muse-spark-open": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"cache_read_input_token_cost_priority": 2.25e-07,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"input_cost_per_token_priority": 1.875e-06,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.25e-06,
|
||||
"source": "https://fireworks.ai/models/fireworks/muse-spark-open",
|
||||
"output_cost_per_token_priority": 6.375e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/pricing",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": {
|
||||
|
|
@ -65823,6 +65865,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_vision": true,
|
||||
"source": "https://deepinfra.com/pricing",
|
||||
"supports_audio_input": true
|
||||
|
|
@ -67292,6 +67335,7 @@
|
|||
"fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"cache_read_input_token_cost_priority": 5.625e-08,
|
||||
"deprecation_date": "2026-11-07",
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"input_cost_per_token_priority": 2.8125e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
|
|
@ -67326,6 +67370,7 @@
|
|||
"fireworks_ai/glm-5p3-flash-us": {
|
||||
"cache_read_input_token_cost": 4.5e-08,
|
||||
"cache_read_input_token_cost_priority": 5.625e-08,
|
||||
"deprecation_date": "2026-11-07",
|
||||
"input_cost_per_token": 2.25e-07,
|
||||
"input_cost_per_token_priority": 2.8125e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
|
|
@ -69262,6 +69307,7 @@
|
|||
"perplexity/google/gemini-3.5-flash-lite": {
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "responses",
|
||||
"supports_sampling_params": false,
|
||||
"supports_web_search": true,
|
||||
"supports_function_calling": true,
|
||||
"input_cost_per_token": 3e-07,
|
||||
|
|
@ -69272,6 +69318,7 @@
|
|||
"perplexity/google/gemini-3.6-flash": {
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "responses",
|
||||
"supports_sampling_params": false,
|
||||
"supports_web_search": true,
|
||||
"supports_function_calling": true,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
|
|
@ -69282,6 +69329,7 @@
|
|||
"perplexity/google/gemini-3.7-flash": {
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "responses",
|
||||
"supports_sampling_params": false,
|
||||
"supports_web_search": true,
|
||||
"supports_function_calling": true,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
|
|
@ -71373,6 +71421,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
|
|
@ -71390,6 +71439,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
|
|
@ -74320,6 +74370,7 @@
|
|||
"perplexity/google/gemini-3.8-flash": {
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "responses",
|
||||
"supports_sampling_params": false,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
"output_cost_per_token": 3.75e-06,
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
|
|
@ -75062,12 +75113,10 @@
|
|||
"gemini-nano-banana-2.1": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"cache_read_input_token_cost_batches": 7.5e-08,
|
||||
"cache_read_input_token_cost_flex": 7.5e-08,
|
||||
"cache_read_input_token_cost_priority": 2.7e-07,
|
||||
"input_cost_per_image": 0.00168,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"input_cost_per_token_batches": 7.5e-07,
|
||||
"input_cost_per_token_flex": 7.5e-07,
|
||||
"input_cost_per_token_priority": 2.7e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 131072,
|
||||
|
|
@ -75078,7 +75127,6 @@
|
|||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"output_cost_per_token_batches": 3.75e-06,
|
||||
"output_cost_per_token_flex": 3.75e-06,
|
||||
"output_cost_per_token_priority": 1.35e-05,
|
||||
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
|
||||
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
|
||||
|
|
@ -75111,12 +75159,10 @@
|
|||
"vertex_ai/gemini-nano-banana-2.1": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"cache_read_input_token_cost_batches": 7.5e-08,
|
||||
"cache_read_input_token_cost_flex": 7.5e-08,
|
||||
"cache_read_input_token_cost_priority": 2.7e-07,
|
||||
"input_cost_per_image": 0.00168,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"input_cost_per_token_batches": 7.5e-07,
|
||||
"input_cost_per_token_flex": 7.5e-07,
|
||||
"input_cost_per_token_priority": 2.7e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 131072,
|
||||
|
|
@ -75127,7 +75173,6 @@
|
|||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"output_cost_per_token_batches": 3.75e-06,
|
||||
"output_cost_per_token_flex": 3.75e-06,
|
||||
"output_cost_per_token_priority": 1.35e-05,
|
||||
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
|
||||
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
|
||||
|
|
@ -76668,14 +76713,14 @@
|
|||
"supports_response_schema": true
|
||||
},
|
||||
"baseten/reflection/Beam-1.0": {
|
||||
"cache_read_input_token_cost": 1.3e-07,
|
||||
"input_cost_per_token": 1.3e-06,
|
||||
"cache_read_input_token_cost": 4e-08,
|
||||
"input_cost_per_token": 4e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.3e-06,
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://inference.baseten.co/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
|
|
@ -76927,5 +76972,83 @@
|
|||
"supported_output_modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash-1017-us": {
|
||||
"cache_read_input_token_cost": 9e-09,
|
||||
"cache_read_input_token_cost_priority": 1.125e-08,
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"input_cost_per_token_priority": 5.625e-07,
|
||||
"litellm_provider": "fireworks_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 393216,
|
||||
"max_tokens": 393216,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-06,
|
||||
"output_cost_per_token_priority": 2.25e-06,
|
||||
"source": "https://docs.fireworks.ai/serverless/us-only-serverless",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"wandb/reflection/Beam-1.0-501B-A23B": {
|
||||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "wandb",
|
||||
"max_input_tokens": 1049000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"wandb/zai-org/GLM-5.3": {
|
||||
"cache_read_input_token_cost": 1.9e-07,
|
||||
"input_cost_per_token": 1.2e-06,
|
||||
"litellm_provider": "wandb",
|
||||
"max_input_tokens": 1049000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-06,
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"azure_ai/MAI-Code-1.1-Flash": {
|
||||
"cache_read_input_token_cost": 2e-08,
|
||||
"deprecation_date": "2026-12-09",
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 256000,
|
||||
"max_output_tokens": 64000,
|
||||
"max_tokens": 64000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing import Literal, TypeAlias
|
|||
|
||||
from pydantic import ConfigDict, Field, model_validator
|
||||
|
||||
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
CredentialSource: TypeAlias = Literal["db", "config"]
|
||||
|
|
@ -101,6 +102,7 @@ class UserConnectionFlowHandle(LiteLLMBaseModel):
|
|||
|
||||
user_id: str
|
||||
credential_name: str
|
||||
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
|
||||
device_code: str = Field(repr=False)
|
||||
interval: int
|
||||
expires_at: float
|
||||
|
|
|
|||
|
|
@ -9,7 +9,14 @@ from litellm.rust_bridge import runtime
|
|||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str, signature
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
native_call,
|
||||
native_call_hook,
|
||||
optional_str,
|
||||
signature,
|
||||
)
|
||||
|
||||
__all__ = ("aocr", "ocr")
|
||||
|
||||
|
|
@ -39,6 +46,7 @@ def _bind_request(
|
|||
|
||||
|
||||
_OCR: Final = signature(_bind_request)
|
||||
_FLAG: Final = AsyncFlag("aocr", sync=_OCR)
|
||||
|
||||
|
||||
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
|
||||
|
|
@ -59,7 +67,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.OCR,
|
||||
request=lambda args, kwargs: _public_request("ocr", args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("aocr") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -73,9 +80,11 @@ def ocr(
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
|
||||
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
|
||||
if _FLAG.requested(kwargs):
|
||||
return aocr(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=NATIVE_OCR,
|
||||
native=native_call_hook,
|
||||
|
|
@ -85,7 +94,7 @@ def ocr(
|
|||
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=NATIVE_AOCR,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -437,6 +437,7 @@ def llm_passthrough_route(
|
|||
params: QueryParamTypes | None = None,
|
||||
cookies: CookieTypes | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
guardrailConfig: object = None,
|
||||
**kwargs,
|
||||
) -> (
|
||||
httpx.Response
|
||||
|
|
@ -547,7 +548,7 @@ def llm_passthrough_route(
|
|||
_request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None)
|
||||
headers, signed_json_body = provider_config.sign_request(
|
||||
headers=headers,
|
||||
litellm_params=litellm_params_dict,
|
||||
litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig},
|
||||
request_data=_request_data,
|
||||
api_base=str(updated_url),
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -1197,6 +1197,8 @@ def should_strip_caller_authorization(
|
|||
admission credential — a virtual key, an IdP JWT, or an SSO / OIDC /
|
||||
session token whose ``api_key`` is ``None`` — and must never reach
|
||||
the upstream, so it is stripped regardless of the ``api_key`` value.
|
||||
- **Every other auth-enabled mode**: strip when ``Authorization`` is the
|
||||
credential admission validated.
|
||||
"""
|
||||
if _authorization_is_admission_jwt(raw_headers, user_api_key_auth):
|
||||
return True
|
||||
|
|
@ -1216,7 +1218,21 @@ def should_strip_caller_authorization(
|
|||
mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream
|
||||
)
|
||||
if not (mcp_server.is_oauth_passthrough or is_delegated_oauth):
|
||||
return False
|
||||
from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
|
||||
master_key,
|
||||
)
|
||||
|
||||
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
|
||||
admission_ran: Final = not is_no_auth_dev_mode(master_key, proxy_general_settings) or (
|
||||
user_api_key_auth is not None and user_api_key_auth.authenticated_by_custom_auth
|
||||
)
|
||||
return admission_ran and (
|
||||
_authorization_is_litellm_admission_credential_with_settings(
|
||||
raw_headers, user_api_key_auth, proxy_general_settings
|
||||
)
|
||||
)
|
||||
|
||||
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
|
||||
if is_delegated_oauth:
|
||||
|
|
@ -1243,16 +1259,39 @@ def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None
|
|||
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
|
||||
|
||||
|
||||
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _authorization_is_litellm_admission_credential(
|
||||
raw_headers: Mapping[str, str] | None,
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
) -> bool:
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
|
||||
)
|
||||
|
||||
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
|
||||
return _authorization_is_litellm_admission_credential_with_settings(
|
||||
raw_headers, user_api_key_auth, proxy_general_settings
|
||||
)
|
||||
|
||||
|
||||
def _authorization_is_litellm_admission_credential_with_settings(
|
||||
raw_headers: Mapping[str, str] | None,
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
proxy_general_settings: Mapping[str, object],
|
||||
) -> bool:
|
||||
"""True when ``Authorization`` carries the LiteLLM key admission validated.
|
||||
|
||||
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
|
||||
same key in both headers.
|
||||
That is the case when no usable admission header was sent, or when the client repeated the same key in both headers.
|
||||
"""
|
||||
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
|
||||
custom_header_name: Final[object | None] = proxy_general_settings.get("litellm_key_header_name")
|
||||
custom_admission_header: Final = (
|
||||
_raw_header_value(raw_headers, custom_header_name.lower())
|
||||
if isinstance(custom_header_name, str) and custom_header_name
|
||||
else None
|
||||
)
|
||||
admission_header: Final = custom_admission_header or _raw_header_value(raw_headers, "x-litellm-api-key")
|
||||
authorization: Final = _raw_header_value(raw_headers, "authorization")
|
||||
if admission_header and authorization:
|
||||
return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer")
|
||||
|
|
@ -4322,6 +4361,7 @@ class MCPServerManager:
|
|||
raw_headers: Mapping[str, str] | None = None,
|
||||
client_ip: str | None = None,
|
||||
protocol_version_override: MCPUpstreamProtocol | None = None,
|
||||
for_tool_listing: bool = False,
|
||||
) -> MCPClient:
|
||||
self.catalog.assert_current(server)
|
||||
record_auth_resolution(server.server_id, AuthResolution.unresolved)
|
||||
|
|
@ -4355,6 +4395,7 @@ class MCPServerManager:
|
|||
else None
|
||||
),
|
||||
upstream_transport=self._upstream_transport,
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
|
||||
_create_mcp_client = create_mcp_client
|
||||
|
|
@ -4501,6 +4542,7 @@ class MCPServerManager:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
raw_headers=raw_headers,
|
||||
client_ip=client_ip,
|
||||
for_tool_listing=True,
|
||||
)
|
||||
|
||||
## HANDLE OPENAPI TOOLS
|
||||
|
|
|
|||
|
|
@ -460,6 +460,8 @@ def validate_static_credential(
|
|||
headers: Mapping[str, str],
|
||||
upstream_token_header: str | None = None,
|
||||
static_header_names: Iterable[str] = (),
|
||||
*,
|
||||
allow_absent: bool = False,
|
||||
) -> Result[None, CredError]:
|
||||
if auth_type not in _STATIC_MODES:
|
||||
return Ok(None)
|
||||
|
|
@ -475,17 +477,23 @@ def validate_static_credential(
|
|||
)
|
||||
)
|
||||
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
|
||||
if allow_absent and not values:
|
||||
return Ok(None)
|
||||
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
|
||||
return Ok(None)
|
||||
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
|
||||
|
||||
|
||||
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
|
||||
async def prepare_mcp_client(server: MCPServer, client: MCPClient, *, for_tool_listing: bool = False) -> MCPClient:
|
||||
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
|
||||
return client
|
||||
request: Final = await client.prepare_request_auth()
|
||||
match validate_static_credential(
|
||||
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
|
||||
server.auth_type,
|
||||
request.headers,
|
||||
server.upstream_token_header,
|
||||
server.static_headers or (),
|
||||
allow_absent=for_tool_listing and server.auth_type == MCPAuth.authorization,
|
||||
):
|
||||
case Error(error):
|
||||
raise_public(error)
|
||||
|
|
|
|||
|
|
@ -208,6 +208,7 @@ async def prepare_upstream_client(
|
|||
sampling_callback: SamplingCallback | None = None,
|
||||
elicitation_callback: ElicitationCallback | None = None,
|
||||
upstream_transport: httpx2.AsyncBaseTransport | None = None,
|
||||
for_tool_listing: bool = False,
|
||||
) -> MCPClient:
|
||||
transport: Final[MCPTransportType] = server.transport or MCPTransport.sse
|
||||
server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server)
|
||||
|
|
@ -272,6 +273,7 @@ async def prepare_upstream_client(
|
|||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
),
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
aws_auth: Final = (
|
||||
MCPSigV4Auth(
|
||||
|
|
@ -316,4 +318,5 @@ async def prepare_upstream_client(
|
|||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
),
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -1,4 +1,4 @@
|
|||
:HL["/litellm-asset-prefix/_next/static/chunks/1kid9zr1--h6y.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/3146e697tym4_.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/0ckv5j_fni3xm.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/06yuul4yvgbd0.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.2bn3s6zvc0dyp.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"tree":{"name":"","param":null,"prefetchHints":4176,"slots":{"children":{"name":"(dashboard)","param":null,"prefetchHints":4192,"slots":{"children":{"name":"__PAGE__","param":null,"prefetchHints":4256,"slots":null}}}}},"staleTime":300,"buildId":"kXnLzJ6ylsRPmgSkCkCKM"}
|
||||
0:{"tree":{"name":"","param":null,"prefetchHints":4176,"slots":{"children":{"name":"(dashboard)","param":null,"prefetchHints":4192,"slots":{"children":{"name":"__PAGE__","param":null,"prefetchHints":4256,"slots":null}}}}},"staleTime":300,"buildId":"u_MWVLx_WA9i-AM2JdCsp"}
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue