mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
Merge branch 'main' into fix/skip-unsatisfiable-tag-fallback-legs-43736
This commit is contained in:
commit
3efb533185
1640 changed files with 241473 additions and 20273 deletions
|
|
@ -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
4
.github/CODEOWNERS
vendored
|
|
@ -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
|
||||
|
|
|
|||
6
.github/workflows/test-linting.yml
vendored
6
.github/workflows/test-linting.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
8
.github/workflows/test-unit.yml
vendored
8
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/token",
|
||||
"/onboarding/",
|
||||
"/audit",
|
||||
"/telemetry/",
|
||||
"/oauth/",
|
||||
"/invitation/",
|
||||
"/jwt/",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Realtime / streaming
|
||||
"/v1/realtime",
|
||||
"/realtime",
|
||||
"/v1/live/sessions",
|
||||
"/live/sessions",
|
||||
# Health & ops
|
||||
"/health",
|
||||
"/metrics",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
441
helm/litellm/tests/component_hosts_tests.yaml
Normal file
441
helm/litellm/tests/component_hosts_tests.yaml
Normal 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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;
|
||||
|
|
@ -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[];
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
4
litellm-rust/Cargo.lock
generated
4
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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!(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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())?);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"[..]])]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
58
litellm-rust/crates/host-python/src/caller.rs
Normal file
58
litellm-rust/crates/host-python/src/caller.rs
Normal 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
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>) {
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)?;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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)))?;
|
||||
|
|
|
|||
770
litellm-rust/crates/host-python/tests/driver.rs
Normal file
770
litellm-rust/crates/host-python/tests/driver.rs
Normal 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"]
|
||||
);
|
||||
});
|
||||
}
|
||||
343
litellm-rust/crates/host-python/tests/handle.rs
Normal file
343
litellm-rust/crates/host-python/tests/handle.rs
Normal 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();
|
||||
});
|
||||
}
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
684
litellm-rust/crates/host-python/tests/support/mod.rs
Normal file
684
litellm-rust/crates/host-python/tests/support/mod.rs
Normal 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()
|
||||
}
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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>>>,
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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}))]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
mod output;
|
||||
mod response;
|
||||
pub mod streaming;
|
||||
pub mod streaming_websocket;
|
||||
|
||||
pub use output::*;
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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!({
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,3 +5,4 @@ pub mod chat;
|
|||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod responses;
|
||||
pub mod sse;
|
||||
|
|
|
|||
|
|
@ -1 +1,2 @@
|
|||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal file
135
litellm-rust/crates/llms/src/base_llm/responses/streaming.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
|
||||
|
|
|
|||
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal file
102
litellm-rust/crates/llms/src/base_llm/sse.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ mod diagnostics;
|
|||
mod errors;
|
||||
mod execution;
|
||||
mod http;
|
||||
mod lifecycle;
|
||||
mod logger;
|
||||
mod marshal;
|
||||
mod preflight;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)?;
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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": (
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue