mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
Merge origin/main into litellm_json_logs_startup_plaintext
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
This commit is contained in:
commit
7d17f126e5
1213 changed files with 53325 additions and 12475 deletions
9
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
9
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -4,6 +4,13 @@
|
|||
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 -->
|
||||
|
||||
<!-- 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 -->
|
||||
|
||||
<details>
|
||||
<summary>Details</summary>
|
||||
|
||||
## TLDR
|
||||
|
||||
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
|
||||
|
|
@ -164,3 +171,5 @@ Example checklists:
|
|||
## Final Attestation
|
||||
|
||||
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
|
||||
|
||||
</details>
|
||||
|
|
|
|||
2
.github/ci-coverage-allowlist.yml
vendored
2
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -5,7 +5,7 @@ description: >-
|
|||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
|
||||
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
|
||||
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
|
||||
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
|
||||
|
|
|
|||
9
.github/workflows/test-unit.yml
vendored
9
.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"
|
||||
|
|
@ -322,6 +323,7 @@ jobs:
|
|||
tests/unit/llms/pass_through
|
||||
tests/unit/llms/perplexity
|
||||
tests/unit/llms/pg_vector
|
||||
tests/unit/llms/pi
|
||||
tests/unit/llms/publicai
|
||||
tests/unit/llms/ragflow
|
||||
tests/unit/llms/recraft
|
||||
|
|
@ -432,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
|
||||
|
|
|
|||
|
|
@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
|
|||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
|
||||
|
||||
### Python SDK - Agents
|
||||
|
||||
|
|
@ -296,7 +296,7 @@ import litellm
|
|||
from litellm import Harness, sandbox
|
||||
|
||||
result = litellm.agent(
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
|
||||
"Find why tests/test_router.py is flaky and fix it.",
|
||||
sandbox=sandbox.local("./repo"),
|
||||
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
|
||||
|
|
@ -305,7 +305,7 @@ result = litellm.agent(
|
|||
print(result.text, result.cost, [f.path for f in result.files])
|
||||
```
|
||||
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
|
||||
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
|
||||
|
||||
|
|
|
|||
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"
|
||||
```
|
||||
|
|
@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper:
|
|||
the tag's value (the mode for that tag).
|
||||
- If no request tag matches, fall back to default mode(s).
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
)
|
||||
from litellm.router_strategy.tag_based_routing import (
|
||||
without_baggage_mapped_tags,
|
||||
)
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
|
|
@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper:
|
|||
|
||||
proxy_server_request = data.get("proxy_server_request", {})
|
||||
|
||||
request_tags = StandardLoggingPayloadSetup.get_request_tags(
|
||||
litellm_params=data,
|
||||
proxy_server_request=proxy_server_request,
|
||||
request_tags = without_baggage_mapped_tags(
|
||||
StandardLoggingPayloadSetup.get_request_tags(
|
||||
litellm_params=data,
|
||||
proxy_server_request=proxy_server_request,
|
||||
),
|
||||
data.get(get_metadata_variable_name_from_kwargs(data)),
|
||||
keep_serving_deployment_tags=True,
|
||||
)
|
||||
|
||||
# Check if any request tag matches a Mode tag key
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -13,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
|
||||
|
|
@ -39,7 +41,12 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
||||
def _error_names_resource(error: Exception, resource_id: str) -> bool:
|
||||
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
|
||||
|
||||
|
||||
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
|
||||
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
|
||||
|
||||
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
|
||||
"completed",
|
||||
|
|
@ -84,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,
|
||||
|
|
@ -109,6 +111,11 @@ class CheckBatchCost:
|
|||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
self.batch_processed_support_confirmed: bool = False
|
||||
self._poll_cursor: tuple[datetime, str] | None = None
|
||||
self._poll_before: datetime = datetime.now(timezone.utc)
|
||||
self._arrival_cursor: tuple[datetime, str] | None = None
|
||||
self._arrival_before: datetime = self._poll_before
|
||||
self._poll_arrivals: bool = False
|
||||
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
|
||||
self._line_item_store_turn: Final = asyncio.Lock()
|
||||
self._poll_idle: Final = asyncio.Event()
|
||||
|
|
@ -161,29 +168,9 @@ class CheckBatchCost:
|
|||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
if not self._has_batch_processed_column:
|
||||
return
|
||||
|
||||
# A row already in a terminal status is never rewritten by the sweep above, so
|
||||
# without this it keeps a poll-page slot forever and starves newer batches.
|
||||
retired: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
"status": {"in": ["complete", "completed"]},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"batch_processed": True},
|
||||
)
|
||||
if retired > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
|
||||
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
|
||||
)
|
||||
|
||||
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
|
||||
"""Query batch jobs without the batch_processed filter (for older schemas)."""
|
||||
return await _managed_object_table(self.prisma_client).find_many(
|
||||
return await self._find_jobs(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {
|
||||
|
|
@ -197,10 +184,59 @@ class CheckBatchCost:
|
|||
]
|
||||
},
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
|
||||
arrivals: Final = self._poll_arrivals
|
||||
page: Final = await self._find_job_page(where, arrivals=arrivals)
|
||||
self._poll_arrivals = not arrivals if page else arrivals
|
||||
if page:
|
||||
return page
|
||||
return await self._find_job_page(where, arrivals=not arrivals)
|
||||
|
||||
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
|
||||
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
|
||||
if cursor is None:
|
||||
if arrivals:
|
||||
self._arrival_before = datetime.now(timezone.utc)
|
||||
else:
|
||||
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
|
||||
before: Final = self._arrival_before if arrivals else self._poll_before
|
||||
page: Final = await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
**where,
|
||||
"created_at": {
|
||||
"lte": before,
|
||||
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
|
||||
},
|
||||
**(
|
||||
{
|
||||
"OR": [
|
||||
{"created_at": {"gt": cursor[0]}},
|
||||
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
|
||||
]
|
||||
}
|
||||
if cursor is not None
|
||||
else {}
|
||||
),
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order=[{"created_at": "asc"}, {"id": "asc"}],
|
||||
)
|
||||
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
|
||||
if arrivals:
|
||||
self._arrival_cursor = next_cursor
|
||||
else:
|
||||
self._poll_cursor = next_cursor
|
||||
return page
|
||||
|
||||
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
|
||||
return {
|
||||
"id": job.id,
|
||||
"updated_at": job.updated_at,
|
||||
**({"batch_processed": False} if self._has_batch_processed_column else {}),
|
||||
}
|
||||
|
||||
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
|
||||
"""
|
||||
Take a row that can never be costed out of the poll page. Leaving it selectable
|
||||
|
|
@ -208,24 +244,19 @@ class CheckBatchCost:
|
|||
once enough such rows accumulate no newer batch is ever reached. Older schemas
|
||||
without batch_processed can only be excluded through the status filter.
|
||||
"""
|
||||
data: Final = (
|
||||
{"batch_processed": True}
|
||||
if self._has_batch_processed_column
|
||||
else {"status": "stale_expired"}
|
||||
)
|
||||
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
|
||||
try:
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
retired: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where=self._unprocessed_snapshot(job),
|
||||
data=data,
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
|
||||
)
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
|
||||
return
|
||||
if retired == 0:
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
|
||||
"so it will no longer be polled"
|
||||
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
|
||||
|
|
@ -312,13 +343,13 @@ class CheckBatchCost:
|
|||
|
||||
from litellm.exceptions import NotFoundError
|
||||
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
|
||||
|
||||
def _batch_deployment_exists(self, model_id: str) -> bool:
|
||||
"""A 404 only proves the batch is gone when it came from the batch's own
|
||||
deployment. Once that deployment leaves the router, default fallbacks can
|
||||
silently send the retrieve to a provider that never saw the batch, so its
|
||||
404 must not retire the row; the staleness sweep bounds it instead."""
|
||||
404 must not retire a potentially recoverable row."""
|
||||
return self.llm_router.get_deployment(model_id=model_id) is not None
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -333,7 +364,7 @@ class CheckBatchCost:
|
|||
|
||||
if not output_file_id:
|
||||
return False
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
|
||||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
|
||||
|
|
@ -355,18 +386,18 @@ class CheckBatchCost:
|
|||
db_batch_object=job,
|
||||
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
|
||||
)
|
||||
update_data: Final[dict] = {
|
||||
update_data: Final[dict[str, object]] = {
|
||||
"status": response.status,
|
||||
"file_object": response.model_dump_json(),
|
||||
**({"batch_processed": True} if self._has_batch_processed_column else {}),
|
||||
}
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where=self._unprocessed_snapshot(job),
|
||||
data=update_data,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
|
||||
)
|
||||
if finalized == 0:
|
||||
return False
|
||||
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
|
||||
return True
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -817,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:
|
||||
|
|
@ -1022,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 {}
|
||||
|
|
@ -1047,6 +1086,7 @@ class CheckBatchCost:
|
|||
"""
|
||||
try:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
prom_logger = PrometheusLogger.get_instance()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
|
||||
|
|
@ -1071,7 +1111,7 @@ class CheckBatchCost:
|
|||
# endpoint may transition a batch to "complete" before
|
||||
# CheckBatchCost runs. The batch_processed=False filter
|
||||
# already prevents reprocessing finished batches.
|
||||
jobs = await _managed_object_table(self.prisma_client).find_many(
|
||||
jobs = await self._find_jobs(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
|
|
@ -1084,8 +1124,6 @@ class CheckBatchCost:
|
|||
]
|
||||
},
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
self.batch_processed_support_confirmed = True
|
||||
except Exception as query_err:
|
||||
|
|
@ -1152,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,12 +61,29 @@ 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):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
id: str | None = None
|
||||
|
||||
|
||||
class _RoutedDeployment(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
model_info: _RoutedModelInfo = _RoutedModelInfo()
|
||||
|
||||
|
||||
_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)
|
||||
|
|
@ -76,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)
|
||||
|
|
@ -94,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,
|
||||
|
|
@ -106,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={},
|
||||
|
|
@ -174,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
|
||||
|
|
@ -183,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)
|
||||
|
|
@ -294,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(
|
||||
|
|
@ -310,12 +398,26 @@ 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]:
|
||||
requested: Final = target.model_metadata.get("model_group")
|
||||
if target.deployment is None or not isinstance(requested, str):
|
||||
return MappingProxyType({})
|
||||
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
|
||||
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
|
||||
)
|
||||
serving_id: Final = target.deployment.model_info.id
|
||||
return MappingProxyType(
|
||||
{"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {}
|
||||
)
|
||||
|
||||
async def _bill_once(
|
||||
|
|
@ -328,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
|
||||
|
|
|
|||
|
|
@ -9,12 +9,17 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast
|
||||
from datetime import datetime
|
||||
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
|
||||
|
|
@ -27,6 +32,12 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
class ManagedObjectRow(Protocol):
|
||||
@property
|
||||
def created_at(self) -> datetime: ...
|
||||
|
||||
@property
|
||||
def updated_at(self) -> datetime: ...
|
||||
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
|
|
@ -63,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:
|
||||
|
|
@ -89,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
|
||||
|
|
@ -96,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,
|
||||
|
|
@ -107,8 +142,15 @@ class ManagedObjectCreatorAttribution:
|
|||
user's alias, because a resolvable name is more useful on the spend row than a null.
|
||||
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
|
||||
to increment organization spend, so leaving it out drops the cost from org accounting.
|
||||
A background response also carries its end user and the key, team, and user budgets, the
|
||||
way the read that billed it before did, and never names a key-scoped creator as its user.
|
||||
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 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:
|
||||
|
|
@ -119,17 +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 {}),
|
||||
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
|
||||
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
|
||||
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
|
||||
}
|
||||
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)
|
||||
),
|
||||
)
|
||||
|
||||
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(
|
||||
|
|
@ -152,23 +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})
|
||||
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)
|
||||
|
|
@ -196,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}"
|
||||
|
||||
|
|
@ -212,7 +354,28 @@ def _organization_id(row: object) -> str | None:
|
|||
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
|
||||
|
||||
|
||||
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
|
||||
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
|
||||
if row is None:
|
||||
return {}
|
||||
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
|
||||
return {
|
||||
spend_key: getattr(row, "spend", None),
|
||||
max_budget_key: getattr(row, "max_budget", None),
|
||||
model_max_budget_key: getattr(row, "model_max_budget", None),
|
||||
}
|
||||
|
||||
|
||||
def _key_budget_fields(key: object) -> Mapping[str, object]:
|
||||
if key is None:
|
||||
return {}
|
||||
tier: Final = getattr(key, "litellm_budget_table", None)
|
||||
return {
|
||||
"user_api_key_spend": getattr(key, "spend", None),
|
||||
"user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"),
|
||||
"user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"),
|
||||
}
|
||||
|
||||
|
||||
def _own_or_tier_budget(key: object, tier: object, field: str) -> object:
|
||||
own: Final = getattr(key, field, None)
|
||||
tier_value: Final = getattr(tier, field, None)
|
||||
return own if tier_value is None or own not in (None, {}) else tier_value
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -551,6 +551,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, ma
|
|||
|
||||
model_cost = get_model_cost_map(url=model_cost_map_url)
|
||||
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
|
||||
apply_search_free_tier: bool = False
|
||||
cost_margin_config: Dict[
|
||||
str, Union[float, Dict[str, float]]
|
||||
] = {} # Provider-specific or global cost margins. Examples:
|
||||
|
|
@ -1504,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 (
|
||||
|
|
@ -1710,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,
|
||||
)
|
||||
|
|
@ -2385,6 +2390,7 @@ _AGENT_EXPORTS: Final = frozenset(
|
|||
"ClaudeCodeOptions",
|
||||
"CodexOptions",
|
||||
"OpenCodeOptions",
|
||||
"PiOptions",
|
||||
"DeepAgentsOptions",
|
||||
"ToolLoopOptions",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -602,9 +602,9 @@ def _build_streaming_logging_obj(
|
|||
litellm_params: dict[str, object] | None,
|
||||
metadata: dict[str, object] | None,
|
||||
proxy_server_request: dict[str, object] | None,
|
||||
start_time: datetime.datetime,
|
||||
) -> Logging:
|
||||
"""Build logging object for streaming A2A requests."""
|
||||
start_time: Final = datetime.datetime.now()
|
||||
model: Final = f"a2a_agent/{agent_name}"
|
||||
|
||||
logging_obj: Final = Logging(
|
||||
|
|
@ -718,6 +718,7 @@ async def asend_message_streaming(
|
|||
if request is None:
|
||||
raise ValueError("request is required")
|
||||
|
||||
called_at: Final = datetime.datetime.now()
|
||||
_raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
|
||||
|
||||
|
|
@ -751,6 +752,7 @@ async def asend_message_streaming(
|
|||
litellm_params=litellm_params,
|
||||
metadata=metadata,
|
||||
proxy_server_request=proxy_server_request,
|
||||
start_time=called_at,
|
||||
)
|
||||
|
||||
verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class A2AStreamingIterator:
|
|||
self.request = request
|
||||
self.logging_obj = logging_obj
|
||||
self.agent_name = agent_name
|
||||
self.start_time = datetime.now()
|
||||
self.start_time = logging_obj.start_time
|
||||
|
||||
# Collect chunks for token counting
|
||||
self.chunks: list[SendStreamingMessageResponse] = []
|
||||
|
|
@ -54,6 +54,8 @@ class A2AStreamingIterator:
|
|||
async def __anext__(self) -> "SendStreamingMessageResponse":
|
||||
try:
|
||||
chunk: Final = await self.stream.__anext__()
|
||||
if self.logging_obj.completion_start_time is None:
|
||||
self.logging_obj.update_completion_start_time(completion_start_time=datetime.now())
|
||||
|
||||
# Store chunk
|
||||
self.chunks.append(chunk)
|
||||
|
|
|
|||
|
|
@ -59,16 +59,6 @@ class PendingBatchRead:
|
|||
previous_access_times: dict[str, float | None]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeclaredBatchRead:
|
||||
"""A ``async_batch_get_cache`` split in two: the memory half done, the Redis half declared on a ``RedisBatch``
|
||||
so it rides that batch's next round trip, resolved later with ``async_resolve_batch_get``."""
|
||||
|
||||
keys: tuple[str, ...]
|
||||
pending: PendingBatchRead
|
||||
result: BatchResult[Mapping[str, object]] | None
|
||||
|
||||
|
||||
def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
|
||||
if future.cancelled():
|
||||
return
|
||||
|
|
@ -382,22 +372,6 @@ class DualCache(BaseCache):
|
|||
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
|
||||
return merged
|
||||
|
||||
async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
|
||||
pending: Final = await self._prepare_batch_get(
|
||||
list(keys),
|
||||
local_only=False,
|
||||
throttle_redis=False,
|
||||
)
|
||||
return DeclaredBatchRead(
|
||||
keys=tuple(keys),
|
||||
pending=pending,
|
||||
result=batch.mget(pending.redis_keys) if pending.redis_keys else None,
|
||||
)
|
||||
|
||||
async def async_resolve_batch_get(self, declared: DeclaredBatchRead) -> list[object | None]:
|
||||
redis_result: Final = None if declared.result is None else await declared.result
|
||||
return await self._apply_batch_get(declared.pending, redis_result)
|
||||
|
||||
async def async_batch_get_cache(
|
||||
self,
|
||||
keys: list,
|
||||
|
|
|
|||
|
|
@ -134,7 +134,12 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
|
|||
"return count"
|
||||
)
|
||||
|
||||
_INCREMENT_IF_PRESENT_LUA: Final = (
|
||||
"if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"
|
||||
)
|
||||
|
||||
_LUA_COUNT: Final = TypeAdapter(int)
|
||||
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
|
||||
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
|
||||
|
||||
|
||||
|
|
@ -675,7 +680,7 @@ class RedisCache(BaseCache):
|
|||
# redis namespaces
|
||||
self.namespace = namespace
|
||||
# for high traffic, we store the redis results in memory and then batch write to redis
|
||||
self.redis_batch_writing_buffer: list = []
|
||||
self.redis_batch_writing_buffer: list[tuple[str, object]] = []
|
||||
if redis_flush_size is None:
|
||||
self.redis_flush_size: int = 100
|
||||
else:
|
||||
|
|
@ -1568,6 +1573,46 @@ class RedisCache(BaseCache):
|
|||
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
|
||||
return _LUA_COUNT.validate_python(count)
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_increment_if_present(
|
||||
self, key: str, value: float, parent_otel_span: Span | None = None
|
||||
) -> float | None:
|
||||
_redis_client: Final = self._async_commands()
|
||||
namespaced_key: Final = self.check_and_fix_namespace(key=key)
|
||||
start_time: Final = time.time()
|
||||
try:
|
||||
total: Final = _LUA_OPTIONAL_TOTAL.validate_python(
|
||||
await _redis_client.eval(_INCREMENT_IF_PRESENT_LUA, 1, namespaced_key, value)
|
||||
)
|
||||
except Exception as e:
|
||||
failed_at: Final = time.time()
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=failed_at - start_time,
|
||||
error=e,
|
||||
call_type="async_increment_if_present",
|
||||
caller=_get_call_stack_info(),
|
||||
start_time=start_time,
|
||||
end_time=failed_at,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
raise
|
||||
end_time: Final = time.time()
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=end_time - start_time,
|
||||
call_type="async_increment_if_present",
|
||||
caller=_get_call_stack_info(),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
return total
|
||||
|
||||
async def flush_cache_buffer(self):
|
||||
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
|
||||
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -256,6 +258,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(
|
||||
|
|
@ -1645,9 +1650,11 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
|
|||
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
|
||||
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"
|
||||
|
|
@ -1811,6 +1818,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"
|
||||
|
|
@ -1854,6 +1862,9 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = get_env_float(
|
|||
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", 300.0)
|
||||
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", 30.0)
|
||||
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = get_env_int("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", 100000)
|
||||
SPEND_LOG_SIZE_CHECK_JOB_NAME: Final = "spend_log_size_check"
|
||||
SPEND_LOG_SIZE_CHECK_INTERVAL: Final = os.getenv("SPEND_LOG_SIZE_CHECK_INTERVAL", "1d")
|
||||
SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES: Final = get_env_int("SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES", 10)
|
||||
SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = get_env_float(
|
||||
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", 5.0
|
||||
)
|
||||
|
|
@ -1975,6 +1986,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
|
|||
"public_model_groups",
|
||||
"public_model_groups_links",
|
||||
"cost_discount_config",
|
||||
"apply_search_free_tier",
|
||||
"cost_margin_config",
|
||||
"block_requests_for_models_without_pricing",
|
||||
"budget_exceeded_throttle_percentage",
|
||||
|
|
@ -2315,6 +2327,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,
|
||||
|
|
@ -1642,10 +1664,10 @@ def completion_cost(
|
|||
*(((model, request_pricing_provider),) if model is not None else ()),
|
||||
)
|
||||
potential_model_names: Final = (
|
||||
raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates)
|
||||
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)
|
||||
|
|
@ -1666,6 +1688,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 +1736,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 +2024,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 +2128,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 +2215,8 @@ def response_cost_calculator(
|
|||
| Response
|
||||
| SearchResponse
|
||||
| DecisionsResponse
|
||||
| OpenAIDecisionResponse,
|
||||
| OpenAIDecisionResponse
|
||||
| ExtractionResponse,
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
call_type: Literal[
|
||||
|
|
@ -2215,6 +2242,8 @@ def response_cost_calculator(
|
|||
"adecisions",
|
||||
"systemone",
|
||||
"asystemone",
|
||||
"extraction",
|
||||
"aextraction",
|
||||
],
|
||||
optional_params: dict,
|
||||
cache_hit: bool | None = None,
|
||||
|
|
@ -2225,6 +2254,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 +2295,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,
|
||||
|
|
@ -427,7 +429,7 @@ def _jsonable_fallback(value: object) -> object:
|
|||
|
||||
def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage:
|
||||
caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python(
|
||||
to_jsonable_python(caller_request, fallback=_jsonable_fallback)
|
||||
to_jsonable_python(caller_request, fallback=_jsonable_fallback, bytes_mode="base64")
|
||||
)
|
||||
sent_fields: Final = json.dumps(
|
||||
dict((name, value) for name, value in caller_fields.items() if value not in (None, [], {}))
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
@ -71,7 +71,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
|
|||
):
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
self.request_body = request_body
|
||||
self.start_time = datetime.now()
|
||||
self.start_time = litellm_logging_obj.start_time
|
||||
self.collected_chunks: list[bytes] = []
|
||||
self.model = model
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
|
|
@ -192,6 +192,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
|
|||
async def __anext__(self):
|
||||
try:
|
||||
chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator)
|
||||
if self.litellm_logging_obj.completion_start_time is None:
|
||||
self.litellm_logging_obj.update_completion_start_time(completion_start_time=datetime.now())
|
||||
self.collected_chunks.append(chunk)
|
||||
return chunk
|
||||
except StopAsyncIteration:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Agent harnesses: run Claude Code, Codex, OpenCode, Deep Agents or Tool Loop on any LiteLLM model.
|
||||
"""Agent harnesses: run Claude Code, Codex, OpenCode, pi, Deep Agents or Tool Loop on any LiteLLM model.
|
||||
|
||||
The entrypoints live on the top-level package:
|
||||
|
||||
|
|
@ -30,6 +30,7 @@ from litellm.harness.options import (
|
|||
CodexOptions,
|
||||
DeepAgentsOptions,
|
||||
OpenCodeOptions,
|
||||
PiOptions,
|
||||
ToolLoopOptions,
|
||||
)
|
||||
from litellm.harness.runtime import (
|
||||
|
|
@ -78,6 +79,7 @@ __all__ = (
|
|||
"OpenCodeOptions",
|
||||
"OptionsMismatch",
|
||||
"OutputInvalid",
|
||||
"PiOptions",
|
||||
"Reasoning",
|
||||
"Result",
|
||||
"SandboxError",
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol
|
|||
|
||||
import httpx
|
||||
import openai
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -172,16 +172,57 @@ def usage_from_mapping(usage: object) -> tuple[int, int]:
|
|||
return _as_int(input_tokens), _as_int(output_tokens)
|
||||
|
||||
|
||||
def usage_from_body(body: object) -> tuple[int, int]:
|
||||
"""Usage from a non-streaming JSON response body."""
|
||||
if not isinstance(body, Mapping):
|
||||
return 0, 0
|
||||
if isinstance(body.get("usage"), Mapping):
|
||||
return usage_from_mapping(body["usage"])
|
||||
response = body.get("response")
|
||||
if isinstance(response, Mapping):
|
||||
return usage_from_mapping(response.get("usage"))
|
||||
return 0, 0
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptCache:
|
||||
"""Prompt tokens read from or written to the provider's cache. Anthropic's usage counts them outside input_tokens,
|
||||
while OpenAI's prompt_tokens and the Responses API's input_tokens already include them."""
|
||||
|
||||
read_tokens: int = 0
|
||||
creation_tokens: int = 0
|
||||
outside_input: bool = False
|
||||
|
||||
def prompt_tokens(self, input_tokens: int) -> int:
|
||||
if self.outside_input:
|
||||
return input_tokens + self.read_tokens + self.creation_tokens
|
||||
return input_tokens
|
||||
|
||||
|
||||
NO_PROMPT_CACHE: Final = PromptCache()
|
||||
|
||||
|
||||
def json_object(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _JSON_OBJECT_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def prompt_cache_from_mapping(usage: object) -> PromptCache:
|
||||
"""Cache token counts from a usage dict using OpenAI, Anthropic, or Responses field names."""
|
||||
fields: Final = json_object(usage)
|
||||
details: Final = json_object(fields.get("prompt_tokens_details", fields.get("input_tokens_details")))
|
||||
return PromptCache(
|
||||
read_tokens=_as_int(fields.get("cache_read_input_tokens") or details.get("cached_tokens")),
|
||||
creation_tokens=_as_int(
|
||||
fields.get("cache_creation_input_tokens")
|
||||
or details.get("cache_write_tokens")
|
||||
or details.get("cache_creation_tokens")
|
||||
),
|
||||
outside_input="prompt_tokens" not in fields and "input_tokens_details" not in fields,
|
||||
)
|
||||
|
||||
|
||||
def usage_mapping_of_body(body: object) -> Mapping[str, object]:
|
||||
fields: Final = json_object(body)
|
||||
usage: Final = json_object(fields.get("usage"))
|
||||
return usage if usage else json_object(json_object(fields.get("response")).get("usage"))
|
||||
|
||||
|
||||
def usage_mapping_of_json(raw: bytes) -> Mapping[str, object]:
|
||||
try:
|
||||
return usage_mapping_of_body(_JSON_VALUE_ADAPTER.validate_python(json.loads(raw)))
|
||||
except ValueError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
class SSEUsageParser:
|
||||
|
|
@ -190,6 +231,7 @@ class SSEUsageParser:
|
|||
def __init__(self) -> None:
|
||||
self.input_tokens = 0
|
||||
self.output_tokens = 0
|
||||
self.prompt_cache = NO_PROMPT_CACHE
|
||||
self._buffer = b""
|
||||
|
||||
def feed(self, chunk: bytes) -> None:
|
||||
|
|
@ -226,21 +268,23 @@ class SSEUsageParser:
|
|||
elif event_type == "response.completed":
|
||||
self._absorb_response_completed(event)
|
||||
elif isinstance(event.get("usage"), Mapping):
|
||||
self._set(*usage_from_mapping(event["usage"]))
|
||||
self._absorb_usage(event["usage"])
|
||||
|
||||
def _absorb_message_start(self, event: Mapping[str, object]) -> None:
|
||||
message = event.get("message")
|
||||
if isinstance(message, Mapping):
|
||||
self._set(*usage_from_mapping(message.get("usage")))
|
||||
self._absorb_usage(json_object(event.get("message")).get("usage"))
|
||||
|
||||
def _absorb_message_delta(self, event: Mapping[str, object]) -> None:
|
||||
# message_delta output_tokens is cumulative for the whole message.
|
||||
self._set(*usage_from_mapping(event.get("usage")))
|
||||
self._absorb_usage(event.get("usage"))
|
||||
|
||||
def _absorb_response_completed(self, event: Mapping[str, object]) -> None:
|
||||
response = event.get("response")
|
||||
if isinstance(response, Mapping):
|
||||
self._set(*usage_from_mapping(response.get("usage")))
|
||||
self._absorb_usage(json_object(event.get("response")).get("usage"))
|
||||
|
||||
def _absorb_usage(self, usage: object) -> None:
|
||||
self._set(*usage_from_mapping(usage))
|
||||
prompt_cache: Final = prompt_cache_from_mapping(usage)
|
||||
if prompt_cache.read_tokens or prompt_cache.creation_tokens:
|
||||
self.prompt_cache = prompt_cache
|
||||
|
||||
def _set(self, input_tokens: int, output_tokens: int) -> None:
|
||||
if input_tokens:
|
||||
|
|
@ -249,13 +293,19 @@ class SSEUsageParser:
|
|||
self.output_tokens = output_tokens
|
||||
|
||||
|
||||
def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
|
||||
"""Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
|
||||
def compute_cost(
|
||||
model: str | None, input_tokens: int, output_tokens: int, prompt_cache: PromptCache = NO_PROMPT_CACHE
|
||||
) -> float:
|
||||
"""Cost from LiteLLM's price map, cached prompt tokens at their cache prices. Never raises; unknown models cost 0.0."""
|
||||
if not model or not (input_tokens or output_tokens):
|
||||
return 0.0
|
||||
try:
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model=model, prompt_tokens=input_tokens, completion_tokens=output_tokens
|
||||
model=model,
|
||||
prompt_tokens=prompt_cache.prompt_tokens(input_tokens),
|
||||
completion_tokens=output_tokens,
|
||||
cache_read_input_tokens=prompt_cache.read_tokens,
|
||||
cache_creation_input_tokens=prompt_cache.creation_tokens,
|
||||
)
|
||||
return float(prompt_cost) + float(completion_cost)
|
||||
except Exception: # accounting must never break a call; the price-map lookup raises bare Exception
|
||||
|
|
@ -565,9 +615,10 @@ class ModelEndpoint:
|
|||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cost: float | None,
|
||||
prompt_cache: PromptCache,
|
||||
) -> None:
|
||||
if cost is None:
|
||||
cost = compute_cost(model, input_tokens, output_tokens)
|
||||
cost = compute_cost(model, input_tokens, output_tokens, prompt_cache)
|
||||
self.usage.add(input_tokens, output_tokens, cost)
|
||||
|
||||
async def _forward(self, request: Request, route: str, body: Mapping[str, object]) -> Response:
|
||||
|
|
@ -617,14 +668,12 @@ class ModelEndpoint:
|
|||
) -> None:
|
||||
if is_sse:
|
||||
parser.close()
|
||||
tokens = (parser.input_tokens, parser.output_tokens)
|
||||
else:
|
||||
try:
|
||||
body: Final[object] = _JSON_VALUE_ADAPTER.validate_python(json.loads(collected))
|
||||
tokens = usage_from_body(body)
|
||||
except ValueError:
|
||||
tokens = (0, 0)
|
||||
self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
|
||||
self._record(
|
||||
model, parser.input_tokens, parser.output_tokens, header_cost(upstream.headers), parser.prompt_cache
|
||||
)
|
||||
return
|
||||
usage: Final = usage_mapping_of_json(collected)
|
||||
self._record(model, *usage_from_mapping(usage), header_cost(upstream.headers), prompt_cache_from_mapping(usage))
|
||||
|
||||
def _sdk_kwargs(
|
||||
self, body: Mapping[str, object]
|
||||
|
|
@ -667,9 +716,9 @@ class ModelEndpoint:
|
|||
return self._responses.StreamingResponse(
|
||||
self._sdk_stream(route, response, model), media_type=SSE_MEDIA_TYPE
|
||||
)
|
||||
data = to_jsonable(response)
|
||||
input_tokens, output_tokens = usage_from_body(data)
|
||||
self._record(model, input_tokens, output_tokens, hidden_cost(response))
|
||||
data: Final = to_jsonable(response)
|
||||
usage: Final = usage_mapping_of_body(data)
|
||||
self._record(model, *usage_from_mapping(usage), hidden_cost(response), prompt_cache_from_mapping(usage))
|
||||
return self._json(data)
|
||||
|
||||
async def _sdk_stream(self, route: str, iterator: AsyncIterable[object], model: str | None) -> AsyncIterator[bytes]:
|
||||
|
|
@ -688,4 +737,4 @@ class ModelEndpoint:
|
|||
yield f"event: error\ndata: {json.dumps(error_body(e, message))}\n\n".encode()
|
||||
finally:
|
||||
parser.close()
|
||||
self._record(model, parser.input_tokens, parser.output_tokens, None)
|
||||
self._record(model, parser.input_tokens, parser.output_tokens, None, parser.prompt_cache)
|
||||
|
|
|
|||
|
|
@ -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 "")),
|
||||
|
|
|
|||
|
|
@ -4,28 +4,37 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
PiThinkingLevel: TypeAlias = Literal["off", "minimal", "low", "medium", "high", "xhigh", "max"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClaudeCodeOptions:
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodexOptions:
|
||||
reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None
|
||||
web_search: bool = False
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OpenCodeOptions:
|
||||
agent: str = "build"
|
||||
config: Mapping[str, object] = field(default_factory=dict)
|
||||
env: Mapping[str, str] = field(default_factory=dict)
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PiOptions:
|
||||
thinking: PiThinkingLevel | None = None
|
||||
config: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
env: Mapping[str, str] = field(default_factory=dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -36,7 +45,9 @@ class DeepAgentsOptions:
|
|||
|
||||
@dataclass(frozen=True)
|
||||
class ToolLoopOptions:
|
||||
completion_kwargs: Mapping[str, object] = field(default_factory=dict)
|
||||
completion_kwargs: Mapping[str, object] = field(default_factory=dict[str, object])
|
||||
|
||||
|
||||
HarnessOptions = ClaudeCodeOptions | CodexOptions | OpenCodeOptions | DeepAgentsOptions | ToolLoopOptions
|
||||
HarnessOptions: TypeAlias = (
|
||||
ClaudeCodeOptions | CodexOptions | OpenCodeOptions | PiOptions | DeepAgentsOptions | ToolLoopOptions
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -411,7 +411,7 @@ def agent(
|
|||
options: HarnessOptions | None = None,
|
||||
install: bool = False,
|
||||
) -> Result | EventStream:
|
||||
"""Run an agent harness (Claude Code, Codex, OpenCode, Deep Agents, Tool Loop) on one prompt.
|
||||
"""Run an agent harness (Claude Code, Codex, OpenCode, pi, Deep Agents, Tool Loop) on one prompt.
|
||||
|
||||
Returns a Result. With stream=True it returns an iterator of events instead.
|
||||
Prefix the model with `litellm_proxy/` to route every model call through your
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ class Harness(Enum):
|
|||
CLAUDE_CODE = "claude_code"
|
||||
CODEX = "codex"
|
||||
OPENCODE = "opencode"
|
||||
PI = "pi"
|
||||
DEEPAGENTS = "deepagents"
|
||||
TOOL_LOOP = "tool_loop"
|
||||
|
||||
|
|
|
|||
|
|
@ -342,6 +342,20 @@ EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints"
|
|||
CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider"
|
||||
|
||||
|
||||
def _accepts_cache_control(part: object) -> bool:
|
||||
match part:
|
||||
case {"type": "thinking" | "redacted_thinking"}:
|
||||
return False
|
||||
case dict():
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _last_index_accepting_cache_control(parts: Sequence[object]) -> int | None:
|
||||
return next((index for index in reversed(range(len(parts))) if _accepts_cache_control(parts[index])), None)
|
||||
|
||||
|
||||
class AnthropicCacheControlHook(CustomPromptManagement):
|
||||
@staticmethod
|
||||
def _request_value(request_kwargs: object, key: str) -> object:
|
||||
|
|
@ -740,10 +754,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
# 1. if string, insert cache control in the message
|
||||
if isinstance(message_content, str):
|
||||
message["cache_control"] = control
|
||||
# 2. list of objects - only apply to last item per Anthropic spec
|
||||
# 2. list of objects - only apply to the last item Claude accepts a breakpoint on
|
||||
elif isinstance(message_content, list):
|
||||
if len(message_content) > 0 and isinstance(message_content[-1], dict):
|
||||
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
|
||||
markable_index: Final = _last_index_accepting_cache_control(message_content)
|
||||
if markable_index is not None:
|
||||
message_content[markable_index]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
|
||||
return message
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -884,7 +884,7 @@ def _set_session_and_user_attrs(span: "Span", kwargs: dict, standard_logging_pay
|
|||
safe_set_attribute(span, "litellm.key_alias", str(key_alias))
|
||||
|
||||
|
||||
def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
|
||||
def _set_response_cost_attr(span: "Span", standard_logging_payload: Mapping[str, object] | None) -> None:
|
||||
"""Emit cost attributes from the StandardLoggingPayload when present.
|
||||
|
||||
Uses the OpenInference `llm.cost.total` key so Arize / Phoenix can
|
||||
|
|
@ -892,15 +892,22 @@ def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
|
|||
single total in `StandardLoggingPayload.response_cost`, so we cannot
|
||||
split it into prompt/completion. We also keep the legacy
|
||||
`llm.response.cost` key for back-compat with any consumer querying it.
|
||||
When the price lookup failed, the call's total is unknown, so no cost is
|
||||
emitted and the backend does not record the call as free. A guardrail fee
|
||||
in `response_cost` stays on the guardrail span, as in OTel v2.
|
||||
"""
|
||||
if not isinstance(standard_logging_payload, dict):
|
||||
return
|
||||
cost: Final = standard_logging_payload.get("response_cost")
|
||||
if cost is None:
|
||||
cost_lookup_failed: Final = (
|
||||
standard_logging_payload.get("response_cost_failure_debug_info") is not None
|
||||
and standard_logging_payload.get("cache_hit") is not True
|
||||
)
|
||||
if not isinstance(cost, (int, float, str)) or cost_lookup_failed:
|
||||
return
|
||||
try:
|
||||
cost_value: Final = float(cost)
|
||||
except (TypeError, ValueError):
|
||||
except ValueError:
|
||||
return
|
||||
safe_set_attribute(span, "llm.cost.total", cost_value)
|
||||
safe_set_attribute(span, "llm.response.cost", cost_value)
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import os
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
|
@ -54,6 +54,7 @@ from litellm.types.utils import (
|
|||
PROMPT_CARRYING_GUARDRAIL_FIELDS,
|
||||
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
|
||||
CallTypes,
|
||||
StandardLoggingMetadata,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
|
|
@ -93,8 +94,12 @@ _ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "ro
|
|||
_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS)
|
||||
|
||||
|
||||
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]:
|
||||
metadata: Final = standard_logging_payload.get("metadata")
|
||||
def _logging_metadata_of(standard_logging_payload: StandardLoggingPayload) -> StandardLoggingMetadata | None:
|
||||
return standard_logging_payload.get("metadata")
|
||||
|
||||
|
||||
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, object]:
|
||||
metadata: Final = _logging_metadata_of(standard_logging_payload)
|
||||
return metadata or _EMPTY_MAPPING
|
||||
|
||||
|
||||
|
|
@ -1057,14 +1062,23 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
# Get budget information from metadata
|
||||
metadata: Final = _metadata_of(standard_logging_payload)
|
||||
logging_metadata: Final = _logging_metadata_of(standard_logging_payload)
|
||||
|
||||
# API key max budget
|
||||
user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget")
|
||||
user_api_key_max_budget: Final = (
|
||||
logging_metadata["user_api_key_max_budget"]
|
||||
if logging_metadata and "user_api_key_max_budget" in logging_metadata
|
||||
else None
|
||||
)
|
||||
if user_api_key_max_budget is not None:
|
||||
spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
|
||||
|
||||
# API key spend
|
||||
user_api_key_spend: Final = metadata.get("user_api_key_spend")
|
||||
user_api_key_spend: Final = (
|
||||
logging_metadata["user_api_key_spend"]
|
||||
if logging_metadata and "user_api_key_spend" in logging_metadata
|
||||
else None
|
||||
)
|
||||
if user_api_key_spend is not None:
|
||||
try:
|
||||
spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
|
|
@ -9,6 +9,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.arize import _utils
|
||||
from litellm.integrations.langfuse.langfuse_otel_attributes import (
|
||||
LangfuseLLMObsOTELAttributes,
|
||||
serialize_langfuse_datetime,
|
||||
)
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
|
|
@ -55,7 +56,7 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj)
|
||||
|
||||
@staticmethod
|
||||
def _extract_langfuse_metadata(kwargs: dict) -> dict:
|
||||
def _extract_langfuse_metadata(kwargs: Mapping[str, Any]) -> dict:
|
||||
"""
|
||||
Extracts Langfuse metadata from the standard LiteLLM kwargs structure.
|
||||
|
||||
|
|
@ -221,7 +222,7 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _set_langfuse_specific_attributes(span: Span, kwargs, response_obj):
|
||||
def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj):
|
||||
"""
|
||||
Sets Langfuse specific metadata attributes onto the OTEL span.
|
||||
|
||||
|
|
@ -255,6 +256,14 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
safe_dumps(messages),
|
||||
)
|
||||
|
||||
completion_start_time: Final = kwargs.get("completion_start_time")
|
||||
if isinstance(completion_start_time, datetime):
|
||||
safe_set_attribute(
|
||||
span,
|
||||
LangfuseSpanAttributes.OBSERVATION_COMPLETION_START_TIME.value,
|
||||
serialize_langfuse_datetime(completion_start_time),
|
||||
)
|
||||
|
||||
LangfuseOtelLogger._set_observation_output(span=span, response_obj=response_obj)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -29,6 +30,10 @@ if TYPE_CHECKING:
|
|||
from opentelemetry.trace import Span
|
||||
|
||||
|
||||
def serialize_langfuse_datetime(value: datetime) -> str:
|
||||
return json.dumps(value.astimezone().isoformat())
|
||||
|
||||
|
||||
def get_output_content_by_type(
|
||||
response_obj: None
|
||||
| dict
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from pydantic import ConfigDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.langfuse.langfuse import PROMPT_CACHE_TTL_ENV, parse_langfuse_debug, whole_number
|
||||
from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -168,7 +169,7 @@ def _string_or_none(value: object) -> str | None:
|
|||
def _serialize_datetime(value: object) -> str | None:
|
||||
"""A datetime the way the SDK's ``EventSerializer`` sends one: a JSON string, naive values read as local time."""
|
||||
if isinstance(value, datetime):
|
||||
return safe_dumps(value.astimezone().isoformat())
|
||||
return serialize_langfuse_datetime(value)
|
||||
return _serialize(value)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ way (`pod_lock`, `budget_reset`, `spend_queue`, `health_check`, `scheduler_queue
|
|||
rather than a bare `redis.set`; `tests/unit/test_internal_context.py` scans every
|
||||
module under `litellm/` and `enterprise/` that calls a shared cache or declares a
|
||||
read or write on the request batch (`reserve_redis_batch_reads`,
|
||||
`declare_batch_get`, `batch.mget`, `batch.set`, `batch.script`) and fails when
|
||||
`batch.mget`, `batch.set`, `batch.script`) and fails when
|
||||
one has no declared family, with the process-local `InMemoryCache` callers listed
|
||||
as the only exemptions. A batch op carries the family that was active when it was
|
||||
declared, so the routing prefetch armed before deployment selection
|
||||
|
|
|
|||
|
|
@ -594,6 +594,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
payload,
|
||||
capture_content=self._capture_span_content(),
|
||||
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
|
||||
completion_start_seconds=call.completion_start_seconds,
|
||||
request_route=request_root_http_route(),
|
||||
request_purpose=call.purpose,
|
||||
trace=call.trace,
|
||||
|
|
|
|||
|
|
@ -12,8 +12,10 @@ the JSON-serialized payloads. ``trace_attributes`` maps the caller's trace contr
|
|||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime
|
||||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import (
|
||||
collect,
|
||||
|
|
@ -45,6 +47,11 @@ class LangfuseMapper:
|
|||
"langfuse.observation.id": lambda d: d.identity.call_id or None,
|
||||
"langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None,
|
||||
"langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None,
|
||||
"langfuse.observation.completion_start_time": lambda d: (
|
||||
serialize_langfuse_datetime(datetime.fromtimestamp(d.completion_start_seconds, tz=timezone.utc))
|
||||
if d.completion_start_seconds is not None
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
# Sub-tables folded into their respective JSON blobs.
|
||||
|
|
|
|||
|
|
@ -258,6 +258,7 @@ class LLMCallEvent:
|
|||
# needs to be reasonable for a span that never gets closed (a leak).
|
||||
provisional_span_name: str
|
||||
time_to_first_chunk_seconds: float | None
|
||||
completion_start_seconds: float | None
|
||||
trace: TraceControls
|
||||
session_id: str | None
|
||||
|
||||
|
|
@ -280,6 +281,7 @@ class LLMCallEvent:
|
|||
cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown
|
||||
provisional_span_name=f"{operation.value} {model}".strip(),
|
||||
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
|
||||
completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")),
|
||||
trace=trace,
|
||||
session_id=caller_session_id(kwargs, trace),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -429,6 +429,7 @@ class LLMCallSpanData:
|
|||
choices_out: tuple[Mapping[str, object], ...] = ()
|
||||
system_fingerprint: str | None = None
|
||||
time_to_first_chunk_seconds: float | None = None
|
||||
completion_start_seconds: float | None = None
|
||||
# The requested output modality, set only on the routes that pin one (image
|
||||
# generation, speech, transcription, OCR), and the litellm route itself, which
|
||||
# keeps routes the convention folds into one operation distinguishable.
|
||||
|
|
@ -450,6 +451,7 @@ class LLMCallSpanData:
|
|||
payload: StandardLoggingPayload,
|
||||
capture_content: bool = False,
|
||||
time_to_first_chunk_seconds: float | None = None,
|
||||
completion_start_seconds: float | None = None,
|
||||
request_route: str | None = None,
|
||||
request_purpose: str | None = None,
|
||||
trace: TraceControls | None = None,
|
||||
|
|
@ -481,6 +483,7 @@ class LLMCallSpanData:
|
|||
embedding_output: Final = (
|
||||
EmbeddingOutput.from_response(response) if operation is GenAIOperation.EMBEDDINGS else None
|
||||
)
|
||||
cost_lookup_failed: Final = payload.get("response_cost_failure_debug_info") is not None and not cache_hit
|
||||
return cls(
|
||||
operation=operation,
|
||||
provider=resolve_provider(as_str(payload.get("custom_llm_provider"))),
|
||||
|
|
@ -491,7 +494,7 @@ class LLMCallSpanData:
|
|||
usage=LLMUsage.from_standard_logging_payload(payload),
|
||||
finish_reasons=finish_reasons,
|
||||
error=_parse_error(payload),
|
||||
response_cost=as_float(payload.get("response_cost")),
|
||||
response_cost=None if cost_lookup_failed else as_float(payload.get("response_cost")),
|
||||
cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))),
|
||||
server=None if cache_hit else ServerInfo.from_api_base(context.api_base),
|
||||
identity=context.identity,
|
||||
|
|
@ -502,6 +505,7 @@ class LLMCallSpanData:
|
|||
choices_out=choices_out if capture_content else (),
|
||||
system_fingerprint=as_str(response.get("system_fingerprint")),
|
||||
time_to_first_chunk_seconds=time_to_first_chunk_seconds,
|
||||
completion_start_seconds=completion_start_seconds,
|
||||
output_type=resolve_output_type(call_type),
|
||||
call_type=call_type or None,
|
||||
request_route=request_route or context.identity.request_route,
|
||||
|
|
|
|||
|
|
@ -214,6 +214,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = {
|
|||
"increment_cache": "incr",
|
||||
"async_increment": "incr",
|
||||
"async_increment_pipeline": "incr",
|
||||
"async_increment_if_present": "incr",
|
||||
"delete_cache": "delete",
|
||||
"async_delete_cache": "delete",
|
||||
"async_rpush": "rpush",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
@ -48,7 +53,7 @@ class BaseInteractionsAPIStreamingIterator:
|
|||
self.finished = False
|
||||
self.interactions_api_config = interactions_api_config
|
||||
self.completed_response: InteractionsAPIStreamingResponse | None = None
|
||||
self.start_time = datetime.now()
|
||||
self.start_time = logging_obj.start_time
|
||||
|
||||
# set request kwargs
|
||||
self.litellm_metadata = litellm_metadata
|
||||
|
|
@ -83,10 +88,12 @@ 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):
|
||||
if self.logging_obj.completion_start_time is None:
|
||||
self.logging_obj.update_completion_start_time(completion_start_time=datetime.now())
|
||||
streaming_response: Final = self.interactions_api_config.transform_streaming_response(
|
||||
model=self.model,
|
||||
parsed_chunk=parsed_chunk,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import posixpath
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import accumulate, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, cast
|
||||
from urllib.parse import quote, unquote
|
||||
|
|
@ -20,15 +21,33 @@ MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
|
|||
_MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({}))
|
||||
|
||||
|
||||
MAX_FILE_ID_DECODE_PASSES: Final = 8
|
||||
|
||||
|
||||
def _unquote_once(current: str, _: str) -> str:
|
||||
return unquote(current)
|
||||
|
||||
|
||||
def _decodings(value: str) -> tuple[str, ...]:
|
||||
return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once))
|
||||
|
||||
|
||||
def is_managed_cloud_storage_uri(file_id: str) -> bool:
|
||||
"""
|
||||
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``).
|
||||
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``),
|
||||
up to ``MAX_FILE_ID_DECODE_PASSES`` layers of percent-encoding deep, or an id
|
||||
encoded deeper than that, which no client produces and which is treated as raw
|
||||
rather than decoded any further.
|
||||
|
||||
These are internal provider artifacts. On the multi-tenant proxy they must be
|
||||
retrieved through their managed unified file id so owner/team access is enforced;
|
||||
a raw URI supplied by a caller bypasses that check.
|
||||
"""
|
||||
return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES)
|
||||
if not isinstance(file_id, str):
|
||||
return False
|
||||
decodings: Final = _decodings(file_id)
|
||||
settled: Final = decodings[-1] == decodings[-2]
|
||||
return decodings[-1].startswith(MANAGED_CLOUD_STORAGE_SCHEMES) or not settled
|
||||
|
||||
|
||||
_SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+")
|
||||
|
|
|
|||
|
|
@ -520,7 +520,7 @@ def reconstruct_model_name(
|
|||
|
||||
# Helper functions used for OTEL logging
|
||||
def get_parent_otel_span_from_kwargs(
|
||||
kwargs: dict[str, object] | None = None,
|
||||
kwargs: Mapping[str, object] | None = None,
|
||||
) -> Span | None:
|
||||
try:
|
||||
if kwargs is None:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,6 +132,7 @@ 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,
|
||||
|
|
@ -337,6 +339,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")
|
||||
|
|
@ -794,6 +797,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
|
||||
|
|
@ -995,6 +999,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 = (
|
||||
|
|
@ -1939,6 +1958,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
used for consistent cost calculation across response headers + logging integrations.
|
||||
"""
|
||||
|
||||
self.model_call_details.pop("response_cost_failure_debug_information", None)
|
||||
|
||||
if cache_hit is None:
|
||||
cache_hit = self.model_call_details.get("cache_hit", False)
|
||||
|
||||
|
|
@ -1988,7 +2009,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
router_model_id = self.get_router_model_id()
|
||||
|
||||
## RESPONSE COST ##
|
||||
custom_pricing: Final = self._custom_pricing_for(priced_result)
|
||||
custom_pricing: Final = self.custom_pricing_for(priced_result)
|
||||
|
||||
prompt: Final = self._prompt_for_cost_calculation()
|
||||
|
||||
|
|
@ -2018,6 +2039,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")
|
||||
|
|
@ -2099,15 +2121,17 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
previous_cost_breakdown: Final = self.cost_breakdown
|
||||
previous_billed_token_rates: Final = self.billed_token_rates
|
||||
debug_key: Final = "response_cost_failure_debug_information"
|
||||
had_failure_debug_info: Final = debug_key in self.model_call_details
|
||||
previous_failure_debug_info: Final = self.model_call_details.get(debug_key)
|
||||
try:
|
||||
cost: Final = self.response_cost_calculator(result=result, cache_hit=False, ignore_stamped_cost=True)
|
||||
return UncachedPrice(cost=cost, cost_breakdown=self.cost_breakdown)
|
||||
finally:
|
||||
self.cost_breakdown = previous_cost_breakdown
|
||||
self.billed_token_rates = previous_billed_token_rates
|
||||
if not had_failure_debug_info:
|
||||
if previous_failure_debug_info is None:
|
||||
self.model_call_details.pop(debug_key, None)
|
||||
else:
|
||||
self.model_call_details[debug_key] = previous_failure_debug_info
|
||||
|
||||
def _record_zero_cost_diagnostic(
|
||||
self,
|
||||
|
|
@ -2160,8 +2184,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
pricing: Final = pricing_entry_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=result,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
custom_pricing=self._custom_pricing_for(result),
|
||||
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(),
|
||||
region_name=_resolve_mantle_region_for_cost(
|
||||
|
|
@ -2186,20 +2210,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:
|
||||
"""
|
||||
|
|
@ -2443,7 +2470,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,
|
||||
|
|
@ -2458,6 +2485,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
|
||||
|
|
@ -2886,6 +2923,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
|
||||
|
|
@ -3485,8 +3523,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,
|
||||
)
|
||||
|
|
@ -3496,8 +3536,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,
|
||||
)
|
||||
|
|
@ -3840,7 +3885,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
|
||||
)
|
||||
##################################
|
||||
|
|
@ -3851,7 +3896,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,
|
||||
|
|
@ -3867,14 +3912,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,
|
||||
)
|
||||
|
|
@ -4050,15 +4095,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()
|
||||
|
|
@ -4069,6 +4120,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
|
||||
|
|
@ -4161,6 +4213,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,
|
||||
|
|
@ -4523,6 +4576,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.
|
||||
|
|
@ -7191,13 +7250,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 = use_custom_pricing_for_model(litellm_params=litellm_params)
|
||||
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)
|
||||
|
||||
|
|
|
|||
|
|
@ -624,7 +624,7 @@ def update_messages_with_model_file_ids(
|
|||
def update_responses_input_with_model_file_ids(
|
||||
input: object,
|
||||
model_id: str | None = None,
|
||||
model_file_id_mapping: dict[str, dict[str, str]] | None = None,
|
||||
model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None,
|
||||
) -> object:
|
||||
"""
|
||||
Updates responses API input with provider-specific file IDs.
|
||||
|
|
@ -665,7 +665,7 @@ def update_responses_input_with_model_file_ids(
|
|||
for content_item in content:
|
||||
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
|
||||
file_id = content_item.get("file_id")
|
||||
if file_id:
|
||||
if isinstance(file_id, str) and file_id:
|
||||
provider_file_id = file_id # Default to original
|
||||
|
||||
# Check if we have a mapping for this file ID
|
||||
|
|
@ -761,7 +761,7 @@ def _decode_vector_store_ids_in_tools(
|
|||
def update_responses_tools_with_model_file_ids(
|
||||
tools: list[dict[str, object]] | None,
|
||||
model_id: str | None = None,
|
||||
model_file_id_mapping: dict[str, dict[str, str]] | None = None,
|
||||
model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None,
|
||||
) -> list[dict[str, object]] | None:
|
||||
"""
|
||||
Updates responses API tools with provider-specific file IDs.
|
||||
|
|
@ -2325,6 +2325,16 @@ def is_unsignable_thinking_block(block: object) -> bool:
|
|||
return not (isinstance(signature, str) and len(signature) > 0)
|
||||
|
||||
|
||||
def has_blank_redacted_data(block: object) -> bool:
|
||||
match block:
|
||||
case {"type": "redacted_thinking", "data": str() as data}:
|
||||
return not data.strip()
|
||||
case {"type": "redacted_thinking"}:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def native_redacted_thinking_data(data: str) -> str:
|
||||
"""Claude's own `data` string for a redacted_thinking block.
|
||||
|
||||
|
|
|
|||
|
|
@ -6,10 +6,10 @@ import json
|
|||
import mimetypes
|
||||
import re
|
||||
import xml.etree.ElementTree as ET
|
||||
from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload
|
||||
from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload
|
||||
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.types.utils import GenericImageParsingChunk
|
|||
from .common_utils import (
|
||||
converse_redacted_content,
|
||||
convert_content_list_to_str,
|
||||
has_blank_redacted_data,
|
||||
infer_content_type_from_url_and_content,
|
||||
is_encrypted_reasoning_block,
|
||||
is_non_content_values_set,
|
||||
|
|
@ -2550,22 +2551,125 @@ def _pair_tool_calls_with_results(messages: Sequence[AllMessageValues]) -> tuple
|
|||
|
||||
def _replayable_thinking_blocks(
|
||||
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
|
||||
*,
|
||||
drop_params: bool = False,
|
||||
) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]:
|
||||
keys_to_drop: Final = _UNSUPPORTED_THINKING_KEYS if drop_params else _STREAM_ONLY_THINKING_KEYS
|
||||
return [
|
||||
_without_stream_index(_with_native_redacted_data(block))
|
||||
_without_keys(_with_native_redacted_data(block), keys_to_drop)
|
||||
for block in thinking_blocks
|
||||
if not is_unsignable_thinking_block(block)
|
||||
if not (is_unsignable_thinking_block(block) or (drop_params and has_blank_redacted_data(block)))
|
||||
]
|
||||
|
||||
|
||||
def _without_stream_index(
|
||||
def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None:
|
||||
match item:
|
||||
case {"type": "thinking" | "redacted_thinking"}:
|
||||
return cast( # cast-ok: the mapping pattern matched a thinking block's type tag
|
||||
ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item
|
||||
)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _thinking_identity(item: object) -> tuple[str | None, ...] | None:
|
||||
match item:
|
||||
case {"type": "redacted_thinking", "data": str() as data}:
|
||||
return ("redacted_thinking", native_redacted_thinking_data(data))
|
||||
case {"type": "thinking", "signature": str() as signature} if signature:
|
||||
return ("thinking", signature)
|
||||
case {"type": str() as kind} if kind in ("thinking", "redacted_thinking"):
|
||||
fields: Final = cast(Mapping[str, object], item) # cast-ok: the mapping pattern matched
|
||||
return _unsigned_thinking_identity(kind, fields.get("thinking"), fields.get("signature"))
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _unsigned_thinking_identity(kind: str, thinking: object, signature: object) -> tuple[str | None, ...] | None:
|
||||
match thinking, signature:
|
||||
case (str() | None as text, str() | None as signed):
|
||||
return (kind, text, signed)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
_Replayed: Final = TypeVar("_Replayed")
|
||||
|
||||
|
||||
def _with_missing_thinking_blocks(
|
||||
content: Iterable[_Replayed], thinking_blocks: Sequence[_Replayed]
|
||||
) -> tuple[_Replayed, ...]:
|
||||
content_items: Final = tuple(content)
|
||||
block_identities: Final = tuple(map(_thinking_identity, thinking_blocks))
|
||||
content_positions: Final[Mapping[tuple[object, ...] | None, int]] = MappingProxyType(
|
||||
{
|
||||
identity: index
|
||||
for index, identity in reversed(tuple(enumerate(map(_thinking_identity, content_items))))
|
||||
if identity is not None
|
||||
}
|
||||
)
|
||||
anchors: Final = tuple(index for index, identity in enumerate(block_identities) if identity in content_positions)
|
||||
if not anchors:
|
||||
return (*thinking_blocks, *content_items)
|
||||
missing_before: Final = MappingProxyType(
|
||||
{
|
||||
content_positions[block_identities[end]]: thinking_blocks[start + 1 : end]
|
||||
for start, end in zip((-1, *anchors), anchors)
|
||||
}
|
||||
)
|
||||
missing_after: Final = MappingProxyType(
|
||||
{content_positions[block_identities[anchors[-1]]]: thinking_blocks[anchors[-1] + 1 :]}
|
||||
)
|
||||
return tuple(_interleave(content_items, missing_before, missing_after))
|
||||
|
||||
|
||||
def _without_repeated_thinking(
|
||||
items: Sequence[_Replayed], identity: Callable[[_Replayed], tuple[object, ...] | None]
|
||||
) -> tuple[_Replayed, ...]:
|
||||
identities: Final = tuple(map(identity, items))
|
||||
first_positions: Final[Mapping[tuple[object, ...], int]] = MappingProxyType(
|
||||
{key: index for index, key in reversed(tuple(enumerate(identities))) if key is not None}
|
||||
)
|
||||
return tuple(
|
||||
item for index, (item, key) in enumerate(zip(items, identities)) if key is None or first_positions[key] == index
|
||||
)
|
||||
|
||||
|
||||
def _interleave(
|
||||
content: Sequence[_Replayed],
|
||||
before: Mapping[int, Sequence[_Replayed]],
|
||||
after: Mapping[int, Sequence[_Replayed]],
|
||||
) -> Iterator[_Replayed]:
|
||||
for index, item in enumerate(content):
|
||||
yield from before.get(index, ())
|
||||
yield item
|
||||
yield from after.get(index, ())
|
||||
|
||||
|
||||
def _thinking_in_replay_order(
|
||||
content: str | Iterable[object] | None, thinking_blocks: Sequence[object] | None
|
||||
) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]:
|
||||
content_items: Final = content if isinstance(content, list) else ()
|
||||
return tuple(
|
||||
block
|
||||
for item in _with_missing_thinking_blocks(content_items, thinking_blocks or ())
|
||||
if (block := _as_thinking_block(item)) is not None
|
||||
)
|
||||
|
||||
|
||||
_STREAM_ONLY_THINKING_KEYS: Final = frozenset({"index"})
|
||||
_UNSUPPORTED_THINKING_KEYS: Final = frozenset({"index", "cache_control"})
|
||||
|
||||
|
||||
def _without_keys(
|
||||
block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
|
||||
keys: frozenset[str],
|
||||
) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock:
|
||||
if "index" not in block:
|
||||
if keys.isdisjoint(block):
|
||||
return block
|
||||
return cast( # cast-ok: dropping one key keeps the block's TypedDict shape
|
||||
return cast( # cast-ok: dropping keys keeps the block's TypedDict shape
|
||||
ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
|
||||
{key: value for key, value in block.items() if key != "index"},
|
||||
{key: value for key, value in block.items() if key not in keys},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2629,6 +2733,7 @@ def anthropic_messages_pt(
|
|||
messages: list[AllMessageValues],
|
||||
model: str,
|
||||
llm_provider: str,
|
||||
drop_params: bool = False,
|
||||
) -> _AnthropicMessageList:
|
||||
"""
|
||||
format messages for anthropic
|
||||
|
|
@ -2816,7 +2921,9 @@ def anthropic_messages_pt(
|
|||
|
||||
_raw_thinking_blocks = assistant_content_block.get("thinking_blocks", None)
|
||||
thinking_blocks = (
|
||||
_replayable_thinking_blocks(_raw_thinking_blocks) if _raw_thinking_blocks is not None else None
|
||||
_replayable_thinking_blocks(_raw_thinking_blocks, drop_params=drop_params)
|
||||
if _raw_thinking_blocks is not None
|
||||
else None
|
||||
)
|
||||
|
||||
# Check if tool_calls contain server tool calls (web search, etc.)
|
||||
|
|
@ -2951,42 +3058,28 @@ def anthropic_messages_pt(
|
|||
# SEQUENTIAL MODE: No server tool calls, or no thinking blocks,
|
||||
# or content is a list. Use the original sequential approach.
|
||||
|
||||
# When content is a list, check if it already contains thinking
|
||||
# blocks inline. If so, skip prepending thinking_blocks to avoid
|
||||
# duplication and preserve the original interleaved order.
|
||||
# Fixes the gap where list-content messages bypass INTERLEAVED
|
||||
# MODE and still get thinking blocks prepended out of order.
|
||||
_content_is_list = "content" in assistant_content_block and isinstance(
|
||||
assistant_content_block["content"], list
|
||||
)
|
||||
_content_list = assistant_content_block.get("content") if _content_is_list else None
|
||||
_list_has_thinking = False
|
||||
if _content_is_list and _content_list is not None:
|
||||
for _item in _content_list:
|
||||
if isinstance(_item, dict) and _item.get("type") in (
|
||||
"thinking",
|
||||
"redacted_thinking",
|
||||
):
|
||||
_list_has_thinking = True
|
||||
break
|
||||
|
||||
if (
|
||||
thinking_blocks is not None and not _list_has_thinking
|
||||
thinking_blocks is not None and not _content_is_list
|
||||
): # IMPORTANT: ADD THIS FIRST, ELSE ANTHROPIC WILL RAISE AN ERROR
|
||||
assistant_content.extend(thinking_blocks)
|
||||
if _content_is_list and _content_list is not None:
|
||||
for m in _content_list:
|
||||
for m in _with_missing_thinking_blocks(_content_list, thinking_blocks or ()):
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
# handle thinking blocks
|
||||
text_block = cast(str, m.get("text", ""))
|
||||
if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m):
|
||||
anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast(
|
||||
ChatCompletionThinkingBlock, m
|
||||
if (thinking_block := _as_thinking_block(m)) is not None:
|
||||
assistant_content.extend(
|
||||
_replayable_thinking_blocks([thinking_block], drop_params=drop_params)
|
||||
)
|
||||
assistant_content.append(anthropic_message)
|
||||
continue
|
||||
text_block = cast(str, m.get("text", ""))
|
||||
# handle text
|
||||
elif (
|
||||
if (
|
||||
m.get("type", "") == "text" and len(text_block) > 0
|
||||
): # don't pass empty text blocks. anthropic api raises errors.
|
||||
anthropic_message = AnthropicMessagesTextParam(type="text", text=text_block)
|
||||
|
|
@ -3058,7 +3151,12 @@ def anthropic_messages_pt(
|
|||
msg_i += 1
|
||||
|
||||
if assistant_content:
|
||||
new_messages.append({"role": "assistant", "content": assistant_content})
|
||||
new_messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": list(_without_repeated_thinking(assistant_content, _thinking_identity)),
|
||||
}
|
||||
)
|
||||
|
||||
## MID-CONVERSATION SYSTEM MESSAGES (placement is the caller's job) ##
|
||||
while msg_i < len(messages) and messages[msg_i]["role"] == "system":
|
||||
|
|
@ -4190,6 +4288,16 @@ def _convert_to_bedrock_tool_call_result(
|
|||
return content_block
|
||||
|
||||
|
||||
def _reasoning_identity(block: BedrockContentBlock) -> tuple[object, ...] | None:
|
||||
match block:
|
||||
case {"reasoningContent": {"redactedContent": redacted}}:
|
||||
return ("redacted_thinking", redacted)
|
||||
case {"reasoningContent": {"reasoningText": {"signature": str() as signature}}} if signature:
|
||||
return ("thinking", signature)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _deduplicate_bedrock_content_blocks(
|
||||
blocks: list[BedrockContentBlock],
|
||||
block_key: str,
|
||||
|
|
@ -4673,6 +4781,7 @@ class BedrockConverseMessagesProcessor:
|
|||
assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
|
||||
deployment_model: str | None = None,
|
||||
base_model: str | None = None,
|
||||
drop_params: bool = False,
|
||||
) -> list[BedrockMessageBlock]:
|
||||
from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
|
||||
|
||||
|
|
@ -4826,33 +4935,21 @@ class BedrockConverseMessagesProcessor:
|
|||
)
|
||||
_assistant_content = assistant_message_block.get("content", None)
|
||||
thinking_blocks = assistant_message_block.get("thinking_blocks")
|
||||
|
||||
if thinking_blocks is not None:
|
||||
converted_thinking_blocks = (
|
||||
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
thinking_blocks, to_claude=to_claude
|
||||
)
|
||||
)
|
||||
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=converted_thinking_blocks,
|
||||
assistant_parts=assistant_content,
|
||||
converted_thinking_blocks = (
|
||||
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
_thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude
|
||||
)
|
||||
)
|
||||
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=converted_thinking_blocks,
|
||||
assistant_parts=assistant_content,
|
||||
)
|
||||
|
||||
if _assistant_content is not None and isinstance(_assistant_content, list):
|
||||
assistants_parts: list[BedrockContentBlock] = []
|
||||
for element in _assistant_content:
|
||||
if isinstance(element, dict):
|
||||
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
|
||||
thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
thinking_blocks=[element], to_claude=to_claude
|
||||
)
|
||||
assistants_parts = (
|
||||
BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=thinking_block,
|
||||
assistant_parts=assistants_parts,
|
||||
)
|
||||
)
|
||||
elif element["type"] == "text":
|
||||
if element["type"] == "text":
|
||||
# Skip completely empty strings to avoid blank content blocks
|
||||
if element.get("text", "").strip():
|
||||
assistants_part = BedrockContentBlock(text=element["text"])
|
||||
|
|
@ -4871,7 +4968,9 @@ class BedrockConverseMessagesProcessor:
|
|||
model=model,
|
||||
deployment_model=deployment_model,
|
||||
)
|
||||
if _cache_point_block is not None:
|
||||
if _cache_point_block is not None and not (
|
||||
drop_params and _as_thinking_block(element) is not None
|
||||
):
|
||||
assistants_parts.append(_cache_point_block)
|
||||
assistant_content.extend(assistants_parts)
|
||||
elif _assistant_content is not None and isinstance(_assistant_content, str):
|
||||
|
|
@ -4899,7 +4998,11 @@ class BedrockConverseMessagesProcessor:
|
|||
|
||||
msg_i += 1
|
||||
|
||||
assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse")
|
||||
assistant_content = list(
|
||||
_without_repeated_thinking(
|
||||
_deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity
|
||||
)
|
||||
)
|
||||
assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content)
|
||||
|
||||
if assistant_content:
|
||||
|
|
@ -5083,6 +5186,7 @@ def _bedrock_converse_messages_pt(
|
|||
assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
|
||||
deployment_model: str | None = None,
|
||||
base_model: str | None = None,
|
||||
drop_params: bool = False,
|
||||
) -> list[BedrockMessageBlock]:
|
||||
"""
|
||||
Converts given messages from OpenAI format to Bedrock format
|
||||
|
|
@ -5244,35 +5348,21 @@ def _bedrock_converse_messages_pt(
|
|||
)
|
||||
_assistant_content = assistant_message_block.get("content", None)
|
||||
thinking_blocks = assistant_message_block.get("thinking_blocks")
|
||||
|
||||
if thinking_blocks is not None:
|
||||
converted_thinking_blocks = (
|
||||
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
thinking_blocks, to_claude=to_claude
|
||||
)
|
||||
)
|
||||
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=converted_thinking_blocks,
|
||||
assistant_parts=assistant_content,
|
||||
converted_thinking_blocks = (
|
||||
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
_thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude
|
||||
)
|
||||
)
|
||||
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=converted_thinking_blocks,
|
||||
assistant_parts=assistant_content,
|
||||
)
|
||||
|
||||
if _assistant_content is not None and isinstance(_assistant_content, list):
|
||||
assistants_parts: list[BedrockContentBlock] = []
|
||||
for element in _assistant_content:
|
||||
if isinstance(element, dict):
|
||||
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
|
||||
thinking_block = (
|
||||
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
thinking_blocks=[element], to_claude=to_claude
|
||||
)
|
||||
)
|
||||
assistants_parts = (
|
||||
BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
|
||||
thinking_blocks=thinking_block,
|
||||
assistant_parts=assistants_parts,
|
||||
)
|
||||
)
|
||||
elif element["type"] == "text":
|
||||
if element["type"] == "text":
|
||||
# AWS Bedrock doesn't allow empty or whitespace-only text content
|
||||
# Skip completely empty strings to avoid blank content blocks
|
||||
if element.get("text", "").strip():
|
||||
|
|
@ -5292,7 +5382,9 @@ def _bedrock_converse_messages_pt(
|
|||
model=model,
|
||||
deployment_model=deployment_model,
|
||||
)
|
||||
if _cache_point_block is not None:
|
||||
if _cache_point_block is not None and not (
|
||||
drop_params and _as_thinking_block(element) is not None
|
||||
):
|
||||
assistants_parts.append(_cache_point_block)
|
||||
assistant_content.extend(assistants_parts)
|
||||
elif _assistant_content is not None and isinstance(_assistant_content, str):
|
||||
|
|
@ -5313,7 +5405,11 @@ def _bedrock_converse_messages_pt(
|
|||
|
||||
msg_i += 1
|
||||
|
||||
assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse")
|
||||
assistant_content = list(
|
||||
_without_repeated_thinking(
|
||||
_deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity
|
||||
)
|
||||
)
|
||||
assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content)
|
||||
|
||||
if assistant_content:
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionUserMessage,
|
||||
)
|
||||
|
||||
from .common_utils import is_unsignable_thinking_block
|
||||
from .common_utils import has_blank_redacted_data, is_unsignable_thinking_block
|
||||
|
||||
CONVERTED_SYSTEM_NOTE: Final = (
|
||||
"Operator note (not from the user): the following was originally a mid-conversation system-role reminder."
|
||||
|
|
@ -267,27 +267,22 @@ def _block_index_by_message(blocks: Sequence[tuple[bool, tuple[int, ...]]]) -> M
|
|||
|
||||
def _thinking_block_renders(block: object) -> bool:
|
||||
"""A thinking block the converter keeps: one Anthropic can verify, so never bridged encrypted reasoning."""
|
||||
return message_field(block, "type") in _THINKING_BLOCK_TYPES and not is_unsignable_thinking_block(block)
|
||||
return message_field(block, "type") in _THINKING_BLOCK_TYPES and not (
|
||||
is_unsignable_thinking_block(block) or has_blank_redacted_data(block)
|
||||
)
|
||||
|
||||
|
||||
def _assistant_part_renders(part: object) -> bool:
|
||||
"""A text part always renders: the converter pads empty text with a placeholder."""
|
||||
part_type: Final = message_field(part, "type")
|
||||
if part_type == "thinking":
|
||||
if part_type in _THINKING_BLOCK_TYPES:
|
||||
return _thinking_block_renders(part)
|
||||
return part_type in _RENDERED_ASSISTANT_PART_TYPES or (
|
||||
isinstance(part_type, str) and part_type.endswith("_tool_result")
|
||||
)
|
||||
|
||||
|
||||
def _separate_thinking_blocks_render(message: object, parts: Sequence[object]) -> bool:
|
||||
"""``thinking_blocks`` reach the wire only when no inline thinking part claims the slot.
|
||||
|
||||
The converter skips the separate blocks as soon as the content list carries a
|
||||
``thinking`` or ``redacted_thinking`` part, whether or not that part itself renders.
|
||||
"""
|
||||
if any(message_field(part, "type") in _THINKING_BLOCK_TYPES for part in parts):
|
||||
return False
|
||||
def _separate_thinking_blocks_render(message: object) -> bool:
|
||||
return any(_thinking_block_renders(block) for block in parts_of(message_field(message, "thinking_blocks")))
|
||||
|
||||
|
||||
|
|
@ -404,7 +399,7 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool:
|
|||
parts: Final = parts_of(content)
|
||||
return (
|
||||
any(_assistant_part_renders(part) for part in parts)
|
||||
or _separate_thinking_blocks_render(message, parts)
|
||||
or _separate_thinking_blocks_render(message)
|
||||
or _tool_calls_render(message, answered)
|
||||
or bool(message_field(message, "function_call"))
|
||||
or bool(message_field(message_field(message, "provider_specific_fields"), "compaction_blocks"))
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from litellm.types.utils import (
|
|||
StreamingChoices,
|
||||
TextChoices,
|
||||
TextCompletionResponse,
|
||||
UnreportedUsage,
|
||||
Usage,
|
||||
)
|
||||
from litellm.utils import print_verbose, token_counter
|
||||
|
|
@ -876,6 +877,8 @@ class ChunkProcessor:
|
|||
|
||||
if isinstance(usage_chunk, dict):
|
||||
return Usage(**usage_chunk)
|
||||
if isinstance(usage_chunk, UnreportedUsage):
|
||||
return None
|
||||
if usage_chunk is None or isinstance(usage_chunk, Usage):
|
||||
return usage_chunk
|
||||
return Usage(**usage_chunk.model_dump())
|
||||
|
|
|
|||
|
|
@ -24,7 +24,12 @@ import litellm
|
|||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, stored_attribute
|
||||
from litellm.litellm_core_utils.hidden_params import (
|
||||
HIDDEN_PARAMS_ATTR,
|
||||
get_hidden_params,
|
||||
set_hidden_param,
|
||||
stored_attribute,
|
||||
)
|
||||
from litellm.litellm_core_utils.model_response_utils import (
|
||||
is_model_response_stream_empty,
|
||||
)
|
||||
|
|
@ -43,6 +48,8 @@ from litellm.types.utils import (
|
|||
ModelResponseStream,
|
||||
PromptTokensDetailsWrapper,
|
||||
StreamingChoices,
|
||||
TextCompletionResponse,
|
||||
UnreportedUsage,
|
||||
Usage,
|
||||
)
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
|
@ -62,6 +69,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
|
|||
_SYNC_ITER_EXHAUSTED: Final = object()
|
||||
|
||||
_GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
|
||||
_CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None
|
||||
_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value})
|
||||
|
||||
|
||||
|
|
@ -335,7 +343,8 @@ class CustomStreamWrapper:
|
|||
self.sent_stream_usage = False
|
||||
self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
|
||||
self.tool_call = False
|
||||
self.chunks: list = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options
|
||||
# keep track of the returned chunks - used for calculating the input/output tokens for stream options
|
||||
self.chunks: list[ModelResponseStream] = []
|
||||
self._repeated_messages_count = 1
|
||||
self.is_function_call = self.check_is_function_call(logging_obj=logging_obj)
|
||||
self.created: int | None = None
|
||||
|
|
@ -544,7 +553,7 @@ class CustomStreamWrapper:
|
|||
last_content: Final = self.chunks[-1].choices[0].delta.content
|
||||
|
||||
if (
|
||||
last_content is None or not isinstance(last_content, str) or len(last_content) <= 2
|
||||
not isinstance(last_content, str) or len(last_content) <= 2
|
||||
): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946
|
||||
self._repeated_messages_count = 1
|
||||
return
|
||||
|
|
@ -1927,6 +1936,51 @@ class CustomStreamWrapper:
|
|||
response.hidden_params["additional_headers"] = {}
|
||||
response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost
|
||||
|
||||
def _builder_failure_response(self, error: Exception) -> ModelResponseStream | None:
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
|
||||
str(error),
|
||||
)
|
||||
try:
|
||||
return self.model_response_creator(chunk={"usage": calculate_total_usage(chunks=self.chunks)})
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _build_complete_streaming_response(self) -> _CompleteStreamingResponse:
|
||||
try:
|
||||
return litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
return self._builder_failure_response(e)
|
||||
|
||||
async def _async_build_complete_streaming_response(self) -> _CompleteStreamingResponse:
|
||||
try:
|
||||
return await asyncify(litellm.stream_chunk_builder)(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
return self._builder_failure_response(e)
|
||||
|
||||
def _set_final_chunk_usage(self, chunk: ModelResponseStream) -> None:
|
||||
reported_usage: Final = _reported_total_usage(chunks=self.chunks)
|
||||
set_hidden_param(
|
||||
chunk,
|
||||
"usage",
|
||||
UnreportedUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
|
||||
if reported_usage is None
|
||||
else reported_usage,
|
||||
)
|
||||
self._last_returned_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior
|
||||
dict[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
|
||||
def __next__(self) -> "ModelResponseStream":
|
||||
if self._closed:
|
||||
raise StopIteration
|
||||
|
|
@ -2007,12 +2061,7 @@ class CustomStreamWrapper:
|
|||
continue
|
||||
# add usage as hidden param
|
||||
if self.sent_last_chunk is True and self.stream_options is None:
|
||||
usage = calculate_total_usage(chunks=self.chunks)
|
||||
response_hidden_params_for_final_chunk = cast( # cast-ok: preserve dynamic mapping behavior
|
||||
dict[str, object], getattr(response, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
response_hidden_params_for_final_chunk["usage"] = usage
|
||||
self._last_returned_hidden_params = response_hidden_params_for_final_chunk
|
||||
self._set_final_chunk_usage(response)
|
||||
# Add MCP metadata to final chunk if present
|
||||
response = self._add_mcp_metadata_to_final_chunk(response)
|
||||
# RETURN RESULT
|
||||
|
|
@ -2020,30 +2069,7 @@ class CustomStreamWrapper:
|
|||
|
||||
except StopIteration:
|
||||
if self.sent_last_chunk is True:
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
# stream_chunk_builder can re-raise (as APIError) on large agentic
|
||||
# streams. The raise originates inside this except-StopIteration block,
|
||||
# so the sibling `except Exception` below does not catch it; it would
|
||||
# escape __next__ and drop the request from SpendLogs. Recover
|
||||
# best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
complete_streaming_response: Final = self._build_complete_streaming_response()
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
|
||||
|
|
@ -2107,9 +2133,8 @@ class CustomStreamWrapper:
|
|||
# last content chunk (stripped there), but never add the synthetic "stop" used when it sent none.
|
||||
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
|
||||
self.chunks.append(processed_chunk)
|
||||
if self.stream_options is None: # add usage as hidden param
|
||||
usage = calculate_total_usage(chunks=self.chunks)
|
||||
processed_chunk.hidden_params["usage"] = usage
|
||||
if self.stream_options is None:
|
||||
self._set_final_chunk_usage(processed_chunk)
|
||||
## LOGGING
|
||||
executor.submit(
|
||||
self.run_success_logging_and_cache_storage,
|
||||
|
|
@ -2224,12 +2249,7 @@ class CustomStreamWrapper:
|
|||
|
||||
# add usage as hidden param
|
||||
if self.sent_last_chunk is True and self.stream_options is None:
|
||||
usage = calculate_total_usage(chunks=self.chunks)
|
||||
processed_chunk_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior
|
||||
dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
processed_chunk_hidden_params["usage"] = usage
|
||||
self._last_returned_hidden_params = processed_chunk_hidden_params
|
||||
self._set_final_chunk_usage(processed_chunk)
|
||||
|
||||
# Call post-call streaming deployment hook for final chunk
|
||||
if self.sent_last_chunk is True:
|
||||
|
|
@ -2285,29 +2305,7 @@ class CustomStreamWrapper:
|
|||
|
||||
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
|
||||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
try:
|
||||
complete_streaming_response = await asyncify(litellm.stream_chunk_builder)(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
# see sync __next__: a raise from stream_chunk_builder inside this
|
||||
# except handler escapes __anext__ and drops the request from SpendLogs.
|
||||
# Recover best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
complete_streaming_response: Final = await self._async_build_complete_streaming_response()
|
||||
response: Final = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
|
||||
|
|
@ -2384,11 +2382,7 @@ class CustomStreamWrapper:
|
|||
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
|
||||
self.chunks.append(processed_chunk)
|
||||
if self.stream_options is None:
|
||||
usage: Final = calculate_total_usage(chunks=self.chunks)
|
||||
processed_chunk_hidden_params: Final = cast( # cast-ok: preserve dynamic mapping behavior
|
||||
dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
processed_chunk_hidden_params["usage"] = usage
|
||||
self._set_final_chunk_usage(processed_chunk)
|
||||
# see sync __next__'s sibling branch: deliberately do NOT restore
|
||||
# here - this chunk is still this call's own data, and restoring
|
||||
# before returning it would corrupt the caller's own log
|
||||
|
|
@ -2659,7 +2653,13 @@ def _coerce_token_details(
|
|||
return details_type(**(raw if isinstance(raw, dict) else raw.model_dump()))
|
||||
|
||||
|
||||
def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
|
||||
def _reported_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage | None:
|
||||
if not any("usage" in chunk and chunk["usage"] is not None for chunk in chunks):
|
||||
return None
|
||||
return calculate_total_usage(chunks=chunks)
|
||||
|
||||
|
||||
def calculate_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage:
|
||||
"""Assume most recent usage chunk has total usage uptil then."""
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
|
||||
attach_cache_creation_token_details,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.llms.base_llm.base_utils import (
|
|||
ResponseFormatSchema,
|
||||
names_a_response_format_property,
|
||||
partial_response_format_answer,
|
||||
response_format_answer_json,
|
||||
response_format_envelope_contents,
|
||||
split_response_format_answer,
|
||||
)
|
||||
|
|
@ -104,17 +105,17 @@ def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatS
|
|||
next_key: Final = _STRING_VALUE_THEN_KEY.match(string_value_start)
|
||||
if next_key is None:
|
||||
return "undecided"
|
||||
return "pass_through" if next_key.group(1) in schema.property_names else "hold"
|
||||
return "pass_through" if next_key.group(1) in schema.declared_names else "hold"
|
||||
|
||||
|
||||
def _is_a_label(text: str, schema: ResponseFormatSchema, tool_name: str) -> bool:
|
||||
return text in schema.property_names or text == tool_name
|
||||
return text in schema.declared_names or text == tool_name
|
||||
|
||||
|
||||
def _label_verdict(label: str, next_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
if _is_a_label(label, schema, tool_name):
|
||||
return "hold"
|
||||
return "pass_through" if next_key in schema.property_names else "undecided"
|
||||
return "pass_through" if next_key in schema.declared_names else "undecided"
|
||||
|
||||
|
||||
def _leading_labels_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
|
|
@ -159,9 +160,9 @@ def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name:
|
|||
return "hold"
|
||||
if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json):
|
||||
return "undecided"
|
||||
if first_key.group(1) == "values" and "values" not in schema.property_names:
|
||||
if first_key.group(1) == "values" and "values" not in schema.declared_names:
|
||||
return "hold"
|
||||
if first_key.group(1) in schema.property_names or not schema.property_names:
|
||||
if first_key.group(1) in schema.declared_names or not schema.declared_names:
|
||||
return "pass_through"
|
||||
return _first_value_verdict(buffered_json[first_key.end() :], schema, tool_name)
|
||||
|
||||
|
|
@ -172,13 +173,13 @@ def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
|
|||
except (json.JSONDecodeError, RecursionError):
|
||||
return held
|
||||
contents: Final = response_format_envelope_contents(payload, schema, RESPONSE_FORMAT_TOOL_NAME)
|
||||
return held if contents is None else json.dumps(contents)
|
||||
return held if contents is None else response_format_answer_json(contents)
|
||||
|
||||
|
||||
def _split_answer_text(blocks: tuple[str, ...], schema: ResponseFormatSchema) -> str:
|
||||
answer: Final = split_response_format_answer(blocks, schema)
|
||||
if answer is not None:
|
||||
return json.dumps(answer)
|
||||
return response_format_answer_json(answer)
|
||||
return "".join(_held_answer(block, schema) for block in blocks)
|
||||
|
||||
|
||||
|
|
@ -1283,7 +1284,7 @@ class ModelResponseIterator:
|
|||
return None
|
||||
self._streamed_answer_keys = streamed | frozenset(fields)
|
||||
separator: Final = ", " if streamed else ""
|
||||
return f"{separator}{json.dumps(fields)[1:-1]}" if fields else ""
|
||||
return f"{separator}{response_format_answer_json(fields)[1:-1]}" if fields else ""
|
||||
|
||||
def _flush_response_format_buffer(self) -> str:
|
||||
schema: Final = self.response_format_schema
|
||||
|
|
@ -1314,6 +1315,17 @@ class ModelResponseIterator:
|
|||
return ""
|
||||
return _split_answer_text(blocks, self.response_format_schema)
|
||||
|
||||
def cut_stream_chunk(self) -> ModelResponseStream | None:
|
||||
unsent: Final = self._flush_response_format_buffer() if self._response_format_buffer.getvalue() else ""
|
||||
text: Final = unsent + self._release_answer_close() + self._drain_split_answer_blocks()
|
||||
if not text:
|
||||
return None
|
||||
return ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=text))],
|
||||
id=self.response_id,
|
||||
model=self.served_model,
|
||||
)
|
||||
|
||||
def _handle_message_delta(
|
||||
self, chunk: dict
|
||||
) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]:
|
||||
|
|
@ -1428,6 +1440,9 @@ class ModelResponseIterator:
|
|||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
cut_stream_chunk = self.cut_stream_chunk()
|
||||
if cut_stream_chunk is not None:
|
||||
return cut_stream_chunk
|
||||
raise StopIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
@ -1474,6 +1489,9 @@ class ModelResponseIterator:
|
|||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
cut_stream_chunk = self.cut_stream_chunk()
|
||||
if cut_stream_chunk is not None:
|
||||
return cut_stream_chunk
|
||||
raise StopAsyncIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from litellm.constants import (
|
|||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
RESPONSE_FORMAT_TOOL_NAME,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
DEFS_MAX_INLINED_BYTES,
|
||||
is_object_schema,
|
||||
|
|
@ -50,6 +50,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
|
|||
from litellm.llms.base_llm.base_utils import (
|
||||
ResponseFormatSchema,
|
||||
as_json_object,
|
||||
response_format_answer_json,
|
||||
response_format_schema,
|
||||
split_response_format_answer,
|
||||
type_to_response_format_param,
|
||||
|
|
@ -163,6 +164,7 @@ _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128
|
|||
_DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool)
|
||||
_TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_TOOLS_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset(
|
||||
("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid")
|
||||
)
|
||||
|
|
@ -2265,6 +2267,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
anthropic_messages_pt,
|
||||
)
|
||||
|
||||
drop_params: Final = litellm.drop_params is True or litellm_params.get("drop_params") is True
|
||||
|
||||
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
|
||||
optional_params["tools"], _ = self._map_tools(add_dummy_tool(custom_llm_provider="anthropic"))
|
||||
|
||||
|
|
@ -2291,7 +2295,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
AnthropicConfig._maybe_drop_speed_param(
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
drop_params=litellm.drop_params or litellm_params.get("drop_params") is True,
|
||||
drop_params=drop_params,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
)
|
||||
|
||||
|
|
@ -2349,6 +2353,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
model=model,
|
||||
messages=list(conversation),
|
||||
llm_provider=self._resolved_provider,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
except Exception as e:
|
||||
raise AnthropicError(
|
||||
|
|
@ -2436,6 +2441,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
drop_params=litellm_params.get("drop_params") is True,
|
||||
output_key="top_k",
|
||||
)
|
||||
output_format: Final = request_params.get("output_format")
|
||||
if (
|
||||
output_format is not None
|
||||
and not isinstance(output_format, Mapping)
|
||||
and drops_unsupported_params(_LITELLM_PARAMS_ADAPTER.validate_python(litellm_params).get("drop_params"))
|
||||
):
|
||||
request_params.pop("output_format")
|
||||
|
||||
data: Final = {
|
||||
"model": model,
|
||||
|
|
@ -3166,7 +3178,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
)
|
||||
if split_answer is not None:
|
||||
return litellm.Message(content=json.dumps(split_answer))
|
||||
return litellm.Message(content=response_format_answer_json(split_answer))
|
||||
try:
|
||||
payload: Final = _DECODED_JSON.validate_python(json.loads(json_mode_content_str))
|
||||
except (json.JSONDecodeError, RecursionError):
|
||||
|
|
@ -3174,7 +3186,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
return litellm.Message(content=json_mode_content_str)
|
||||
if response_format_schema is not None:
|
||||
return litellm.Message(
|
||||
content=json.dumps(
|
||||
content=response_format_answer_json(
|
||||
unwrap_response_format_envelope(payload, response_format_schema, RESPONSE_FORMAT_TOOL_NAME)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
Anthropic Token Counter implementation using the CountTokens API.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
|
@ -10,6 +12,11 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountFields,
|
||||
ChatCountTranslationError,
|
||||
anthropic_count_fields,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
|
|
@ -38,6 +45,8 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
context_management: CountContextManagement | None = None,
|
||||
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
count_options: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using Anthropic's CountTokens API.
|
||||
|
|
@ -57,6 +66,30 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
if not messages:
|
||||
return None
|
||||
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
match translated := await anthropic_count_fields(
|
||||
model_to_use, messages, system, tools, messages_format, AnthropicConfig(), chat_params
|
||||
):
|
||||
case ChatCountTranslationError():
|
||||
return translated.response(request_model, model_to_use, "anthropic_api")
|
||||
case AnthropicCountFields():
|
||||
return await self.count_tokens(
|
||||
model_to_use=model_to_use,
|
||||
messages=[dict(message) for message in translated.messages],
|
||||
contents=contents,
|
||||
deployment=deployment,
|
||||
request_model=request_model,
|
||||
tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None,
|
||||
system=translated.system,
|
||||
messages_format="anthropic",
|
||||
output_format=output_format,
|
||||
context_management=context_management,
|
||||
count_options=translated.options,
|
||||
)
|
||||
case None:
|
||||
pass
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
api_base: Final = litellm_params.get("api_base")
|
||||
|
|
@ -89,6 +122,7 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
api_base=api_base,
|
||||
tools=tools,
|
||||
system=system,
|
||||
optional_params=count_options,
|
||||
)
|
||||
|
||||
if result is not None:
|
||||
|
|
|
|||
|
|
@ -4,25 +4,37 @@ Anthropic CountTokens API transformation logic.
|
|||
This module handles the transformation of requests to Anthropic's CountTokens API format.
|
||||
"""
|
||||
|
||||
import copy
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import (
|
||||
Final,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
cast, # noqa: TID251 # caller chat messages go to chat inference translation as-is
|
||||
)
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
|
||||
anthropic_system_blocks,
|
||||
split_leading_system_run,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
|
||||
from litellm.llms.anthropic.wif import resolve_anthropic_base
|
||||
from litellm.types.llms.openai import ChatCompletionImageObject
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue])
|
||||
_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject)
|
||||
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
|
||||
_PLAIN_MESSAGE_KEYS: Final = frozenset({"role", "content"})
|
||||
|
||||
|
||||
def _count_image(block: JsonValue) -> JsonValue:
|
||||
|
|
@ -74,6 +86,213 @@ def _lift_leading_system(
|
|||
return conversation, lifted or system
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AnthropicCountFields:
|
||||
messages: tuple[Mapping[str, JsonValue], ...]
|
||||
system: JsonValue
|
||||
tools: tuple[Mapping[str, JsonValue], ...] | None
|
||||
options: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChatCountTranslationError:
|
||||
message: str
|
||||
|
||||
def response(self, request_model: str, model_used: str, tokenizer_type: str) -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_used,
|
||||
tokenizer_type=tokenizer_type,
|
||||
error=True,
|
||||
error_message=self.message,
|
||||
status_code=400,
|
||||
rejected_before_provider_call=True,
|
||||
)
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="ignore"))
|
||||
class _AnthropicChatRequest(TypedDict):
|
||||
messages: ReadOnly[tuple[Mapping[str, JsonValue], ...]]
|
||||
system: NotRequired[ReadOnly[JsonValue]]
|
||||
tools: NotRequired[ReadOnly[tuple[Mapping[str, JsonValue], ...]]]
|
||||
thinking: NotRequired[ReadOnly[JsonValue]]
|
||||
tool_choice: NotRequired[ReadOnly[JsonValue]]
|
||||
output_config: NotRequired[ReadOnly[JsonValue]]
|
||||
output_format: NotRequired[ReadOnly[JsonValue]]
|
||||
|
||||
|
||||
_ANTHROPIC_CHAT_REQUEST: Final = TypeAdapter(_AnthropicChatRequest)
|
||||
_ChatParams: TypeAlias = dict[str, object]
|
||||
_ChatMessages: TypeAlias = list[AllMessageValues]
|
||||
|
||||
|
||||
class _ChatRequestTranslator(Protocol):
|
||||
def map_openai_params(
|
||||
self, non_default_params: _ChatParams, optional_params: _ChatParams, model: str, drop_params: bool
|
||||
) -> _ChatParams: ...
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: _ChatMessages,
|
||||
optional_params: _ChatParams,
|
||||
litellm_params: _ChatParams,
|
||||
headers: _ChatParams,
|
||||
) -> _ChatParams: ...
|
||||
|
||||
|
||||
def _needs_chat_translation(message: Mapping[str, object]) -> bool:
|
||||
return message.get("role") == "tool" or not _PLAIN_MESSAGE_KEYS.issuperset(message.keys())
|
||||
|
||||
|
||||
def _is_chat_tool(tool: Mapping[str, object]) -> bool:
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import is_flat_responses_tool
|
||||
|
||||
return is_flat_responses_tool(tool) or (tool.get("type") == "function" and "function" in tool)
|
||||
|
||||
|
||||
def _has_non_chat_user_block(message: Mapping[str, object]) -> bool:
|
||||
content: Final = message.get("content")
|
||||
if message.get("role") != "user" or not isinstance(content, list):
|
||||
return False
|
||||
blocks: Final = _SYSTEM_BLOCKS.validate_python(content)
|
||||
return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks)
|
||||
|
||||
|
||||
def _gets_placeholder_name(tool: Mapping[str, object]) -> bool:
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import is_provider_native_tool_dict
|
||||
|
||||
name: Final = tool.get("name")
|
||||
return not (_is_chat_tool(tool) or is_provider_native_tool_dict(tool) or (isinstance(name, str) and name.strip()))
|
||||
|
||||
|
||||
def _untranslatable_input(
|
||||
messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]]
|
||||
) -> ChatCountTranslationError | None:
|
||||
invalid_index: Final = next(
|
||||
(index for index, message in enumerate(messages) if _has_non_chat_user_block(message)), None
|
||||
)
|
||||
if invalid_index is not None:
|
||||
return ChatCountTranslationError(
|
||||
message=f"Invalid user message at index {invalid_index}. "
|
||||
"Please ensure all user messages are valid OpenAI chat completion messages."
|
||||
)
|
||||
unnamed_tool: Final = next(filter(_gets_placeholder_name, tools), None)
|
||||
if unnamed_tool is not None:
|
||||
return ChatCountTranslationError(
|
||||
message=f"Invalid request for token counting: tool type {unnamed_tool.get('type')!r} "
|
||||
"has no Chat Completions equivalent"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _system_messages(system: object) -> tuple[Mapping[str, object], ...]:
|
||||
message: Final[Mapping[str, object]] = {"role": "system", "content": system}
|
||||
return (message,) if isinstance(system, (str, list)) and system else ()
|
||||
|
||||
|
||||
def _count_options(request: _AnthropicChatRequest) -> Mapping[str, JsonValue]:
|
||||
"""Structured output reaches chat inference as the beta ``output_format``, which count_tokens
|
||||
refuses; its GA ``output_config.format`` form counts the same schema the same way."""
|
||||
output_config: Final = request.get("output_config")
|
||||
output_format: Final = request.get("output_format")
|
||||
counted_config: Final = {
|
||||
**(output_config if isinstance(output_config, dict) else {}),
|
||||
**({} if output_format is None else {"format": output_format}),
|
||||
}
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("thinking", request.get("thinking")),
|
||||
("tool_choice", request.get("tool_choice")),
|
||||
("output_config", counted_config or None),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _translated_chat_request(
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
system: object,
|
||||
tools: Sequence[Mapping[str, object]],
|
||||
chat_params: Mapping[str, object],
|
||||
translator: _ChatRequestTranslator,
|
||||
) -> _ChatParams:
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for
|
||||
|
||||
chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(chat_completion_tools_for, tools))]
|
||||
openai_params: Final = {name: value for name, value in chat_params.items() if name != "drop_params"}
|
||||
mapped_params: Final = {**openai_params, **({"tools": chat_tools} if chat_tools else {})}
|
||||
chat_messages: Final = cast( # cast-ok: chat-shaped request messages, read the way chat inference reads them
|
||||
_ChatMessages, copy.deepcopy([*_system_messages(system), *messages])
|
||||
)
|
||||
return translator.transform_request(
|
||||
model=model,
|
||||
messages=chat_messages,
|
||||
optional_params=translator.map_openai_params(
|
||||
mapped_params, {}, model, drop_params=chat_params.get("drop_params") is True
|
||||
)
|
||||
if mapped_params
|
||||
else {},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
def anthropic_fields_for_chat_tool_use(
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
system: object,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
translator: _ChatRequestTranslator,
|
||||
chat_params: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> AnthropicCountFields | ChatCountTranslationError | None:
|
||||
counted_tools: Final = tools or ()
|
||||
if (
|
||||
not chat_params
|
||||
and not any(map(_is_chat_tool, counted_tools))
|
||||
and not any(map(_needs_chat_translation, messages))
|
||||
):
|
||||
return None
|
||||
untranslatable: Final = _untranslatable_input(messages, counted_tools)
|
||||
if untranslatable is not None:
|
||||
return untranslatable
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
|
||||
try:
|
||||
request: Final = _translated_chat_request(model, messages, system, counted_tools, chat_params, translator)
|
||||
except (AnthropicError, BadRequestError) as e:
|
||||
return ChatCountTranslationError(message=str(e.message))
|
||||
except (ValueError, KeyError) as e:
|
||||
return ChatCountTranslationError(message=f"Invalid request for token counting: {e}")
|
||||
translated: Final = _ANTHROPIC_CHAT_REQUEST.validate_python(request)
|
||||
return AnthropicCountFields(
|
||||
messages=translated["messages"],
|
||||
system=translated.get("system"),
|
||||
tools=translated.get("tools"),
|
||||
options=_count_options(translated),
|
||||
)
|
||||
|
||||
|
||||
async def anthropic_count_fields(
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
system: object,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
translator: _ChatRequestTranslator,
|
||||
chat_params: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> AnthropicCountFields | ChatCountTranslationError | None:
|
||||
if messages_format != "openai":
|
||||
return None
|
||||
return await asyncify(anthropic_fields_for_chat_tool_use)(model, messages, system, tools, translator, chat_params)
|
||||
|
||||
|
||||
class AnthropicCountTokensConfig:
|
||||
"""
|
||||
Configuration and transformation logic for Anthropic CountTokens API.
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import re
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
|
|
@ -153,6 +154,8 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat
|
|||
await self._handle_streaming_logging(chunks)
|
||||
raise StopAsyncIteration
|
||||
chunk: Final = self.chunks[self.current_index]
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = datetime.now()
|
||||
self.current_index += 1
|
||||
return chunk
|
||||
|
||||
|
|
|
|||
|
|
@ -506,7 +506,7 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
):
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
self.request_body = request_body
|
||||
self.start_time = datetime.now()
|
||||
self.start_time = litellm_logging_obj.start_time
|
||||
self.completion_start_time: datetime | None = None
|
||||
|
||||
async def _handle_streaming_logging(self, collected_chunks: list[bytes], *, stream_teardown: bool = False):
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_beta_header_supported_by_bundled_config
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import drop_non_string_tool_result_texts
|
||||
from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards
|
||||
|
|
@ -627,6 +628,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
):
|
||||
anthropic_messages_optional_request_params.pop("safeguards") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict
|
||||
|
||||
output_format: Final[object] = anthropic_messages_optional_request_params.get("output_format") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the request params arrive as a bare dict
|
||||
if (
|
||||
output_format is not None
|
||||
and not isinstance(output_format, Mapping)
|
||||
and drops_unsupported_params(litellm_params.drop_params)
|
||||
):
|
||||
anthropic_messages_optional_request_params.pop("output_format") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict
|
||||
|
||||
self._translate_reasoning_effort_to_anthropic(
|
||||
model=model,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Azure AI Anthropic CountTokens API handler.
|
|||
Uses httpx for HTTP requests with Azure authentication.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -36,6 +37,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: object = None,
|
||||
optional_params: Mapping[str, JsonValue] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Handle a CountTokens request using httpx with Azure authentication.
|
||||
|
|
@ -66,6 +68,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
|
|||
messages=messages,
|
||||
tools=tools,
|
||||
system=system,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
|
||||
verbose_logger.debug("Transformed request: %s", request_body)
|
||||
|
|
|
|||
|
|
@ -3,12 +3,19 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
|
|||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.anthropic.count_tokens.transformation import (
|
||||
AnthropicCountFields,
|
||||
ChatCountTranslationError,
|
||||
anthropic_count_fields,
|
||||
)
|
||||
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
|
||||
AzureAIAnthropicCountTokensHandler,
|
||||
)
|
||||
|
|
@ -55,6 +62,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
context_management: CountContextManagement | None = None,
|
||||
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
count_options: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using Azure AI Anthropic's CountTokens API.
|
||||
|
|
@ -74,6 +83,30 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
if not messages:
|
||||
return None
|
||||
|
||||
from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig
|
||||
|
||||
match translated := await anthropic_count_fields(
|
||||
model_to_use, messages, system, tools, messages_format, AzureAnthropicConfig(), chat_params
|
||||
):
|
||||
case ChatCountTranslationError():
|
||||
return translated.response(request_model, model_to_use, "azure_ai_anthropic_api")
|
||||
case AnthropicCountFields():
|
||||
return await self.count_tokens(
|
||||
model_to_use=model_to_use,
|
||||
messages=[dict(message) for message in translated.messages],
|
||||
contents=contents,
|
||||
deployment=deployment,
|
||||
request_model=request_model,
|
||||
tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None,
|
||||
system=translated.system,
|
||||
messages_format="anthropic",
|
||||
output_format=output_format,
|
||||
context_management=context_management,
|
||||
count_options=translated.options,
|
||||
)
|
||||
case None:
|
||||
pass
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
|
||||
|
|
@ -101,6 +134,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
litellm_params=litellm_params,
|
||||
tools=tools,
|
||||
system=system,
|
||||
optional_params=count_options,
|
||||
)
|
||||
|
||||
if result is not None:
|
||||
|
|
|
|||
|
|
@ -5,9 +5,11 @@ Utility functions for base LLM classes.
|
|||
import copy
|
||||
import itertools
|
||||
import json
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals
|
||||
|
|
@ -39,6 +41,7 @@ class BaseTokenCounter(ABC):
|
|||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
context_management: CountContextManagement | None = None,
|
||||
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
|
||||
) -> TokenCountResponse | None:
|
||||
pass
|
||||
|
||||
|
|
@ -172,7 +175,7 @@ _convert_tool_response_to_message = convert_tool_response_to_message
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponseFormatSchema:
|
||||
property_names: frozenset[str]
|
||||
declared_names: frozenset[str]
|
||||
required_names: frozenset[str] = frozenset()
|
||||
allows_undeclared_keys: bool = False
|
||||
|
||||
|
|
@ -216,9 +219,10 @@ def response_format_schema(value: object) -> ResponseFormatSchema | None:
|
|||
properties: Final = as_json_object(document.get("properties", {}))
|
||||
if properties is None:
|
||||
return None
|
||||
required_names: Final = _names(document.get("required", ()))
|
||||
return ResponseFormatSchema(
|
||||
property_names=frozenset(properties),
|
||||
required_names=_names(document.get("required", ())),
|
||||
declared_names=frozenset(properties) | required_names,
|
||||
required_names=required_names,
|
||||
allows_undeclared_keys=_allows_undeclared_keys(document),
|
||||
)
|
||||
|
||||
|
|
@ -244,8 +248,9 @@ def _json_object_or_encoded_object(value: JsonValue) -> Mapping[str, JsonValue]
|
|||
def _label_keys(payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> frozenset[str]:
|
||||
if not schema.required_names:
|
||||
return frozenset()
|
||||
declared: Final = schema.property_names | schema.required_names
|
||||
return frozenset(key for key, value in payload_object.items() if isinstance(value, str) and key not in declared)
|
||||
return frozenset(
|
||||
key for key, value in payload_object.items() if isinstance(value, str) and key not in schema.declared_names
|
||||
)
|
||||
|
||||
|
||||
def _envelope_layer(
|
||||
|
|
@ -256,7 +261,7 @@ def _envelope_layer(
|
|||
)
|
||||
label_keys: Final = name_keys or _label_keys(payload_object, schema)
|
||||
wrapper_keys: Final = tuple(key for key in payload_object if key not in label_keys)
|
||||
if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.property_names:
|
||||
if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.declared_names:
|
||||
return None
|
||||
wrapped: Final = payload_object[wrapper_keys[0]]
|
||||
inner: Final = _json_object_or_encoded_object(wrapped) if name_keys else as_json_object(wrapped)
|
||||
|
|
@ -275,15 +280,16 @@ def _envelope_layers(
|
|||
def _schema_object_inside_envelope(
|
||||
payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
layers: Final = _envelope_layers(payload_object, schema, tool_name)
|
||||
candidates: Final = layers if schema.property_names else itertools.islice(layers, 1)
|
||||
return next((layer.inner for layer in candidates if _holds_the_schema_object(layer, schema)), None)
|
||||
walk: Final = _envelope_layers(payload_object, schema, tool_name)
|
||||
layers: Final = tuple(walk if schema.declared_names else itertools.islice(walk, 1))
|
||||
schema_object: Final = next((layer.inner for layer in layers if _shares_a_schema_key(layer.inner, schema)), None)
|
||||
if schema_object is not None:
|
||||
return schema_object
|
||||
return next((layer.inner for layer in layers if layer.names_the_tool), None)
|
||||
|
||||
|
||||
def _holds_the_schema_object(layer: _EnvelopeLayer, schema: ResponseFormatSchema) -> bool:
|
||||
if layer.names_the_tool:
|
||||
return True
|
||||
return not schema.accepts_a_lone_undeclared_key and not schema.property_names.isdisjoint(layer.inner)
|
||||
def _shares_a_schema_key(candidate: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> bool:
|
||||
return not schema.accepts_a_lone_undeclared_key and not schema.declared_names.isdisjoint(candidate)
|
||||
|
||||
|
||||
def response_format_envelope_contents(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object | None:
|
||||
|
|
@ -293,7 +299,15 @@ def response_format_envelope_contents(payload: object, schema: ResponseFormatSch
|
|||
schema_object: Final = _schema_object_inside_envelope(payload_object, schema, tool_name)
|
||||
if schema_object is not None:
|
||||
return schema_object
|
||||
return None if "values" in schema.property_names else payload_object.get("values")
|
||||
return None if "values" in schema.declared_names else payload_object.get("values")
|
||||
|
||||
|
||||
_LONE_SURROGATE: Final = re.compile("[\ud800-\udfff]")
|
||||
|
||||
|
||||
def response_format_answer_json(value: object) -> str:
|
||||
text: Final = json.dumps(value, ensure_ascii=False)
|
||||
return json.dumps(value) if _LONE_SURROGATE.search(text) else text
|
||||
|
||||
|
||||
def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object:
|
||||
|
|
@ -305,7 +319,7 @@ def _field_names(
|
|||
blocks: tuple[Mapping[str, JsonValue], ...], key: str, schema: ResponseFormatSchema
|
||||
) -> tuple[str, ...] | None:
|
||||
values: Final = (block[key] for block in blocks)
|
||||
names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.property_names)
|
||||
names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.declared_names)
|
||||
return names if len(frozenset(names)) == len(blocks) else None
|
||||
|
||||
|
||||
|
|
@ -324,7 +338,7 @@ def _split_answer_fields(
|
|||
blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
keys: Final = frozenset(blocks[0])
|
||||
if not keys.isdisjoint(schema.property_names) or any(frozenset(block) != keys for block in blocks):
|
||||
if not keys.isdisjoint(schema.declared_names) or any(frozenset(block) != keys for block in blocks):
|
||||
return None
|
||||
named: Final = tuple((key, names) for key in keys if (names := _field_names(blocks, key, schema)) is not None)
|
||||
if len(named) != 1:
|
||||
|
|
@ -340,7 +354,7 @@ def _merged_partial_answers(
|
|||
blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
keys: Final = frozenset(itertools.chain.from_iterable(blocks))
|
||||
if not keys <= schema.property_names or len(keys) != sum(len(block) for block in blocks):
|
||||
if not keys <= schema.declared_names or len(keys) != sum(len(block) for block in blocks):
|
||||
return None
|
||||
return dict(itertools.chain.from_iterable(block.items() for block in blocks))
|
||||
|
||||
|
|
@ -358,14 +372,14 @@ def split_response_format_answer(
|
|||
|
||||
def partial_response_format_answer(encoded_block: str, schema: ResponseFormatSchema) -> Mapping[str, JsonValue] | None:
|
||||
block: Final = _json_object_or_encoded_object(encoded_block)
|
||||
return block if block is not None and frozenset(block) <= schema.property_names else None
|
||||
return block if block is not None and frozenset(block) <= schema.declared_names else None
|
||||
|
||||
|
||||
def names_a_response_format_property(encoded_block: str, schema: ResponseFormatSchema) -> bool:
|
||||
block: Final = _json_object_or_encoded_object(encoded_block)
|
||||
if block is None:
|
||||
return False
|
||||
return any(isinstance(value, str) and value in schema.property_names for value in block.values())
|
||||
return any(isinstance(value, str) and value in schema.declared_names for value in block.values())
|
||||
|
||||
|
||||
def _dict_to_response_format_helper(response_format: dict, ref_template: str | None = None) -> dict:
|
||||
|
|
|
|||
|
|
@ -428,7 +428,7 @@ class BaseConfig(ABC):
|
|||
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"],
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None = False,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
pass
|
||||
|
||||
def get_model_response_iterator_for_request(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
@ -18,6 +18,7 @@ def _responses_request_keys() -> frozenset[str]:
|
|||
def responses_batch_body_to_chat_body(
|
||||
openai_request_body: Mapping[str, object],
|
||||
custom_llm_provider: str | None = None,
|
||||
drop_params: bool | None = None,
|
||||
) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict
|
||||
"""
|
||||
Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body.
|
||||
|
|
@ -50,5 +51,6 @@ def responses_batch_body_to_chat_body(
|
|||
input=chat_input,
|
||||
responses_api_request=responses_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
drop_params=drop_params,
|
||||
metadata=openai_request_body.get("metadata"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -70,6 +70,22 @@ def structured_output_instruction(schema: Mapping[str, object]) -> str:
|
|||
)
|
||||
|
||||
|
||||
def build_instructions(ctx: SessionContext) -> str | None:
|
||||
"""System instructions plus the structured-output instruction, for runtimes that take one prompt file."""
|
||||
schema_part: Final = (
|
||||
structured_output_instruction(ctx.output.model_json_schema()) if ctx.output is not None else None
|
||||
)
|
||||
sections: Final = tuple(section for section in (ctx.instructions, schema_part) if section)
|
||||
return "\n\n".join(sections) if sections else None
|
||||
|
||||
|
||||
def turn_prompt(ctx: SessionContext, prompt: str) -> str:
|
||||
"""Repeat the schema instruction in the user turn; system instructions alone are too weak."""
|
||||
if ctx.output is None:
|
||||
return prompt
|
||||
return f"{prompt}\n\n{structured_output_instruction(ctx.output.model_json_schema())}"
|
||||
|
||||
|
||||
def strict_json_schema(schema: JSONValue, depth: int = 0) -> JSONValue:
|
||||
"""Make a JSON schema acceptable to OpenAI strict structured outputs.
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -80,6 +80,12 @@ class _AssumeRoleParams(TypedDict):
|
|||
Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]]
|
||||
|
||||
|
||||
class _StsClientKwargs(TypedDict):
|
||||
verify: ReadOnly[bool | str]
|
||||
endpoint_url: ReadOnly[NotRequired[str]]
|
||||
region_name: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
def _assume_role_params(
|
||||
aws_role_name: str,
|
||||
aws_session_name: str,
|
||||
|
|
@ -275,7 +281,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
"aws_session_tags",
|
||||
]
|
||||
|
||||
def _get_ssl_verify(self, ssl_verify: bool | str | None = None):
|
||||
def _get_ssl_verify(self, ssl_verify: bool | str | None = None) -> bool | str:
|
||||
"""
|
||||
Get SSL verification setting for boto3 clients.
|
||||
|
||||
|
|
@ -875,15 +881,18 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
) -> dict:
|
||||
) -> _StsClientKwargs:
|
||||
"""STS client kwargs with aligned endpoint_url and region_name (SigV4)."""
|
||||
kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)}
|
||||
if aws_sts_endpoint is not None:
|
||||
kwargs["endpoint_url"] = aws_sts_endpoint
|
||||
sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name)
|
||||
if sts_region is not None:
|
||||
kwargs["region_name"] = sts_region
|
||||
return kwargs
|
||||
verify: Final = self._get_ssl_verify(ssl_verify)
|
||||
match (aws_sts_endpoint, self._resolve_sts_region(aws_sts_endpoint, aws_region_name)):
|
||||
case (None, None):
|
||||
return _StsClientKwargs(verify=verify)
|
||||
case (None, region):
|
||||
return _StsClientKwargs(verify=verify, region_name=region)
|
||||
case (endpoint, None):
|
||||
return _StsClientKwargs(verify=verify, endpoint_url=endpoint)
|
||||
case (endpoint, region):
|
||||
return _StsClientKwargs(verify=verify, endpoint_url=endpoint, region_name=region)
|
||||
|
||||
def get_aws_region_name_for_non_llm_api_calls(
|
||||
self,
|
||||
|
|
@ -1813,13 +1822,13 @@ async def run_aws_signing(
|
|||
)
|
||||
|
||||
|
||||
async def sign_request_off_loop_if_aws(
|
||||
async def run_off_loop_if_aws(
|
||||
provider_config: object,
|
||||
sign_request: Callable[_SignParams, _SignedRequest],
|
||||
call: Callable[_SignParams, _SignedRequest],
|
||||
/,
|
||||
*args: _SignParams.args,
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped call's signature
|
||||
) -> _SignedRequest:
|
||||
if isinstance(provider_config, SignsRequestsWithAWS):
|
||||
return await run_aws_signing(sign_request, *args, **kwargs)
|
||||
return sign_request(*args, **kwargs)
|
||||
return await run_aws_signing(call, *args, **kwargs)
|
||||
return call(*args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -385,7 +385,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
@staticmethod
|
||||
def _requires_min_max_tokens(model: str) -> bool:
|
||||
return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None
|
||||
return re.search(r"openai\.gpt-\d|xai\.grok-|moonshotai\.kimi-k3", model) is not None
|
||||
|
||||
def _is_nova_2_model(self, model: str) -> bool:
|
||||
"""
|
||||
|
|
@ -2127,6 +2127,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
llm_provider="bedrock_converse",
|
||||
user_continue_message=litellm_params.pop("user_continue_message", None),
|
||||
base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)),
|
||||
drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True,
|
||||
),
|
||||
capability_model,
|
||||
)
|
||||
|
|
@ -2207,6 +2208,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
llm_provider="bedrock_converse",
|
||||
user_continue_message=litellm_params.pop("user_continue_message", None),
|
||||
base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)),
|
||||
drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True,
|
||||
),
|
||||
capability_model,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -857,6 +857,9 @@ class AWSEventStreamDecoder:
|
|||
undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers)
|
||||
if undecoded_stream_error is not None:
|
||||
raise undecoded_stream_error
|
||||
stream_end_chunk: Final = self._stream_end_chunk()
|
||||
if stream_end_chunk is not None:
|
||||
yield stream_end_chunk
|
||||
|
||||
async def aiter_bytes(
|
||||
self, iterator: AsyncIterator[bytes], *, response_headers: Mapping[str, str] | None = None
|
||||
|
|
@ -878,6 +881,12 @@ class AWSEventStreamDecoder:
|
|||
undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers)
|
||||
if undecoded_stream_error is not None:
|
||||
raise undecoded_stream_error
|
||||
stream_end_chunk: Final = self._stream_end_chunk()
|
||||
if stream_end_chunk is not None:
|
||||
yield stream_end_chunk
|
||||
|
||||
def _stream_end_chunk(self) -> ModelResponseStream | None:
|
||||
return None
|
||||
|
||||
def _response_stream_shape(self) -> "Shape":
|
||||
response_stream_shape: Final = get_bedrock_response_stream_shape()
|
||||
|
|
@ -956,6 +965,9 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
|
|||
def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
|
||||
return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
|
||||
|
||||
def _stream_end_chunk(self) -> ModelResponseStream | None:
|
||||
return self.anthropic_model_response_iterator.cut_stream_chunk()
|
||||
|
||||
|
||||
class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder):
|
||||
def __init__(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.pass_through.messages.transformation import (
|
||||
DEFAULT_ANTHROPIC_API_VERSION,
|
||||
AnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
|
@ -22,14 +23,14 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
|
|||
|
||||
def validate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: dict[str, str], # mutable-ok: overrides the base signature
|
||||
model: str,
|
||||
messages: list[Any],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
messages: list[object], # mutable-ok: overrides the base signature
|
||||
optional_params: dict[str, object], # mutable-ok: overrides the base signature
|
||||
litellm_params: dict[str, object], # mutable-ok: overrides the base signature
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]:
|
||||
) -> tuple[dict[str, str], str | None]: # mutable-ok: overrides the base signature
|
||||
workspace_id: Final = self._get_workspace_id(optional_params, litellm_params)
|
||||
if workspace_id is None:
|
||||
raise litellm.AuthenticationError(
|
||||
|
|
@ -66,6 +67,27 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
|
|||
api_base,
|
||||
)
|
||||
|
||||
async def avalidate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict[str, str], # mutable-ok: mirrors validate_anthropic_messages_environment
|
||||
model: str,
|
||||
messages: list[object], # mutable-ok: mirrors validate_anthropic_messages_environment
|
||||
optional_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
|
||||
litellm_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors validate_anthropic_messages_environment
|
||||
return await run_aws_signing(
|
||||
self.validate_anthropic_messages_environment,
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue