Merge branch 'main' into fix/skip-unsatisfiable-tag-fallback-legs-43736

This commit is contained in:
Quinn Xu 2026-10-16 00:55:13 +08:00 • committed by GitHub
commit 3efb533185
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1640 changed files with 241473 additions and 20273 deletions

View file

@ -201,7 +201,7 @@ if [ "$isolated_workers" = false ]; then
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ] || [ "$suite" = mcp ]; then
if [ "$suite" = management ] || [ "$suite" = mcp ] || [ "$suite" = accounting ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"

4
.github/CODEOWNERS vendored
View file

@ -1,2 +1,2 @@
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri

View file

@ -600,6 +600,9 @@ jobs:
- name: check_unsafe_enterprise_import
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
- name: check_telemetry_imports
run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py
- name: ban_copy_deepcopy_kwargs
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
@ -615,6 +618,9 @@ jobs:
- name: check_migrations_no_data_rewrites
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
- name: check_migrations_lock_timeout
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_lock_timeout.py
- name: check_no_publicly_known_master_key
run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py

View file

@ -233,6 +233,7 @@ jobs:
tests/unit/llms/base_llm
tests/unit/llms/baseten
tests/unit/llms/black_forest_labs
tests/unit/llms/brave
tests/unit/llms/bytez
tests/unit/llms/cerebras
tests/unit/llms/chat
@ -268,6 +269,7 @@ jobs:
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
tests/unit/llms/google_pse
tests/unit/llms/github_copilot
tests/unit/llms/gradient_ai
tests/unit/llms/groq
@ -300,6 +302,7 @@ jobs:
tests/unit/llms/nebius
tests/unit/llms/neosantara
tests/unit/llms/nimble
tests/unit/llms/nlp_cloud
tests/unit/llms/novita
tests/unit/llms/nscale
tests/unit/llms/nvidia_nim
@ -310,6 +313,7 @@ jobs:
tests/unit/llms/oobabooga
tests/unit/llms/openai_like
tests/unit/llms/opencode
tests/unit/llms/opensandbox
tests/unit/llms/openrouter
tests/unit/llms/ovhcloud
tests/unit/llms/parallel_ai
@ -342,10 +346,12 @@ jobs:
tests/unit/llms/tinyfish
tests/unit/llms/together_ai
tests/unit/llms/tool_loop
tests/unit/llms/topaz
tests/unit/llms/triton
tests/unit/llms/v0
tests/unit/llms/valkey
tests/unit/llms/vercel_ai_gateway
tests/unit/llms/vllm
tests/unit/llms/volcengine
tests/unit/llms/voyage
tests/unit/llms/wandb
@ -412,6 +418,7 @@ jobs:
tests/unit/rag
tests/unit/realtime_api
tests/unit/rerank_api
tests/unit/search
tests/unit/vector_stores
tests/unit/videos
python-version: "3.12"
@ -446,6 +453,7 @@ jobs:
tests/unit/repositories
tests/unit/secret_managers
tests/unit/skills/test_skills_main.py
tests/unit/telemetry tests/unit/proxy/telemetry
tests/unit/tracing
python-version: "3.12"
workers: 2

View file

@ -23,7 +23,7 @@ When adding new features, add meaningful tests. Don't add tests that don't check
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
Never test structure or implementation details of code, only function of it. A test should catch a change that breaks behavior by accident. If the only way it can fail is someone deliberately undoing the change, don't write it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent". Assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken

View file

@ -290,7 +290,7 @@ test-rust-extension:
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
litellm.rust_bridge._native && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
rust-sqlx-prepare:

View file

@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/token",
"/onboarding/",
"/audit",
"/telemetry/",
"/oauth/",
"/invitation/",
"/jwt/",

View file

@ -322,7 +322,12 @@ def create_migration(
# Write the SQL to migration.sql
migration_file = migration_dir / "migration.sql"
migration_file.write_text(diff_sql)
migration_sql: Final = (
diff_sql
if "CONCURRENTLY" in diff_sql.upper()
else "SET LOCAL lock_timeout = '10s';\n\n" + diff_sql
)
migration_file.write_text(migration_sql)
print(f"Created migration in {migration_dir}")
return True

View file

@ -24,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
resolve_resource_owner_id,
)
from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
from litellm.proxy._types import (
CallTypes,
LiteLLM_ManagedFileTable,
@ -102,6 +103,7 @@ from litellm.types.utils import (
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
LlmProviders,
LLMResponseTypes,
SpecialEnums,
)
@ -194,6 +196,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
def _stored_json_text(raw_file_object: object) -> str:
"""JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
if isinstance(raw_file_object, BaseModel):
return raw_file_object.model_dump_json()
return json.dumps(raw_file_object)
def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
"""Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
if not size_bytes or file_object.bytes:
return file_object
return file_object.model_copy(update={"bytes": size_bytes})
def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
"""Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
if stored_object is None:
return file_object
return _with_measured_size(file_object, stored_object.bytes).model_copy(
update={"created_at": stored_object.created_at}
)
def _proxy_llm_router() -> Router | None:
import litellm.proxy.proxy_server as proxy_server_module
@ -252,6 +280,23 @@ def _has_provider_file_retrieve_route(
return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
def _is_xai_batch_output(
*,
llm_router: Router | None,
model_id: str | None,
model_name: str | None,
provider_file_id: str,
) -> bool:
"""Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
if not is_xai_batch_results_id(provider_file_id):
return False
router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
model_name
)
return provider == LlmProviders.XAI.value
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
@ -845,6 +890,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
litellm_parent_otel_span=None,
)
async def _save_xai_refreshed_file_object(
self,
stored: LiteLLM_ManagedFileTable,
file_object: OpenAIFileObject,
) -> OpenAIFileObject | None:
"""Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
sized_file_object: Final = _over_stored_object(file_object, stored_object)
if not await ManagedFileRepository(self.prisma_client).update_file_object(
stored.unified_file_id,
sized_file_object,
if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
):
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
)
return None
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id,
value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
litellm_parent_otel_span=None,
)
return sized_file_object
async def _save_batch_output_file_object(
self,
stored: LiteLLM_ManagedFileTable,
file_object: OpenAIFileObject,
size_bytes: int | None,
) -> None:
"""Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
current_row: LiteLLM_ManagedFileTable | None = stored # rebind-ok: the row as last read
details: OpenAIFileObject = file_object # rebind-ok: a concurrent refresh's details replace the job's copy
for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
if current_row is None:
return
saved = await self._save_xai_refreshed_file_object( # rebind-ok: per attempt
current_row, _with_measured_size(details, size_bytes)
)
if saved is not None or not size_bytes:
return
latest_row = await ManagedFileRepository( # rebind-ok: per attempt
self.prisma_client, use_writer=True
).table.find_first(where={"unified_file_id": stored.unified_file_id})
current_row = (
None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
)
latest_object = None if current_row is None else current_row.file_object # rebind-ok: per attempt
if latest_object is not None and not latest_object.litellm_details_fallback:
details = latest_object
async def store_batch_output_file(
self,
*,
@ -885,6 +981,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
return
if stored_object is not None and not stored_object.litellm_details_fallback:
if (
stored_file is not None
and size_bytes
and not stored_object.bytes
and _is_xai_batch_output(
llm_router=_proxy_llm_router(),
model_id=model_id,
model_name=model_name,
provider_file_id=provider_file_id,
)
):
await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
return
fallback_written_recently: Final = (
@ -924,13 +1032,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
)
xai_batch_output: Final = _is_xai_batch_output(
llm_router=_proxy_llm_router(),
model_id=model_id,
model_name=model_name,
provider_file_id=provider_file_id,
)
if stored_file is not None and xai_batch_output:
await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
return
if stored_file is not None:
await self._save_refreshed_file_object(stored_file, file_object)
return
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
litellm_parent_otel_span=litellm_parent_otel_span,
model_mappings={model_id: provider_file_id} if model_id else {},
user_api_key_dict=owner,
@ -1943,8 +2060,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if refreshed_file_object is None:
return _public_file_object(file_object, file_id)
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
return _public_file_object(refreshed_file_object, file_id)
if not _is_xai_batch_output(
llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
model_id=model_id,
model_name=model_id,
provider_file_id=provider_file_id,
):
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
return _public_file_object(refreshed_file_object, file_id)
saved_file_object: Final = await self._save_xai_refreshed_file_object(
stored_file_object, refreshed_file_object
)
return _public_file_object(
saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
)
except Exception as error:
verbose_logger.warning(
"Failed to refresh batch file object for "

View file

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

View file

@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
# Realtime / streaming
"/v1/realtime",
"/realtime",
"/v1/live/sessions",
"/live/sessions",
# Health & ops
"/health",
"/metrics",

View file

@ -440,10 +440,10 @@ db:
# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
# the pod NotReady and blocks only until in-flight requests actually finish
# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
# waiting the worst-case duration. The drain runs once (the preStop hook and
# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
# SIGKILL.
# waiting the worst-case duration. After SIGTERM the server waits up to
# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
# them, so set terminationGracePeriodSeconds a few seconds above twice
# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
#
# /health/drain is off by default; enable it with
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop

View file

@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr
Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured.
```yaml
ingress:
enabled: true
host: ""
componentHosts:
ui: ui.example.com
backend: backend.example.com
gateway: gateway.example.com
extraPaths: []
ui:
publicUrl: https://ui.example.com
backendUrl: https://backend.example.com
backend:
gatewayUrl: https://gateway.example.com
```

View file

@ -22,12 +22,14 @@ Reminders:
- database.writer.{host,port,dbname} (writer connection pieces)
- database.writer.passwordSecret.{name,usernameKey,passwordKey}
(Secret holding the writer DB username + password)
- database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret)
- database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and
omits DATABASE_PASSWORD / DATABASE_URL so the proxy
mints the URL from an IAM token at startup)
- database.reader.host (optional — enables read-replica routing; reader
.passwordSecret.name is required when set, unless
.useIAMAuth is true)
- database.reader.host / database.reader.passwordSecret.{hostKey,urlKey}
(optional, enables read-replica routing. The Secret
name is required for hostKey, urlKey, and for host
with password auth)
- database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true —
chart emits DATABASE_*_READ_REPLICA env vars and
omits DATABASE_PASSWORD_READ_REPLICA /
@ -35,6 +37,8 @@ Reminders:
reader URL from an IAM token at startup)
- redis.passwordSecret.name (optional — set when redis.host is provided and the
cache requires auth)
- redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret)
(the Secret name is required when either key is set)
- redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from
redis.host / redis.port so the proxy's Cache()
constructs a RedisClusterCache; the cluster client

View file

@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates.
{{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}}
{{- end -}}
{{- define "litellm.ui.publicOrigin" -}}
{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.ui.backendOrigin" -}}
{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.backend.gatewayOrigin" -}}
{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.commonLabels" -}}
app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's
Sensitive values (master key, DB username + password, Redis password) come
only from referenced Secrets; the chart never accepts inline values for them.
The chart never assembles DATABASE_URL itself. It emits only the discrete
The chart never assembles DATABASE_URL itself. When
database.writer.passwordSecret.urlKey is set, the chart passes the operator's
URL through from the Secret. Otherwise it emits only the discrete
DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth)
vars; the proxy's entrypoint (DatabaseURLSettings in
litellm/proxy/db/db_url_settings.py) builds the URL from them and
@ -218,11 +244,11 @@ token instead of a static password. `database.writer.useAzureEntraAuth: true`
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
for Azure Database for PostgreSQL. The two are mutually exclusive.
The read replica is opt-in via `database.reader.host`. The chart emits
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
password auth. When `database.reader.useIAMAuth: true` (or
The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`.
The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete
DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus
DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for password auth.
When `database.reader.useIAMAuth: true` (or
`database.reader.useAzureEntraAuth: true`) it omits
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
same way. Reader token auth only takes effect when the writer uses the same
@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
value: {{ $component.logLevel | quote }}
{{- end }}
{{- with $root.Values.database.writer }}
{{- $hostKey := .passwordSecret.hostKey }}
{{- $urlKey := .passwordSecret.urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (not .host) (not $hostKey) (not $urlKey) }}
{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }}
{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: DATABASE_URL
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
- name: DATABASE_HOST
value: {{ required "database.writer.host is required" .host | quote }}
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT
value: {{ .port | default 5432 | quote }}
- name: DATABASE_USER
@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
key: {{ .passwordSecret.passwordKey | default "password" }}
{{- end }}
{{- end }}
{{- end }}
{{- with $root.Values.database.reader }}
{{- if .host }}
{{- $readerSecret := .passwordSecret | default dict }}
{{- $hostKey := $readerSecret.hostKey }}
{{- $urlKey := $readerSecret.urlKey }}
{{- $readerEnabled := or .host $hostKey $urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }}
{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $readerEnabled }}
{{- if $urlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
{{- end }}
@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
{{- end }}
- name: DATABASE_HOST_READ_REPLICA
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT_READ_REPLICA
value: {{ .port | default 5432 | quote }}
- name: DATABASE_NAME_READ_REPLICA
value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }}
value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }}
{{- if .schema }}
- name: DATABASE_SCHEMA_READ_REPLICA
value: {{ .schema | quote }}
@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- end }}
{{- end }}
{{- end }}
{{- end }}
{{/*
The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of
`prisma migrate deploy`. Without this, every gateway/backend pod also runs
@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods.
tracking, pod lock manager) via its REDIS_* env fallback. An explicit
`general_settings.coordination_redis` block in proxy_config takes
precedence over anything emitted here. */}}
{{- if $root.Values.redis.host }}
{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }}
{{- $hostKey := $redisSecret.hostKey }}
{{- $urlKey := $redisSecret.urlKey }}
{{- if and $root.Values.redis.host $hostKey }}
{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }}
{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }}
{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: REDIS_URL
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else if or $root.Values.redis.host $hostKey }}
- name: REDIS_HOST
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ $root.Values.redis.host | quote }}
{{- end }}
- name: REDIS_PORT
value: {{ $root.Values.redis.port | quote }}
{{- if $root.Values.redis.passwordSecret.name }}
@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods.
configured endpoint; the cluster client discovers the remaining nodes from
CLUSTER SLOTS at startup. */}}
- name: REDIS_CLUSTER_NODES
{{- if $hostKey }}
value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }}
{{- else }}
value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
{{- end }}
{{- end }}
{{- end }}
{{- with $component.extraEnv }}

View file

@ -1,4 +1,6 @@
{{- if .Values.backend.enabled }}
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
apiVersion: apps/v1
kind: Deployment
metadata:
@ -58,7 +60,17 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
{{- if .Values.backend.gatewayUrl }}
- name: LITELLM_GATEWAY_URL
value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }}
{{- end }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml

View file

@ -56,6 +56,10 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH

View file

@ -1,3 +1,36 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if $hasComponentHosts }}
{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }}
{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }}
{{- end }}
{{- if not .Values.gateway.enabled }}
{{- fail "ingress.componentHosts requires gateway.enabled=true" }}
{{- end }}
{{- if not .Values.backend.enabled }}
{{- fail "ingress.componentHosts requires backend.enabled=true" }}
{{- end }}
{{- if not .Values.ui.enabled }}
{{- fail "ingress.componentHosts requires ui.enabled=true" }}
{{- end }}
{{- if .Values.ingress.host }}
{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if .Values.ingress.extraPaths }}
{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if not (trim .Values.ui.backendUrl) }}
{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.ui.backendOrigin" . }}
{{- if not (trim .Values.backend.gatewayUrl) }}
{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.backend.gatewayOrigin" . }}
{{- if not (trim .Values.ui.publicUrl) }}
{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- end }}
{{- if .Values.ingress.enabled -}}
{{- $gatewayName := include "litellm.gateway.fullname" . -}}
{{- $backendName := include "litellm.backend.fullname" . -}}
@ -90,6 +123,66 @@
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
{{- if $hasComponentHosts }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
name: {{ include "litellm.fullname" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
{{- with .Values.ingress.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
{{- with .Values.ingress.className }}
ingressClassName: {{ . | quote }}
{{- end }}
{{- with .Values.ingress.tls }}
tls:
{{- toYaml . | nindent 4 }}
{{- end }}
rules:
- host: {{ $componentHosts.ui | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- host: {{ $componentHosts.backend | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $backendName }}
port:
number: {{ $backendPort }}
- host: {{ $componentHosts.gateway | quote }}
http:
paths:
{{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }}
- path: /lens-ingest
pathType: Prefix
backend:
service:
name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }}
port:
number: {{ .Values.lensWorker.service.port }}
{{- end }}
- path: /
pathType: Prefix
backend:
service:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
{{- else }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -217,3 +310,4 @@ spec:
port:
number: {{ $backendPort }}
{{- end }}
{{- end }}

View file

@ -1,3 +1,5 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if .Values.ui.enabled }}
apiVersion: apps/v1
kind: Deployment
@ -56,9 +58,9 @@ spec:
- name: LITELLM_LOG
value: {{ .Values.ui.logLevel | quote }}
{{- end }}
{{- if .Values.ui.backendUrl }}
{{- if and .Values.ui.backendUrl $hasComponentHosts }}
- name: LITELLM_BACKEND_URL
value: {{ .Values.ui.backendUrl | quote }}
value: {{ include "litellm.ui.backendOrigin" . | quote }}
{{- end }}
{{- with .Values.ui.extraEnv }}
{{- toYaml . | nindent 12 }}

View file

@ -0,0 +1,441 @@
suite: ingress component hosts
templates:
- ingress.yaml
- backend/deployment.yaml
- gateway/deployment.yaml
- gateway/configmap.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: routes each hostname to its component and injects the external URLs
template: ingress.yaml
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
lensWorker.mode: disabled
ingress.tls:
- hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- equal:
path: spec.rules[0].host
value: ui.example.com
- equal:
path: spec.tls[0]
value:
hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
- equal:
path: spec.rules[0].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- equal:
path: spec.rules[1].host
value: backend.example.com
- equal:
path: spec.rules[1].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- equal:
path: spec.rules[2].host
value: gateway.example.com
- equal:
path: spec.rules[2].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: keeps Lens ingestion ahead of the gateway catch-all
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
lensWorker.mode: external
lensWorker.externalServiceName: existing-lens
lensWorker.externalUrl: https://lens.example.com
lensWorker.publicUrl: https://lens.example.com
lensWorker.serviceTokenSecret.name: existing-lens-signing
lensWorker.gateway.secretName: existing-lens-signing
asserts:
- equal:
path: spec.rules[2].http.paths[0].path
value: /lens-ingest
- equal:
path: spec.rules[2].http.paths[0].backend.service.name
value: existing-lens
- equal:
path: spec.rules[2].http.paths[1].path
value: /
template: ingress.yaml
- it: rejects a partial component host configuration
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together
template: ingress.yaml
- it: rejects component hosts when the gateway is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
gateway.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires gateway.enabled=true
template: ingress.yaml
- it: rejects component hosts when the backend is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
backend.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires backend.enabled=true
template: ingress.yaml
- it: rejects component hosts when the UI is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
ui.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires ui.enabled=true
template: ingress.yaml
- it: rejects a legacy host alongside component hosts
set:
ingress.enabled: true
ingress.host: legacy.example.com
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
asserts:
- failedTemplate:
errorMessage: ingress.host must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: rejects extra paths alongside component hosts
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ingress.extraPaths:
- path: /custom
asserts:
- failedTemplate:
errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the public UI origin in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.publicUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: exposes the gateway URL in the backend Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: " https://gateway.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: https://gateway.example.com
- it: preserves the backend gateway URL verbatim in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
backend.gatewayUrl: " https://gateway.example.com/api/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: " https://gateway.example.com/api/ "
- it: exposes the backend URL in the UI Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: " https://backend.example.com/ "
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
template: ui/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
value: https://backend.example.com
- it: omits the backend URL from the UI Deployment in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
ui.backendUrl: https://litellm.example.com/api
template: ui/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
any: true
- it: enables credentialed CORS from the public UI origin on the backend
set:
ui.publicUrl: https://ui.example.com
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: lets backend extraEnv override generated CORS values
set:
ui.publicUrl: https://ui.example.com
backend.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://backend-override.example.com
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "false"
template: backend/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[1].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[1].value
value: "true"
- equal:
path: spec.template.spec.containers[0].env[10].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[10].value
value: https://backend-override.example.com
- equal:
path: spec.template.spec.containers[0].env[11].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[11].value
value: "false"
- it: exposes the public UI origin to the gateway without enabling credentials
set:
ui.publicUrl: https://ui.example.com
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: trims the public UI origin before setting backend CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: trims the public UI origin before setting gateway CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: rejects a UI public URL with a path
set:
ui.publicUrl: https://ui.example.com/app
template: gateway/deployment.yaml
asserts:
- failedTemplate:
errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path
- it: lets gateway extraEnv override the generated CORS origin
set:
ui.publicUrl: https://ui.example.com
gateway.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://gateway-override.example.com
template: gateway/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[9].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[9].value
value: https://gateway-override.example.com
---
suite: component host UI backend origin validation
templates:
- ingress.yaml
values:
- ./values/required.yaml
tests:
- it: rejects a pathful UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com/api
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set
- it: rejects a pathful backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com/v1
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
- it: rejects a backend gateway URL without a scheme in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: gateway.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set

View file

@ -191,3 +191,400 @@ tests:
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- it: writer hostKey reads the host from the Secret and supports IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.schema: ""
database.writer.sslMode: ""
database.writer.sslRootCert: ""
database.writer.passwordSecret.urlKey: connection-url
database.connectionPool.enabled: true
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: backend/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: backend/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: migrations-job.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: migrations-job.yaml
- it: writer rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.writer.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set"
- it: writer rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.schema: public
database.writer.sslMode: require
database.writer.sslRootCert: /etc/ssl/ca.crt
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with Entra auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer requires an endpoint
template: gateway/deployment.yaml
set:
database.writer.host: ""
asserts:
- failedTemplate:
errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required"
- it: writer hostKey requires a Secret name
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set"
- it: reader renders without replica env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
database.reader.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
any: true
- it: reader hostKey enables replica routing and reads its host from the Secret
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.hostKey: hostname
database.reader.dbname: litellm
database.reader.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
value: litellm
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: reader urlKey emits only the replica URL
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA_READ_REPLICA
any: true
- it: reader rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.reader.host: reader.example.com
database.reader.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set"
- it: reader rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.dbname: litellm
database.reader.schema: public
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader rejects urlKey with token auth
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.urlKey: connection-url
database.reader.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader urlKey requires a Secret name
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set"

View file

@ -21,6 +21,27 @@ tests:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway renders without redis env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway emits host, port and password when redis is configured
template: gateway/deployment.yaml
set:
@ -68,6 +89,27 @@ tests:
name: redis-secret
key: redis-password
- it: backend renders without redis env vars when passwordSecret is null
template: backend/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway omits REDIS_PASSWORD for an auth-less redis
template: gateway/deployment.yaml
set:
@ -107,3 +149,174 @@ tests:
content:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway gets redis host from hostKey and keeps port and password
template: gateway/deployment.yaml
set:
redis.host: ""
redis.port: 6380
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.passwordKey: redis-password
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6380"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: redis-password
- it: backend gets redis host from hostKey with the same coordination env
template: backend/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6379"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: password
- it: gateway expands the redis cluster seed from hostKey at runtime
template: gateway/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.cluster: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
value: '[{"host":"$(REDIS_HOST)","port":6379}]'
- equal:
path: spec.template.spec.containers[0].env[8].name
value: REDIS_HOST
- equal:
path: spec.template.spec.containers[0].env[11].name
value: REDIS_CLUSTER_NODES
- it: gateway gets REDIS_URL from urlKey without discrete Redis env
template: gateway/deployment.yaml
set:
redis.passwordSecret.name: redis-secret
redis.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
valueFrom:
secretKeyRef:
name: redis-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
any: true
- it: redis rejects host together with hostKey
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set"
- it: redis urlKey rejects host
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects hostKey
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects cluster mode
template: gateway/deployment.yaml
set:
redis.cluster: true
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis hostKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
- it: redis urlKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.urlKey: connection-url
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"

View file

@ -23,7 +23,11 @@ ingress:
# wildcard pathType, so that rule could never match there.
controller: alb
annotations: {}
host: "" # optional; if set, becomes the rule's host
host: "" # optional; if set, becomes the single ingress rule's host
componentHosts:
ui: ""
backend: ""
gateway: ""
tls: []
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
# UI / gateway / backend path is still rendered, these entries are placed
@ -208,6 +212,8 @@ database:
name: litellm-writer-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the writer database host
urlKey: "" # Secret key containing the complete writer database URL
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
# sslRootCert on its own implies sslMode verify-full
@ -229,6 +235,8 @@ database:
name: litellm-reader-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the reader database host
urlKey: "" # Secret key containing the complete reader database URL
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
@ -274,6 +282,8 @@ redis:
passwordSecret:
name: "" # Leave empty for auth-less Redis
passwordKey: password
hostKey: "" # Secret key containing the Redis host
urlKey: "" # Secret key containing the complete Redis URL
# ---------- gateway (LLM data plane) ----------
gateway:
@ -493,6 +503,7 @@ gateway:
backend:
enabled: true
logLevel: INFO
gatewayUrl: ""
extraEnv: []
envConfigMaps: []
envSecrets: []
@ -575,9 +586,10 @@ ui:
service:
type: ClusterIP
port: 3000
# The dashboard expects to know where to reach the backend API. Set this to
# the externally-routable URL (typically the ingress host + /api or similar).
# The backend's public origin, passed to the UI only when ingress.componentHosts is configured.
backendUrl: ""
# Public origin of the UI, used for backend and gateway CORS.
publicUrl: ""
resources:
requests:
cpu: 500m

View file

@ -1,5 +1,4 @@
import random
import re
import time
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -10,12 +9,11 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
migration_ddl_lock_timeout,
migration_lock_timeout,
)
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
if TYPE_CHECKING:
import psycopg
@ -30,29 +28,12 @@ def _migration_url(database_url: str, direct_url: "str | None") -> str:
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
def _with_ddl_lock_timeout(url: str) -> str:
split: Final = urlsplit(url)
pairs: Final = tuple(parse_qsl(split.query))
if any(key == "pgbouncer" and value == "true" for key, value in pairs):
return url
existing: Final = next((value for key, value in pairs if key == "options"), "")
if _LOCK_TIMEOUT_PINNED_RE.search(existing):
return url
timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
return urlunsplit(split._replace(query=urlencode(rewritten)))
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
database_url: Final = environment.get("DATABASE_URL")
if not database_url:
return environment
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
return {
**environment,
"DATABASE_URL": _with_ddl_lock_timeout(migration_url),
}
return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
@dataclass(frozen=True, slots=True)
@ -75,10 +56,15 @@ class MigrationCoordinator:
def check_connection(self) -> None:
self.connection.execute("SELECT 1")
def release_statement_snapshot(self) -> None:
# a parameterized statement leaves its portal snapshot registered on the backend until the
# next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
self.connection.execute("SELECT 1")
def acquire_prisma_lock(self) -> None:
deadline: Final = time.monotonic() + migration_lock_timeout()
while time.monotonic() < deadline:
if _try_lock(self.connection, 72707369):
if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
return
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
raise RuntimeError(
@ -115,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
)
@contextmanager
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
transaction pooler can never strand that lock on an idle backend after the command exits."""
with migration_lock(database_url) as coordinator:
coordinator.acquire_prisma_lock()
coordinator.release_statement_snapshot()
yield
@contextmanager
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit

View file

@ -0,0 +1,14 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_TelemetryReport" (
"id" TEXT NOT NULL,
"window_start" DOUBLE PRECISION NOT NULL,
"window_end" DOUBLE PRECISION NOT NULL,
"report" JSONB NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_TelemetryReport_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_window_end_idx" ON "LiteLLM_TelemetryReport"("window_end");
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_created_at_idx" ON "LiteLLM_TelemetryReport"("created_at");

View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;

View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[];

View file

@ -60,13 +60,11 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
BOOTSTRAP_ARG = "--version"
PRISMA_CONSOLE_SCRIPT = "prisma"
@ -113,11 +111,6 @@ def migration_lock_timeout() -> float:
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
def migration_ddl_lock_timeout() -> float:
"""Seconds one migration statement may wait on a table lock before Postgres cancels it."""
return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
def prisma_bootstrap_timeout() -> float:
"""Seconds the one-time Node toolchain install may run for."""
return _timeout_from_env(

View file

@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
mcp_tool_permissions Json?
access_agent_ids String[] @default([])
access_skill_names String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])
@ -2045,3 +2047,14 @@ model LiteLLM_LensTraceSignal {
@@id([trace_id, trace_ref])
}
model LiteLLM_TelemetryReport {
id String @id
window_start Float
window_end Float
report Json
created_at DateTime @default(now())
@@index([window_end])
@@index([created_at])
}

View file

@ -9,6 +9,7 @@ import tempfile
import threading
import time
from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
@ -18,7 +19,6 @@ from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
ensure_prisma_toolchain,
@ -1022,7 +1022,11 @@ class ProxyExtrasDBManager:
def _setup_database_v2(use_migrate: bool) -> bool:
if not use_migrate:
return ProxyExtrasDBManager._run_database_v2(False)
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
from litellm_proxy_extras.migration_lock import (
migration_environment,
migration_lock,
prisma_migration_lock,
)
from litellm_proxy_extras.migration_recovery import (
baseline_current_schema,
recover_completed_migration,
@ -1056,7 +1060,9 @@ class ProxyExtrasDBManager:
migration_environment(_get_prisma_env()),
)
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
while not ProxyExtrasDBManager._run_database_v2(
True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
):
continue
return True
@ -1065,6 +1071,7 @@ class ProxyExtrasDBManager:
use_migrate: bool,
recover_completed: Callable[[str], bool] = lambda name: False,
baseline_existing: "Callable[[str], None] | None" = None,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> bool:
"""
v2 migration resolver (opt-in via --use_v2_migration_resolver).
@ -1127,11 +1134,12 @@ class ProxyExtrasDBManager:
while not budget.exhausted:
attempt_started = time.monotonic()
try:
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
with deploy_lock():
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
return True
@ -1156,6 +1164,7 @@ class ProxyExtrasDBManager:
e,
budget,
time.monotonic() - attempt_started,
deploy_lock=deploy_lock,
)
if next_budget.attempts_left < budget.attempts_left:
@ -1165,10 +1174,9 @@ class ProxyExtrasDBManager:
raise RuntimeError(
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
"Check database connectivity, load, long-running transactions on the migrated tables, "
"and _prisma_migrations ledger state. Raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
"Check database connectivity, load and _prisma_migrations ledger state. "
"If a migration hit SET LOCAL lock_timeout, check for long-running transactions on the migrated tables. "
f"Raise {PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
)
finally:
os.chdir(original_dir)
@ -1189,15 +1197,19 @@ class ProxyExtrasDBManager:
return match.group(1) if match else None
@staticmethod
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
def _v2_roll_back_migration_best_effort(
migration_name: str,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> None:
from litellm_proxy_extras.migration_lock import migration_environment
try:
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
with deploy_lock():
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass
@ -1206,6 +1218,7 @@ class ProxyExtrasDBManager:
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
attempt_seconds: float = 0.0,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
@ -1227,14 +1240,14 @@ class ProxyExtrasDBManager:
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
logger.info(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
logger.info(
@ -1278,17 +1291,17 @@ class ProxyExtrasDBManager:
"rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
logger.warning(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
"Migration %s timed out waiting for a table lock after hitting SET LOCAL lock_timeout, rolling "
"its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
raise RuntimeError(

View file

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

View file

@ -3406,6 +3406,7 @@ dependencies = [
name = "litellm-auth-gcp"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"gcp_auth",
"google-cloud-auth",
"http 1.4.2",
@ -3922,13 +3923,16 @@ version = "0.1.0"
dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-host",
"parking_lot",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
]

View file

@ -80,6 +80,7 @@ http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
parking_lot = "0.12.5"
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }

View file

@ -11,6 +11,7 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"]
[dependencies]
litellm-auth-types.workspace = true
base64.workspace = true
moka.workspace = true
serde.workspace = true
serde_json.workspace = true

View file

@ -1,4 +1,4 @@
use std::{future::Future, path::Path, pin::Pin, sync::Arc};
use std::{future::Future, pin::Pin, sync::Arc};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
use crate::{
config::{
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env,
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
is_base64_encoded_json, names_a_credentials_file, non_empty_env,
},
constants::{
CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
@ -176,14 +177,7 @@ impl VertexProviderLoader for GcpProviderLoader {
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
let service_account = if Path::new(configured).is_file() {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)?;
Arc::new(service_account)
Arc::new(load_trusted_service_account(configured.expose())?)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
@ -240,6 +234,23 @@ enum CredentialCacheKey {
Adc,
}
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
Set it to a readable file path, or to the credentials JSON itself.";
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
BASE64_CREDENTIALS.into(),
)));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)
}
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
"Vertex AI credentials",
@ -264,6 +275,73 @@ mod tests {
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
}
#[rstest::rstest]
#[case::missing_file(
"/nonexistent/vertex-key.json",
"failed to open application credentials file"
)]
#[case::malformed_json(
" {\"type\": \"service_account\",",
"failed to deserialize ApplicationCredentials"
)]
#[case::quoted_json(
"'{\"type\": \"service_account\"}'",
"failed to deserialize ApplicationCredentials"
)]
#[tokio::test]
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
#[case] configured: &str,
#[case] expected: &str,
) {
let Err(error) = GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
else {
panic!("{configured:?} should not load");
};
assert!(error.to_string().contains(expected), "{error}");
}
#[rstest::rstest]
#[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
#[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
#[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
#[tokio::test]
async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
assert!(
trusted_load_error(configured)
.await
.contains("is base64-encoded JSON, not the JSON itself")
);
}
async fn trusted_load_error(configured: &str) -> String {
match GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => error.to_string(),
}
}
#[rstest::rstest]
#[tokio::test]
async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let absolute_path = std::fs::canonicalize(key.path()).unwrap();
assert_eq!(
trusted_load_error(relative_name).await,
trusted_load_error(absolute_path.to_str().unwrap()).await
);
}
#[rstest::rstest]
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
assert_ne!(

View file

@ -1,5 +1,9 @@
use std::{collections::BTreeMap, path::Path};
use base64::{
Engine, alphabet,
engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
};
use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced};
use serde_json::{Map, Value};
@ -80,9 +84,10 @@ pub fn get_vertex_ai_project_from_credentials(
CredentialSource::Inline(configured) => configured.expose().to_string(),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
match std::fs::read_to_string(configured) {
Ok(contents) if Path::new(configured).is_file() => contents,
_ => configured.to_string(),
if names_a_credentials_file(configured) {
std::fs::read_to_string(configured).ok()?
} else {
configured.to_string()
}
}
CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?,
@ -105,6 +110,32 @@ pub fn get_vertex_ai_location(
.or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup))
}
const INLINE_JSON_LEADING_NOISE: [char; 7] = [' ', '\t', '\r', '\n', '\u{feff}', '"', '\''];
fn is_inline_json(configured: &str) -> bool {
configured
.trim_start_matches(INLINE_JSON_LEADING_NOISE)
.starts_with('{')
}
pub(crate) fn names_a_credentials_file(configured: &str) -> bool {
!is_inline_json(configured) || Path::new(configured).is_file()
}
const BASE64_ANY_PADDING: GeneralPurpose = GeneralPurpose::new(
&alphabet::STANDARD,
GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
);
pub(crate) fn is_base64_encoded_json(configured: &str) -> bool {
let compact: String = configured.split_whitespace().collect();
let quoted = compact.trim_matches(['"', '\'']);
let encoded = quoted.strip_prefix("base64:").unwrap_or(quoted);
BASE64_ANY_PADDING
.decode(encoded.replace('-', "+").replace('_', "/"))
.is_ok_and(|decoded| serde_json::from_slice::<Map<String, Value>>(&decoded).is_ok())
}
#[derive(Clone, Debug)]
pub enum CredentialSource {
Inline(SecretValue),

View file

@ -93,3 +93,22 @@ fn project_is_read_from_the_credentials_that_would_authenticate() {
None
);
}
#[rstest]
fn project_is_read_from_a_relative_key_file_named_like_json() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let from_file = VertexConfig::from_params(&VertexParams {
vertex_credentials: Some(relative_name.into()),
..VertexParams::default()
});
assert_eq!(
get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(),
Some("file-project")
);
}

View file

@ -9,7 +9,7 @@ use litellm_host::{
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing, epoch_seconds},
};
use litellm_host_python::{HookStep, from_py, missing_state, to_py};
use litellm_host_python::{Error as HostError, HookStep, from_py, to_py};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
@ -160,7 +160,7 @@ impl LegacyLogging {
self.response
.as_ref()
.map(|response| HookStep::Ready(response.clone_ref(py)))
.ok_or_else(missing_state)
.ok_or_else(|| HostError::MissingState.into())
}
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
@ -209,7 +209,7 @@ impl LegacyLogging {
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or_else(missing_state)?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
py,
(
@ -525,7 +525,7 @@ impl LegacyLogging {
return self.stream_failure(py);
}
if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
let error = self.error.as_ref().ok_or_else(missing_state)?;
let error = self.error.as_ref().ok_or(HostError::MissingState)?;
return Ok(HookStep::Await(
DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
Self::resume_deployment_failure,
@ -536,7 +536,7 @@ impl LegacyLogging {
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
if self.stream_billing().is_none() {
return Err(missing_state());
return Err(HostError::MissingState.into());
}
if let Some(key) = &self.cache_key {
head.bind(py).set_item("cache_key", key)?;
@ -551,7 +551,7 @@ impl LegacyLogging {
}
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
let stream = self.stream.as_mut().ok_or(HostError::MissingState)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
}

View file

@ -3,7 +3,7 @@
Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks
- R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/execution.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
- R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime
- R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states
- R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for

View file

@ -1,6 +1,6 @@
use std::str;
use bytes::{Buf, BufMut, BytesMut};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use tokio_util::codec::{Decoder, Encoder};
use crate::SseError;
@ -29,38 +29,107 @@ impl Decoder for SseCodec {
return Ok(None);
}
while let Some(end) = block_end(src) {
let block = src.split_to(end);
let pending = lines(&block)
.map(|(line, _)| line)
.take_while(|line| !line.is_empty())
.try_fold(Pending::default(), Pending::apply)?;
if let Some(event) = pending.dispatch() {
if let Some(event) = event_in(&src.split_to(end))? {
return Ok(Some(event));
}
}
Ok(None)
}
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
Ok(None)
/// An event the stream ended inside of is dropped, as the SSE spec requires.
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
self.decode(pending)
}
}
pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
let mut codec = SseCodec::default();
let mut buffer = BytesMut::from(wire);
let mut events = Vec::new();
while let Some(event) = codec.decode_eof(&mut buffer)? {
events.push(event);
}
Ok(events)
}
impl SseCodec {
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
if self.past_bom {
return true;
}
if src.starts_with(BOM) {
src.advance(BOM.len());
} else if BOM.starts_with(src) {
return false;
}
self.past_bom = true;
true
skip_bom(&mut self.past_bom, src)
}
}
fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
Ok(lines(block)
.map(|(line, _)| line)
.take_while(|line| !line.is_empty())
.try_fold(Pending::default(), Pending::apply)?
.dispatch())
}
/// One event together with the exact wire bytes that carried it, for relays that forward a
/// provider's bytes untouched but still need to see each event go by.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SseBlock {
pub bytes: Bytes,
pub event: SseEvent,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct SseBlockCodec {
past_bom: bool,
}
impl Decoder for SseBlockCodec {
type Item = SseBlock;
type Error = SseError;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
self.decode_block(src, false)
}
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
self.decode_block(pending, true)
}
}
impl SseBlockCodec {
fn decode_block(
&mut self,
src: &mut BytesMut,
eof: bool,
) -> Result<Option<SseBlock>, SseError> {
if !skip_bom(&mut self.past_bom, src) {
return Ok(None);
}
while let Some(end) = block_end(src) {
if !eof && end == src.len() && src.ends_with(b"\r") {
return Ok(None);
}
let block = src.split_to(end);
if let Some(event) = event_in(&block)? {
return Ok(Some(SseBlock {
bytes: block.freeze(),
event,
}));
}
}
Ok(None)
}
}
fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
if *past_bom {
return true;
}
if src.starts_with(BOM) {
src.advance(BOM.len());
} else if BOM.starts_with(src) {
return false;
}
*past_bom = true;
true
}
fn block_end(bytes: &[u8]) -> Option<usize> {
lines(bytes)
.find(|(line, _)| line.is_empty())
@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
Ok(())
}
}
#[cfg(test)]
mod tests {
use futures_util::{StreamExt, TryStreamExt, stream};
use rstest::rstest;
use super::*;
use crate::frames;
fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
wire.chunks(size).map(Bytes::copy_from_slice).collect()
}
async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
frames(
stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
SseBlockCodec::default(),
)
.try_collect()
.await
}
#[rstest]
#[case::lf(
b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
b"event: delta\ndata: a\ndata: b\n\n"
)]
#[case::crlf(
b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
)]
#[case::cr(
b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
b"event: delta\rdata: a\rdata: b\r\r"
)]
#[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
#[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
#[case::extensions(
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
)]
#[tokio::test]
async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
#[case] wire: &[u8],
#[case] expected: &[u8],
#[values(1, 3, 4096)] size: usize,
) {
let blocks = blocks(wire, size).await.unwrap();
let expected_events = decode_all(expected).unwrap();
assert_eq!(expected_events.len(), 1);
assert_eq!(
blocks,
vec![SseBlock {
bytes: Bytes::copy_from_slice(expected),
event: expected_events[0].clone(),
}]
);
assert_eq!(expected_events, decode_all(wire).unwrap());
}
#[rstest]
#[case::empty(b"", vec![])]
#[case::unterminated(b"data: partial\n", vec![])]
#[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
#[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
#[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
#[tokio::test]
async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
#[case] wire: &[u8],
#[case] expected: Vec<&[u8]>,
) {
let mut blocks = Box::pin(frames(
stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
SseBlockCodec::default(),
));
for expected in expected {
assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
}
assert!(blocks.next().await.is_none());
}
#[rstest]
fn bulk_decode_preserves_every_field() {
let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
assert_eq!(
events,
vec![SseEvent {
event: Some("delta".into()),
data: "a\nb".into(),
id: Some("1".into()),
retry: Some(20),
}]
);
}
#[rstest]
#[case::only_a_tail(b"data: partial\n", vec![])]
#[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
#[case::no_terminator_at_all(b"data: partial", vec![])]
fn bulk_decode_drops_the_event_a_stream_ends_inside(
#[case] wire: &[u8],
#[case] expected: Vec<&str>,
) {
let data: Vec<String> = decode_all(wire)
.unwrap()
.into_iter()
.map(|e| e.data)
.collect();
assert_eq!(data, expected);
}
#[rstest]
#[case::trailing_blank_lines(b"data: x\n\n\n\n")]
#[case::trailing_comment(b"data: x\n\n: bye\n\n")]
#[case::lone_cr_at_eof(b"data: x\r\r")]
fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
assert_eq!(decode_all(wire).unwrap().len(), 1);
}
#[rstest]
fn bulk_decode_rejects_invalid_utf8() {
assert!(matches!(
decode_all(b"data: \xff\n\n"),
Err(SseError::InvalidUtf8(_))
));
}
#[rstest]
#[tokio::test]
async fn block_framing_preserves_transport_errors() {
let input = stream::iter([
Ok(Bytes::from_static(b"data: first\n\n")),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
]);
let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
assert_eq!(
blocks.next().await.unwrap().unwrap().bytes,
Bytes::from_static(b"data: first\n\n")
);
let Some(Err(SseError::Body(error))) = blocks.next().await else {
panic!("expected body error");
};
let cause = error
.get_ref()
.unwrap()
.downcast_ref::<std::io::Error>()
.unwrap();
assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
assert!(blocks.next().await.is_none());
}
#[rstest]
#[tokio::test]
async fn invalid_utf8_stops_block_framing_after_prior_events() {
let mut blocks = Box::pin(frames(
stream::iter(
pieces(b"data: first\n\ndata: \xff\n\n", 1)
.into_iter()
.map(Ok::<_, std::io::Error>),
),
SseBlockCodec::default(),
));
assert_eq!(
blocks.next().await.unwrap().unwrap().bytes,
Bytes::from_static(b"data: first\n\n")
);
assert!(matches!(
blocks.next().await,
Some(Err(SseError::InvalidUtf8(_)))
));
assert!(blocks.next().await.is_none());
}
}

View file

@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
}
#[rstest]
#[case::unterminated_single(b"data: partial\n", vec![])]
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
#[tokio::test]
async fn eof_dispatches_only_terminated_events(
async fn eof_between_events_ends_the_stream_cleanly(
#[case] wire: &[u8],
#[case] expected: Vec<SseEvent>,
) {
@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
);
}
#[rstest]
#[case::unterminated_single(b"data: partial\n", vec![])]
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
#[case::partial_bom(b"\xEF\xBB", vec![])]
#[tokio::test]
async fn eof_inside_an_event_drops_it_after_the_terminated_events(
#[case] wire: &[u8],
#[case] expected: Vec<SseEvent>,
) {
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
}
#[rstest]
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]

View file

@ -6,7 +6,7 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines
`src/native.rs` belongs here: it runs a generic machine through the Python runtime and owns its pending execution and abort handle. Keep provider selection, request projection and public exception policy out of it. A rename to `machine_runner.rs` is optional and must not change behavior
The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
`run_call` returns the unstarted `Execution` and never names a Python stream class or driver. The consumer drives it from Python and supplies the step classes through `StepTypes` rather than have this crate import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy
@ -28,6 +28,7 @@ Terminal notifications share the selected response or exception. An ordinary not
- `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
- A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is
- A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed
- Consumers import this crate's error as `use litellm_host_python::Error as HostError`, so it never shadows their own `Error`
- Use standard PyO3 ownership and conversion APIs
- Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention. Binding/unbinding does not copy payloads
- Use `pythonize` for selected Serde data, never a JSON-text round trip. Share conversion with `Pythonized<T>`
@ -37,7 +38,7 @@ Terminal notifications share the selected response or exception. An ordinary not
- Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release
- Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal
- Keep coroutine driving in the shared Python driver and the native handle
- Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
- Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
- Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract
- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html)
- [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html)

View file

@ -6,11 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-coroutine.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
parking_lot.workspace = true
serde.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
pyo3.workspace = true

View file

@ -0,0 +1,58 @@
use std::ops::ControlFlow;
use litellm_coroutine::{Abandoned, Co, Reply};
use pyo3::prelude::*;
use crate::Error;
use crate::handle::Suspension;
pub(crate) type Answer = PyResult<Py<PyAny>>;
pub(crate) struct Yielded {
pub(crate) suspension: Suspension,
pub(crate) reply: Reply<Answer>,
}
/// Hands suspensions to the Python caller and receives its answers.
#[derive(Clone)]
pub(crate) struct PythonCaller {
co: Co<Yielded>,
}
impl PythonCaller {
/// Connects the call body to its coroutine's yield handle.
pub(crate) fn new(co: Co<Yielded>) -> Self {
Self { co }
}
/// Asks the caller to await Python work and return its value or exception.
pub(crate) async fn await_python(&self, awaitable: Py<PyAny>) -> Answer {
self.suspend(Suspension::Await(awaitable)).await
}
/// Hands the caller a stream head and waits for demand or closure.
pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::StreamOpen(head)).await
}
/// Delivers a chunk and waits for the caller to request another or close.
pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::StreamChunk(chunk)).await
}
/// Maps stream demand to continuation and any error to detachment.
async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> {
match self.suspend(suspension).await {
Ok(_) => ControlFlow::Continue(()),
Err(_) => ControlFlow::Break(()),
}
}
/// Yields a suspension with its reply handle and waits for the answer.
async fn suspend(&self, suspension: Suspension) -> Answer {
match self.co.yield_(|reply| Yielded { suspension, reply }).await {
Ok(answer) => answer,
Err(Abandoned) => Err(Error::MissingState.into()),
}
}
}

File diff suppressed because it is too large Load diff

View file

@ -1,20 +1,58 @@
use pyo3::prelude::*;
use pyo3::{
PyErr,
exceptions::{PyRuntimeError, PyTypeError},
};
/// A failure host-python raises itself, surfacing in Python as the exception `PyErr::from` picks.
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("missing native call state")]
MissingState,
#[error("execution is already running")]
ExecutionRunning,
#[error("execution is closed")]
ExecutionClosed,
#[error("can't send non-None value to a just-started execution")]
SendBeforeStart,
#[error("execution requires start before resume and can only start once")]
ExecutionOutOfOrder,
#[error("execution was closed while running")]
ExecutionClosedWhileRunning,
#[error("native execution closed")]
NativeExecutionClosed,
#[error("synchronous native routes cannot run from a Tokio context; use the async route")]
SyncInsideTokio,
#[error("the context body produced no value")]
ContextBodyProducedNothing,
#[error("file read must return bytes or str, got {got}")]
FileContentType { got: String },
}
impl From<Error> for PyErr {
fn from(error: Error) -> Self {
match error {
Error::SendBeforeStart | Error::FileContentType { .. } => {
PyTypeError::new_err(error.to_string())
}
Error::MissingState
| Error::ExecutionRunning
| Error::ExecutionClosed
| Error::ExecutionOutOfOrder
| Error::ExecutionClosedWhileRunning
| Error::NativeExecutionClosed
| Error::SyncInsideTokio
| Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()),
}
}
}
/// Why a custom operation the host answered did not produce a result: the route's own code
/// rejected it, which the route classifies like any other native failure, or Python code
/// raised, which reaches the caller as it was raised.
#[derive(Debug)]
#[derive(Debug, thiserror::Error)]
pub enum InvokeError<E> {
#[error("{0}")]
Native(E),
Python(PyErr),
}
impl<E> From<PyErr> for InvokeError<E> {
fn from(error: PyErr) -> Self {
Self::Python(error)
}
}
pub fn missing_state() -> PyErr {
pyo3::exceptions::PyRuntimeError::new_err("missing native call state")
#[error(transparent)]
Python(#[from] PyErr),
}

View file

@ -3,13 +3,14 @@
use bytes::Bytes;
use pyo3::{
exceptions::PyTypeError,
gc::{PyTraverseError, PyVisit},
prelude::*,
pybacked::PyBackedBytes,
types::{PyBytes, PyString},
};
use crate::Error;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FileContent {
pub bytes: Bytes,
@ -50,10 +51,10 @@ impl PythonFileReader {
} else if value.is_instance_of::<PyBytes>() {
py_bytes(&value)?
} else {
return Err(PyTypeError::new_err(format!(
"file read must return bytes or str, got {}",
value.get_type(),
)));
return Err(Error::FileContentType {
got: value.get_type().to_string(),
}
.into());
};
Ok(FileContent {
bytes,

View file

@ -3,7 +3,9 @@ use std::sync::{
atomic::{AtomicU64, Ordering},
};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
use pyo3::prelude::*;
use crate::Error;
/// The caller's `contextvars` context, captured at the Python entry point so blocking Python
/// work started from Rust sees the same request-local values as the Python caller.
@ -49,7 +51,7 @@ impl PythonContext {
.lock()
.expect("context value slot poisoned")
.take()
.ok_or_else(|| PyRuntimeError::new_err("the context body produced no value"))
.ok_or_else(|| Error::ContextBodyProducedNothing.into())
}
}

View file

@ -1,19 +1,30 @@
use std::panic::{AssertUnwindSafe, catch_unwind};
use crate::panic_to_pyerr;
use pyo3::exceptions::{PyBaseException, PyRuntimeError};
use crate::{Error, panic_to_pyerr};
use pyo3::exceptions::PyBaseException;
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
pub type StepTypes = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
pub enum ExecutionStep {
Return(Py<PyAny>),
/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
/// boundary the caller's stream reads past.
pub enum Suspension {
Await(Py<PyAny>),
/// The call streams: the caller gets a stream over this execution carrying this head,
/// and the execution stays suspended until the stream asks for a chunk.
Open(Py<PyAny>),
Yield(Py<PyAny>),
StreamOpen(Py<PyAny>),
StreamChunk(Py<PyAny>),
}
enum Resume {
Send(Py<PyAny>),
Throw(PyErr),
}
pub enum ExecutionStep {
Suspended(Suspension),
Return(Py<PyAny>),
}
pub trait ExecutionBody: Send + Sync {
@ -31,77 +42,59 @@ enum ExecutionState {
#[pyclass]
pub struct Execution {
state: ExecutionState,
lifecycle: PythonLifecycle,
steps: StepTypes,
}
impl Execution {
pub fn new(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
pub fn new(body: impl ExecutionBody + 'static, steps: StepTypes) -> Self {
Self {
state: ExecutionState::Created(Box::new(body)),
lifecycle,
steps,
}
}
pub fn into_coroutine(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
let binding = (self.lifecycle)(py)?;
let execution = Py::new(py, self)?;
binding.getattr("drive")?.call1((execution,))
}
pub(crate) fn into_sync_stream(
self,
py: Python<'_>,
head: Py<PyAny>,
) -> PyResult<Bound<'_, PyAny>> {
(self.lifecycle)(py)?
.getattr("SyncStream")?
.call1((Py::new(py, self)?, head))
}
/// An execution already started elsewhere and now waiting for its next input.
pub fn suspended(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
Self {
state: ExecutionState::Suspended(Box::new(body)),
lifecycle,
}
}
fn advance(
slf: &Bound<'_, Self>,
py: Python<'_>,
result: Option<PyResult<Py<PyAny>>>,
) -> PyResult<Py<PyAny>> {
let mut body = {
fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
let (mut body, started) = {
let mut execution = slf.borrow_mut();
match (&execution.state, result.is_some()) {
(ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
let started = match (&execution.state, &input) {
(ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
(ExecutionState::Created(_), Resume::Send(_)) => {
return Err(Error::SendBeforeStart.into());
}
(ExecutionState::Suspended(_), _) => true,
(ExecutionState::Running, _) => {
return Err(PyRuntimeError::new_err("execution is already running"));
return Err(Error::ExecutionRunning.into());
}
(ExecutionState::Closed, _) => {
return Err(PyRuntimeError::new_err("execution is closed"));
return Err(Error::ExecutionClosed.into());
}
_ => {
return Err(PyRuntimeError::new_err(
"execution requires start before resume and can only start once",
));
(ExecutionState::Created(_), Resume::Throw(_)) => {
return Err(Error::ExecutionOutOfOrder.into());
}
}
};
match std::mem::replace(&mut execution.state, ExecutionState::Running) {
ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
_ => unreachable!(),
}
};
let lifecycle = slf.borrow().lifecycle;
let result = match input {
Resume::Send(value) => started.then_some(Ok(value)),
Resume::Throw(error) => Some(Err(error)),
};
let steps = slf.borrow().steps;
let outcome = catch_unwind(AssertUnwindSafe(|| {
let step = body.resume(result)?;
let (tag, value, suspended) = match step {
ExecutionStep::Await(value) => ("Await", value, true),
ExecutionStep::Open(head) => ("Open", head, true),
ExecutionStep::Yield(value) => ("Yield", value, true),
ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
("StreamOpen", head, true)
}
ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
("StreamChunk", chunk, true)
}
ExecutionStep::Return(value) => ("Complete", value, false),
};
let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();
let step = steps(py)?.getattr(tag)?.call1((value,))?.unbind();
Ok((step, suspended))
}))
.map_err(panic_to_pyerr)
@ -116,9 +109,7 @@ impl Execution {
drop(body);
outcome.and_then(|(step, suspended)| {
if suspended {
Err(PyRuntimeError::new_err(
"execution was closed while running",
))
Err(Error::ExecutionClosedWhileRunning.into())
} else {
Ok(step)
}
@ -130,24 +121,16 @@ impl Execution {
#[pymethods]
impl Execution {
fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, None)
fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Resume::Send(value))
}
fn resume_value(
slf: &Bound<'_, Self>,
py: Python<'_>,
value: Py<PyAny>,
) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Some(Ok(value)))
}
fn resume_error(
fn throw(
slf: &Bound<'_, Self>,
py: Python<'_>,
error: Bound<'_, PyBaseException>,
) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
}
fn close(slf: &Bound<'_, Self>) {

View file

@ -6,7 +6,7 @@ use pyo3::{
types::PyDict,
};
use crate::{HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, missing_state};
use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned};
pub(super) enum ChainStep<T> {
Ready(T),
@ -110,7 +110,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Py<PyDict>>> {
let Some(Continuation::Arguments(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Arguments))
@ -132,7 +132,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Box<WireRequest>>> {
let Some(Continuation::Wire(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Wire))
@ -154,7 +154,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Py<PyAny>>> {
let Some(Continuation::Response(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Response))
@ -171,7 +171,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<()>> {
let Some(Continuation::Event(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Event))

View file

@ -10,9 +10,7 @@ use pyo3::{
types::PyDict,
};
use crate::{
HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime, missing_state,
};
use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
use super::adapter::{ChainHooks, ChainStep, HookAdapter};
@ -72,7 +70,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Py<PyDict>>> {
let (index, context) = self.arguments.take().ok_or_else(missing_state)?;
let (index, context) = self.arguments.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_arguments(py, result)?;
self.arguments_step(py, index, step, context)
@ -113,7 +111,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let (index, context) = self.wire.take().ok_or_else(missing_state)?;
let (index, context) = self.wire.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_wire(py, result)?;
self.wire_step(py, index, step, &context)
@ -154,7 +152,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
let (index, context) = self.response.take().ok_or_else(missing_state)?;
let (index, context) = self.response.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_response(py, result)?;
self.response_step(py, index, step, context)
@ -195,7 +193,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, ()>> {
let (index, event) = self.event.take().ok_or_else(missing_state)?;
let (index, event) = self.event.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let result = self.hooks[index].resume_event(py, result);
let step = notification_result(py, is_terminal(&event), result)?;

View file

@ -5,6 +5,7 @@
mod argument;
mod binding;
mod caller;
mod conversion_cache;
mod driver;
mod error;
@ -23,11 +24,11 @@ pub use argument::{effective_py_args, lookup, present};
pub use binding::PythonBinding;
pub use conversion_cache::{FromPythonCache, ToPythonCache};
pub use driver::{CallOptions, run_call};
pub use error::{InvokeError, missing_state};
pub use error::{Error, InvokeError};
pub use file_reader::{FileContent, PythonFileReader, py_bytes};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle};
pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
pub use marshal::{
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,

View file

@ -5,11 +5,10 @@ use futures_util::future::{AbortHandle, Abortable};
use litellm_host::call::HostedCompletion;
use litellm_host::machine::{HostFailure, Machine, MachineStep};
use litellm_host::protocol::Protocol;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use tokio::sync::Mutex;
use crate::missing_state;
use crate::Error;
use crate::runtime::{poll_async_value, run_async_value, run_sync_value};
type NativeResult<M> = Result<
@ -65,7 +64,7 @@ where
py: Python<'_>,
interruption: Option<HostFailure<<M::Protocol as Protocol>::Error>>,
) -> PyResult<NativePoll<NativeResult<M>>> {
let state = Arc::clone(self.state.as_ref().ok_or_else(missing_state)?);
let state = Arc::clone(self.state.as_ref().ok_or(Error::MissingState)?);
let future = async move {
let mut state = state.lock().await;
let result = match interruption {
@ -90,7 +89,7 @@ where
run_async_value(py, async move {
Abortable::new(future, registration)
.await
.map_err(|_| PyRuntimeError::new_err("native execution closed"))?
.map_err(|_| Error::NativeExecutionClosed)?
})?
.unbind(),
))
@ -103,12 +102,12 @@ where
pub(super) fn take_result(&self) -> PyResult<NativeResult<M>> {
self.state
.as_ref()
.ok_or_else(missing_state)?
.ok_or(Error::MissingState)?
.try_lock()
.map_err(|_| missing_state())?
.map_err(|_| Error::MissingState)?
.result
.take()
.ok_or_else(missing_state)
.ok_or_else(|| Error::MissingState.into())
.map(|result| {
result.map(|step| match step {
MachineStep::Suspended(op) => MachineStep::Suspended(op),

View file

@ -5,7 +5,7 @@ use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use crate::{Error, Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
@ -112,9 +112,7 @@ where
F: Future<Output = PyResult<T>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
return Err(Error::SyncInsideTokio.into());
}
release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?
}
@ -131,9 +129,7 @@ where
F: Future<Output = Result<T, E>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
return Err(Error::SyncInsideTokio.into());
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;

View file

@ -0,0 +1,770 @@
mod support;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_host::{
hooks::CallHooks,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
};
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
PyValueError, asyncio::CancelledError,
},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use rstest::rstest;
use support::{
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
let options = CallOptions {
asynchronous: false,
steps: lifecycle_module,
observers: Some(sender),
};
(options, move || {
std::iter::from_fn(|| receiver.try_recv().ok()).collect()
})
}
#[rstest]
#[case::success(OpScript::AwaitAnswer)]
#[case::failure(OpScript::AwaitFailure)]
fn asynchronous_host_operations_resume_the_same_machine(
#[from(interpreter)] _python: Interpreter,
#[case] op: OpScript,
) {
Python::attach(|py| {
let machine: MachineFactory = Box::new(|_| {
CallMachine::<Synthetic>::new(None, |host| {
Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
})
});
let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
match op {
OpScript::AwaitAnswer => {
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
}
OpScript::AwaitFailure => {
assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
1
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
}
_ => unreachable!(),
}
});
}
#[rstest]
fn composed_hooks_share_prepared_arguments_and_one_execution(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let chain = HookChain::new()
.with(hooks(&log, HookScript::ObserveArguments))
.with(hooks(&log, HookScript::RewriteArguments))
.with(hooks(&log, HookScript::ReplaceResponse));
let result = handed(
py,
run_call(
py,
move |_, _, request| Ok(success_machine()(request)),
binding(&log, OpScript::Answer),
chain,
PyDict::new(py).unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let response = settled(py, result, asynchronous).unwrap();
assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
let entries = log.entries();
assert!(entries.iter().any(|entry| entry == "adopted:hook"));
assert_eq!(
entries.iter().filter(|entry| *entry == "project").count(),
1
);
assert_eq!(
entries.iter().filter(|entry| *entry == "op:sign").count(),
1
);
assert_eq!(
entries
.iter()
.filter(|entry| *entry == "succeeded:replaced")
.count(),
3
);
assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
});
}
#[rstest]
#[case::draining(4, false)]
#[case::full(1, false)]
#[case::closed(4, true)]
fn observation_delivery_preserves_callback_order_and_response(
#[from(interpreter)] _python: Interpreter,
#[case] capacity: usize,
#[case] closed: bool,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (expected, expected_log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
if closed {
receiver.close();
}
let log = Log::default();
let (actual, actual_log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::ReplaceResponse),
std::convert::identity,
CallOptions {
asynchronous,
steps: lifecycle_module,
observers: Some(sender.clone()),
},
);
assert_eq!(
actual.unwrap().extract::<String>(py).unwrap(),
expected.unwrap().extract::<String>(py).unwrap()
);
assert_eq!(actual_log, expected_log);
let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
if !closed && capacity >= 3 {
let [
CallEvent::Started { start_time },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
CallEvent::Succeeded { timing, .. },
] = events.as_slice()
else {
panic!("expected start, provider response and success: {events:?}");
};
assert_eq!(raw.body, "raw");
assert_eq!(*start_time, timing.start_time);
assert!(timing.end_time >= timing.start_time);
}
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, false, 1)]
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] cancelled: bool,
#[case] failure_callbacks: usize,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (options, mut events) = observed(4);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, script),
std::convert::identity,
CallOptions {
asynchronous,
..options
},
);
let error = result.unwrap_err();
assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
let events = events();
assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
assert_eq!(
matches!(events.last(), Some(CallEvent::Cancelled { .. })),
cancelled
);
assert_eq!(
matches!(
events.last(),
Some(CallEvent::Failed {
origin: FailureOrigin::Host,
..
})
),
!cancelled
);
let terminal = events.iter().filter(|event| {
matches!(
event,
CallEvent::Failed { .. }
| CallEvent::Succeeded { .. }
| CallEvent::Cancelled { .. }
)
});
assert_eq!(terminal.count(), 1);
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
failure_callbacks
);
assert!(
!log.iter()
.any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
);
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
fn startup_failures_release_the_factory_without_constructing_a_machine(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] op: OpScript,
#[values(false, true)] asynchronous: bool,
) {
struct Release(Arc<AtomicBool>);
impl Drop for Release {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
Python::attach(|py| {
let constructed = Arc::new(AtomicBool::new(false));
let did_construct = constructed.clone();
let released = Arc::new(AtomicBool::new(false));
let release = Release(released.clone());
let machine: MachineFactory = Box::new(move |request| {
let _release = release;
did_construct.store(true, Ordering::SeqCst);
success_machine()(request)
});
let (result, log) = run_scripted(py, machine, op, script, asynchronous);
assert!(result.is_err());
assert!(!constructed.load(Ordering::SeqCst));
assert!(released.load(Ordering::SeqCst));
assert_eq!(
log.iter()
.filter(|event| event.starts_with("failed:"))
.count(),
1
);
assert_eq!(
log.iter().filter(|event| *event == "adapter.close").count(),
1
);
assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
});
}
#[rstest]
fn success_runs_every_step_in_order_and_returns_the_public_response(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:1|sign|rewritten"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"before_provider_request",
"response:raw",
"complete",
"after_success",
"succeeded:project:1|sign|rewritten",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (options, mut events) = observed(3);
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
CallOptions {
asynchronous,
..options
},
);
if asynchronous {
for method in ["__anext__", "aclose", "aclose"] {
let stop = stream
.call_method0(method)
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
assert!(stop.is_instance_of::<PyStopIteration>(py));
}
} else {
assert_eq!(
stream
.call_method0("__next__")
.unwrap()
.extract::<String>()
.unwrap(),
"first"
);
stream.call_method0("close").unwrap();
stream.call_method0("close").unwrap();
}
assert!(matches!(
events().as_slice(),
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
));
assert_eq!(
log.entries(),
[
"started",
"begin",
"opened",
"delivered",
"succeeded:None",
"adapter.close"
]
);
});
}
#[rstest]
fn a_stream_carries_its_head_before_the_first_chunk(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
call_options(asynchronous),
);
let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
stream.getattr("_hidden_params").unwrap().extract().unwrap();
assert_eq!(
hidden["additional_headers"],
std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
);
assert_eq!(log.entries(), ["started", "begin", "opened"]);
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
});
}
struct Classification {
provider_fails: bool,
op: OpScript,
classifier_fails: bool,
message: &'static str,
native_context: Option<&'static str>,
steps: &'static [&'static str],
}
#[rstest]
#[case::provider_failure(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: false,
message: "classified: provider exploded",
native_context: None,
steps: &["project", "classify:provider exploded"],
})]
#[case::host_op_rejection(Classification {
provider_fails: false,
op: OpScript::RejectNatively,
classifier_fails: false,
message: "classified: op rejected",
native_context: None,
steps: &["project", "op:sign", "classify:op rejected"],
})]
#[case::host_op_python_exception(Classification {
provider_fails: false,
op: OpScript::RaisePython,
classifier_fails: false,
message: "op failed",
native_context: None,
steps: &["project", "op:sign"],
})]
#[case::failing_classifier(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: true,
message: "classifier failed",
native_context: Some("provider exploded"),
steps: &["project", "classify:provider exploded"],
})]
fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
#[from(interpreter)] _python: Interpreter,
#[case] case: Classification,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let Classification {
provider_fails,
op,
classifier_fails,
message,
native_context,
steps,
} = case;
let log = Log::default();
let binding = if classifier_fails {
binding(&log, op).with_failing_classifier()
} else {
binding(&log, op)
};
let machine = if provider_fails {
failing_machine()
} else {
success_machine()
};
let (result, log) = run_composed(
py,
machine,
binding,
hooks(&log, HookScript::Plain),
std::convert::identity,
call_options(asynchronous),
);
let error = result.unwrap_err();
assert_eq!(error.value(py).to_string(), message);
match native_context {
Some(context) => {
assert!(error.is_instance_of::<PyTypeError>(py));
let chained = error.context(py).unwrap();
assert!(chained.is_instance_of::<PyRuntimeError>(py));
assert_eq!(chained.value(py).to_string(), context);
}
None => assert!(error.is_instance_of::<PyValueError>(py)),
}
let expected: Vec<String> = ["started", "begin"]
.iter()
.chain(steps)
.map(|step| step.to_string())
.chain([
format!("failed:Call:{message}"),
"adapter.close".into(),
"host.close".into(),
])
.collect();
assert_eq!(log, expected);
});
}
#[rstest]
fn begin_failures_are_host_failures_without_provider_mapping(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailBegin,
false,
);
assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:begin failed",
"adapter.close",
"host.close"
]
);
});
}
enum ArgumentPolicy {
Inherit,
Reject(Py<PyBaseException>),
}
impl CallHooks<PythonRuntime> for ArgumentPolicy {
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
match self {
Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
}
}
}
impl PythonOwned for ArgumentPolicy {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::Inherit => Ok(()),
Self::Reject(error) => visit.call(error),
}
}
}
#[rstest]
fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[values(false, true)] fail_setup: bool,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let error = PyValueError::new_err("setup failed").into_value(py);
let setup_error = error.clone_ref(py);
let arguments = PyDict::new(py);
arguments.set_item("api_key", "original").unwrap();
let result = handed(
py,
run_call(
py,
move |py, prepared, request| {
let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
setup_log.push(format!("setup:{source}:{key}"));
if fail_setup {
return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
}
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
HookChain::new()
.with(hooks(&log, HookScript::RewriteArguments))
.with(ArgumentPolicy::Inherit),
arguments.clone().unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let settled = settled(py, result, asynchronous);
assert_eq!(
arguments
.get_item("api_key")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"original"
);
assert_eq!(
&log.entries()[..4],
["started", "begin", "project", "setup:hook:inherited"]
);
if fail_setup {
assert!(settled.unwrap_err().value(py).is(error.bind(py)));
assert_eq!(
log.entries(),
[
"started",
"begin",
"project",
"setup:hook:inherited",
"failed:Host:setup failed",
"adapter.close",
"host.close",
]
);
} else {
assert!(settled.is_ok());
assert!(
log.entries()
.iter()
.any(|entry| entry.starts_with("succeeded:"))
);
}
});
}
#[rstest]
fn closing_an_unstarted_call_never_initializes_resources(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let pending = handed(
py,
run_call(
py,
move |_, _, request| {
setup_log.push("setup");
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
PyDict::new(py).unbind(),
call_options(true),
),
true,
)
.unwrap();
assert!(log.entries().is_empty());
pending.call_method0(py, "close").unwrap();
drop(pending);
assert_eq!(log.entries(), ["adapter.close", "host.close"]);
});
}
#[rstest]
fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let raised = PyValueError::new_err("over budget").into_value(py);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| {
HookChain::new()
.with(hooks)
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
},
call_options(asynchronous),
);
assert!(result.unwrap_err().value(py).is(raised.bind(py)));
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:over budget",
"adapter.close",
"host.close"
]
);
});
}
#[rstest]
fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (result, _) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
call_options(asynchronous),
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:2|sign|rewritten"
);
});
}
#[rstest]
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailAfterSuccess,
asynchronous,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"after_success failed"
);
assert_eq!(
&log[log.len() - 4..],
[
"after_success",
"failed:Host:after_success failed",
"adapter.close",
"host.close"
]
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
});
}
#[rstest]
fn cancellation_ends_the_call_without_terminal_dispatch(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::CancelRequest,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(!error.is_instance_of::<PyException>(py));
assert_eq!(
log,
["started", "begin", "project", "adapter.close", "host.close"]
);
});
}

View file

@ -0,0 +1,343 @@
mod support;
use litellm_host::{hooks::CallHooks, lifecycle::Timing};
use litellm_host_python::{
Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
Suspension, run_call,
};
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyCFunction, PyDict},
};
use rstest::rstest;
use support::{
Interpreter, Log, OpScript, binding, call_options, handed, install_lifecycle_module,
interpreter, lifecycle_module, success_machine,
};
struct RetainingHost {
retained: Option<Py<PyAny>>,
}
impl ExecutionBody for RetainingHost {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.retained)
}
}
#[pyfunction]
fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
Py::new(
py,
Execution::new(
RetainingHost {
retained: Some(retained),
},
lifecycle_module,
),
)
}
struct AwaitBody(Option<Py<PyAny>>);
impl ExecutionBody for AwaitBody {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
match self.0.take() {
Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
None => result
.expect("selected await completed")
.map(ExecutionStep::Return),
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn await_execution(awaitable: Py<PyAny>) -> Execution {
Execution::new(AwaitBody(Some(awaitable)), lifecycle_module)
}
struct CallingBody(Py<PyAny>);
impl ExecutionBody for CallingBody {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn calling_execution(callback: Py<PyAny>) -> Execution {
Execution::new(CallingBody(callback), lifecycle_module)
}
struct ErrorBody(Option<Py<PyBaseException>>);
impl ExecutionBody for ErrorBody {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| {
Err(PyErr::from_value(
self.0.take().unwrap().into_bound(py).into_any(),
))
})
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
Execution::new(ErrorBody(Some(error.unbind())), lifecycle_module)
}
/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
/// response or exception, with the call paused in the hook that is told about it.
#[derive(Clone, Copy, Debug)]
enum Holder {
Hooks,
Response,
Exception,
}
struct PausingHooks {
retained: Option<Py<PyAny>>,
pause: Option<Py<PyAny>>,
holder: Holder,
}
impl CallHooks<PythonRuntime> for PausingHooks {
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
_: Timing,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
match self.holder {
Holder::Hooks => Ok(HookStep::Ready(response)),
Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
Holder::Exception => Err(PyErr::from_value(
self.retained.take().unwrap().into_bound(py),
)),
}
}
fn on_event(
&mut self,
_: Python<'_>,
event: PythonCallEvent<'_>,
) -> PyResult<HookStep<Self, ()>> {
let pauses_here = match self.holder {
Holder::Hooks => true,
Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
};
Ok(match self.pause.take_if(|_| pauses_here) {
Some(pause) => {
HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
}
None => HookStep::Ready(()),
})
}
}
impl PythonOwned for PausingHooks {
fn close(&mut self, _: Python<'_>) {
self.retained = None;
self.pause = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.retained)?;
visit.call(&self.pause)
}
}
fn paused_call(
py: Python<'_>,
retained: Py<PyAny>,
pause: Py<PyAny>,
holder: Holder,
) -> PyResult<Py<PyAny>> {
let machine = success_machine();
handed(
py,
run_call(
py,
move |_, _, request| Ok(machine(request)),
binding(&Log::default(), OpScript::Answer),
PausingHooks {
retained: Some(retained),
pause: Some(pause),
holder,
},
PyDict::new(py).unbind(),
call_options(true),
),
true,
)
}
#[rstest]
fn every_suspension_is_awaited_inline_on_the_callers_task(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let module = install_lifecycle_module(py);
let locals = PyDict::new(py);
locals
.set_item("drive", module.getattr("drive").unwrap())
.unwrap();
locals
.set_item(
"await_execution",
wrap_pyfunction!(await_execution, py).unwrap(),
)
.unwrap();
locals
.set_item(
"calling_execution",
wrap_pyfunction!(calling_execution, py).unwrap(),
)
.unwrap();
let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
py.run(&probe, Some(&locals), Some(&locals)).unwrap();
});
}
#[rstest]
fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
let locals = PyDict::new(py);
locals
.set_item(
"error_execution",
wrap_pyfunction!(error_execution, py).unwrap(),
)
.unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained:
pass
def cycle():
retained = Retained()
try:
raise ValueError('retained traceback')
except ValueError as error:
retained.owner = error_execution(error)
return weakref.ref(retained)
reference = cycle()
gc.collect()
assert reference() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}
#[rstest]
fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
let locals = PyDict::new(py);
locals
.set_item(
"retaining_coroutine",
wrap_pyfunction!(retaining_coroutine, py).unwrap(),
)
.unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained:
pass
def cycle():
retained = Retained()
coroutine = retaining_coroutine(retained)
retained.coroutine = coroutine
return weakref.ref(retained)
retained_ref = cycle()
gc.collect()
assert retained_ref() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}
#[rstest]
#[case::held_by_the_hooks(Holder::Hooks)]
#[case::selected_as_the_response(Holder::Response)]
#[case::selected_as_the_exception(Holder::Exception)]
fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
#[from(interpreter)] _python: Interpreter,
#[case] holder: Holder,
) {
Python::attach(|py| {
let locals = PyDict::new(py);
let call = PyCFunction::new_closure(py, None, None, move |args, _| {
let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
paused_call(args.py(), retained, pause, holder)
})
.unwrap();
locals.set_item("paused_call", call).unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained(Exception):
pass
class Pause:
def __await__(self):
yield
def cycle():
retained = Retained()
coroutine = paused_call(retained, Pause())
coroutine.send(None)
retained.coroutine = coroutine
return weakref.ref(retained)
reference = cycle()
gc.collect()
assert reference() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}

View file

@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
}
HookStep::Await(awaitable, resume) => {
self.resume = Some(resume);
Ok(litellm_host_python::ExecutionStep::Await(awaitable))
Ok(litellm_host_python::ExecutionStep::Suspended(
litellm_host_python::Suspension::Await(awaitable),
))
}
}
})
@ -518,9 +520,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
let source =
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py"))
.unwrap();
PyModule::from_code(py, &source, c"lifecycle.py", c"hook_chain_test_lifecycle")
PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle")
}
#[rstest]
@ -556,9 +558,16 @@ second = Hooks('b', invoke)
},
lifecycle,
);
locals
.set_item("call", execution.into_coroutine(py).unwrap())
let module = lifecycle(py).unwrap();
let call = module
.getattr("drive")
.unwrap()
.call1((
Py::new(py, execution).unwrap(),
module.getattr("AsyncStream").unwrap(),
))
.unwrap();
locals.set_item("call", call).unwrap();
py.run(
c"
async def exercise():

View file

@ -36,11 +36,17 @@ async def exercise():
releaser = asyncio.create_task(release())
execution = await_execution(CustomAwaitable())
try:
execution.resume_value(None)
execution.throw(ValueError("early"))
except RuntimeError:
pass
else:
raise AssertionError("resumed an unstarted execution")
try:
execution.send(1)
except TypeError:
pass
else:
raise AssertionError("sent a value into an unstarted execution")
wrapped = drive(execution)
try:
wrapped.send(1)
@ -129,7 +135,7 @@ async def exercise():
def reenter():
try:
active.start()
active.send(None)
except RuntimeError as error:
observed.append(str(error))
return result
@ -151,7 +157,7 @@ async def exercise():
execution = calling_execution(callback)
callback.owner = execution
if started:
assert execution.start().value is result
assert execution.send(None).value is result
return weakref.ref(callback)
for started in (False, True):

View file

@ -0,0 +1,684 @@
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
hooks::CallHooks,
interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
lifecycle::{ExecutionEvent, Timing},
machine::{CallMachine, MachineFault},
protocol::{Protocol, Reply},
};
use litellm_host_python::{
CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
PythonHostCalls, PythonOwned, PythonRuntime, run_call,
};
use pyo3::{
exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyString},
};
use rstest::fixture;
static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
/// One initialized interpreter with the real lifecycle protocol installed, held for the
/// length of a test so tests sharing the interpreter run one at a time.
pub struct Interpreter(MutexGuard<'static, ()>);
#[fixture]
pub fn interpreter() -> Interpreter {
let guard = PYTHON_GLOBALS
.lock()
.unwrap_or_else(PoisonError::into_inner);
Python::initialize();
Python::attach(|py| {
py.import("asyncio").unwrap();
install_lifecycle_module(py);
});
Interpreter(guard)
}
pub fn lifecycle_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
py.import("host_test_lifecycle")
}
/// Hands a call's execution to its caller the way the product entrypoints do: an async call
/// gets the coroutine from `drive`, a sync call the value or stream `run` settles to.
pub fn handed(
py: Python<'_>,
execution: PyResult<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let execution = execution?;
let module = lifecycle_module(py)?;
let handed = if asynchronous {
module.getattr("drive")?.call1((execution,))?
} else {
module
.getattr("run")?
.call1((execution, module.getattr("SyncStream")?))?
};
Ok(handed.unbind())
}
pub fn call_options(asynchronous: bool) -> CallOptions {
CallOptions::new(asynchronous, lifecycle_module)
}
pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
let source = std::ffi::CString::new(concat!(
include_str!("../../../../../litellm/rust_bridge/execution.py"),
"\n_drive = drive\n\n\nasync def drive(execution):\n return await _drive(execution, AsyncStream)\n",
))
.unwrap();
PyModule::from_code(
py,
&source,
pyo3::ffi::c_str!("execution.py"),
pyo3::ffi::c_str!("host_test_lifecycle"),
)
.unwrap()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Error(pub String);
impl std::fmt::Display for Error {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self(format!("{fault:?}"))
}
}
pub struct Synthetic;
impl Protocol for Synthetic {
type Response = String;
type Error = Error;
type Request = String;
type HostCall = (&'static str, Reply<String>);
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
pub fn wire() -> WireRequest {
WireRequest {
url: "https://example.invalid".into(),
headers: Vec::new(),
body: serde_json::json!({}),
}
}
pub fn context() -> RequestContext {
RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: serde_json::json!({}),
secret_fields: Vec::new(),
api_key: None,
}
}
#[derive(Clone, Default)]
pub struct Log(Arc<Mutex<Vec<String>>>);
impl Log {
pub fn push(&self, entry: impl Into<String>) {
self.0.lock().unwrap().push(entry.into());
}
pub fn entries(&self) -> Vec<String> {
self.0.lock().unwrap().clone()
}
}
#[derive(Clone, Copy, Debug)]
pub enum OpScript {
Answer,
RaisePython,
RejectNatively,
RejectRequestNatively,
RaiseRequestPython,
CancelRequest,
AwaitAnswer,
AwaitFailure,
}
/// The fake route's public exception, kept as a value so a test sees what `classify`
/// produced before the driver raises it.
#[derive(Debug, PartialEq, Eq)]
pub struct Classified(pub String);
impl From<Classified> for PyErr {
fn from(classified: Classified) -> Self {
PyValueError::new_err(format!("classified: {}", classified.0))
}
}
pub struct SyntheticBinding {
log: Log,
op: OpScript,
classifier_fails: bool,
pending_reply: Option<Reply<String>>,
}
pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
SyntheticBinding {
log: log.clone(),
op,
classifier_fails: false,
pending_reply: None,
}
}
impl SyntheticBinding {
pub fn with_failing_classifier(self) -> Self {
Self {
classifier_fails: true,
..self
}
}
fn answer(&self, value: impl FnOnce() -> String) -> Result<String, InvokeError<Error>> {
match self.op {
OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
OpScript::RaisePython | OpScript::RaiseRequestPython => {
Err(PyValueError::new_err("op failed").into())
}
OpScript::RejectNatively | OpScript::RejectRequestNatively => {
Err(InvokeError::Native(Error("op rejected".into())))
}
OpScript::CancelRequest => Err(CancelledError::new_err(()).into()),
}
}
}
impl PythonBinding for SyntheticBinding {
type Protocol = Synthetic;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<String, InvokeError<Error>> {
self.log.push("project");
match self.op {
OpScript::RejectRequestNatively
| OpScript::RaiseRequestPython
| OpScript::CancelRequest => self.answer(String::new),
_ => Ok(format!("project:{}", arguments.len())),
}
}
fn encode_stream_head(
&mut self,
_: Python<'_>,
head: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match head {}
}
fn encode_chunk(
&mut self,
_: Python<'_>,
chunk: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match chunk {}
}
fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
self.log.push("complete");
Ok(PyString::new(py, &response).into_any().unbind())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
self.log.push(format!("classify:{error}"));
if self.classifier_fails {
return Err(PyTypeError::new_err("classifier failed"));
}
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Synthetic> for SyntheticBinding {
fn handle_host_call(
&mut self,
_: Python<'_>,
(op, reply): (&'static str, Reply<String>),
) -> Result<(), InvokeError<Error>> {
self.log.push(format!("op:{op}"));
self.answer(|| op.to_string())
.map(|answer| reply.send(answer))
}
fn begin_host_call(
&mut self,
py: Python<'_>,
(op, reply): (&'static str, Reply<String>),
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
return self.handle_host_call(py, (op, reply)).map(|()| None);
}
self.pending_reply = Some(reply);
let module = PyModule::from_code(
py,
pyo3::ffi::c_str!(
"async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n"
),
pyo3::ffi::c_str!("host_op.py"),
pyo3::ffi::c_str!("host_op"),
)?;
Ok(Some(
module
.getattr("answer")?
.call1((matches!(self.op, OpScript::AwaitFailure),))?
.unbind(),
))
}
fn resume_host_call(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
let answer = result?.extract::<String>(py)?;
self.pending_reply.take().unwrap().send(answer);
Ok(None)
}
}
impl PythonOwned for SyntheticBinding {
fn close(&mut self, _: Python<'_>) {
self.log.push("host.close");
}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
#[derive(Clone, Copy, Debug)]
pub enum HookScript {
Plain,
RewriteArguments,
ObserveArguments,
FailBegin,
ReplaceResponse,
FailAfterSuccess,
CancelTerminal,
FailTerminal,
}
pub struct SyntheticHooks {
log: Log,
script: HookScript,
}
pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
SyntheticHooks {
log: log.clone(),
script,
}
}
impl CallHooks<PythonRuntime> for SyntheticHooks {
fn prepare_arguments(
&mut self,
_: Python<'_>,
arguments: Py<PyDict>,
_: f64,
) -> PyResult<HookStep<Self, Py<PyDict>>> {
self.log.push("begin");
if matches!(self.script, HookScript::FailBegin) {
return Err(PyValueError::new_err("begin failed"));
}
if matches!(self.script, HookScript::RewriteArguments) {
let prepared = Python::attach(|py| -> PyResult<Py<PyDict>> {
let prepared = arguments.bind(py).copy()?;
prepared.set_item("prepared", "hook")?;
prepared.set_item("api_key", "hook-key")?;
Ok(prepared.unbind())
})?;
return Ok(HookStep::Ready(prepared));
}
Ok(HookStep::Ready(arguments))
}
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
if matches!(self.script, HookScript::ObserveArguments) {
let value: String = arguments
.bind(py)
.get_item("prepared")?
.unwrap()
.extract()?;
self.log.push(format!("adopted:{value}"));
}
Ok(())
}
fn before_provider_request(
&mut self,
_: Python<'_>,
wire: Box<WireRequest>,
_: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
self.log.push("before_provider_request");
Ok(HookStep::Ready(Box::new(WireRequest {
url: "rewritten".into(),
..*wire
})))
}
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
_: Timing,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.log.push("after_success");
match self.script {
HookScript::ReplaceResponse => Ok(HookStep::Ready(
"replaced".into_pyobject(py)?.into_any().unbind(),
)),
HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
_ => Ok(HookStep::Ready(response)),
}
}
fn on_event(
&mut self,
py: Python<'_>,
event: PythonCallEvent<'_>,
) -> PyResult<HookStep<Self, ()>> {
let terminal = matches!(
event,
PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
);
if terminal && matches!(self.script, HookScript::CancelTerminal) {
return Err(CancelledError::new_err("callback cancelled"));
}
if terminal && matches!(self.script, HookScript::FailTerminal) {
return Err(PyValueError::new_err("callback failed"));
}
self.log.push(match event {
PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
PythonCallEvent::Started { .. } => "started".into(),
PythonCallEvent::Cancelled { .. } => "cancelled".into(),
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
PythonCallEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}
PythonCallEvent::Failed { origin, error, .. } => {
format!("failed:{origin:?}:{}", error.value(py))
}
});
Ok(HookStep::Ready(()))
}
fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("opened");
Ok(())
}
fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("delivered");
Ok(())
}
}
impl PythonOwned for SyntheticHooks {
fn close(&mut self, _: Python<'_>) {
self.log.push("adapter.close");
}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
pub type MachineFactory = Box<dyn FnOnce(String) -> CallMachine<Synthetic> + Send + Sync>;
/// Answers to projection, to the route op and to `before_provider_request` all reach the
/// response, so a driver that misroutes a reply changes what the call returns.
pub fn success_machine() -> MachineFactory {
Box::new(move |projected| {
CallMachine::new(None, move |host| {
Box::pin(async move {
let signed = host.services.call(|reply| ("sign", reply)).await?;
let wire = host
.interceptors
.before_provider_request(wire(), context())
.await?;
host.interceptors
.after_provider_response(RawResponse { body: "raw".into() })
.await?;
Ok(format!("{projected}|{signed}|{}", wire.url))
})
})
})
}
pub fn failing_machine() -> MachineFactory {
Box::new(move |_| {
CallMachine::new(None, |_| {
Box::pin(async move { Err(Error("provider exploded".into())) })
})
})
}
/// Drives an asynchronous call's coroutine to its settled value from a sync test. A
/// synchronous call is already settled.
pub fn settled(
py: Python<'_>,
result: PyResult<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
if !asynchronous {
return result;
}
let coroutine = result?;
let completed = coroutine
.call_method1(py, "send", (py.None(),))
.unwrap_err();
if !completed.is_instance_of::<pyo3::exceptions::PyStopIteration>(py) {
return Err(completed);
}
completed.value(py).getattr("value").map(Bound::unbind)
}
pub fn run_scripted(
py: Python<'_>,
machine: MachineFactory,
op: OpScript,
script: HookScript,
asynchronous: bool,
) -> (PyResult<Py<PyAny>>, Vec<String>) {
let log = Log::default();
run_composed(
py,
machine,
binding(&log, op),
hooks(&log, script),
std::convert::identity,
call_options(asynchronous),
)
}
pub fn run_composed<L: PythonCallHooks + 'static>(
py: Python<'_>,
machine: MachineFactory,
binding: SyntheticBinding,
hooks: SyntheticHooks,
compose: impl FnOnce(SyntheticHooks) -> L,
options: CallOptions,
) -> (PyResult<Py<PyAny>>, Vec<String>) {
let asynchronous = options.asynchronous;
let log = binding.log.clone();
let arguments = PyDict::new(py);
arguments.set_item("model", "m").unwrap();
let result = handed(
py,
run_call(
py,
move |_, _, request| Ok(machine(request)),
binding,
compose(hooks),
arguments.unbind(),
options,
),
asynchronous,
);
(settled(py, result, asynchronous), log.entries())
}
pub struct Streaming;
impl Protocol for Streaming {
type Response = ();
type Error = Error;
type Request = ();
type HostCall = std::convert::Infallible;
type Chunk = &'static str;
type StreamHead = Vec<(&'static str, &'static str)>;
}
pub struct StreamingBinding;
impl PythonBinding for StreamingBinding {
type Protocol = Streaming;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
_: &Bound<'_, PyDict>,
) -> Result<(), InvokeError<Error>> {
Ok(())
}
fn encode_stream_head(
&mut self,
py: Python<'_>,
head: Vec<(&'static str, &'static str)>,
) -> PyResult<Py<PyAny>> {
let headers = PyDict::new(py);
for (name, value) in head {
headers.set_item(name, value)?;
}
let hidden = PyDict::new(py);
hidden.set_item("additional_headers", headers)?;
Ok(hidden.into_any().unbind())
}
fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
Ok(PyString::new(py, chunk).into_any().unbind())
}
fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
Ok(py.None())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Streaming> for StreamingBinding {
fn handle_host_call(
&mut self,
_: Python<'_>,
op: std::convert::Infallible,
) -> Result<(), InvokeError<Error>> {
match op {}
}
}
impl PythonOwned for StreamingBinding {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine<Streaming> + Send + Sync {
move |()| {
hosted_call((), None, |(), _, _, _observations| async {
Ok(CallOutput::Stream {
head: vec![("request-id", "req_1")],
chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
})
})
}
}
/// Starts a streaming call and returns the caller's stream object.
pub fn open_stream(
py: Python<'_>,
hooks: SyntheticHooks,
options: CallOptions,
) -> Bound<'_, PyAny> {
let asynchronous = options.asynchronous;
let handed = handed(
py,
run_call(
py,
|_, _, request| Ok(streaming_machine()(request)),
StreamingBinding,
hooks,
PyDict::new(py).unbind(),
options,
),
asynchronous,
)
.unwrap();
if asynchronous {
let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
stop.value(py).getattr("value").unwrap()
} else {
handed.into_bound(py)
}
}
/// Drives an `AsyncStream` or `SyncStream` to completion from a sync test.
pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
if !asynchronous {
return stream
.try_iter()
.unwrap()
.map(|chunk| chunk.unwrap().extract().unwrap())
.collect();
}
std::iter::from_fn(|| {
let stop = stream
.call_method0("__anext__")
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
return None;
}
assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
})
.collect()
}

View file

@ -140,7 +140,7 @@ impl Default for HttpSettings {
proxies: EnvironmentProxies::default(),
connect_timeout: Duration::from_secs(10),
tcp_keepalive: None,
pool_idle_timeout: Duration::from_secs(120),
pool_idle_timeout: Duration::from_secs(15),
}
}
}
@ -330,7 +330,7 @@ mod tests {
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
))]);
assert_eq!(settings.tcp_keepalive, None);
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15));
}
fn proxy_flags(

View file

@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
messages::{
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
streaming::{self, ByteStream, StreamDecoder},
transformation::BaseMessagesConfig,
},
};
@ -245,7 +245,7 @@ fn decoded_chunks(
let Some(event) = events.try_next().await? else {
return Ok(None);
};
let chunk = encode_anthropic_sse(&event)?;
let chunk = streaming::encode(&event)?;
log_chunk(provider, "client_response", &chunk);
Ok(Some((chunk, events)))
})
@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
#[cfg(test)]
mod tests {
use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
use rstest::rstest;
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
@ -300,7 +299,7 @@ mod tests {
.await
.unwrap();
let MessagesCallResponse::Stream { mut chunks, .. } =
streaming_response(response, Some(anthropic_sse_event_stream), "test")
streaming_response(response, Some(streaming::event_stream), "test")
else {
panic!("a streaming response returns chunks");
};

View file

@ -2,7 +2,6 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
use futures_util::{SinkExt, StreamExt};
use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{
Message,
@ -12,17 +11,6 @@ use tokio_tungstenite::tungstenite::{
use super::Error;
pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
matches!(
event_type,
ResponsesWsEventType::ResponseCreated
| ResponsesWsEventType::ResponseCompleted
| ResponsesWsEventType::ResponseFailed
| ResponsesWsEventType::ResponseIncomplete
| ResponsesWsEventType::Error
)
}
#[derive(Clone)]
pub struct ResponsesWebSocketConnection {
socket: Arc<Mutex<Option<UpstreamWebSocket>>>,

View file

@ -33,16 +33,82 @@ pub struct ChatCompletionStreamingChoice {
pub logprobs: Option<Value>,
}
/// Compatible hosts omit `id`, `created` and `object` on chunks, and Python's
/// `ModelResponseStream` fills them in, so none of the envelope fields is required here.
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Default)]
pub struct ChatCompletionChunk {
pub id: String,
pub created: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub created: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub object: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub object: Option<String>,
#[serde(default)]
pub choices: Vec<ChatCompletionStreamingChoice>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<ChatCompletionsUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
#[rstest]
#[case::openai(json!({
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi", "future": 1}, "finish_reason": null, "logprobs": null}],
"usage": null
}))]
#[case::usage_only(json!({
"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [],
"usage": {
"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3,
"prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": 1}
}
}))]
#[case::bare_host(json!({"choices": [{"index": 0, "delta": {"content": "hi"}}]}))]
#[case::host_extensions(json!({"id": "x", "choices": [], "citations": ["a"], "future": {"k": null}}))]
fn chunks_round_trip_without_losing_fields(#[case] wire: Value) {
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
assert_eq!(serde_json::to_value(chunk).unwrap(), without_nulls(wire));
}
/// Typed optional fields serialize by omission; extension fields keep their nulls.
fn without_nulls(value: Value) -> Value {
let Value::Object(fields) = value else {
return value;
};
Value::Object(
fields
.into_iter()
.filter(|(_, value)| !value.is_null())
.map(|(key, value)| match (key.as_str(), value) {
("choices", Value::Array(choices)) => (
key,
Value::Array(choices.into_iter().map(without_nulls).collect()),
),
(_, value) => (key, value),
})
.collect(),
)
}
#[rstest]
#[case::choices_not_a_list(json!({"choices": {}}))]
#[case::choice_without_delta(json!({"choices": [{"index": 0}]}))]
#[case::created_not_a_number(json!({"created": "now", "choices": []}))]
#[case::non_object(json!([]))]
fn malformed_chunks_are_rejected(#[case] wire: Value) {
assert!(serde_json::from_value::<ChatCompletionChunk>(wire).is_err());
}
}

View file

@ -109,7 +109,9 @@ pub struct MessagesStreamError {
}
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(strum::AsRefStr)]
#[serde(tag = "type", rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum MessagesStreamEvent {
MessageStart {
message: MessagesStreamMessage,
@ -175,6 +177,23 @@ mod tests {
assert_eq!(serde_json::to_value(event).unwrap(), wire);
}
#[rstest]
#[case::start(json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {}}}))]
#[case::block_start(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}))]
#[case::block_delta(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "a"}}))]
#[case::block_stop(json!({"type": "content_block_stop", "index": 0}))]
#[case::delta(json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))]
#[case::stop(json!({"type": "message_stop"}))]
#[case::ping(json!({"type": "ping"}))]
#[case::error(json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}))]
fn the_event_name_is_the_serialized_type_tag(#[case] wire: Value) {
let event: MessagesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
assert_eq!(
Some(event.as_ref()),
wire.get("type").and_then(Value::as_str)
);
}
#[rstest]
#[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))]
#[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))]

View file

@ -1,7 +1,7 @@
# rules
- Responses originated at OpenAI and is now offered by other hosts and gateways. Types here describe the format, not OpenAI's deployment of it
- Served over HTTP and over the WebSocket mode. `ResponsesWsEvent` is a wire event and lives here
- Served over HTTP and over the WebSocket mode. `ResponsesStreamEvent` and `ResponsesEventType` in `streaming.rs` are the wire event for both
- `ResponsesWsTransformResult` wraps a provider transformation's output, so it lives in `llms/src/base_llm/responses/`
- Hosts that accept only part of the format, or add output item types, are handled in `llms/src/<provider>/responses/`. Keep unknown output items passing through instead of rejecting them
- Responses-to-Chat emulation and model-specific parameter rewriting are transformations, not format rules

View file

@ -1,5 +1,6 @@
mod output;
mod response;
pub mod streaming;
pub mod streaming_websocket;
pub use output::*;

View file

@ -0,0 +1,305 @@
use serde_json::{Map, Value};
/// Every event name on https://developers.openai.com/api/reference/resources/responses/streaming-events
/// as of 2026-10-11, plus `response.create` from the WebSocket mode. Names OpenAI adds later ride in
/// `Other` untouched.
#[derive(
Clone,
Debug,
PartialEq,
Eq,
strum::AsRefStr,
strum::EnumIter,
strum::EnumString,
strum::Display,
serde_with::DeserializeFromStr,
serde_with::SerializeDisplay,
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
pub enum ResponsesEventType {
#[strum(serialize = "response.create")]
ResponseCreate,
#[strum(serialize = "response.created")]
ResponseCreated,
#[strum(serialize = "response.queued")]
ResponseQueued,
#[strum(serialize = "response.in_progress")]
ResponseInProgress,
#[strum(serialize = "response.completed")]
ResponseCompleted,
#[strum(serialize = "response.failed")]
ResponseFailed,
#[strum(serialize = "response.incomplete")]
ResponseIncomplete,
#[strum(serialize = "error")]
Error,
#[strum(serialize = "response.audio.delta")]
AudioDelta,
#[strum(serialize = "response.audio.done")]
AudioDone,
#[strum(serialize = "response.audio.transcript.delta")]
AudioTranscriptDelta,
#[strum(serialize = "response.audio.transcript.done")]
AudioTranscriptDone,
#[strum(serialize = "response.code_interpreter_call.completed")]
CodeInterpreterCallCompleted,
#[strum(serialize = "response.code_interpreter_call.in_progress")]
CodeInterpreterCallInProgress,
#[strum(serialize = "response.code_interpreter_call.interpreting")]
CodeInterpreterCallInterpreting,
#[strum(serialize = "response.code_interpreter_call_code.delta")]
CodeInterpreterCallCodeDelta,
#[strum(serialize = "response.code_interpreter_call_code.done")]
CodeInterpreterCallCodeDone,
#[strum(serialize = "response.compaction.compacting")]
CompactionCompacting,
#[strum(serialize = "response.content_part.added")]
ContentPartAdded,
#[strum(serialize = "response.content_part.done")]
ContentPartDone,
#[strum(serialize = "response.custom_tool_call_input.delta")]
CustomToolCallInputDelta,
#[strum(serialize = "response.custom_tool_call_input.done")]
CustomToolCallInputDone,
#[strum(serialize = "response.file_search_call.completed")]
FileSearchCallCompleted,
#[strum(serialize = "response.file_search_call.in_progress")]
FileSearchCallInProgress,
#[strum(serialize = "response.file_search_call.searching")]
FileSearchCallSearching,
#[strum(serialize = "response.function_call_arguments.delta")]
FunctionCallArgumentsDelta,
#[strum(serialize = "response.function_call_arguments.done")]
FunctionCallArgumentsDone,
#[strum(serialize = "response.image_generation_call.completed")]
ImageGenerationCallCompleted,
#[strum(serialize = "response.image_generation_call.generating")]
ImageGenerationCallGenerating,
#[strum(serialize = "response.image_generation_call.in_progress")]
ImageGenerationCallInProgress,
#[strum(serialize = "response.image_generation_call.partial_image")]
ImageGenerationCallPartialImage,
#[strum(serialize = "response.mcp_call.completed")]
McpCallCompleted,
#[strum(serialize = "response.mcp_call.failed")]
McpCallFailed,
#[strum(serialize = "response.mcp_call.in_progress")]
McpCallInProgress,
#[strum(serialize = "response.mcp_call_arguments.delta")]
McpCallArgumentsDelta,
#[strum(serialize = "response.mcp_call_arguments.done")]
McpCallArgumentsDone,
#[strum(serialize = "response.mcp_list_tools.completed")]
McpListToolsCompleted,
#[strum(serialize = "response.mcp_list_tools.failed")]
McpListToolsFailed,
#[strum(serialize = "response.mcp_list_tools.in_progress")]
McpListToolsInProgress,
#[strum(serialize = "response.output_item.added")]
OutputItemAdded,
#[strum(serialize = "response.output_item.done")]
OutputItemDone,
#[strum(serialize = "response.output_text.annotation.added")]
OutputTextAnnotationAdded,
#[strum(serialize = "response.output_text.delta")]
OutputTextDelta,
#[strum(serialize = "response.output_text.done")]
OutputTextDone,
#[strum(serialize = "response.reasoning_summary_part.added")]
ReasoningSummaryPartAdded,
#[strum(serialize = "response.reasoning_summary_part.done")]
ReasoningSummaryPartDone,
#[strum(serialize = "response.reasoning_summary_text.delta")]
ReasoningSummaryTextDelta,
#[strum(serialize = "response.reasoning_summary_text.done")]
ReasoningSummaryTextDone,
#[strum(serialize = "response.reasoning_text.delta")]
ReasoningTextDelta,
#[strum(serialize = "response.reasoning_text.done")]
ReasoningTextDone,
#[strum(serialize = "response.refusal.delta")]
RefusalDelta,
#[strum(serialize = "response.refusal.done")]
RefusalDone,
#[strum(serialize = "response.shell_call_command.added")]
ShellCallCommandAdded,
#[strum(serialize = "response.shell_call_command.delta")]
ShellCallCommandDelta,
#[strum(serialize = "response.shell_call_command.done")]
ShellCallCommandDone,
#[strum(serialize = "response.shell_call_output_content.delta")]
ShellCallOutputContentDelta,
#[strum(serialize = "response.shell_call_output_content.done")]
ShellCallOutputContentDone,
#[strum(serialize = "response.web_search_call.completed")]
WebSearchCallCompleted,
#[strum(serialize = "response.web_search_call.in_progress")]
WebSearchCallInProgress,
#[strum(serialize = "response.web_search_call.searching")]
WebSearchCallSearching,
#[strum(default, transparent)]
Other(String),
}
impl ResponsesEventType {
/// The last event of a response. `response.completed`, `response.incomplete` and
/// `response.failed` each close the response and carry its final body; `error` aborts it.
pub fn is_terminal(&self) -> bool {
matches!(
self,
Self::ResponseCompleted | Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
)
}
/// A terminal event after which the response is not a success.
pub fn is_failure(&self) -> bool {
matches!(
self,
Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
)
}
}
#[macro_rules_attribute::apply(crate::wire_type)]
pub struct ResponsesStreamEvent {
#[serde(rename = "type")]
pub event_type: ResponsesEventType,
#[serde(flatten)]
pub data: Map<String, Value>,
}
impl ResponsesStreamEvent {
pub fn model(&self) -> Option<&str> {
let model = self.data.get("model").and_then(Value::as_str);
if model.is_some() {
return model;
}
self.data
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("model"))
.and_then(Value::as_str)
}
pub fn is_response_create(&self) -> bool {
self.event_type == ResponsesEventType::ResponseCreate
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use strum::IntoEnumIterator;
use super::*;
#[rstest]
fn every_named_event_type_round_trips() {
let named: Vec<_> = ResponsesEventType::iter()
.filter(|event_type| !matches!(event_type, ResponsesEventType::Other(_)))
.collect();
assert!(named.len() > 50);
for event_type in named {
let wire = event_type.to_string();
assert_eq!(wire, event_type.as_ref());
assert_eq!(serde_json::to_value(&event_type).unwrap(), json!(wire));
assert_eq!(wire.parse::<ResponsesEventType>().unwrap(), event_type);
assert_eq!(
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
event_type
);
}
}
#[rstest]
#[case::unknown("future.event", ResponsesEventType::Other("future.event".into()))]
#[case::empty("", ResponsesEventType::Other(String::new()))]
#[case::case_sensitive(
"Response.Completed",
ResponsesEventType::Other("Response.Completed".into())
)]
#[case::escaped("future\"\\\n", ResponsesEventType::Other("future\"\\\n".into()))]
fn unknown_event_types_round_trip_through_other(
#[case] wire: &str,
#[case] expected: ResponsesEventType,
) {
assert_eq!(
serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
expected
);
assert_eq!(serde_json::to_value(&expected).unwrap(), json!(wire));
assert_eq!(expected.as_ref(), wire);
assert_eq!(expected.to_string(), wire);
}
/// Python's `ResponsesAPIStreamingIterator` stores the response on `response.completed`,
/// `response.incomplete` and `response.failed`, and raises on `response.failed` and `error`.
#[rstest]
#[case::completed(ResponsesEventType::ResponseCompleted, true, false)]
#[case::incomplete(ResponsesEventType::ResponseIncomplete, true, true)]
#[case::failed(ResponsesEventType::ResponseFailed, true, true)]
#[case::error(ResponsesEventType::Error, true, true)]
#[case::create(ResponsesEventType::ResponseCreate, false, false)]
#[case::created(ResponsesEventType::ResponseCreated, false, false)]
#[case::in_progress(ResponsesEventType::ResponseInProgress, false, false)]
#[case::delta(ResponsesEventType::OutputTextDelta, false, false)]
#[case::item_done(ResponsesEventType::OutputItemDone, false, false)]
#[case::unknown(ResponsesEventType::Other("response.future".into()), false, false)]
#[case::tool_failed(ResponsesEventType::McpCallFailed, false, false)]
fn terminal_and_failure_events_match_the_python_iterator(
#[case] event_type: ResponsesEventType,
#[case] terminal: bool,
#[case] failure: bool,
) {
assert_eq!(event_type.is_terminal(), terminal);
assert_eq!(event_type.is_failure(), failure);
}
#[rstest]
#[case::number("17")]
#[case::boolean("true")]
#[case::null("null")]
#[case::array("[]")]
#[case::object("{}")]
fn event_types_reject_non_strings(#[case] wire: &str) {
assert!(serde_json::from_str::<ResponsesEventType>(wire).is_err());
}
#[rstest]
#[case::text(json!({"type":"response.output_text.delta","delta":"hello","sequence_number":3,"future":null}))]
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1","output_index":0}))]
#[case::unknown(json!({"type":"future.event","payload":{"items":[1,null]},"future":true}))]
#[case::terminal(json!({"type":"response.completed","response":{"id":"r","output":[],"future":null}}))]
fn events_preserve_every_payload_field(#[case] wire: Value) {
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
assert_eq!(serde_json::to_value(event).unwrap(), wire);
}
#[rstest]
#[case::missing_type(json!({"delta":"hello"}))]
#[case::numeric_type(json!({"type":7}))]
#[case::null_type(json!({"type":null}))]
#[case::non_object(json!([]))]
fn malformed_events_are_rejected(#[case] wire: Value) {
assert!(serde_json::from_value::<ResponsesStreamEvent>(wire).is_err());
}
#[rstest]
#[case::flat(json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
#[case::nested(json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
#[case::none(json!({"type":"response.created","response":{}}), None)]
fn model_reads_flat_and_nested_shapes(#[case] wire: Value, #[case] expected: Option<&str>) {
let event: ResponsesStreamEvent = serde_json::from_value(wire).unwrap();
assert_eq!(event.model(), expected);
}
#[cfg(feature = "schema")]
#[rstest]
fn event_type_schema_accepts_unknown_strings() {
let schema = schemars::schema_for!(ResponsesEventType).to_value();
assert_eq!(schema.get("type"), Some(&json!("string")));
assert!(schema.get("enum").is_none());
}
}

View file

@ -2,68 +2,6 @@ use super::ResponsesOutputItem;
use crate::recognized::Recognized;
use serde_json::{Map, Value};
#[derive(
Clone,
Debug,
PartialEq,
Eq,
strum::AsRefStr,
strum::EnumString,
strum::Display,
serde_with::DeserializeFromStr,
serde_with::SerializeDisplay,
)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "schema", schemars(with = "String"))]
pub enum ResponsesWsEventType {
#[strum(serialize = "response.create")]
ResponseCreate,
#[strum(serialize = "response.created")]
ResponseCreated,
#[strum(serialize = "response.completed")]
ResponseCompleted,
#[strum(serialize = "response.failed")]
ResponseFailed,
#[strum(serialize = "response.incomplete")]
ResponseIncomplete,
#[strum(serialize = "error")]
Error,
#[strum(default, transparent)]
Other(String),
}
impl ResponsesWsEventType {
pub fn as_str(&self) -> &str {
self.as_ref()
}
}
#[macro_rules_attribute::apply(crate::wire_type)]
pub struct ResponsesWsEvent {
#[serde(rename = "type")]
pub event_type: ResponsesWsEventType,
#[serde(flatten)]
pub data: Map<String, Value>,
}
impl ResponsesWsEvent {
pub fn model(&self) -> Option<&str> {
let model = self.data.get("model").and_then(Value::as_str);
if model.is_some() {
return model;
}
self.data
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("model"))
.and_then(Value::as_str)
}
pub fn is_response_create(&self) -> bool {
self.event_type == ResponsesWsEventType::ResponseCreate
}
}
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Eq)]
pub struct ResponsesErrorFrame {
@ -112,8 +50,7 @@ mod tests {
use crate::formats::responses::ResponsesOutputItem;
use crate::{
formats::responses::streaming_websocket::{ResponsesEventResponse, ResponsesWsEventType},
recognized::Recognized,
formats::responses::streaming_websocket::ResponsesEventResponse, recognized::Recognized,
test_support::*,
};
@ -134,66 +71,6 @@ mod tests {
);
}
#[rstest]
#[case::flat(serde_json::json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
#[case::nested(serde_json::json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
fn model_reads_flat_and_nested_create_shapes(
#[case] payload: serde_json::Value,
#[case] expected: Option<&str>,
) {
let event: ResponsesWsEvent = serde_json::from_value(payload).expect("valid event");
assert_eq!(event.model(), expected);
}
#[rstest]
#[case::create("response.create", ResponsesWsEventType::ResponseCreate)]
#[case::created("response.created", ResponsesWsEventType::ResponseCreated)]
#[case::completed("response.completed", ResponsesWsEventType::ResponseCompleted)]
#[case::failed("response.failed", ResponsesWsEventType::ResponseFailed)]
#[case::incomplete("response.incomplete", ResponsesWsEventType::ResponseIncomplete)]
#[case::error("error", ResponsesWsEventType::Error)]
#[case::unknown(
"response.output_text.delta",
ResponsesWsEventType::Other("response.output_text.delta".to_string())
)]
#[case::empty("", ResponsesWsEventType::Other(String::new()))]
#[case::case_sensitive(
"Response.Completed",
ResponsesWsEventType::Other("Response.Completed".into())
)]
#[case::escaped("future\"\\\n", ResponsesWsEventType::Other("future\"\\\n".into()))]
fn websocket_event_type_round_trips(
#[case] wire: &str,
#[case] expected: ResponsesWsEventType,
) {
let serialized = serde_json::to_string(&expected).unwrap();
assert_eq!(serialized, serde_json::to_string(wire).unwrap());
assert_eq!(
serde_json::from_str::<ResponsesWsEventType>(&serialized).unwrap(),
expected
);
}
#[rstest]
#[case::number("17")]
#[case::boolean("true")]
#[case::null("null")]
#[case::array("[]")]
#[case::object("{}")]
fn websocket_event_type_rejects_non_strings(#[case] wire: &str) {
assert!(serde_json::from_str::<ResponsesWsEventType>(wire).is_err());
}
#[cfg(feature = "schema")]
#[rstest]
fn websocket_event_type_schema_is_open_string() {
let schema = schemars::schema_for!(ResponsesWsEventType);
assert_eq!(
schema.to_value().get("type"),
Some(&serde_json::json!("string"))
);
}
#[rstest]
fn nested_event_response_exposes_typed_output_and_preserves_extensions() {
let wire = json!({

View file

@ -32,6 +32,13 @@
- Provider-wide docs (auth, errors, regions): `src/<provider>/`
- One host's endpoint docs: `src/<provider>/<format>/`
## Streaming
- Each format's `src/base_llm/<format>/streaming.rs` exposes the same three names: `event_stream` (bytes to typed events), `decode` (a whole stream to typed events) and `encode` (one typed event to the SSE frame LiteLLM's proxy sends clients). Callers use the module path, never a vendor prefix
- SSE framing, `[DONE]` and the generic helpers live in `src/base_llm/sse.rs`. Messages frames are `event:` plus `data:`; Chat Completions and Responses frames are bare `data:`, as Python's proxy writes them
- Decoding is strict: a malformed event, or `[DONE]` on a Messages stream, is `InvalidResponse` and ends the stream. Python never fails a stream this way; the divergence is deliberate and `src/base_llm/sse.rs` lists what Python does per format. A stream cut inside an event drops that event, as the SSE spec requires
- What ends or fails a stream belongs on the format's event type (`is_terminal`, `is_failure`). The cache's success check in `inference/src/caching.rs` still matches type strings; move it onto the typed events rather than adding new string matches
## Python pairs
- Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure

View file

@ -30,7 +30,7 @@ use crate::{
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
},
},
messages::streaming::anthropic_sse_event_stream,
messages::streaming::event_stream as messages_event_stream,
},
};
@ -196,7 +196,7 @@ impl BaseConfig for AnthropicConfig {
fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
Some(ChatStream::new(
anthropic_sse_event_stream,
messages_event_stream,
ModelResponseIterator::new(shape),
))
}

View file

@ -1,15 +1,32 @@
use std::collections::HashMap;
use bytes::Bytes;
use futures_util::{StreamExt, stream::BoxStream};
use litellm_llms_types::formats::chat_completions::ChatCompletionChunk;
use crate::{
Error,
base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
base_llm::{
base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
sse::{self, is_done},
},
};
pub type ChatChunkStream = BoxStream<'static, Result<ChatCompletionChunk, Error>>;
/// OpenAI-compatible SSE: bare `data:` frames closed by `data: [DONE]`.
pub fn event_stream(bytes: ByteStream) -> ChatChunkStream {
sse::event_stream(bytes, is_done)
}
pub fn decode(wire: &[u8]) -> Result<Vec<ChatCompletionChunk>, Error> {
sse::decode(wire, is_done)
}
pub fn encode(chunk: &ChatCompletionChunk) -> Result<Bytes, Error> {
sse::encode_data(chunk)
}
/// What Python's `map_openai_params` decides about the stream and `completion`
/// hands to `ModelResponseIterator`: it is settled while the request is built,
/// never re-derived from the body.
@ -52,3 +69,109 @@ impl ChatStream {
(self.run)(bytes)
}
}
#[cfg(test)]
mod tests {
use futures_util::{StreamExt, TryStreamExt, stream};
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
stream::iter(pieces.into_iter().map(Ok)).boxed()
}
#[rstest]
#[case::role(json!({"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}))]
#[case::text(json!({"id":"c","choices":[{"index":0,"delta":{"content":"hi\nthere"}}]}))]
#[case::tool(json!({"id":"c","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":"{\"x\":"}}]}}]}))]
#[case::finish(json!({"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}))]
#[case::usage(json!({"id":"c","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":0,"cache_creation_tokens":0,"text_tokens":1}}}))]
#[case::host_extension(json!({"id":"c","choices":[],"citations":[{"url":"u"}]}))]
fn encoded_chunks_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
let encoded = encode(&chunk).unwrap();
assert!(encoded.starts_with(b"data: "));
assert!(encoded.ends_with(b"\n\n"));
assert_eq!(
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
wire
);
assert_eq!(decode(&encoded).unwrap(), vec![chunk]);
}
#[rstest]
#[case::lf("\n")]
#[case::crlf("\r\n")]
#[case::cr("\r")]
#[tokio::test]
async fn streams_split_anywhere_skip_comments_and_done(
#[case] newline: &str,
#[values(1, 5, 4096)] size: usize,
) {
let wire = format!(
": keepalive{newline}{newline}data: {{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"a\"}}}}]}}{newline}{newline}data:{{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}{newline}{newline}data: [DONE]{newline}{newline}"
);
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
assert_eq!(streamed.len(), 2);
assert_eq!(streamed[0].choices[0].delta.content.as_deref(), Some("a"));
assert_eq!(
streamed[1].choices[0].finish_reason.as_deref(),
Some("stop")
);
}
#[rstest]
#[case::spaced(b"data: [DONE]\n\n")]
#[case::unspaced(b"data:[DONE]\n\n")]
#[case::trailing_space(b"data: [DONE] \n\n")]
fn every_done_spelling_closes_the_stream_without_a_chunk(#[case] wire: &[u8]) {
assert_eq!(decode(wire).unwrap(), vec![]);
}
#[rstest]
#[case::invalid_json(b"data: {\n\n")]
#[case::choices_not_a_list(b"data: {\"choices\":{}}\n\n")]
#[case::invalid_utf8(b"data: \xff\n\n")]
#[tokio::test]
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
assert!(matches!(
streamed.last(),
Some(Err(Error::InvalidResponse(_))) | None
));
}
#[rstest]
#[tokio::test]
async fn a_body_error_ends_the_stream_after_delivered_chunks() {
let mut chunks = event_stream(
stream::iter([
Ok(Bytes::from_static(
b"data: {\"id\":\"c\",\"choices\":[]}\n\n",
)),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
])
.boxed(),
);
assert_eq!(
chunks.next().await.unwrap().unwrap().id.as_deref(),
Some("c")
);
assert!(matches!(
chunks.next().await,
Some(Err(Error::InvalidResponse(_)))
));
assert!(chunks.next().await.is_none());
}
}

View file

@ -1,38 +1,32 @@
use bytes::Bytes;
use futures_util::{StreamExt, stream::BoxStream};
use litellm_framer::{frames, sse::SseCodec};
use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent;
use crate::Error;
pub use crate::base_llm::base_model_iterator::ByteStream;
use crate::{Error, base_llm::sse};
pub type EventStream = BoxStream<'static, Result<MessagesStreamEvent, Error>>;
pub type EventStream = sse::EventStream<MessagesStreamEvent>;
pub type StreamDecoder = fn(ByteStream) -> EventStream;
pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream {
Box::pin(frames(bytes, SseCodec::default()).map(|event| {
let event = event.map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
})?;
serde_json::from_str(&event.data).map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
})
}))
fn no_control_frames(_: &litellm_framer::sse::SseEvent) -> bool {
false
}
pub fn encode_anthropic_sse(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
let data = serde_json::to_value(event).map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
})?;
let name = data
.get("type")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| {
Error::InvalidResponse(
"Anthropic stream event is invalid: stream event has no type".into(),
)
})?;
Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n")))
pub fn event_stream(bytes: ByteStream) -> EventStream {
sse::event_stream(bytes, no_control_frames)
}
pub fn decode(wire: &[u8]) -> Result<Vec<MessagesStreamEvent>, Error> {
sse::decode(wire, no_control_frames)
}
/// Anthropic names every frame after the event's `type`, and LiteLLM's proxy does the same
/// when it re-emits translated events.
pub fn encode(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
let data = serde_json::to_string(event).map_err(sse::invalid)?;
Ok(Bytes::from(format!(
"event: {}\ndata: {data}\n\n",
event.as_ref()
)))
}
#[cfg(test)]
@ -41,6 +35,7 @@ mod tests {
use litellm_llms_types::formats::messages::streaming::{
MessagesContentBlockDelta, MessagesStreamUsage,
};
use rstest::rstest;
use serde_json::json;
use super::*;
@ -48,51 +43,11 @@ mod tests {
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
fn in_pieces(wire: &[u8]) -> ByteStream {
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
stream::iter(pieces.into_iter().map(Ok)).boxed()
}
#[tokio::test]
async fn sse_frames_split_anywhere_decode_into_typed_events() {
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![MessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: MessagesContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
#[tokio::test]
async fn decodes_citations_delta_events() {
let wire = concat!(
"event: content_block_delta\n",
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
"\n\n",
);
let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert!(matches!(
events.as_slice(),
[MessagesStreamEvent::ContentBlockDelta {
delta: MessagesContentBlockDelta::Citations { .. },
..
}]
));
}
fn events() -> Vec<MessagesStreamEvent> {
vec![
MessagesStreamEvent::Ping,
@ -110,27 +65,80 @@ mod tests {
]
}
#[tokio::test]
async fn encoded_events_decode_back_to_themselves() {
let wire = events()
.iter()
.map(encode_anthropic_sse)
.collect::<Result<Vec<_>, _>>()
.unwrap();
fn wire(newline: &str) -> String {
format!(
"event: content_block_delta{newline}data: {TEXT_DELTA}{newline}{newline}: keepalive{newline}{newline}event: message_stop{newline}data: {{\"type\":\"message_stop\"}}{newline}{newline}"
)
}
let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed())
#[rstest]
#[case::lf("\n")]
#[case::crlf("\r\n")]
#[case::cr("\r")]
#[tokio::test]
async fn frames_split_anywhere_decode_into_typed_events(
#[case] newline: &str,
#[values(1, 3, 4096)] size: usize,
) {
let wire = wire(newline);
let expected = vec![
MessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: MessagesContentBlockDelta::TextDelta {
text: "hello".into(),
},
},
MessagesStreamEvent::MessageStop { usage: None },
];
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(streamed, expected);
assert_eq!(decode(wire.as_bytes()).unwrap(), expected);
}
assert_eq!(decoded, events());
#[tokio::test]
async fn decodes_citations_delta_events() {
let wire = concat!(
"event: content_block_delta\n",
r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
"\n\n",
);
let events = event_stream(in_pieces(wire.as_bytes(), 3))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert!(matches!(
events.as_slice(),
[MessagesStreamEvent::ContentBlockDelta {
delta: MessagesContentBlockDelta::Citations { .. },
..
}]
));
}
#[tokio::test]
async fn encoded_events_decode_back_to_themselves() {
let wire: Vec<u8> = events()
.iter()
.map(encode)
.collect::<Result<Vec<_>, _>>()
.unwrap()
.into_iter()
.flatten()
.collect();
let streamed = event_stream(in_pieces(&wire, 3))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(streamed, events());
assert_eq!(decode(&wire).unwrap(), events());
}
#[test]
fn an_event_is_named_by_its_type() {
let encoded =
encode_anthropic_sse(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
let encoded = encode(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
assert_eq!(
encoded,
Bytes::from(format!(
@ -139,4 +147,68 @@ mod tests {
))
);
}
#[rstest]
#[case::json(b"data: {\n\n")]
#[case::unknown_event(b"data: {\"type\":\"future_event\"}\n\n")]
#[case::malformed_tool(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":3}}\n\n")]
#[case::utf8(b"data: \xff\n\n")]
#[case::done_is_not_an_anthropic_frame(b"data: [DONE]\n\n")]
#[tokio::test]
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
assert!(matches!(
streamed.last(),
Some(Err(Error::InvalidResponse(_))) | None
));
}
#[rstest]
fn decoding_keeps_tool_deltas_and_block_extensions() {
let wire = concat!(
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"lookup\",\"input\":{},\"future\":null}}\n\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n",
);
let decoded = decode(wire.as_bytes()).unwrap();
let [
MessagesStreamEvent::ContentBlockStart { content_block, .. },
MessagesStreamEvent::ContentBlockDelta {
delta: MessagesContentBlockDelta::InputJsonDelta { partial_json },
..
},
] = decoded.as_slice()
else {
panic!("expected tool events");
};
assert_eq!(
content_block.extra.get("future"),
Some(&serde_json::Value::Null)
);
assert_eq!(partial_json, "{\"x\":");
}
#[rstest]
#[tokio::test]
async fn a_body_error_ends_the_stream_after_delivered_events() {
let mut events = event_stream(
stream::iter([
Ok(Bytes::from_static(b"data: {\"type\":\"ping\"}\n\n")),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
])
.boxed(),
);
assert_eq!(
events.next().await.unwrap().unwrap(),
MessagesStreamEvent::Ping
);
assert!(matches!(
events.next().await,
Some(Err(Error::InvalidResponse(_)))
));
assert!(events.next().await.is_none());
}
}

View file

@ -5,3 +5,4 @@ pub mod chat;
pub mod messages;
pub mod ocr;
pub mod responses;
pub mod sse;

View file

@ -1 +1,2 @@
pub mod streaming;
pub mod transformation;

View file

@ -0,0 +1,135 @@
use bytes::Bytes;
use litellm_llms_types::formats::responses::streaming::ResponsesStreamEvent;
pub use crate::base_llm::base_model_iterator::ByteStream;
use crate::{
Error,
base_llm::sse::{self, is_done},
};
pub type EventStream = sse::EventStream<ResponsesStreamEvent>;
pub fn event_stream(bytes: ByteStream) -> EventStream {
sse::event_stream(bytes, is_done)
}
pub fn decode(wire: &[u8]) -> Result<Vec<ResponsesStreamEvent>, Error> {
sse::decode(wire, is_done)
}
/// A bare `data:` frame. Python's proxy serializes every Responses event this way and never
/// writes an `event:` line, so clients see one shape whichever runtime served them.
pub fn encode(event: &ResponsesStreamEvent) -> Result<Bytes, Error> {
sse::encode_data(event)
}
#[cfg(test)]
mod tests {
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
stream::iter(pieces.into_iter().map(Ok)).boxed()
}
#[rstest]
#[case::created(json!({"type":"response.created","response":{"id":"r","output":[]}}))]
#[case::text(json!({"type":"response.output_text.delta","delta":"hello\nworld","sequence_number":3}))]
#[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1"}))]
#[case::unknown(json!({"type":"future.event","payload":[1,null],"future":{"nested":true}}))]
#[case::completed(json!({"type":"response.completed","response":{"id":"r","usage":{"input_tokens":3}}}))]
#[case::failed(json!({"type":"response.failed","response":{"error":{"code":"failure"}}}))]
#[case::incomplete(json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"limit"}}}))]
#[case::error(json!({"type":"error","error":{"message":"boom"}}))]
fn encoded_events_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
let encoded = encode(&event).unwrap();
assert!(encoded.starts_with(b"data: "));
assert!(encoded.ends_with(b"\n\n"));
assert_eq!(
serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
wire
);
assert_eq!(decode(&encoded).unwrap(), vec![event]);
}
#[rstest]
#[case::lf("\n")]
#[case::crlf("\r\n")]
#[case::cr("\r")]
#[tokio::test]
async fn streams_split_anywhere_skip_comments_and_done(
#[case] newline: &str,
#[values(1, 5, 4096)] size: usize,
) {
let wire = format!(
"\u{feff}event: response.created{newline}data: {{\"type\":\"response.created\"}}{newline}{newline}: keepalive{newline}{newline}event: future.event{newline}data: {{\"type\":\"future.event\",\"payload\":null}}{newline}{newline}data: [DONE]{newline}{newline}"
);
let streamed = event_stream(in_pieces(wire.as_bytes(), size))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
assert_eq!(streamed.len(), 2);
assert_eq!(streamed[0].event_type, ResponsesEventType::ResponseCreated);
assert_eq!(
streamed[1].event_type,
ResponsesEventType::Other("future.event".into())
);
assert_eq!(streamed[1].data.get("payload"), Some(&Value::Null));
}
#[rstest]
#[case::spaced(b"data: [DONE]\n\n")]
#[case::unspaced(b"data:[DONE]\n\n")]
#[case::trailing_space(b"data: [DONE] \n\n")]
fn every_done_spelling_closes_the_stream_without_an_event(#[case] wire: &[u8]) {
assert_eq!(decode(wire).unwrap(), vec![]);
}
#[rstest]
#[case::invalid_json(b"data: {\n\n")]
#[case::missing_type(b"data: {}\n\n")]
#[case::numeric_type(b"data: {\"type\":3}\n\n")]
#[case::invalid_utf8(b"data: \xff\n\n")]
#[tokio::test]
async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
assert!(matches!(
streamed.last(),
Some(Err(Error::InvalidResponse(_))) | None
));
}
#[rstest]
#[tokio::test]
async fn a_body_error_ends_the_stream_after_delivered_events() {
let mut events = event_stream(
stream::iter([
Ok(Bytes::from_static(
b"data: {\"type\":\"response.created\"}\n\n",
)),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
])
.boxed(),
);
assert_eq!(
events.next().await.unwrap().unwrap().event_type,
ResponsesEventType::ResponseCreated
);
assert!(matches!(
events.next().await,
Some(Err(Error::InvalidResponse(_)))
));
assert!(events.next().await.is_none());
}
}

View file

@ -1,5 +1,5 @@
use litellm_llms_types::formats::responses::{
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
ResponsesApiResponse, streaming::ResponsesStreamEvent,
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
@ -38,11 +38,11 @@ pub trait BaseResponsesApiConfig: Sync {
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ResponsesWsTransformResult {
pub events: Vec<ResponsesWsEvent>,
pub events: Vec<ResponsesStreamEvent>,
}
impl ResponsesWsTransformResult {
pub fn passthrough(event: ResponsesWsEvent) -> Self {
pub fn passthrough(event: ResponsesStreamEvent) -> Self {
Self {
events: vec![event],
}
@ -58,18 +58,18 @@ pub trait ResponsesWebSocketProviderConfig: Sync {
fn transform_ws_request(
&self,
event: &ResponsesWsEvent,
event: &ResponsesStreamEvent,
model: &str,
) -> Result<ResponsesWsTransformResult, Error>;
fn transform_ws_response(
&self,
event: &ResponsesWsEvent,
event: &ResponsesStreamEvent,
model: &str,
) -> Result<ResponsesWsTransformResult, Error>;
}
pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent {
pub fn enforce_model(event: &ResponsesStreamEvent, model: &str) -> ResponsesStreamEvent {
if !event.is_response_create() {
return event.clone();
}
@ -103,7 +103,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent
mod tests {
use super::*;
fn event(value: serde_json::Value) -> ResponsesWsEvent {
fn event(value: serde_json::Value) -> ResponsesStreamEvent {
serde_json::from_value(value).expect("valid event")
}

View file

@ -0,0 +1,102 @@
//! SSE decoding and encoding shared by every format's `streaming` module.
//!
//! DIVERGENCE: decoding is strict. An event whose `data` is not a valid event of the format is
//! `InvalidResponse` and ends the stream, and so is `[DONE]` on a Messages stream. Python
//! never fails a stream on a malformed event:
//! - Chat Completions (`BaseModelResponseIterator`) and Responses
//! (`BaseResponsesAPIStreamingIterator`) skip the event and keep going
//! - Messages (`BaseAnthropicMessagesStreamingIterator`) forwards the provider's bytes to the
//! client without parsing them
//! - Anthropic to Chat (`anthropic/chat/handler.py` `ModelResponseIterator`) buffers the invalid
//! lines on the guess that one payload was split across them
//!
//! Skipping loses content or a tool-call fragment with nothing to show for it, so Rust reports it.
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream};
use litellm_framer::{
frames,
sse::{SseCodec, SseEvent, decode_all},
};
use serde::{Serialize, de::DeserializeOwned};
use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream};
pub type EventStream<E> = BoxStream<'static, Result<E, Error>>;
/// Which frames are control frames rather than events of the format.
pub type IsControl = fn(&SseEvent) -> bool;
/// OpenAI's end-of-stream marker on Chat Completions and Responses streams. It closes the
/// stream and is not an event; Anthropic Messages streams never carry it.
pub const DONE_FRAME: &[u8] = b"data: [DONE]\n\n";
pub fn is_done(event: &SseEvent) -> bool {
event.data.trim() == "[DONE]"
}
pub fn event_stream<E: DeserializeOwned + Send + 'static>(
bytes: ByteStream,
is_control: IsControl,
) -> EventStream<E> {
frames(bytes, SseCodec::default())
.map_err(framing)
.try_filter(move |event| ready(!is_control(event)))
.and_then(|event| ready(parse(&event)))
.boxed()
}
/// Decodes a whole stream at once. An event the stream ends inside is dropped.
pub fn decode<E: DeserializeOwned>(wire: &[u8], is_control: IsControl) -> Result<Vec<E>, Error> {
decode_all(wire)
.map_err(framing)?
.iter()
.filter(|event| !is_control(event))
.map(parse)
.collect()
}
/// A bare `data:` frame, which is what LiteLLM's proxy sends to clients on OpenAI formats.
pub fn encode_data<E: Serialize>(event: &E) -> Result<Bytes, Error> {
let data = serde_json::to_string(event).map_err(invalid)?;
Ok(Bytes::from(format!("data: {data}\n\n")))
}
fn parse<E: DeserializeOwned>(event: &SseEvent) -> Result<E, Error> {
serde_json::from_str(&event.data).map_err(invalid)
}
pub fn invalid(error: serde_json::Error) -> Error {
Error::InvalidResponse(ErrorDetail::invalid("stream event", error))
}
fn framing(error: litellm_framer::SseError) -> Error {
Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::exact("[DONE]", true)]
#[case::padded(" [DONE] ", true)]
#[case::json("{\"type\":\"x\"}", false)]
#[case::mentions("[DONE] was sent", false)]
#[case::empty("", false)]
fn done_is_an_exact_marker(#[case] data: &str, #[case] done: bool) {
let event = SseEvent {
data: data.into(),
..SseEvent::default()
};
assert_eq!(is_done(&event), done);
}
#[rstest]
fn the_done_frame_we_send_decodes_as_a_control_frame() {
let events: Vec<serde_json::Value> = decode(DONE_FRAME, is_done).unwrap();
assert!(events.is_empty());
}
}

View file

@ -102,7 +102,7 @@ mod tests {
use litellm_llms_types::formats::messages::streaming::MessagesContentBlockDelta;
use super::*;
use crate::base_llm::messages::streaming::anthropic_sse_event_stream;
use crate::base_llm::messages::streaming;
fn in_pieces(wire: &[u8]) -> ByteStream {
let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
@ -147,7 +147,7 @@ mod tests {
.try_collect::<Vec<_>>()
.await
.unwrap();
let from_sse = anthropic_sse_event_stream(in_pieces(sse.as_bytes()))
let from_sse = streaming::event_stream(in_pieces(sse.as_bytes()))
.try_collect::<Vec<_>>()
.await
.unwrap();

View file

@ -332,7 +332,7 @@ mod tests {
use litellm_auth_aws::constants::DEFAULT_BEDROCK_REGION;
use super::*;
use crate::base_llm::messages::streaming::encode_anthropic_sse;
use crate::base_llm::messages::streaming;
fn event(value: Value) -> MessagesStreamEvent {
serde_json::from_value(value).unwrap()
@ -493,7 +493,7 @@ mod tests {
.boxed();
let sse = bedrock_anthropic_messages_event_stream(bytes)
.map_ok(|event| encode_anthropic_sse(&event).unwrap())
.map_ok(|event| streaming::encode(&event).unwrap())
.try_collect::<Vec<_>>()
.await
.unwrap()
@ -508,7 +508,7 @@ mod tests {
)),
]
.iter()
.flat_map(|event| encode_anthropic_sse(event).unwrap())
.flat_map(|event| streaming::encode(event).unwrap())
.collect();
assert_eq!(sse, expected);
}

View file

@ -1,5 +1,5 @@
use litellm_llms_types::formats::responses::{
ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
ResponsesApiResponse, streaming::ResponsesStreamEvent,
};
use serde_json::{Map, Value};
@ -34,7 +34,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
fn transform_ws_request(
&self,
event: &ResponsesWsEvent,
event: &ResponsesStreamEvent,
model: &str,
) -> Result<ResponsesWsTransformResult, Error> {
Ok(ResponsesWsTransformResult::passthrough(enforce_model(
@ -44,7 +44,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
fn transform_ws_response(
&self,
event: &ResponsesWsEvent,
event: &ResponsesStreamEvent,
_model: &str,
) -> Result<ResponsesWsTransformResult, Error> {
Ok(ResponsesWsTransformResult::passthrough(event.clone()))
@ -230,7 +230,7 @@ mod tests {
#[test]
fn openai_config_is_native_and_enforces_model() {
let event: ResponsesWsEvent =
let event: ResponsesStreamEvent =
serde_json::from_value(serde_json::json!({"type":"response.create"}))
.expect("valid event");
let result = OPENAI_RESPONSES_WS_CONFIG

View file

@ -48,8 +48,8 @@ Each step needs focused regression tests in the owning crate and Python integrat
- Preserve provider-specific upload/submission/poll observation and encoding boundaries. Signed/build-captured bytes must not be silently reserialized
- Conversion errors and every failure after the call starts are terminal
- Disabled/unavailable native execution may select legacy once. Callback exceptions never authorize fallback or replay
- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle and call driver in `litellm-host-python`
- Contract: `start`, `resume_value`, `resume_error`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
- Use one ordinary inline `async def` driver in `litellm/rust_bridge/execution.py`, with the native handle and call driver in `litellm-host-python`
- Contract: generator-style `send` (`send(None)` starts), `throw`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
- Validate Created/Running/Suspended/Closed protocol states. The machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy
- Defer effectful setup/context reads/timestamps until start. Unstarted-handle destruction releases inputs independently of Python `finally`
- Catch only the selected await's errors. Start/resume errors propagate, `GeneratorExit` closes without further awaits

View file

@ -6,7 +6,6 @@ mod diagnostics;
mod errors;
mod execution;
mod http;
mod lifecycle;
mod logger;
mod marshal;
mod preflight;

View file

@ -1,10 +0,0 @@
use litellm_host_python::CallOptions;
use pyo3::prelude::*;
pub(crate) fn binding(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
py.import("litellm.rust_bridge.streams")
}
pub(crate) fn call_options(asynchronous: bool) -> CallOptions {
CallOptions::new(asynchronous, binding)
}

View file

@ -12,9 +12,9 @@ pub(crate) mod traces;
use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall};
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{
HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
};
use litellm_host_python::{missing_state, present};
use litellm_host_python::{Error as HostError, present};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
@ -64,7 +64,9 @@ impl RequestView {
}
fn dicts(&self) -> PyResult<&RequestDicts> {
self.0.as_ref().ok_or_else(missing_state)
self.0
.as_ref()
.ok_or_else(|| HostError::MissingState.into())
}
pub(crate) fn resolve(
@ -72,7 +74,7 @@ impl RequestView {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<()> {
let dicts = self.0.as_mut().ok_or_else(missing_state)?;
let dicts = self.0.as_mut().ok_or(HostError::MissingState)?;
dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind();
Ok(())
}
@ -127,6 +129,10 @@ fn call_hooks(
))
}
fn execution_steps(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
py.import("litellm.rust_bridge.execution")
}
fn run_public_call<H, M>(
py: Python<'_>,
arguments: Py<PyDict>,
@ -157,7 +163,7 @@ where
.with(hooks)
.with(crate::preflight::SdkPolicy),
arguments,
crate::lifecycle::call_options(asynchronous),
CallOptions::new(asynchronous, execution_steps),
)
}

View file

@ -1,5 +1,5 @@
use litellm_auth::ResolvedCredential;
use litellm_host_python::{InvokeError, PythonBinding, missing_state};
use litellm_host_python::{Error as HostError, InvokeError, PythonBinding};
use litellm_host_python::{PythonHostCalls, PythonOwned};
use litellm_inference_ocr::route::{Ocr, OcrCall, OcrOp};
use litellm_llms::base_llm::ocr::error::Error;
@ -42,7 +42,7 @@ impl OcrPythonHost {
fn handles(&self) -> PyResult<&OcrHostHandles> {
match &self.data {
OcrHostData::Projected(handles) => Ok(handles),
_ => Err(missing_state()),
_ => Err(HostError::MissingState.into()),
}
}
@ -50,13 +50,13 @@ impl OcrPythonHost {
self.handles()?
.azure_ad_token_provider
.as_ref()
.ok_or_else(missing_state)?
.ok_or(HostError::MissingState)?
.acquire(py)
}
fn projection(&mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<OcrCall> {
let OcrHostData::Unprojected = self.data else {
return Err(missing_state());
return Err(HostError::MissingState.into());
};
self.view.resolve(py, arguments)?;
let (request, handles) = project_request(self.view.request(py)?, arguments)?;

View file

@ -441,7 +441,6 @@ autorouter_presets_url: str = os.getenv(
"LITELLM_AUTOROUTER_PRESETS_URL",
"https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json",
)
whats_new_url: str | None = os.getenv("LITELLM_WHATS_NEW_URL")
suppress_debug_info: bool = False
dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None
@ -1711,6 +1710,12 @@ if TYPE_CHECKING:
from .llms.hosted_vllm.decisions.transformation import (
HostedVLLMDecisionsConfig as HostedVLLMDecisionsConfig,
)
from .llms.laya.decisions.transformation import (
LayaDecisionsConfig as LayaDecisionsConfig,
)
from .llms.bespoke.decisions.transformation import (
BespokeDecisionsConfig as BespokeDecisionsConfig,
)
from .llms.openai.decisions.transformation import (
OpenAIDecisionsConfig as OpenAIDecisionsConfig,
)
@ -2282,12 +2287,14 @@ if TYPE_CHECKING:
get_response_string: Callable[..., str]
supports_function_calling: Callable[..., bool]
supports_web_search: Callable[..., bool]
supports_tool_search: Callable[..., bool]
supports_url_context: Callable[..., bool]
supports_response_schema: Callable[..., bool]
supports_parallel_function_calling: Callable[..., bool]
supports_vision: Callable[..., bool]
supports_audio_input: Callable[..., bool]
supports_video_input: Callable[..., bool]
supports_box_questions: Callable[..., bool]
supports_audio_output: Callable[..., bool]
supports_system_messages: Callable[..., bool]
supports_reasoning: Callable[..., bool]

View file

@ -31,12 +31,14 @@ UTILS_NAMES: Final = (
"create_tokenizer",
"supports_function_calling",
"supports_web_search",
"supports_tool_search",
"supports_url_context",
"supports_response_schema",
"supports_parallel_function_calling",
"supports_vision",
"supports_audio_input",
"supports_video_input",
"supports_box_questions",
"supports_audio_output",
"supports_system_messages",
"supports_reasoning",
@ -166,6 +168,8 @@ LLM_CONFIG_NAMES: Final = (
"StrandsDeciderDecisionsConfig",
"DatabricksDecisionsConfig",
"HostedVLLMDecisionsConfig",
"LayaDecisionsConfig",
"BespokeDecisionsConfig",
"OpenAIDecisionsConfig",
"AzureAIDecisionsConfig",
"NvidiaNimRerankConfig",
@ -475,6 +479,7 @@ _UTILS_IMPORT_MAP: Final = {
"create_tokenizer": (".utils", "create_tokenizer"),
"supports_function_calling": (".utils", "supports_function_calling"),
"supports_web_search": (".utils", "supports_web_search"),
"supports_tool_search": (".utils", "supports_tool_search"),
"supports_url_context": (".utils", "supports_url_context"),
"supports_response_schema": (".utils", "supports_response_schema"),
"supports_parallel_function_calling": (
@ -484,6 +489,7 @@ _UTILS_IMPORT_MAP: Final = {
"supports_vision": (".utils", "supports_vision"),
"supports_audio_input": (".utils", "supports_audio_input"),
"supports_video_input": (".utils", "supports_video_input"),
"supports_box_questions": (".utils", "supports_box_questions"),
"supports_audio_output": (".utils", "supports_audio_output"),
"supports_system_messages": (".utils", "supports_system_messages"),
"supports_reasoning": (".utils", "supports_reasoning"),
@ -736,6 +742,8 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.hosted_vllm.decisions.transformation",
"HostedVLLMDecisionsConfig",
),
"LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"),
"BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"),
"OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
"AzureAIDecisionsConfig": (".llms.azure_ai.decisions.transformation", "AzureAIDecisionsConfig"),
"NvidiaNimRerankConfig": (

View file

@ -675,7 +675,10 @@ def _get_redis_client_logic(**env_overrides):
if redis_kwargs.get("startup_nodes")
else redis.Redis
)
return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
coerced_kwargs: Final = _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
if coerced_kwargs.get("socket_timeout") is None:
coerced_kwargs["socket_timeout"] = REDIS_SOCKET_TIMEOUT
return coerced_kwargs
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
@ -722,7 +725,6 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:
@ -755,7 +757,6 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:

View file

@ -152,6 +152,11 @@ class ServiceLogging(CustomLogger):
except Exception as e:
verbose_logger.exception("Error dispatching service event - %s", e)
def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
prometheus_logger: Final[PrometheusServicesLogger | None] = getattr(self, "prometheusServicesLogger", None)
if prometheus_logger is not None:
prometheus_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
def service_success_hook(
self,
service: ServiceTypes,

View file

@ -117,7 +117,7 @@ class BedrockAgentCoreA2AHandler:
response.raise_for_status()
# Check content type — AgentCore may return JSON instead of SSE
content_type: Final = response.headers.get("content-type", "").lower()
content_type: Final = ", ".join(response.headers.get_list("content-type")).lower()
if "application/json" in content_type:
# Single JSON response fallback (not SSE)

View file

@ -768,7 +768,7 @@ def _get_response_from_batch_job_output_file(
def _batch_response_was_successful(
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
batch_job_output_file: Mapping[str, object], custom_llm_provider: str = "openai"
) -> bool:
"""
Check if the batch job response was successful

View file

@ -65,7 +65,7 @@ class AzureBlobCache(BaseCache):
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes: Final = self.container_client.download_blob(key).readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = json.loads(as_str)
cached_response: Final[object] = json.loads(as_str)
verbose_logger.debug(
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
@ -86,7 +86,7 @@ class AzureBlobCache(BaseCache):
blob: Final = await self.async_container_client.download_blob(key)
as_bytes: Final = await blob.readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = json.loads(as_str)
cached_response: Final[object] = json.loads(as_str)
verbose_logger.debug(
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
key,

View file

@ -709,7 +709,7 @@ class Cache:
cache_key = self.get_cache_key(**kwargs)
if cache_key is not None:
if isinstance(result, BaseModel):
result = result.model_dump_json()
result = result.model_dump_json(exclude_unset=True)
## DEFAULT TTL ##
if self.ttl is not None:

View file

@ -1088,6 +1088,7 @@ class LLMCachingHandler:
kwargs=kwargs,
start_time=self.start_time,
end_time=datetime.datetime.now(),
cache_hit=True,
)
return cached_result
@ -1180,7 +1181,7 @@ class LLMCachingHandler:
)
)
else:
result_json: Final = result.model_dump_json()
result_json: Final = result.model_dump_json(exclude_unset=True)
create_cache_write_task(
lambda: cache.async_add_cache(
result_json,

View file

@ -51,6 +51,7 @@ class InMemoryCache(BaseCache):
self.ttl_dict: dict = {}
self.expiration_heap: list[tuple[float, str]] = []
self._increment_lock = threading.Lock()
self._write_lock = threading.Lock()
self._clock = clock if clock is not None else lambda: time.time()
def check_value_size(self, value: Any):
@ -96,6 +97,9 @@ class InMemoryCache(BaseCache):
"""
return key in self.ttl_dict and self._clock() > self.ttl_dict[key]
def _cache_keys(self) -> tuple[str, ...]:
return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state
def _remove_key(self, key: str) -> None:
"""
Remove a key from both cache_dict and ttl_dict
@ -133,13 +137,23 @@ class InMemoryCache(BaseCache):
# Case 3: Entry is valid and not expired
break
if len(self.expiration_heap) > 2 * self.max_size_in_memory:
self.expiration_heap = [entry for entry in self.expiration_heap if self.ttl_dict.get(entry[1]) == entry[0]]
heapq.heapify(self.expiration_heap)
# Step 2: Evict if cache is still full
while len(self.cache_dict) >= self.max_size_in_memory:
while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory:
expiration_time, key = heapq.heappop(self.expiration_heap)
# Skip if key was removed or updated
if self.ttl_dict.get(key) == expiration_time:
self._remove_key(key)
if not self.expiration_heap:
untracked_keys: Final = self._cache_keys()
overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1
for untracked_key in untracked_keys[: max(overflow, 0)]:
self._remove_key(untracked_key)
# de-reference the removed item
# https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/
# One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used.
@ -150,7 +164,7 @@ class InMemoryCache(BaseCache):
Check if ttl is set for a key
"""
ttl_time: Final = self.ttl_dict.get(key)
if ttl_time is None or float(ttl_time) < self._clock(): # if ttl is not set, allow override
if ttl_time is None or ttl_time < self._clock(): # if ttl is not set, allow override
return True
else:
return False
@ -160,21 +174,22 @@ class InMemoryCache(BaseCache):
if self.max_size_in_memory == 0:
return # Don't cache anything if max size is 0
# Always prune expired/outdated heap roots before inserting.
# This keeps expiration_heap bounded even when the live cache stays
# below max_size_in_memory and keys are reinserted after TTL expiry.
self.evict_cache()
if not self.check_value_size(value):
return
with self._write_lock:
# Always prune expired/outdated heap roots before inserting.
# This keeps expiration_heap bounded even when the live cache stays
# below max_size_in_memory and keys are reinserted after TTL expiry.
self.evict_cache()
if not self.check_value_size(value):
return
self.cache_dict[key] = value
if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl
if "ttl" in kwargs and kwargs["ttl"] is not None:
self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
else:
self.ttl_dict[key] = self._clock() + self.default_ttl
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
self.cache_dict[key] = value
if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl
if "ttl" in kwargs and kwargs["ttl"] is not None:
self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
else:
self.ttl_dict[key] = self._clock() + self.default_ttl
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
async def async_set_cache(self, key, value, **kwargs):
self.set_cache(key=key, value=value, **kwargs)
@ -258,14 +273,15 @@ class InMemoryCache(BaseCache):
return results
def flush_cache(self):
self.cache_dict.clear()
self.ttl_dict.clear()
self.expiration_heap.clear()
with self._write_lock:
self.cache_dict.clear()
self.ttl_dict.clear()
self.expiration_heap.clear()
async def disconnect(self):
pass
def delete_cache(self, key):
def delete_cache(self, key: str) -> None:
self._remove_key(key)
async def async_get_ttl(self, key: str) -> int | None:

View file

@ -18,7 +18,9 @@ import json
import logging
import threading
import time
from collections.abc import Awaitable, Callable, Iterator, Sequence
import weakref
from collections import Counter
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
@ -51,12 +53,12 @@ from .base_cache import BaseCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prometheus_client import Counter as _PromCounter
from prometheus_client import Gauge as _PromGauge
from redis.asyncio import Redis, RedisCluster
from redis.asyncio.client import Pipeline
from redis.asyncio.cluster import ClusterPipeline
from litellm._service_logger import ServiceLogging
pipeline = Pipeline
cluster_pipeline = ClusterPipeline
async_redis_client = Redis
@ -100,8 +102,7 @@ _GENERIC_CALLER_MODULES: Final = frozenset(
"litellm.caching.redis_batch",
"litellm.caching.dual_cache",
"litellm.caching.caching",
"litellm.rust_bridge.lifecycle",
"litellm.rust_bridge.streams",
"litellm.rust_bridge.execution",
"contextlib",
}
)
@ -243,7 +244,9 @@ class RedisCircuitBreaker:
recovery_timeout: int,
enabled: bool = True,
timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
service_logger: "ServiceLogging | None" = None,
) -> None:
self._service_logger = service_logger
self.failure_threshold = failure_threshold
self.recovery_timeout = recovery_timeout
self.enabled = enabled
@ -255,7 +258,12 @@ class RedisCircuitBreaker:
self._opened_at: float | None = None
self._state = self.CLOSED
self._generation = 0
_breaker_metrics().record_state_change(None, self._state)
_LIVE_BREAKERS.add(self)
self._report(transitions_to=None, failure_class=None)
@property
def state(self) -> str:
return self._state
@property
def generation(self) -> int:
@ -301,7 +309,7 @@ class RedisCircuitBreaker:
self._timeout_count = 0
self._timeout_streak_started_at = None
self._opened_at = now
_breaker_metrics().record_failure("timeout" if is_timeout else "connectivity")
self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity")
if self._should_open(now):
if self._state != self.OPEN:
verbose_logger.warning(
@ -327,10 +335,13 @@ class RedisCircuitBreaker:
def _set_state(self, state: str) -> None:
if state == self._state:
return
_breaker_metrics().record_transition(state)
_breaker_metrics().record_state_change(self._state, state)
self._state = state
self._generation += 1
self._report(transitions_to=state, failure_class=None)
def _report(self, transitions_to: str | None, failure_class: str | None) -> None:
if self._service_logger is not None:
self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
_RedisCallResult = TypeVar("_RedisCallResult")
@ -400,12 +411,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]:
``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from
py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own
class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither.
``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have
no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a
dead Redis, like the standalone pool's chained wait timeout.
"""
try:
from redis.exceptions import TimeoutError as RedisTimeoutError
from redis import exceptions as redis_exceptions
except ImportError:
return (TimeoutError, asyncio.TimeoutError)
return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError)
max_connections_error: Final = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
"type[Exception]",
redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # types-redis 4.6 stubs predate it
)
return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error)
_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20
@ -430,59 +448,6 @@ def is_redis_timeout_failure(exc: BaseException) -> bool:
return any(isinstance(link, timeout_types) for link in _explicit_causes(exc))
class _BreakerMetrics:
"""Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent.
Registered lazily on the default registry (which /metrics serves) via the module-level
``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers.
"""
def __init__(self) -> None:
self._state_gauge: _PromGauge | None = None
self._transitions: _PromCounter | None = None
self._failures: _PromCounter | None = None
try:
from prometheus_client import Counter as PromCounter
from prometheus_client import Gauge
except ImportError:
return
self._state_gauge = Gauge(
"litellm_redis_circuit_breaker_state",
"Number of Redis circuit breakers currently in each state",
labelnames=("state",),
)
self._transitions = PromCounter(
"litellm_redis_circuit_breaker_transitions",
"Redis circuit breaker state transitions",
labelnames=("state",),
)
self._failures = PromCounter(
"litellm_redis_circuit_breaker_failures",
"Redis health failures counted by the circuit breaker",
labelnames=("failure_class",),
)
def record_state_change(self, old_state: str | None, new_state: str) -> None:
if self._state_gauge is None:
return
if old_state is not None:
self._state_gauge.labels(old_state).dec()
self._state_gauge.labels(new_state).inc()
def record_transition(self, state: str) -> None:
if self._transitions is not None:
self._transitions.labels(state).inc()
def record_failure(self, failure_class: str) -> None:
if self._failures is not None:
self._failures.labels(failure_class).inc()
@functools.lru_cache(maxsize=1)
def _breaker_metrics() -> _BreakerMetrics:
return _BreakerMetrics()
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
"""Record a Redis failure that the calling method is about to swallow.
@ -498,6 +463,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
def redis_breaker_state_counts() -> Mapping[str, int]:
"""How many live Redis circuit breakers sit in each state, for the Prometheus services logger."""
return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS)))
class RedisCircuitBreakerOpenError(Exception):
pass
@ -637,9 +610,14 @@ def _redis_circuit_breaker_guard(method):
"""
@functools.wraps(method)
async def wrapper(self, *args, **kwargs):
async def wrapper(self: "RedisCache", *args, **kwargs):
async def guarded() -> object:
return await method(self, *args, **kwargs)
return await _run_under_circuit_breaker(
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator
wrapper.__name__,
lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator
)
return wrapper
@ -684,12 +662,9 @@ class RedisCache(BaseCache):
redis_kwargs["socket_timeout"] = socket_timeout
### HEALTH MONITORING OBJECT ###
if kwargs.get("service_logger_obj", None) is not None and isinstance(
kwargs["service_logger_obj"], ServiceLogging
):
self.service_logger_obj: ServiceLogging = ServiceLogging()
if isinstance(kwargs.get("service_logger_obj"), ServiceLogging):
self.service_logger_obj = kwargs.pop("service_logger_obj")
else:
self.service_logger_obj = ServiceLogging()
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
@ -716,6 +691,7 @@ class RedisCache(BaseCache):
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
service_logger=self.service_logger_obj,
)
self._setup_health_pings()
@ -725,6 +701,9 @@ class RedisCache(BaseCache):
else:
super().__init__() # defaults to 60s
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
return call()
def _setup_health_pings(self):
"""Setup async and sync health pings for Redis."""
# ASYNC HEALTH PING
@ -1105,7 +1084,9 @@ class RedisCache(BaseCache):
litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
return await executor(keys=keys, args=args, client=client)
return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute)
return await _run_under_circuit_breaker(
self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute)
)
return run_script

View file

@ -5,9 +5,17 @@ Key differences:
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
"""
from typing import TYPE_CHECKING, Any, Final
import asyncio
import contextvars
import os
import weakref
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
from litellm._logging import verbose_logger
from litellm.caching.redis_cache import RedisCache
from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -23,9 +31,100 @@ else:
Span = Any
_T: Final = TypeVar("_T")
def _connection_limit(max_connections: object) -> int | None:
try:
limit: Final = int(max_connections) # pyright: ignore[reportArgumentType] # config value, str or int
except (TypeError, ValueError):
return None
return limit if limit > 0 else None
_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore]
@dataclass(slots=True)
class _SemaphoreRegistry:
"""Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a
closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired."""
_per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field(
default_factory=weakref.WeakKeyDictionary, repr=False
)
def get(self, client_key: str, limit: int) -> asyncio.Semaphore:
per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary())
existing: Final = per_loop.get(client_key)
if existing is not None:
return existing
created: Final = asyncio.Semaphore(limit)
per_loop[client_key] = created
return created
_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"]
_SEMAPHORES: Final = _SemaphoreRegistry()
_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar(
"redis_cluster_gate_admitted", default=frozenset()
)
@dataclass(frozen=True, slots=True)
class ClusterConnectionGate:
"""The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises
``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes
guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no
slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error
handling and the breaker."""
limit: int
wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT
@classmethod
def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None":
limit: Final = _connection_limit(max_connections)
return None if limit is None else cls(limit=limit)
async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T:
admission: Final[_Admission] = (client_key, asyncio.current_task())
if admission in _ADMITTED.get():
return await call()
semaphore: Final = _SEMAPHORES.get(client_key, self.limit)
try:
await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout)
except asyncio.TimeoutError:
verbose_logger.debug(
"Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway",
self.wait_timeout,
self.limit,
)
return await call()
token: Final = _ADMITTED.set(_ADMITTED.get() | {admission})
try:
return await call()
finally:
_ADMITTED.reset(token)
semaphore.release()
class RedisClusterCache(RedisCache):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs):
pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections}
super().__init__(*args, **pool_kwargs, **kwargs)
self._connection_gate: Final = (
connection_gate
if isinstance(connection_gate, ClusterConnectionGate)
else ClusterConnectionGate.from_max_connections(
os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections
)
)
def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]:
if self._connection_gate is None:
return call()
return self._connection_gate.run(call, self._get_async_client_cache_key())
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""

View file

@ -2,8 +2,14 @@
Handler for transforming /chat/completions api requests to litellm.responses requests
"""
from collections.abc import AsyncIterable, Coroutine, Iterable
from typing import TYPE_CHECKING, Any, Final, Union
from collections.abc import AsyncIterable, Coroutine, Iterable, Mapping
from typing import (
TYPE_CHECKING,
Any,
Final,
Union,
cast, # noqa: TID251 # optional_params reaches the bridge as an untyped dict
)
from typing_extensions import TypedDict
@ -245,10 +251,15 @@ class ResponsesToCompletionBridgeHandler:
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(result, model, custom_llm_provider)
from .transformation import is_structured_output_request
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=True,
json_mode=kwargs.get("json_mode"),
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,
@ -343,10 +354,15 @@ class ResponsesToCompletionBridgeHandler:
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(result, model, custom_llm_provider)
from .transformation import is_structured_output_request
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=False,
json_mode=kwargs.get("json_mode"),
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,

View file

@ -20,7 +20,7 @@ from openai.types.responses.response_input_param import (
from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm import ModelResponse
@ -35,6 +35,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.bridges.completion_transformation import (
CompletionTransformationBridge,
)
from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
from litellm.responses.sse_output_recovery import (
parse_sse_json_chunk,
record_output_item_chunk,
@ -304,6 +305,136 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[
)
_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"})
_TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_ITEM_TYPES})
_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
def _json_field(value: object, key: str) -> object:
return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON
def _tool_search_replay_item(item: object) -> Mapping[str, object] | None:
raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item
if not isinstance(raw, dict):
return None
fields: Final = cast(dict[str, object], raw) # cast-ok: provider JSON item, keys forwarded verbatim
if fields.get("type") not in _TOOL_SEARCH_REPLAY_ITEM_TYPES:
return None
return {key: value for key, value in fields.items() if value is not None}
def _tool_search_items_from_output_items(output_items: object) -> tuple[Mapping[str, object], ...]:
if not isinstance(output_items, list):
return ()
entries: Final = cast(Sequence[object], output_items) # cast-ok: each entry is re-validated
items: Final = tuple(item for item in (_tool_search_replay_item(entry) for entry in entries) if item is not None)
last_search: Final = max(
(index for index, item in enumerate(items) if item.get("type") in _TOOL_SEARCH_ITEM_TYPES), default=-1
)
return items[: last_search + 1]
def _tool_search_provider_fields(
items: Sequence[Mapping[str, object]],
) -> dict[str, object] | None: # mutable-ok: Message.provider_specific_fields is typed as a dict
return {_TOOL_SEARCH_ITEMS_KEY: [dict(item) for item in items]} if items else None
def _tool_search_input_items(msg: object) -> tuple[Mapping[str, object], ...]:
return _tool_search_items_from_output_items(
_json_field(_json_field(msg, "provider_specific_fields"), _TOOL_SEARCH_ITEMS_KEY)
)
def carries_tool_search_items(provider_specific_fields: object) -> bool:
return bool(_json_field(provider_specific_fields, _TOOL_SEARCH_ITEMS_KEY))
def _tool_definition_name(tool: object) -> tuple[str, ...]:
name: Final = _json_field(tool, "server_label" if _json_field(tool, "type") == "mcp" else "name")
return (name,) if isinstance(name, str) else ()
def _namespace_members(tool: object) -> Sequence[object]:
members: Final = _json_field(tool, "tools")
if _json_field(tool, "type") != "namespace" or not isinstance(members, list):
return (tool,)
return cast(Sequence[object], members) # cast-ok: untyped JSON, each member re-checked
def _loaded_tools(item: Mapping[str, object]) -> Sequence[object]:
tools: Final = item.get("tools")
if item.get("type") != "tool_search_output" or not isinstance(tools, list):
return ()
return cast(Sequence[object], tools) # cast-ok: untyped JSON, each tool re-checked
def _json_list(value: object) -> Sequence[object] | None:
return cast(Sequence[object], value) if isinstance(value, list) else None # cast-ok: untyped JSON
def _message_loaded_tools(message: object) -> Iterator[object]:
loaded: Final = chain.from_iterable(_loaded_tools(item) for item in _tool_search_input_items(message))
return chain.from_iterable(_namespace_members(tool) for tool in loaded)
def replayed_tool_search_tools(request: object) -> tuple[object, ...]:
messages: Final = _json_list(_json_field(request, "messages")) or ()
return tuple(chain.from_iterable(_message_loaded_tools(message) for message in messages))
def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
return any(name in names for name in loaded_names)
def _without_tool_search_items(message: object) -> object:
fields: Final = _json_field(message, "provider_specific_fields")
if not isinstance(fields, dict):
return message
kept: Final = {
key: value
for key, value in cast(Mapping[str, object], fields).items() # cast-ok: untyped JSON
if key != _TOOL_SEARCH_ITEMS_KEY
}
return {**cast(Mapping[str, object], message), "provider_specific_fields": kept} # cast-ok: dict checked above
def without_tool_search_items_loading(
request: object, names: frozenset[str]
) -> list[object] | None: # mutable-ok: request messages stay a list for every downstream handler
messages: Final = _json_list(_json_field(request, "messages"))
if messages is None:
return None
return [_without_tool_search_items(message) if _loads_any_tool(message, names) else message for message in messages]
def _effective_tools(request: object) -> Sequence[object]:
override: Final = _json_list(_json_field(_json_field(request, "extra_body"), "tools"))
return override if override is not None else (_json_list(_json_field(request, "tools")) or ())
def _requests_tool_search(request: object) -> bool:
return any(_json_field(tool, "type") == "tool_search" for tool in _effective_tools(request))
def _reasoning_ids(items: Sequence[Mapping[str, object]]) -> frozenset[str]:
ids: Final = (item.get("id") for item in items if item.get("type") == "reasoning")
return frozenset(item_id for item_id in ids if isinstance(item_id, str))
def _assistant_replay_items(msg: "AllMessageValues", replay_tool_search: bool) -> tuple[Mapping[str, object], ...]:
searched: Final = _tool_search_input_items(msg) if replay_tool_search else ()
searched_reasoning_ids: Final = _reasoning_ids(searched)
reasoning: Final = (item for item in _reasoning_input_items(msg) if item.get("id") not in searched_reasoning_ids)
return (*searched, *reasoning)
def _as_chat_reasoning_items(
reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem],
) -> list[ChatCompletionReasoningItem] | None:
@ -346,14 +477,21 @@ def _fold_choices_into_one(
tool_calls: Sequence[Mapping[str, object]],
trailing_reasoning_content: str | None,
trailing_reasoning_item: _BuiltReasoningItem | None,
tool_search_items: Sequence[Mapping[str, object]],
structured_output: bool = False,
) -> "Choices":
"""Clients that only read choices[0] must still see every text and tool call."""
"""Clients that only read choices[0] must still see every text and tool call.
Under a structured ``response_format`` each message item is a complete document, so
``content`` is the first one, the same item the openai SDK returns as ``output_parsed``.
"""
from litellm.types.utils import Choices, Message
messages: Final = tuple(choice.message for choice in choices)
texts: Final = tuple(str(message.content or "") for message in messages)
content_messages: Final = messages[:1] if structured_output else messages
texts: Final = tuple(str(message.content or "") for message in content_messages)
offsets: Final = (0, *accumulate(len(text) for text in texts))
annotations: Final = tuple(_shifted_annotations(messages, offsets))
annotations: Final = tuple(_shifted_annotations(content_messages, offsets))
reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content)
reasoning_items: Final = (
*chain.from_iterable(_message_reasoning_items(m) for m in messages),
@ -367,12 +505,45 @@ def _fold_choices_into_one(
tool_calls=list(tool_calls) or None,
reasoning_content=" ".join(value for value in reasoning_contents if value) or None,
reasoning_items=_as_chat_reasoning_items(reasoning_items),
provider_specific_fields=_tool_search_provider_fields(tool_search_items),
),
finish_reason="tool_calls" if tool_calls else "stop",
index=0,
)
def _responses_event_type(chunk: Mapping[str, object]) -> object:
event_type: Final = chunk.get("type")
return event_type.value if isinstance(event_type, ResponsesAPIStreamEvents) else event_type
class _TypedPayload(BaseModel):
"""The ``type`` discriminator of a response_format dict or a Responses output item."""
model_config = ConfigDict(frozen=True, extra="ignore")
type: str | None = None
_TYPED_PAYLOAD: Final = TypeAdapter(_TypedPayload)
def _payload_type(payload: object) -> str | None:
try:
return _TYPED_PAYLOAD.validate_python(payload).type
except ValidationError:
return None
def _is_message_item_done(chunk: Mapping[str, object]) -> bool:
if _responses_event_type(chunk) != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
return False
return _payload_type(chunk.get("item")) == "message"
def is_structured_output_request(optional_params: Mapping[str, object]) -> bool:
return _payload_type(optional_params.get("response_format")) in ("json_schema", "json_object")
_ToolChoiceT = TypeVar("_ToolChoiceT")
@ -419,7 +590,9 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
provider_specific_fields: Mapping[str, object]
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
def tool_call_dict_from_output_item(
item: Mapping[str, Any], index: int, *, include_namespace: bool = False
) -> _ChatToolCallDict:
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
string payload in ``input`` rather than ``arguments``; both map to
@ -448,12 +621,24 @@ def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Cha
provider_specific_fields = vars(raw_provider_fields)
else:
provider_specific_fields = None
if provider_specific_fields:
tool_call_dict["provider_specific_fields"] = provider_specific_fields
function_chunk["provider_specific_fields"] = provider_specific_fields
namespace: Final = item.get("namespace") if include_namespace else None
merged_provider_fields: Final = {
**(provider_specific_fields or {}),
**({"namespace": namespace} if isinstance(namespace, str) else {}),
}
if merged_provider_fields:
tool_call_dict["provider_specific_fields"] = merged_provider_fields
function_chunk["provider_specific_fields"] = merged_provider_fields
return tool_call_dict
def _tool_call_namespace(tool_call: object) -> str | None:
top_level: Final = _json_field(_json_field(tool_call, "provider_specific_fields"), "namespace")
function_fields: Final = _json_field(_json_field(tool_call, "function"), "provider_specific_fields")
nested: Final = _json_field(function_fields, "namespace")
return top_level if isinstance(top_level, str) else nested if isinstance(nested, str) else None
def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam:
if choice_type == "custom":
return ToolChoiceCustomParam(type="custom", name=name)
@ -479,6 +664,50 @@ def _reasoning_item_to_response_input(
}
def _requested_defer_loading(tool: Mapping[str, object], function_tool: Mapping[str, object]) -> object:
tool_level: Final = tool.get("defer_loading")
return tool_level if tool_level is not None else function_tool.get("defer_loading")
def _defer_loading_field(
tool: Mapping[str, object],
function_tool: Mapping[str, object],
*,
tool_search: Literal["hosted", "client", "absent"],
drop_params: bool,
) -> Mapping[str, bool]:
defer_loading: Final = _requested_defer_loading(tool, function_tool)
if defer_loading is None or (defer_loading is False and tool_search == "absent"):
return {}
if tool_search == "client":
verbose_logger.debug(
"Chat provider: dropping defer_loading=%r beside a client-executed tool_search", defer_loading
)
return {}
if isinstance(defer_loading, bool) and tool_search == "hosted":
return {"defer_loading": defer_loading}
if drop_params:
verbose_logger.debug("Chat provider: dropping defer_loading=%r under drop_params", defer_loading)
return {}
problem: Final = (
'defer_loading needs a {"type": "tool_search"} tool in the same request'
if isinstance(defer_loading, bool)
else f"defer_loading must be a boolean, got {defer_loading!r}"
)
raise litellm.UnsupportedParamsError(
message=f"{problem}. To drop it, set `litellm.drop_params=True` or pass `drop_params=True`."
)
def _without_dropped_defer_loading(
tool: Mapping[str, object], *, tool_search: Literal["hosted", "client", "absent"], drop_params: bool
) -> Mapping[str, object]:
if tool.get("type") != "function" or "defer_loading" not in tool:
return tool
kept: Final = _defer_loading_field(tool, {}, tool_search=tool_search, drop_params=drop_params)
return {**{key: value for key, value in tool.items() if key != "defer_loading"}, **kept}
class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"""
Handler for transforming /chat/completions api requests to litellm.responses requests
@ -559,9 +788,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
*,
drop_params: bool = False,
keep_prompt_cache_breakpoints: bool = False,
replay_tool_search: bool = False,
include_namespace: bool = False,
) -> tuple[list[object], str | None]:
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
messages, drop_params=drop_params
messages,
drop_params=drop_params,
replay_tool_search=replay_tool_search,
include_namespace=include_namespace,
)
return (
converted_input_items
@ -571,7 +805,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
def _convert_chat_completion_messages_to_responses_input(
self, messages: list["AllMessageValues"], *, drop_params: bool = False
self,
messages: list["AllMessageValues"],
*,
drop_params: bool = False,
replay_tool_search: bool = False,
include_namespace: bool = False,
) -> tuple[list[object], str | None]:
input_items: Final[list[object]] = []
instructions: str | None = None
@ -653,7 +892,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
)
elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
input_items.extend(_reasoning_input_items(msg))
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
if content:
input_items.append(
{
@ -676,6 +915,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
input_tool_call["name"] = function["name"]
if "arguments" in function:
input_tool_call["arguments"] = function["arguments"]
if include_namespace and (
namespace := _tool_call_namespace(cast(object, tool_call)) # cast-ok: untyped
):
input_tool_call["namespace"] = namespace
input_items.append(input_tool_call)
elif isinstance(custom, dict):
input_items.append(
@ -690,7 +933,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
raise ValueError(f"tool call not supported: {tool_call}")
elif content is not None:
if role == "assistant":
input_items.extend(_reasoning_input_items(msg))
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
input_items.append(
{
"type": "message",
@ -701,7 +944,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
}
)
elif role == "assistant":
input_items.extend(_reasoning_input_items(msg))
input_items.extend(_assistant_replay_items(msg, replay_tool_search))
return input_items, instructions
@ -709,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
self,
optional_params: dict,
responses_api_request: "ResponsesAPIOptionalRequestParams",
drop_params: bool = False,
) -> None:
"""Map optional_params into responses_api_request (mutates in place)."""
for key, value in optional_params.items():
@ -718,7 +962,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
responses_api_request["max_output_tokens"] = value
elif key == "tools" and value is not None:
responses_api_request["tools"] = self._convert_tools_to_responses_format(
cast(list[dict[str, object]], value)
cast(list[dict[str, object]], value), drop_params=drop_params
)
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
@ -807,6 +1051,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
litellm_logging_obj: "LiteLLMLoggingObj",
client: object | None = None,
) -> dict:
optional_params = self._extract_extra_body_params(optional_params)
provider: Final = _served_provider(litellm_params)
base_model: Final = litellm_params.get("base_model")
supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or (
@ -818,6 +1063,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
messages,
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast
include_namespace=True,
)
# OpenAI's Responses API rejects an empty input. For a system-only
# request, carry the system message as a system-role input item instead
@ -842,8 +1089,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
instructions: Final = None if is_system_only_request else converted_instructions
optional_params = self._extract_extra_body_params(optional_params)
# Build responses API request using the reverse transformation logic
responses_api_request: Final = ResponsesAPIOptionalRequestParams()
@ -851,7 +1096,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if instructions:
responses_api_request["instructions"] = instructions
self._map_optional_params_to_responses_api_request(optional_params, responses_api_request)
self._map_optional_params_to_responses_api_request(
optional_params,
responses_api_request,
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
)
stream: Final = optional_params.get("stream") or litellm_params.get("stream", False)
verbose_logger.debug("Chat provider: Stream parameter: %s", stream)
@ -897,6 +1146,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _convert_response_output_to_choices(
output_items: Sequence[object],
handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None,
structured_output: bool = False,
) -> list["Choices"]:
"""
Convert Responses API output items to chat completion choices.
@ -904,6 +1154,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
Args:
output_items: List of items from ResponsesAPIResponse.output
handle_raw_dict_callback: Optional callback for handling raw dict items
structured_output: The chat request carried a json response_format, so only the
first message item becomes content
Returns:
List of Choices objects
@ -1010,7 +1262,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Tool calls accumulate into the single trailing tool_calls choice
# like the typed branches above; a choice per call would hide every
# call after choices[0] from chat clients
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
accumulated_tool_calls.append(
tool_call_dict_from_output_item(raw_item, tool_call_index, include_namespace=True)
)
tool_call_index += 1
elif handle_raw_dict_callback is not None:
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@ -1019,21 +1273,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
else:
pass # don't fail request if item in list is not supported
tool_search_items: Final = _tool_search_items_from_output_items(output_items)
if not choices and not accumulated_tool_calls:
return choices
if not tool_search_items:
return choices
return [
LiteLLMResponsesTransformationHandler._build_empty_choice(
output_items, "stop", tool_search_items=tool_search_items
)
]
return [
_fold_choices_into_one(
choices=choices,
tool_calls=accumulated_tool_calls,
trailing_reasoning_content=reasoning_content,
trailing_reasoning_item=pending_reasoning_item,
tool_search_items=tool_search_items,
structured_output=structured_output,
)
]
@staticmethod
def _build_empty_incomplete_choice(
def _build_empty_choice(
output_items: Sequence[object],
finish_reason: Literal["length", "content_filter"],
finish_reason: Literal["stop", "length", "content_filter"],
tool_search_items: Sequence[Mapping[str, object]] = (),
) -> "Choices":
from litellm.types.utils import Choices, Message
@ -1045,9 +1309,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if summary_block.get("text")
)
message: Final = Message(
content="",
content=None if tool_search_items else "",
reasoning_content=reasoning_content if reasoning_content else None,
reasoning_items=_as_chat_reasoning_items(reasoning_items),
provider_specific_fields=_tool_search_provider_fields(tool_search_items),
)
return Choices(message=message, finish_reason=finish_reason, index=0)
@ -1157,6 +1422,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
choices: Final = self._convert_response_output_to_choices(
output_items=output_items,
handle_raw_dict_callback=self._handle_raw_dict_response_item,
structured_output=is_structured_output_request(
cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
),
)
response_is_incomplete: Final = raw_response.status == "incomplete" or (
@ -1171,7 +1439,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
)
if len(choices) == 0:
choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
choices.append(self._build_empty_choice(output_items, incomplete_finish_reason))
else:
for choice in choices:
choice.finish_reason = incomplete_finish_reason
@ -1219,8 +1487,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
sync_stream: bool,
json_mode: bool | None = False,
structured_output: bool = False,
) -> BaseModelResponseIterator:
return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
return OpenAiResponsesToChatCompletionStreamIterator(
streaming_response, sync_stream, json_mode, structured_output=structured_output
)
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
if role in ("user", "system", "developer", "tool"):
@ -1369,14 +1640,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return result
def _convert_tools_to_responses_format(
self, tools: list[dict[str, object]]
self, tools: list[dict[str, object]], drop_params: bool = False
) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
"""Convert chat completion tools to responses API tools format"""
tool_search: Final[Literal["hosted", "client", "absent"]] = (
"client"
if any(tool.get("type") == "tool_search" and tool.get("execution") == "client" for tool in tools)
else "hosted"
if any(tool.get("type") == "tool_search" or is_anthropic_tool_search_chat_tool(tool) for tool in tools)
else "absent"
)
responses_tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = []
for tool in tools:
# convert function tool from chat completion to responses API format
if tool.get("type") == "function" and isinstance(tool.get("function"), dict):
function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function"))
if tool.get("type") == "function" and isinstance(function_payload := tool.get("function"), dict):
function_tool = cast(ChatCompletionToolParamFunctionChunk, function_payload) # cast-ok: client tool
responses_tools.append(
FunctionToolParam(
name=function_tool["name"],
@ -1384,6 +1662,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
strict=function_tool.get("strict"),
type="function",
description=function_tool.get("description"),
**_defer_loading_field(tool, function_tool, tool_search=tool_search, drop_params=drop_params),
)
)
elif tool.get("type") == "custom" and isinstance(custom_payload := tool.get("custom"), dict):
@ -1398,7 +1677,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
flat_custom["format"] = convert_custom_tool_format_to_responses_shape(custom_payload["format"])
responses_tools.append(flat_custom)
else:
responses_tools.append(tool)
responses_tools.append(
_without_dropped_defer_loading(tool, tool_search=tool_search, drop_params=drop_params)
)
return cast(list["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools)
@ -1587,11 +1868,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
sync_stream: bool,
json_mode: bool | None = False,
structured_output: bool = False,
):
super().__init__(streaming_response, sync_stream, json_mode)
self._chat_completion_id: str | None = None
self._served_service_tier: str | None = None
self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state
self._structured_output: Final = structured_output
self._first_message_item_done: bool = False
def _handle_string_chunk(
self, str_line: Union[str, "BaseModel"]
@ -1692,7 +1976,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
# New output item added
output_item = parsed_chunk.get("item", {})
if output_item.get("type") in ("function_call", "custom_tool_call"):
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
converted: Final = tool_call_dict_from_output_item(
output_item, parsed_chunk.get("output_index", 0), include_namespace=True
)
provider_specific_fields: Final = converted.get("provider_specific_fields")
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@ -1705,14 +1991,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index(
tool_call_index_map, parsed_chunk.get("output_index", 0)
)
tool_call_chunk: Final = ChatCompletionToolCallChunk(
tool_call_chunk: Final = _ChatToolCallDict(
id=converted["id"],
index=tool_call_index,
type="function",
function=function_chunk,
**({"provider_specific_fields": provider_specific_fields} if provider_specific_fields else {}),
)
if provider_specific_fields:
tool_call_chunk.provider_specific_fields = provider_specific_fields
return ModelResponseStream(
choices=[
@ -1768,7 +2053,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
delta=Delta(
tool_calls=(
tool_call_dict_from_output_item(
output_item, parsed_chunk.get("output_index", 0)
output_item, parsed_chunk.get("output_index", 0), include_namespace=True
),
)
),
@ -1863,6 +2148,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
delta=Delta(
content="",
reasoning_items=terminal_reasoning_items_typed,
provider_specific_fields=_tool_search_provider_fields(
_tool_search_items_from_output_items(cast(object, output_items)) # cast-ok: untyped
),
),
finish_reason=finish_reason,
)
@ -1903,13 +2191,31 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
"""
verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk)
self._remember_served_service_tier(chunk)
return self._with_served_service_tier(
self._with_stream_scoped_id(
OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
chunk, tool_call_index_map=self._tool_call_index_map
)
translated: Final = (
self._empty_content_chunk()
if self._is_text_delta_after_structured_answer(chunk)
else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
chunk, tool_call_index_map=self._tool_call_index_map
)
)
if self._structured_output and _is_message_item_done(chunk):
self._first_message_item_done = True
return self._with_served_service_tier(self._with_stream_scoped_id(translated))
def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool:
"""Under a structured response_format the first message item is the whole answer, so
text from any later message item is dropped, matching the non-stream fold."""
return (
self._structured_output
and self._first_message_item_done
and _responses_event_type(chunk) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value
)
@staticmethod
def _empty_content_chunk() -> "ModelResponseStream":
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason=None)])
def _remember_served_service_tier(self, chunk: dict[str, object]) -> None:
response_payload: Final = chunk.get("response")

File diff suppressed because it is too large Load diff

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