Merge remote-tracking branch 'origin/main' into litellm_key_budget_duration_tier_auth

This commit is contained in:
yucheng-berri 2026-10-20 16:18:32 -07:00
commit e9fc0d6c2c
891 changed files with 25193 additions and 8525 deletions

View file

@ -4,8 +4,6 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
## Description
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->

View file

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -433,6 +434,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge

View file

@ -852,20 +852,26 @@ class CheckBatchCost:
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job, original_status, original_file_object)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -1186,7 +1192,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.78"
version = "0.1.79"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.78"
version = "0.1.79"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -1,7 +1,7 @@
"""Path allowlist for the gateway component.
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/extraction",
"/extraction",
"/v1/ocr",
"/ocr",
"/v1/rag/",

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.110"
version = "0.4.111"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.110"
version = "0.4.111"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -4,7 +4,7 @@ mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
use litellm_cache::Error;
use pyo3::{

View file

@ -279,7 +279,7 @@ pub(in crate::cache) fn configured(
)> {
let handle = native_handle(configured)?.ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err(
"the configured cache changed to a Python cache after native admission",
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
)
})?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;

View file

@ -76,21 +76,6 @@ fn selected_cache<'py>(
Ok(Some(configured))
}
pub(crate) fn admit_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<()> {
if let Some(configured) = selected_cache(py, kwargs, call_type)?
&& native::v2::native_handle(&configured)?.is_none()
{
return Err(crate::errors::RustBridgeDeclined::new_err(
"the configured cache requires Python inference",
));
}
Ok(())
}
pub(crate) fn configured_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,

View file

@ -177,7 +177,6 @@ where
} else {
R::SYNC_CALL_TYPE
};
crate::cache::admit_native(py, &call.kwargs, call_type)?;
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
super::run_public_call(
py,

View file

@ -8,59 +8,19 @@ use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::present;
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use serde_json::Value;
pub(crate) use websocket::ResponsesWebSocketConnection;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let resolved = call.resolved()?;
if let Some(reason) = py
.import(ROUTE_HOST_MODULE)?
.getattr("decline_reason")?
.call1((&resolved,))?
.extract::<Option<String>>()?
{
return Err(RustBridgeDeclined::new_err(reason));
}
let argument = |name: &str| present(&call.kwargs, &call.base, name);
let model = argument("model")?
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
.extract::<String>()?;
let provider = argument("custom_llm_provider")?
.map(|value| value.extract::<String>())
.transpose()?;
if provider
.as_deref()
.is_some_and(|provider| provider != "openai")
|| model
.strip_prefix("openai/")
.unwrap_or(&model)
.contains('/')
{
return Err(RustBridgeDeclined::new_err(
"native HTTP responses provider",
));
}
if argument("stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost::new(host))
}

View file

@ -1505,6 +1505,7 @@ from .rag.main import *
from .sandbox.main import *
from .decisions.main import *
from .systemone.main import *
from .extraction.main import *
from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop
from .search.main import *
from .realtime_api.main import (
@ -1711,6 +1712,9 @@ if TYPE_CHECKING:
from .llms.cloudflare.decisions.transformation import (
CloudflareDecisionsConfig as CloudflareDecisionsConfig,
)
from .llms.cloudflare.extraction.transformation import (
CloudflareExtractionConfig as CloudflareExtractionConfig,
)
from .llms.strands_decider.decisions.transformation import (
StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig,
)

View file

@ -166,6 +166,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterDecisionsConfig",
"FireworksAIDecisionsConfig",
"CloudflareDecisionsConfig",
"CloudflareExtractionConfig",
"StrandsDeciderDecisionsConfig",
"DatabricksDecisionsConfig",
"HostedVLLMDecisionsConfig",
@ -735,6 +736,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
"FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),
"CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"),
"CloudflareExtractionConfig": (".llms.cloudflare.extraction.transformation", "CloudflareExtractionConfig"),
"StrandsDeciderDecisionsConfig": (
".llms.strands_decider.decisions.transformation",
"StrandsDeciderDecisionsConfig",

View file

@ -1,15 +1,19 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext
from litellm.rust_bridge.chat_completions import detectors as route_detectors
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
)
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
@ -44,6 +48,7 @@ def _python_acompletion() -> PythonAcompletion:
_PYTHON_COMPLETION: Final = _python_completion()
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION)
_PYTHON_ACOMPLETION: Final = _python_acompletion()
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
@ -61,11 +66,32 @@ def _public_request(
return native_call(legacy, args, kwargs)
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
Feature.MOCK_RESPONSE: route_detectors.chat_mock_response,
Feature.AGENTIC_LOOP_HOOK: route_detectors.chat_agentic_loop_hook,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.STREAMING: detectors.streaming,
Feature.DROP_PARAMS: detectors.drop_params,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.EXTRA_BODY: detectors.extra_body,
Feature.RETRIES_FALLBACKS: route_detectors.chat_retries_fallbacks,
Feature.PROMPT_MANAGEMENT: detectors.prompt_management,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.MODEL_PARAM_GATING: route_detectors.chat_model_param_gating,
Feature.MODEL_LIST: route_detectors.model_list,
Feature.MODEL_ALIAS_MAP: route_detectors.model_alias_map,
}
)
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.CHAT_COMPLETIONS,
provider=optional_str(request.resolved.get("custom_llm_provider")),
model=str(request.resolved["model"]),
features=detectors.detected(_FEATURES, request),
)
@ -73,7 +99,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("acompletion") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -88,9 +113,11 @@ def completion(
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
) -> ChatResult | Coroutine[object, object, ChatResult]:
python: Final = _PYTHON_COMPLETION
if _FLAG.requested(kwargs):
return acompletion(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_COMPLETION,
native=native_call_hook,
@ -101,7 +128,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
python: Final = _PYTHON_ACOMPLETION
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_ACOMPLETION,
native=native_call_hook,

View file

@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
@ -256,6 +258,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range(
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize
)
MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3)
MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float(
@ -1812,6 +1817,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token"
CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1."
CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS: Final = 600
CLI_SESSION_KEY_PREFIX: Final = "cli-session"
@ -2320,6 +2326,7 @@ HARNESS_STDERR_TAIL_LINES: Final = 40
HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024
HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024
HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS: Final = 30.0
HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
{
".git",

View file

@ -112,6 +112,7 @@ from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_toke
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage
from litellm.types.extraction import ExtractionResponse, ExtractionUsage
from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
@ -1171,6 +1172,8 @@ def get_usage_object(
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
if isinstance(completion_response, ExtractionResponse):
return _extraction_usage(completion_response.usage)
if isinstance(completion_response, HttpxBinaryResponseContent):
return completion_response.usage
if usage_obj is None:
@ -1227,6 +1230,14 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage:
)
def _extraction_usage(usage: ExtractionUsage) -> Usage:
return Usage(
prompt_tokens=usage.input_tokens,
completion_tokens=usage.output_tokens,
total_tokens=usage.input_tokens + usage.output_tokens,
)
def _is_known_usage_objects(usage_obj):
"""Returns True if the usage obj is a known Usage type"""
return (
@ -1234,6 +1245,7 @@ def _is_known_usage_objects(usage_obj):
or isinstance(usage_obj, ResponseAPIUsage)
or isinstance(usage_obj, DecisionsUsage)
or isinstance(usage_obj, OpenAIDecisionUsage)
or isinstance(usage_obj, ExtractionUsage)
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
)
@ -1263,6 +1275,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj
return "send_message"
elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return "adecisions"
elif isinstance(completion_response, ExtractionResponse):
return "aextraction"
return call_type
@ -1511,6 +1525,7 @@ def completion_cost(
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
pricing_provider: str | None = None,
### SERVICE TIER ###
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
@ -1531,6 +1546,8 @@ def completion_cost(
"""
try:
call_type = infer_call_type(call_type, completion_response) or "completion"
if pricing_provider is not None:
custom_llm_provider = pricing_provider # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups
if call_type == CallTypes.aresponses_websocket.value and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
@ -1556,6 +1573,7 @@ def completion_cost(
optional_params=optional_params,
custom_pricing=custom_pricing,
base_model=base_model,
pricing_provider=pricing_provider,
standard_built_in_tools_params=standard_built_in_tools_params,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
@ -1649,7 +1667,7 @@ def completion_cost(
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
)
for idx, (model, pricing_provider) in enumerate(potential_model_names):
for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names):
try:
if verbose_logger.isEnabledFor(logging.DEBUG):
verbose_logger.debug("selected model name for cost calculation: %s", model)
@ -1670,6 +1688,8 @@ def completion_cost(
)
if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
_usage = _decisions_usage(usage_obj).model_dump()
elif isinstance(usage_obj, ExtractionUsage):
_usage = _extraction_usage(usage_obj).model_dump()
elif usage_obj is None:
_usage = {}
elif isinstance(usage_obj, BaseModel):
@ -1716,7 +1736,9 @@ def completion_cost(
hidden_params = getattr(completion_response, "_hidden_params", None)
if hidden_params is not None:
custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None)
custom_llm_provider = pricing_provider or hidden_params.get(
"custom_llm_provider", custom_llm_provider or None
)
region_name = hidden_params.get("region_name", region_name)
else:
@ -2002,7 +2024,7 @@ def completion_cost(
prompt_tokens=prompt_tokens or 0,
completion_tokens=completion_tokens or 0,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
candidate_provider=candidate_pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
@ -2106,7 +2128,7 @@ def completion_cost(
_token_type_breakdown = get_token_type_cost_breakdown(
model=model,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
candidate_provider=candidate_pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
@ -2193,7 +2215,8 @@ def response_cost_calculator(
| Response
| SearchResponse
| DecisionsResponse
| OpenAIDecisionResponse,
| OpenAIDecisionResponse
| ExtractionResponse,
model: str,
custom_llm_provider: str | None,
call_type: Literal[
@ -2219,6 +2242,8 @@ def response_cost_calculator(
"adecisions",
"systemone",
"asystemone",
"extraction",
"aextraction",
],
optional_params: dict,
cache_hit: bool | None = None,
@ -2229,6 +2254,7 @@ def response_cost_calculator(
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
pricing_provider: str | None = None,
### SERVICE TIER ###
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
@ -2269,6 +2295,7 @@ def response_cost_calculator(
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
litellm_logging_obj=litellm_logging_obj,
pricing_provider=pricing_provider,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,

View file

@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import (
NATIVE_AEMBEDDING,
NATIVE_EMBEDDING,
)
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
native_call_hook,
optional_str,
signature,
)
from litellm.types.utils import EmbeddingResponse
__all__ = ("aembedding", "embedding")
@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi
PythonAembedding, main.aembedding
)
_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE)
def _public_request(
@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.EMBEDDINGS,
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aembedding") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -64,9 +72,11 @@ def embedding(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public embedding call shape
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
if _FLAG.requested(kwargs):
return aembedding(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_EMBEDDING,
binding=NATIVE_EMBEDDING,
native=native_call_hook,
@ -76,7 +86,7 @@ def embedding(
async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_AEMBEDDING,
binding=NATIVE_AEMBEDDING,
native=native_call_hook,

View file

@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_MAX_SSE_EVENT_SIZE_BYTES,
MCP_NPM_CACHE_DIR,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
@ -292,6 +293,7 @@ class _MCPResponseStream(httpx2.AsyncByteStream):
class _MCPHTTPClient(httpx2.AsyncClient):
cleanup_scope: anyio.CancelScope | None = None
cleanup_errors: tuple[Exception, ...] = ()
post_failure: asyncio.Future[Exception] | None = None
def _record_cleanup_error(self, error: Exception) -> None:
if self.cleanup_scope is not None and self.cleanup_scope.shield:
@ -319,7 +321,12 @@ class _MCPHTTPClient(httpx2.AsyncClient):
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
if request.method == "POST" and response.is_error and response.status_code != 404:
await response.aclose()
response.raise_for_status()
try:
response.raise_for_status()
except httpx2.HTTPStatusError as error:
if self.post_failure is not None and not self.post_failure.done():
self.post_failure.set_result(error)
raise
if stream:
response.stream = _MCPResponseStream(
cast(httpx2.AsyncByteStream, response.stream), self._record_cleanup_error
@ -641,6 +648,7 @@ class MCPClient:
transport_ctx: Final = streamable_http_client(
url=self.server_url,
http_client=http_client,
max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES,
)
return transport_ctx, http_client
@ -757,6 +765,8 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
if isinstance(http_client, _MCPHTTPClient):
http_client.post_failure = stream_error # rebind-ok: operation-scoped stream error
_notify_stream_error(on_stream_error, stream_error)
session_kwargs: Final = (
@ -823,6 +833,7 @@ class MCPClient:
if isinstance(http_client, _MCPHTTPClient):
http_client.cleanup_errors = ()
http_client.cleanup_scope = None
http_client.post_failure = None # rebind-ok: clear operation-scoped stream error
await anyio.lowlevel.checkpoint_if_cancelled()
if cleanup_scope.cancel_called:
raise (

View file

@ -0,0 +1,3 @@
from litellm.extraction.main import aextraction, extraction
__all__ = ["aextraction", "extraction"]

351
litellm/extraction/main.py Normal file
View file

@ -0,0 +1,351 @@
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.decisions.call import log_caller_request
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.extraction import (
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
ExtractionAnswer,
ExtractionBase64Image,
ExtractionField,
ExtractionOptions,
ExtractionRequestBody,
ExtractionResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders, is_litellm_owned_kwarg
from litellm.utils import (
ProviderConfigManager,
client, # pyright: ignore[reportUnknownVariableType] # client is an untyped decorator
get_model_cost_key,
)
ExtractionFields: TypeAlias = Mapping[str, ExtractionField | Mapping[str, object]]
ExtractionImages: TypeAlias = Sequence[str | ExtractionBase64Image | Mapping[str, object]]
ExtractionOptionsParam: TypeAlias = ExtractionOptions | Mapping[str, object]
_REQUEST_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody)
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_HANDLER: Final = BaseLLMHTTPHandler()
_LOCAL_KWARGS: Final = frozenset(
(
*GenericLiteLLMParams.model_fields,
"extra_headers",
"base_url",
"default_headers",
"deployment_id",
"additional_drop_params",
)
)
@dataclass(frozen=True, slots=True, repr=False)
class _ExtractionCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseExtractionConfig
body: Mapping[str, object]
api_base: str
api_key: str | None
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object]
client: HTTPHandler | AsyncHTTPHandler | None
def supported_extraction_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_extraction_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseExtractionConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_extraction_config(model, provider)
)
if provider_config is None:
raise litellm.BadRequestError(
message=(
f"Unknown Extraction provider '{custom_llm_provider}'. "
f"Supported providers: {', '.join(supported_extraction_providers())}"
),
model=model,
llm_provider=custom_llm_provider,
)
return provider_config
def _request_fields(
*,
state: str | None,
fields: ExtractionFields | None,
images: ExtractionImages | None,
reference_date: str | None,
options: ExtractionOptionsParam | None,
) -> Mapping[str, object]:
sent: Final[Mapping[str, object | None]] = MappingProxyType(
{"state": state, "fields": fields, "images": images, "reference_date": reference_date, "options": options}
)
return MappingProxyType({name: value for name, value in sent.items() if value is not None})
def _is_unsupported(name: str, value: object, policy: DropPolicy) -> bool:
return (
value is not None
and value is not False
and name not in _LOCAL_KWARGS
and not is_litellm_owned_kwarg(name)
and not policy.drops(name)
)
def _unsupported_params(kwargs: Mapping[str, object]) -> tuple[str, ...]:
policy: Final = drop_policy(kwargs)
return tuple(sorted(name for name, value in kwargs.items() if _is_unsupported(name, value, policy)))
def _prepare_call(
*,
model: str,
request_fields: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> _ExtractionCall:
log_caller_request(kwargs, request_fields)
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(model, provider)
unsupported: Final = _unsupported_params(kwargs)
if unsupported:
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: {list(unsupported)}, for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
try:
request: Final = _REQUEST_ADAPTER.validate_python(request_fields)
except ValidationError as error:
raise litellm.BadRequestError(
message=f"Invalid Extraction request: {error}", model=model, llm_provider=provider
) from error
canonical_model: Final = provider_config.canonical_model(upstream_model)
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=provider_config.missing_api_base_message(provider), model=model, llm_provider=provider
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None:
raise litellm.BadRequestError(
message=f"Missing API key for Extraction provider '{provider}'", model=model, llm_provider=provider
)
logging_obj: Final = kwargs.get("litellm_logging_obj")
http_client: Final = kwargs.get("client")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging takes untyped dicts
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return _ExtractionCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
body=provider_config.transform_extraction_request(canonical_model, request),
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=MappingProxyType(dict(extra_headers or {})),
timeout=timeout,
litellm_params=kwargs,
client=http_client if isinstance(http_client, HTTPHandler | AsyncHTTPHandler) else None,
)
def _state_rewritten_by_guardrail(litellm_params: Mapping[str, object]) -> bool:
metadata: Final = litellm_params.get(get_metadata_variable_name_from_kwargs(litellm_params))
return (
isinstance(metadata, Mapping)
and _METADATA_ADAPTER.validate_python(metadata).get(EXTRACTION_STATE_REWRITTEN_METADATA_KEY) is True
)
def _without_spans(answers: Mapping[str, ExtractionAnswer]) -> Mapping[str, ExtractionAnswer]:
return {name: answer.model_copy(update={"span": None, "spans": None}) for name, answer in answers.items()}
def _formatted_response(response: ExtractionResponse, call: _ExtractionCall) -> ExtractionResponse:
spans_index_hidden_text: Final = response.state_text is None and _state_rewritten_by_guardrail(call.litellm_params)
formatted: Final = response.model_copy(
update={
"model": call.requested_model,
**({"answers": _without_spans(response.answers)} if spans_index_hidden_text else {}),
}
)
formatted.set_hidden_params(
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": (
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
or f"{call.custom_llm_provider}/{call.model}"
),
}
)
return formatted
def _upstream_exception(error: Exception, call: _ExtractionCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
return litellm.APIConnectionError(
message=f"{call.custom_llm_provider.capitalize()}Exception - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
async def _asend(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = await _HANDLER.aextraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, AsyncHTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
def _send(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = _HANDLER.extraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, HTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
@client
async def aextraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = await asyncio.to_thread(
partial(
_prepare_call,
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
)
return await _asend(call)
@client
def extraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = _prepare_call(
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
return _send(call)
__all__ = ["aextraction", "extraction"]

View file

@ -12,17 +12,27 @@ from __future__ import annotations
import asyncio
import os
from collections import deque
from collections.abc import AsyncIterator, Sequence
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from litellm._logging import verbose_logger
from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES
from litellm.constants import (
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS,
HARNESS_STDERR_TAIL_LINES,
HARNESS_STREAM_READ_CHUNK_BYTES,
)
from litellm.harness.context import SessionContext
from litellm.harness.errors import HarnessInstallFailed, SandboxError
from litellm.harness.errors import HarnessError, HarnessInstallFailed, SandboxError
from litellm.harness.handlers.base import BaseHarnessHandler
from litellm.harness.sandbox.base import Process, Sandbox
from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox
from litellm.harness.types import Event
from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup
from litellm.llms.base_llm.harness.transformation import (
BaseCLIHarnessConfig,
HarnessSessionCopy,
HarnessSessionSetup,
HarnessTurnError,
)
from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files
# Link <private_dir>/<dir> to a LiteLLM-owned cache dir so a later session can resume.
@ -81,6 +91,52 @@ async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> N
)
@dataclass(frozen=True, slots=True)
class ResumeFailed:
reason: str
async def bounded_resume_run(
sandbox: Sandbox, cmd: Sequence[str], env: Mapping[str, str]
) -> CompletedRun | ResumeFailed:
try:
run: Final = await sandbox.run(cmd, env=env, timeout=HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS)
except SandboxError as exc:
return ResumeFailed(str(exc))
return run if run.exit_code == 0 else ResumeFailed(run.stderr.strip() or f"exit code {run.exit_code}")
async def import_session_copy(sandbox: Sandbox, copy: HarnessSessionCopy, env: Mapping[str, str]) -> str | ResumeFailed:
try:
await sandbox.write(copy.path, copy.data)
except SandboxError as exc:
return ResumeFailed(str(exc))
imported: Final = await bounded_resume_run(sandbox, copy.import_command, env)
return imported if isinstance(imported, ResumeFailed) else copy.session_id
async def resumable_session_id(
config: BaseCLIHarnessConfig[object, object],
sandbox: Sandbox,
setup_env: Mapping[str, str],
private_dir: str,
native_session_id: str,
) -> str | ResumeFailed:
export_cmd: Final = config.transform_resume_export_command(native_session_id)
if export_cmd is None:
return native_session_id
exported: Final = await bounded_resume_run(sandbox, export_cmd, setup_env)
if isinstance(exported, ResumeFailed):
return exported
try:
copy: Final = config.transform_resume_copy(exported.stdout, private_dir, sandbox.workdir)
except HarnessError as exc:
return ResumeFailed(str(exc))
if copy is None:
return native_session_id
return await import_session_copy(sandbox, copy, setup_env)
async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None:
for skill in skills:
name = os.path.basename(os.path.realpath(os.fspath(skill)))
@ -89,13 +145,14 @@ async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str)
class CLIHarnessHandler(BaseHarnessHandler):
config: BaseCLIHarnessConfig
config: BaseCLIHarnessConfig[object, object]
def __init__(self, config: BaseCLIHarnessConfig) -> None:
super().__init__(config)
self._private_dir: str | None = None
self._setup: HarnessSessionSetup | None = None
self._native_id: str | None = None
self._pending_resume: str | None = None
self._proc: Process | None = None
async def start(self, ctx: SessionContext) -> None:
@ -119,6 +176,9 @@ class CLIHarnessHandler(BaseHarnessHandler):
async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
if self._setup is None or self._private_dir is None:
raise RuntimeError("CLIHarnessHandler.turn() called before start()")
if self._pending_resume is not None:
self._native_id = await self._prepared_session_id(ctx, self._setup, self._private_dir, self._pending_resume)
self._pending_resume = None
request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id)
argv: Final = list(request.argv) # mutable-ok: Sandbox.exec takes list[str]
proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd)
@ -158,4 +218,16 @@ class CLIHarnessHandler(BaseHarnessHandler):
return self._native_id
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
self._native_id = native_session_id
if native_session_id != self._native_id:
self._native_id = native_session_id
self._pending_resume = native_session_id
async def _prepared_session_id(
self, ctx: SessionContext, setup: HarnessSessionSetup, private_dir: str, saved_id: str
) -> str:
target: Final = await resumable_session_id(self.config, ctx.sandbox, setup.env, private_dir, saved_id)
match target:
case ResumeFailed(reason=reason):
raise HarnessTurnError(f"could not prepare session {saved_id} for resume: {reason}")
case session_id:
return session_id

View file

@ -14,7 +14,7 @@ from typing import Final
from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox.base import CompletedRun
from litellm.harness.sandbox.base import CompletedRun, Process
from litellm.harness.sandbox.snapshot import snapshot_local
_SECRET_PREFIXES: Final = (
@ -103,7 +103,7 @@ class SubprocessHandle:
await self._proc.wait()
async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
async def _read_all(handle: Process) -> tuple[bytes, bytes, int]:
if handle.stdin is not None:
handle.stdin.close()
stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
@ -111,13 +111,16 @@ async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
return stdout, stderr, exit_code
async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
"""Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout."""
async def collect_output(handle: Process, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
"""Close stdin, read stdout/stderr to EOF; kill on cancel, and kill and raise SandboxError on timeout."""
try:
stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout)
except asyncio.TimeoutError:
await handle.kill()
raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}")
except asyncio.CancelledError:
await handle.kill()
raise
return CompletedRun(
stdout=stdout.decode("utf-8", errors="replace"),
stderr=stderr.decode("utf-8", errors="replace"),
@ -180,14 +183,15 @@ class LocalSandbox:
self._check_open()
if not cmd:
raise SandboxError("exec() needs a non-empty command")
work_dir: Final = self._resolve_cwd(cwd)
try:
proc = await asyncio.create_subprocess_exec(
*cmd,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self._resolve_cwd(cwd),
env=self.child_env(env),
cwd=work_dir,
env=self.child_env(MappingProxyType({**(env or {}), "PWD": work_dir})),
start_new_session=True,
)
except (FileNotFoundError, PermissionError) as exc:

View file

@ -77,6 +77,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"otpm",
"use_xai_oauth",
"github_copilot_auth_type",
"github_copilot_host",
"github_copilot_user_session",
"fireworks_forward_user_id",
"responses_api_version",

View file

@ -4,6 +4,7 @@ Helper functions for health check calls.
import base64
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from litellm.llms.base_llm.ocr.transformation import DocumentType
@ -25,6 +26,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
)
ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16
EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
{"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})}
)
def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None:
@ -215,6 +219,7 @@ class HealthCheckHelpers:
"anthropic_messages",
"ocr",
"decisions",
"extraction",
],
Callable,
]:
@ -237,7 +242,11 @@ class HealthCheckHelpers:
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_for_health_check,
)
from litellm.litellm_core_utils.health_check_utils import DECISIONS_CALL_PARAMS, _filter_model_params
from litellm.litellm_core_utils.health_check_utils import (
DECISIONS_CALL_PARAMS,
EXTRACTION_CALL_PARAMS,
_filter_model_params,
)
from litellm.realtime_api.main import realtime_health_check
return {
@ -322,4 +331,13 @@ class HealthCheckHelpers:
}
)
),
"extraction": lambda: litellm.aextraction(
**EXTRACTION_CALL_PARAMS.validate_python(
{
"state": prompt or "health check",
"fields": EXTRACTION_HEALTH_CHECK_FIELDS,
**_filter_model_params(model_params=model_params),
}
)
),
}

View file

@ -8,8 +8,10 @@ from typing import Final
from pydantic import TypeAdapter
from litellm.types.decisions import DecisionsCallParams
from litellm.types.extraction import ExtractionCallParams
DECISIONS_CALL_PARAMS: Final[TypeAdapter[DecisionsCallParams]] = TypeAdapter(DecisionsCallParams)
EXTRACTION_CALL_PARAMS: Final[TypeAdapter[ExtractionCallParams]] = TypeAdapter(ExtractionCallParams)
OPTIONAL_STR: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)

View file

@ -117,7 +117,8 @@ from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request
from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_custom_logger,
redact_message_input_output_from_logging,
redact_streaming_responses_for_custom_logger,
redacted_model_call_details_for_custom_logger,
redacted_response_obj_for_custom_logger,
should_redact_message_logging,
)
from litellm.litellm_core_utils.served_output_texts import (
@ -131,6 +132,7 @@ from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.containers.main import ContainerObject
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
from litellm.types.extraction import ExtractionResponse
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
from litellm.types.interactions import (
InteractionsAPIResponse,
@ -337,6 +339,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
_METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object])
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
("user_api_key_auth", "user_api_key_budget_reservation")
@ -794,6 +797,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.defer_async_logging: bool = False
self.enqueue_deferred_logging: Callable[[], None] | None = None
self._async_success_scheduled: bool = False
self._async_success_enqueued: bool = False
self._sync_success_scheduled: bool = False
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
@ -995,6 +999,21 @@ class Logging(LiteLLMLoggingBaseClass):
return model_id
return None
def get_router_pricing_provider(self) -> str | None:
"""The deployment's ``model_info.pricing_provider``, the provider its cost-map lookups run under."""
litellm_params: Final = cast( # cast-ok: callback metadata is caller-provided
Mapping[str, object] | None, getattr(self, "litellm_params", None)
)
if not litellm_params:
return None
candidates: Final = (
_METADATA_MAPPING.validate_python(
_METADATA_MAPPING.validate_python(litellm_params.get(key) or {}).get("model_info") or {}
).get("pricing_provider")
for key in ("litellm_metadata", "metadata")
)
return next((provider for provider in candidates if isinstance(provider, str) and provider), None)
def _routed_deployment_model(self) -> str | None:
litellm_params: Final = self._model_call_details_view.get("litellm_params")
metadata: Final = (
@ -2020,6 +2039,7 @@ class Logging(LiteLLMLoggingBaseClass):
"prompt": prompt,
"standard_built_in_tools_params": self.standard_built_in_tools_params,
"router_model_id": router_model_id,
"pricing_provider": self.get_router_pricing_provider(),
"litellm_logging_obj": self,
"data_residency": (
self.litellm_params.get("data_residency")
@ -2164,7 +2184,7 @@ class Logging(LiteLLMLoggingBaseClass):
pricing: Final = pricing_entry_for_cost_calc(
model=model,
completion_response=result,
custom_llm_provider=custom_llm_provider,
custom_llm_provider=self.get_router_pricing_provider() or custom_llm_provider,
custom_pricing=self.custom_pricing_for(result),
base_model=get_base_model_from_metadata(model_call_details=self.model_call_details),
router_model_id=router_model_id or self.get_router_model_id(),
@ -2190,20 +2210,23 @@ class Logging(LiteLLMLoggingBaseClass):
usage=usage,
)
def custom_pricing_for(self, result: object) -> bool:
litellm_params: Final = getattr(self, "litellm_params", None)
def spilled_over_ptu_for(self, result: object) -> bool:
result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({})
additional_headers: Final = (
result_hidden_params.get("additional_headers")
if isinstance(result_hidden_params, dict)
else getattr(result_hidden_params, "additional_headers", None)
)
spilled_over: Final = is_spilled_over_ptu_request(
model_info=_deployment_model_info(litellm_params),
return is_spilled_over_ptu_request(
model_info=_deployment_model_info(getattr(self, "litellm_params", None)),
response_headers=self.model_call_details.get("response_headers"),
additional_headers=additional_headers,
)
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
def custom_pricing_for(self, result: object) -> bool:
if self.spilled_over_ptu_for(result):
return False
return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None))
def _prompt_for_cost_calculation(self) -> str:
"""
@ -2447,7 +2470,7 @@ class Logging(LiteLLMLoggingBaseClass):
await self.async_failure_handler(exception, traceback_exception)
finally:
if self._should_run_sync_failure_callbacks_for_async_calls():
executor.submit(self.failure_handler, exception, traceback_exception)
executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True)
def should_run_logging(
self,
@ -2462,6 +2485,16 @@ class Logging(LiteLLMLoggingBaseClass):
except Exception:
return True
@property
def async_success_log_enqueued(self) -> bool:
"""True once the success log task was actually created, not merely claimed, so a later failure raised
around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second
time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False."""
return self._async_success_enqueued
def mark_async_success_log_enqueued(self) -> None:
self._async_success_enqueued = True
def claim_async_success_log(self) -> bool:
"""One async success log per request. The innermost @client wrapper always exits first,
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
@ -2890,6 +2923,7 @@ class Logging(LiteLLMLoggingBaseClass):
or isinstance(logging_result, SearchResponse) # Search API
or isinstance(logging_result, DecisionsResponse)
or isinstance(logging_result, OpenAIDecisionResponse)
or isinstance(logging_result, ExtractionResponse)
or (
isinstance(logging_result, InteractionsAPIResponse)
and logging_result.usage is not None
@ -3489,8 +3523,10 @@ class Logging(LiteLLMLoggingBaseClass):
): # custom logger class
if self.stream and complete_streaming_response is None:
callback.log_stream_event(
kwargs=call_details,
response_obj=result,
kwargs=redacted_model_call_details_for_custom_logger(
model_call_details=call_details, custom_logger=callback
),
response_obj=redacted_response_obj_for_custom_logger(result, callback),
start_time=start_time,
end_time=end_time,
)
@ -3500,8 +3536,13 @@ class Logging(LiteLLMLoggingBaseClass):
result = call_details_view["complete_response"]
callback.log_success_event(
kwargs=callback.redact_standard_logging_payload_from_model_call_details(call_details),
response_obj=result,
kwargs=redacted_model_call_details_for_custom_logger(
model_call_details=callback.redact_standard_logging_payload_from_model_call_details(
call_details
),
custom_logger=callback,
),
response_obj=redacted_response_obj_for_custom_logger(result, callback),
start_time=start_time,
end_time=end_time,
)
@ -3844,7 +3885,7 @@ class Logging(LiteLLMLoggingBaseClass):
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
model_call_details=model_call_details
)
model_call_details = redact_streaming_responses_for_custom_logger(
model_call_details = redacted_model_call_details_for_custom_logger(
model_call_details=model_call_details, custom_logger=callback
)
##################################
@ -3855,7 +3896,7 @@ class Logging(LiteLLMLoggingBaseClass):
kwargs=model_call_details,
response_obj=model_call_details["async_complete_streaming_response"]
if self.stream
else result,
else redacted_response_obj_for_custom_logger(result, callback),
start_time=start_time,
end_time=end_time,
guardrail_snapshot=self.shadow_eval_request_snapshot,
@ -3871,14 +3912,14 @@ class Logging(LiteLLMLoggingBaseClass):
else:
await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function
kwargs=model_call_details,
response_obj=result,
response_obj=redacted_response_obj_for_custom_logger(result, callback),
start_time=start_time,
end_time=end_time,
)
else:
await callback.async_log_success_event(
kwargs=model_call_details,
response_obj=result,
response_obj=redacted_response_obj_for_custom_logger(result, callback),
start_time=start_time,
end_time=end_time,
)
@ -4054,15 +4095,21 @@ class Logging(LiteLLMLoggingBaseClass):
traceback_exception: str,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
*,
for_async_request: bool = False,
) -> None:
"""Restores trace_id/session_id contextvars once this attempt's own failure
logging (including any nested calls its callbacks trigger) is fully done."""
logging (including any nested calls its callbacks trigger) is fully done.
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
try:
return self._failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
for_async_request=for_async_request,
)
finally:
self._restore_correlation_context()
@ -4073,6 +4120,7 @@ class Logging(LiteLLMLoggingBaseClass):
traceback_exception: str,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
for_async_request: bool = False,
) -> None:
verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback)
if not self.should_run_logging(event_type="sync_failure"): # prevent double logging
@ -4165,6 +4213,7 @@ class Logging(LiteLLMLoggingBaseClass):
isinstance(callback, CustomLogger)
and is_sync_request
and self.call_type != CallTypes.pass_through.value
and not (for_async_request and self._async_handler_logs_to(callback))
): # custom logger class
callback.log_failure_event(
start_time=start_time,
@ -4527,6 +4576,12 @@ class Logging(LiteLLMLoggingBaseClass):
cb_name: Final = self._get_callback_name(cb)
return any(prefix in cb_name for prefix in INTERNAL_PREFIXES)
def _async_handler_logs_to(self, logger: CustomLogger) -> bool:
"""True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``."""
if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped
return False
return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor
def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list:
"""
Removes internal custom logger callbacks from the list.
@ -7195,13 +7250,16 @@ def get_standard_logging_object_payload(
saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0
## Get model cost information ##
base_model = get_base_model_from_metadata(model_call_details=kwargs)
# The router overrides completion_response.model to the model-group alias before
# this payload is built, so cost-map lookup via that alias always misses.
# Fall back to the actual deployment model set by the router in metadata.
if base_model is None:
base_model = metadata.get("deployment")
custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj)
configured_base_model: Final = get_base_model_from_metadata(model_call_details=kwargs)
spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj)
# The router overrides completion_response.model to the model-group alias before this payload is
# built, so cost-map lookup via that alias misses and the deployment model stands in for base_model.
# A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates,
# which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there.
base_model: Final = (
metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model
)
custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj)
raw_response_cost: Final = kwargs.get("response_cost")
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))

View file

@ -4,6 +4,7 @@ from functools import reduce
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.constants import LITELLM_DETAILED_TIMING
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
@ -101,6 +102,32 @@ def response_timing_metrics(
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object])
def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]:
try:
return _HIDDEN_PARAMS_DICT.validate_python(hidden_params)
except ValidationError:
return {}
def _existing_provider_response_model(result: object) -> str | None:
"""The pricing model a response already carries, like the decisions path sets, wins over the response model."""
hidden_params: Final[object] = getattr(result, "_hidden_params", None)
carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr(
hidden_params, "provider_response_model", None
)
return carried if isinstance(carried, str) and carried else None
def _provider_response_model(result: object) -> str | None:
"""The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the
client asked for and cost-map lookups for the standard logging payload would otherwise miss."""
model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None
return model if isinstance(model, str) and model else None
class ResponseMetadata:
"""
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
@ -147,6 +174,8 @@ class ResponseMetadata:
preserve_litellm_internal_headers=True,
),
"litellm_model_name": model,
"provider_response_model": _existing_provider_response_model(self.result)
or _provider_response_model(self.result),
}
self._update_hidden_params(new_params)

View file

@ -10,10 +10,12 @@
import asyncio
import copy
import inspect
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from dataclasses import replace
from typing import TYPE_CHECKING, Any, Final
from pydantic import TypeAdapter
import litellm
from litellm.constants import REDACTED_BY_LITELLM
from litellm.integrations.custom_logger import CustomLogger
@ -40,36 +42,79 @@ else:
LiteLLMLoggingObject = Any
def _opted_out_of_message_logging(custom_logger: CustomLogger) -> bool:
return hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True
def redact_message_input_output_from_custom_logger(
litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger
):
if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True:
) -> object:
if _opted_out_of_message_logging(custom_logger):
return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False)
return result
def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
def redacted_model_call_details_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
"""
Returns a copy of model_call_details whose streaming response entries are redacted deepcopies
when the custom logger has opted out of message logging. The shared model_call_details is left
untouched so other callbacks still receive the unredacted response.
Returns a copy of model_call_details with the request messages, instructions and streaming responses
redacted when the custom logger has opted out of message logging. The shared model_call_details is left
untouched so other callbacks and the spend log still receive the unredacted values.
"""
if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True):
if not _opted_out_of_message_logging(custom_logger):
return model_call_details
redacted_entries: Final = {
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
if model_call_details.get(streaming_key) is not None
return {
**{key: value for key, value in model_call_details.items() if key != SERVED_OUTPUT_TEXTS_KEY},
**_redacted_entries_for_opted_out_logger(model_call_details),
}
if not redacted_entries:
return model_call_details
return {**model_call_details, **redacted_entries}
def _redacted_streaming_response_copy(streaming_response):
redacted_response: Final = copy.deepcopy(streaming_response)
_redact_streaming_response(redacted_response)
return redacted_response
def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, object]) -> Mapping[str, object]:
optional_params: Final = _as_mapping(model_call_details.get("optional_params"))
litellm_params: Final = _as_mapping(model_call_details.get("litellm_params"))
standard_logging_object: Final = _as_mapping(model_call_details.get("standard_logging_object"))
return {
"messages": [{"role": "user", "content": REDACTED_BY_LITELLM}],
"prompt": "",
"input": "",
**({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}),
**({"optional_params": _with_redacted_instructions(optional_params)} if optional_params is not None else {}),
**({"litellm_params": redacted_litellm_params(litellm_params)} if litellm_params is not None else {}),
**(
{"standard_logging_object": _redact_standard_logging_object(standard_logging_object)}
if standard_logging_object is not None
else {}
),
**{
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
for streaming_key in (
"complete_streaming_response",
"async_complete_streaming_response",
"complete_response",
)
if model_call_details.get(streaming_key) is not None
},
}
def redacted_response_obj_for_custom_logger(response_obj: object, custom_logger: CustomLogger) -> object:
if not _opted_out_of_message_logging(custom_logger) or not _is_redactable_response(response_obj):
return response_obj
return perform_redaction({}, response_obj)
def _is_redactable_response(response_obj: object) -> bool:
return isinstance(
response_obj, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse)
) or (isinstance(response_obj, dict) and ("choices" in response_obj or "output" in response_obj))
def _redacted_streaming_response_copy(streaming_response: object) -> object:
redacted_response: Final = (
streaming_response
if isinstance(streaming_response, litellm.ResponsesAPIResponse)
else copy.deepcopy(streaming_response)
)
return _redact_streaming_response(redacted_response)
def _redact_streaming_response(streaming_response):
@ -77,10 +122,13 @@ def _redact_streaming_response(streaming_response):
for choice in streaming_response.choices:
_redact_choice_content(choice)
redact_vertex_ai_metadata_from_logged_object(streaming_response)
elif isinstance(streaming_response, litellm.ResponsesAPIResponse):
return perform_redaction({}, streaming_response)
elif hasattr(streaming_response, "output"):
_redact_responses_api_output(streaming_response.output)
if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None:
streaming_response.reasoning = None
return streaming_response
def _redact_tool_calls(tool_calls) -> None:
@ -177,31 +225,75 @@ def _redact_responses_api_output(output_items):
output_item.input = REDACTED_BY_LITELLM
def _redact_responses_api_output_dict(output_items, redacted_str: str):
"""Helper to redact ResponsesAPIResponse output items in dict form."""
for output_item in output_items:
if not isinstance(output_item, dict):
continue
_RESPONSE_ITEMS: Final = TypeAdapter(tuple[object, ...])
if output_item.get("text") is not None:
output_item["text"] = redacted_str
if isinstance(output_item.get("content"), list):
for content_item in output_item["content"]:
if isinstance(content_item, dict) and content_item.get("text") is not None:
content_item["text"] = redacted_str
if isinstance(content_item, dict) and content_item.get("refusal") is not None:
content_item["refusal"] = redacted_str
def _redacted_responses_api_output_part(part: object, fields: tuple[str, ...], redacted_str: str) -> object:
if (part_mapping := _as_mapping(part)) is None:
return part
return {key: redacted_str if key in fields and value is not None else value for key, value in part_mapping.items()}
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
for summary_item in output_item["summary"]:
if isinstance(summary_item, dict) and summary_item.get("text") is not None:
summary_item["text"] = redacted_str
if output_item.get("type") == "function_call" and "arguments" in output_item:
output_item["arguments"] = redacted_str
if output_item.get("type") == "custom_tool_call" and "input" in output_item:
output_item["input"] = redacted_str
def _redacted_responses_api_output_item(item: object, redacted_str: str) -> object:
if (item_mapping := _as_mapping(item)) is None:
return item
content: Final = item_mapping.get("content")
summary: Final = item_mapping.get("summary")
return {
**item_mapping,
**({"text": redacted_str} if item_mapping.get("text") is not None else {}),
**(
{
"content": [
_redacted_responses_api_output_part(part, ("text", "refusal"), redacted_str)
for part in _RESPONSE_ITEMS.validate_python(content)
]
}
if isinstance(content, list)
else {}
),
**(
{
"summary": [
_redacted_responses_api_output_part(part, ("text",), redacted_str)
for part in _RESPONSE_ITEMS.validate_python(summary)
]
}
if item_mapping.get("type") == "reasoning" and isinstance(summary, list)
else {}
),
**(
{"arguments": redacted_str}
if item_mapping.get("type") == "function_call" and "arguments" in item_mapping
else {}
),
**(
{"input": redacted_str}
if item_mapping.get("type") == "custom_tool_call" and "input" in item_mapping
else {}
),
}
def _redact_responses_api_output_dict(output_items: Sequence[object], redacted_str: str) -> Sequence[object]:
return [_redacted_responses_api_output_item(item, redacted_str) for item in output_items]
def _with_redacted_instructions(params: Mapping[str, object]) -> Mapping[str, object]:
return params if params.get("instructions") is None else {**params, "instructions": REDACTED_BY_LITELLM}
def _redacted_responses_api_response(
response: Mapping[str, object], *, redact_instructions: bool = True
) -> Mapping[str, object]:
output: Final = response.get("output")
redacted: Final = {
**response,
"output": _redact_responses_api_output_dict(_RESPONSE_ITEMS.validate_python(output), REDACTED_BY_LITELLM)
if isinstance(output, list)
else output,
}
return _with_redacted_instructions(redacted) if redact_instructions else redacted
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:
@ -215,20 +307,25 @@ def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[
return _redact_standard_logging_object(payload)
def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]:
def _redact_standard_logging_object(
payload: Mapping[str, object], *, redact_instructions: bool = True
) -> dict[str, object]:
standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload))
redacted_str: Final = REDACTED_BY_LITELLM
if standard_logging_object.get("messages") is not None:
standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}]
model_parameters: Final = _as_mapping(standard_logging_object.get("model_parameters"))
if model_parameters is not None and redact_instructions:
standard_logging_object["model_parameters"] = _with_redacted_instructions(model_parameters)
response: Final = standard_logging_object.get("response")
if response is not None:
if isinstance(response, dict) and "output" in response:
# ResponsesAPIResponse format - redact content in output items
if isinstance(response.get("output"), list):
_redact_responses_api_output_dict(response["output"], redacted_str)
redact_vertex_ai_metadata_from_logged_object(response)
standard_logging_object["response"] = _redacted_responses_api_response(
response, redact_instructions=redact_instructions
)
redact_vertex_ai_metadata_from_logged_object(standard_logging_object["response"])
elif isinstance(response, dict) and "choices" in response:
# ModelResponse dict format - redact content in choices
if isinstance(response.get("choices"), list):
@ -320,11 +417,11 @@ def _redact_hidden_provider_specific_fields(params: Mapping[str, object]) -> Non
def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body: bool = True) -> dict[str, object]:
"""
Copy of litellm_params without classifier audit fields or baseline route request parameters.
redact_request_body also redacts the nested proxy_server_request body.
redact_request_body also redacts the request instructions and the nested proxy_server_request body.
"""
request: Final = _as_mapping(params.get("proxy_server_request"))
return {
**params,
**(_with_redacted_instructions(params) if redact_request_body else params),
**{
key: _redacted_baseline_metadata(value)
for key, value in params.items()
@ -344,8 +441,9 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
redact_streaming_responses=False skips the in-place redaction of the shared streaming
response entries; per-callback redaction hands each opted-out callback its own redacted
copy via redact_streaming_responses_for_custom_logger instead. It also leaves the shared
proxy request body and hidden provider fields as they are, since the spend log reads them.
copy via redacted_model_call_details_for_custom_logger instead. It also leaves the shared
request instructions, proxy request body and hidden provider fields as they are, since the
spend log and the other callbacks read them.
"""
# Redact model_call_details
params: Final = _as_mapping(model_call_details.get("litellm_params"))
@ -358,16 +456,31 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
model_call_details["prompt"] = ""
model_call_details["input"] = ""
optional_params: Final = _as_mapping(model_call_details.get("optional_params"))
if redact_streaming_responses:
model_call_details.update(
{
**({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}),
**({"optional_params": _with_redacted_instructions(optional_params)} if optional_params else {}),
}
)
model_call_details.pop(SERVED_OUTPUT_TEXTS_KEY, None)
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if isinstance(standard_logging_object, Mapping):
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
model_call_details["standard_logging_object"] = _redact_standard_logging_object(
standard_logging_object, redact_instructions=redact_streaming_responses
)
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
# Redact streaming response
if redact_streaming_responses and model_call_details.get("stream", False) is True:
for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"):
_redact_streaming_response(model_call_details.get(_streaming_key))
model_call_details.update(
{
streaming_key: _redact_streaming_response(model_call_details[streaming_key])
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
if model_call_details.get(streaming_key) is not None
}
)
# Redact result
if result is not None:
@ -381,10 +494,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
# For async objects, return a simple redacted response without deepcopy
return {"text": REDACTED_BY_LITELLM}
if not (
isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse))
or (isinstance(result, dict) and ("choices" in result or "output" in result))
):
if not _is_redactable_response(result):
return {"text": REDACTED_BY_LITELLM}
_result: Final = copy.deepcopy(result)
@ -399,14 +509,18 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
_redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM)
redact_vertex_ai_metadata_from_logged_object(_result)
elif isinstance(_result, dict) and "output" in _result:
if isinstance(_result.get("output"), list):
_redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM)
return _redacted_responses_api_response(_result, redact_instructions=redact_streaming_responses)
elif isinstance(_result, litellm.ResponsesAPIResponse):
if hasattr(_result, "output"):
_redact_responses_api_output(_result.output)
# Redact reasoning field in ResponsesAPIResponse
if hasattr(_result, "reasoning") and _result.reasoning is not None:
_result.reasoning = None
return _result.model_copy(
update={"instructions": REDACTED_BY_LITELLM}
if redact_streaming_responses and _result.instructions is not None
else {}
)
elif isinstance(_result, litellm.EmbeddingResponse):
if hasattr(_result, "data") and _result.data is not None:
_result.data = []

View file

@ -13,7 +13,7 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping, MutableSequence, Sequence
from collections.abc import Iterable, Mapping, MutableSequence, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
@ -27,8 +27,10 @@ from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.pass_through.adapters.transformation import (
OPENAI_MAX_TOOL_NAME_LENGTH,
LiteLLMAnthropicMessagesAdapter,
is_provider_native_tool_dict,
truncate_tool_name,
)
from litellm.llms.anthropic.pass_through.convertible_blocks import (
keep_non_string_tool_result_texts,
@ -202,6 +204,60 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool:
)
def _tool_use_name(block: object) -> str | None:
fields: Final = as_str_mapping(block)
if fields is None or fields.get("type") != "tool_use":
return None
return anthropic_tool_name(fields)
def _tool_use_names(message: object) -> tuple[str, ...]:
return tuple(name for name in map(_tool_use_name, list_field(message, "content")) if name is not None)
@dataclass(frozen=True, slots=True)
class _OriginalToolNames:
"""Caller tool names over the guardrail's limit, keyed by the shortened name the guardrail sees."""
declared: Mapping[str, str]
used: Mapping[str, str]
def _restorable_names(names: Iterable[str]) -> Mapping[str, str]:
listed: Final = tuple(names)
taken: Final = frozenset(listed)
shortened: Final = tuple(
(truncate_tool_name(name), name) for name in listed if len(name) > OPENAI_MAX_TOOL_NAME_LENGTH
)
return MappingProxyType({alias: name for alias, name in shortened if alias not in taken})
def _original_tool_names(request: Mapping[str, object]) -> _OriginalToolNames:
return _OriginalToolNames(
declared=_restorable_names(chain.from_iterable(map(anthropic_tool_names, list_field(request, "tools")))),
used=_restorable_names(chain.from_iterable(map(_tool_use_names, list_field(request, "messages")))),
)
def _restore_tool_use_names(messages: Sequence[object], original_names: Mapping[str, str]) -> None:
for message in messages:
for block in list_field(message, "content"):
if (name := _tool_use_name(block)) is not None and name in original_names and isinstance(block, dict):
block["name"] = original_names[name]
def _with_original_name(tool: ChatCompletionToolParam, original_names: Mapping[str, str]) -> ChatCompletionToolParam:
if "function" not in tool:
return tool
function: Final = tool["function"]
fields: Final = as_str_mapping(function)
name: Final = None if fields is None else fields.get("name")
original: Final = original_names.get(name) if isinstance(name, str) else None
if original is None:
return tool
return {**tool, "function": {**function, "name": original}}
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
if not isinstance(system, list):
return
@ -629,6 +685,9 @@ class AnthropicMessagesHandler(BaseTranslation):
model: Final = data.get("model")
if model:
inputs["model"] = model
original_tool_names: Final = _original_tool_names(
cast(Mapping[str, object], data) # cast-ok: Anthropic request body
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
@ -643,7 +702,9 @@ class AnthropicMessagesHandler(BaseTranslation):
anthropic_config: Final = AnthropicConfig()
anthropic_tools: Final[list[AllAnthropicToolsValues]] = []
for tool in guardrailed_tools:
converted_tool, mcp_server = anthropic_config.map_tool_helper(tool)
converted_tool, mcp_server = anthropic_config.map_tool_helper(
_with_original_name(tool, original_tool_names.declared)
)
if converted_tool is not None:
anthropic_tools.append(converted_tool)
# Note: MCP servers are handled separately in the main transformation
@ -676,6 +737,7 @@ class AnthropicMessagesHandler(BaseTranslation):
),
hoisted_system_message=hoisted_system_message,
preserve_system_messages=has_midturn_system_message,
original_tool_names=original_tool_names.used,
)
else:
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
@ -840,6 +902,8 @@ class AnthropicMessagesHandler(BaseTranslation):
structured_messages: list, # mutable-ok: API message payload
hoisted_system_message: object = None,
preserve_system_messages: bool = False,
*,
original_tool_names: Mapping[str, str],
) -> None:
"""Write a guardrail's structured-message rewrite back without losing corrections."""
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -903,6 +967,10 @@ class AnthropicMessagesHandler(BaseTranslation):
for block in content:
if isinstance(block, dict) and block.get("type") == "thinking":
block.pop("cache_control", None)
_restore_tool_use_names(
cast(Sequence[object], converted), # cast-ok: Anthropic message payload
original_tool_names,
)
data["messages"] = keep_non_string_tool_result_texts(converted, data.get("messages"))
@staticmethod

View file

@ -15,7 +15,7 @@ How it works:
"""
import uuid
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import litellm
@ -48,7 +48,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
def can_handle(
self,
tools: list[dict] | None,
tools: Sequence[Mapping[str, object]] | None,
custom_llm_provider: str | None,
) -> bool:
if not tools:

View file

@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping, Sequence
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
@ -21,7 +21,7 @@ class MessagesInterceptor(ABC):
@abstractmethod
def can_handle(
self,
tools: list[dict] | None,
tools: Sequence[Mapping[str, object]] | None,
custom_llm_provider: str | None,
) -> bool:
"""Return True if this interceptor should handle the request."""

View file

@ -0,0 +1,3 @@
from .transformation import BaseExtractionConfig
__all__ = ["BaseExtractionConfig"]

View file

@ -0,0 +1,86 @@
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.extraction import ExtractionRequestBody, ExtractionResponse
_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
_RESPONSE_ADAPTER: Final[TypeAdapter[ExtractionResponse]] = TypeAdapter(ExtractionResponse)
_RESERVED_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "content-type"})
class BaseExtractionConfig(ABC):
api_key_env: tuple[str, ...] = ()
api_base_env: tuple[str, ...] = ()
def get_default_api_base(self) -> str | None:
return None
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return f"api_base is required for Extraction provider '{custom_llm_provider}'"
def resolve_api_base(self, api_base: str | None) -> str | None:
return api_base or self._first_secret(self.api_base_env) or self.get_default_api_base()
def resolve_api_key(self, api_key: str | None) -> str | None:
return api_key or self._first_secret(self.api_key_env)
@staticmethod
def _first_secret(names: tuple[str, ...]) -> str | None:
return next((value for value in (get_secret_str(name) for name in names) if value), None)
def canonical_model(self, model: str) -> str:
return model
def request_model(self, model: str) -> str:
return model
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
return {
**{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS},
**({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}),
"Content-Type": "application/json",
}
@abstractmethod
def get_complete_url(self, api_base: str, model: str) -> str: ...
def transform_extraction_request(self, model: str, request: ExtractionRequestBody) -> Mapping[str, object]:
return {**request.model_dump(mode="json", exclude_unset=True), "model": self.request_model(model)}
def unwrap_response(self, payload: object) -> object:
return payload
def transform_extraction_response(
self,
model: str,
custom_llm_provider: str,
raw_response: httpx.Response,
) -> ExtractionResponse:
payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content)
try:
return _RESPONSE_ADAPTER.validate_python(self.unwrap_response(payload))
except ValidationError as error:
raise BaseLLMException(
status_code=500,
message=f"Extraction provider '{custom_llm_provider}' returned an unexpected response: {error}",
) from error
def get_error_class(
self,
error_message: str,
status_code: int,
headers: Mapping[str, str] | httpx.Headers,
) -> BaseLLMException:
return BaseLLMException(status_code=status_code, message=error_message, headers=httpx.Headers(headers))

View file

@ -68,6 +68,16 @@ class HarnessTurnRequest:
cwd: str | None = None
@dataclass(frozen=True)
class HarnessSessionCopy:
"""A saved session under fresh ids: written to `path` in the sandbox, then loaded by `import_command`."""
session_id: str
path: str
data: bytes
import_command: Sequence[str]
@dataclass(frozen=True)
class HarnessTurnResponse:
"""What the runtime produced for one turn, after the process exited."""
@ -128,6 +138,18 @@ class BaseCLIHarnessConfig(BaseHarnessConfig[OptionsT], Generic[OptionsT, Stream
) -> HarnessTurnRequest:
"""argv / env / stdin for one turn. native_session_id is set after the first turn."""
def transform_resume_export_command(self, native_session_id: str) -> Sequence[str] | None:
"""Command that prints a saved session before resuming one this handler did not run.
None resumes the saved session as is. A non-zero exit or a timeout fails the turn, and the next turn
tries again.
"""
return None
def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None:
"""The copy to import and continue, from the export command's stdout, or None to continue the original."""
return None
@abstractmethod
def create_stream_state(self) -> StreamStateT:
"""Fresh per-turn parser state."""

View file

@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None:
return headers.get(name, None)
def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
try:
guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
except ValidationError as e:
@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
guardrail_headers: Final = (
()
if raw_guardrail_config is None
else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
)
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)

View file

@ -14,6 +14,9 @@ from litellm.constants import AWS_ERROR_TYPE_HEADER
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
bedrock_invoke_guardrail_headers,
)
from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
from litellm.types.utils import ModelResponseStream
@ -42,6 +45,9 @@ if TYPE_CHECKING:
BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"})
_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"})
_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset(
{"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"}
)
def _bedrock_endpoint_action(endpoint: str) -> str:
@ -123,6 +129,31 @@ def _with_invoke_request_metadata(
return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body
def _converse_guardrail_config(action: str, guardrail_config: object) -> object:
fields: Final = as_json_object(guardrail_config)
if action == "converse-stream" or fields is None:
return guardrail_config
return {key: value for key, value in fields.items() if key != "streamProcessingMode"}
def _with_deployment_guardrail(
action: str,
headers: Mapping[str, str],
body: Mapping[str, object],
guardrail_config: object,
) -> tuple[Mapping[str, str], Mapping[str, object]]:
if guardrail_config is None:
return headers, body
if action in _BEDROCK_CONVERSE_ACTIONS:
return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)}
if action in _BEDROCK_INVOKE_ACTIONS:
guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items())
return merge_bedrock_invoke_headers(
dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS
), body
return headers, body
def _is_cache_point(entry: object) -> bool:
fields: Final = as_json_object(entry)
return fields is not None and "cachePoint" in fields
@ -453,11 +484,14 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body,
litellm_params=litellm_params,
)
guarded_headers, guarded_body = _with_deployment_guardrail(
action, signed_headers, signed_body, litellm_params.get("guardrailConfig")
)
return self._sign_request(
service_name="bedrock",
headers=dict(signed_headers),
headers=dict(guarded_headers),
optional_params=optional_params,
request_data=dict(signed_body),
request_data=dict(guarded_body),
api_base=api_base,
model=model,
api_key=self.get_bedrock_bearer_token(optional_params),

View file

@ -256,12 +256,11 @@ def merge_bedrock_invoke_headers(
"""
Merge the ``X-Amzn-*`` headers the Invoke paths derive from params.
``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is
the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's
headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry
proxy-authenticated identity into an AWS billing record that the caller must not be able to
write. Names are compared case-insensitively so a caller cannot leave a second spelling in
the dict and let the transport pick the winner.
``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped
from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names
the caller must not be able to write, such as proxy-authenticated identity headed for an AWS
billing record or a deployment's guardrail. Names are compared case-insensitively so a caller
cannot leave a second spelling in the dict and let the transport pick the winner.
"""
if not caller_owned and not proxy_owned and not proxy_owned_names:
return headers

View file

@ -0,0 +1,53 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.secret_managers.main import (
get_secret_str,
normalize_nonempty_secret_str,
)
CLOUDFLARE_MISSING_API_BASE_MESSAGE: Final = (
"Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
)
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def workers_ai_default_api_base() -> str | None:
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
if account_id is None:
return None
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
def workers_ai_canonical_model(model: str) -> str:
if model.startswith("@cf/"):
return model
return f"@cf/cloudflare/{model}"
def workers_ai_request_model(model: str) -> str:
return model.rsplit("/", maxsplit=1)[-1]
def workers_ai_run_url(api_base: str, model: str) -> str:
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai/v1"):
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
if normalized_api_base.endswith("/ai/run"):
return f"{normalized_api_base}/{model}"
return f"{normalized_api_base}/ai/run/{model}"
def unwrap_workers_ai_result(payload: object) -> object:
try:
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
except ValidationError:
return payload
if "answers" in response_mapping:
return response_mapping
try:
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
except ValidationError:
return response_mapping

View file

@ -1,16 +1,13 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.secret_managers.main import (
get_secret_str,
normalize_nonempty_secret_str,
from litellm.llms.cloudflare.common_utils import (
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
unwrap_workers_ai_result,
workers_ai_canonical_model,
workers_ai_default_api_base,
workers_ai_request_model,
workers_ai_run_url,
)
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
class CloudflareDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("CLOUDFLARE_API_KEY",)
@ -23,38 +20,19 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig):
inlines_remote_images = True
def get_default_api_base(self) -> str | None:
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
if account_id is None:
return None
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
return workers_ai_default_api_base()
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
def canonical_model(self, model: str) -> str:
if model.startswith("@cf/"):
return model
return f"@cf/cloudflare/{model}"
return workers_ai_canonical_model(model)
def request_model(self, model: str) -> str:
return model.rsplit("/", maxsplit=1)[-1]
return workers_ai_request_model(model)
def get_complete_url(self, api_base: str, model: str) -> str:
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai/v1"):
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
if normalized_api_base.endswith("/ai/run"):
return f"{normalized_api_base}/{model}"
return f"{normalized_api_base}/ai/run/{model}"
return workers_ai_run_url(api_base, model)
def unwrap_response(self, payload: object) -> object:
try:
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
except ValidationError:
return payload
if "answers" in response_mapping:
return response_mapping
try:
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
except ValidationError:
return response_mapping
return unwrap_workers_ai_result(payload)

View file

@ -0,0 +1,11 @@
from typing import Final
from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler
from litellm.types.utils import CallTypes
guardrail_translation_mappings: Final = {
CallTypes.extraction: ExtractionHandler,
CallTypes.aextraction: ExtractionHandler,
}
__all__ = ["ExtractionHandler", "guardrail_translation_mappings"]

View file

@ -0,0 +1,276 @@
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias, TypeVar
from pydantic import Field, JsonValue, TypeAdapter
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.extraction import (
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
ExtractionAnswer,
ExtractionBase64Image,
ExtractionResponse,
)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter(
Mapping[str, Mapping[str, object]]
)
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_IMAGES_ADAPTER: Final[TypeAdapter[tuple[ExtractionBase64Image | JsonValue, ...]]] = TypeAdapter(
tuple[Annotated[ExtractionBase64Image | JsonValue, Field(union_mode="left_to_right")], ...]
)
_AnswerKey: TypeAlias = Literal["value", "guess", "mean", "median", "p10", "p90", "quantiles"]
_ESTIMATE_KEYS: Final[tuple[_AnswerKey, ...]] = ("value", "guess", "mean", "median", "p10", "p90", "quantiles")
_STATISTIC_KEYS: Final[frozenset[_AnswerKey]] = frozenset(("mean", "median", "p10", "p90", "quantiles"))
_NULLED_ANSWER: Final[Mapping[str, None]] = MappingProxyType(
dict.fromkeys(("value", "guess", "mean", "median", "p10", "p90", "quantiles", "steps", "span", "spans"), None)
)
@dataclass(frozen=True, slots=True)
class _AnswerText:
name: str
key: _AnswerKey
item: int | None
text: str
_SPEC_TEXT_KEYS: Final = ("instructions", "unit")
@dataclass(frozen=True, slots=True)
class _SpecText:
name: str
key: str
text: str
def _spec_texts_of(name: str, spec: Mapping[str, object]) -> tuple[_SpecText, ...]:
return tuple(
_SpecText(name=name, key=key, text=text) for key in _SPEC_TEXT_KEYS if isinstance(text := spec.get(key), str)
)
def _spec_texts(fields: Mapping[str, Mapping[str, object]]) -> tuple[_SpecText, ...]:
return tuple(chain.from_iterable(_spec_texts_of(name, spec) for name, spec in fields.items()))
_NamedText = TypeVar("_NamedText", _AnswerText, _SpecText)
def _scans_by_name(scanned: Iterable[tuple[_NamedText, str]]) -> Mapping[str, tuple[tuple[_NamedText, str], ...]]:
return {name: tuple(group) for name, group in groupby(scanned, key=lambda pair: pair[0].name)}
def _guardrailed_spec(spec: Mapping[str, object], scanned: Sequence[tuple[_SpecText, str]]) -> Mapping[str, object]:
return {**spec, **{text.key: guardrailed for text, guardrailed in scanned}}
def _metadata_key(data: Mapping[str, object]) -> Literal["metadata", "litellm_metadata"]:
return get_metadata_variable_name_from_kwargs(data)
def _metadata(data: Mapping[str, object]) -> Mapping[str, object]:
return _METADATA_ADAPTER.validate_python(data.get(_metadata_key(data)) or {})
def _is_text_answer(answer: ExtractionAnswer, spec: Mapping[str, object] | None) -> bool:
return answer.type == "text" or (answer.type == "list" and (spec is None or spec.get("items") == "text"))
def _image_urls(images: object) -> tuple[str, ...]:
if not isinstance(images, list):
return ()
return tuple(
f"data:{image.content_type};base64,{image.base64}" if isinstance(image, ExtractionBase64Image) else image
for image in _IMAGES_ADAPTER.validate_python(images)
if isinstance(image, (str, ExtractionBase64Image))
)
def _scalar_text(value: JsonValue, whole: bool) -> str | None:
if isinstance(value, bool):
return None
if isinstance(value, float) and whole:
return format(value, ".0f")
return str(value) if isinstance(value, (str, int, float)) else None
def _item_texts(items: Sequence[JsonValue], whole: bool) -> tuple[tuple[int, str], ...]:
return tuple((index, text) for index, item in enumerate(items) if (text := _scalar_text(item, whole)) is not None)
def _value_texts(name: str, key: _AnswerKey, value: JsonValue) -> tuple[_AnswerText, ...]:
whole: Final = key in _STATISTIC_KEYS
if isinstance(value, list):
return tuple(
_AnswerText(name=name, key=key, item=index, text=text) for index, text in _item_texts(value, whole)
)
text: Final = _scalar_text(value, whole)
return () if text is None else (_AnswerText(name=name, key=key, item=None, text=text),)
def _step_texts(name: str, step: JsonValue) -> tuple[_AnswerText, ...]:
if not isinstance(step, dict):
return ()
return tuple(chain.from_iterable(_value_texts(name, key, step.get(key)) for key in _ESTIMATE_KEYS))
def _answer_texts(name: str, answer: ExtractionAnswer) -> tuple[_AnswerText, ...]:
return (
*_value_texts(name, "value", answer.value),
*_value_texts(name, "guess", answer.guess),
*_value_texts(name, "mean", answer.mean),
*_value_texts(name, "median", answer.median),
*_value_texts(name, "p10", answer.p10),
*_value_texts(name, "p90", answer.p90),
*_value_texts(name, "quantiles", list(answer.quantiles or ())),
*chain.from_iterable(_step_texts(name, step) for step in answer.steps or ()),
)
def _guardrailed_value(value: JsonValue, scanned: Mapping[int | None, str]) -> JsonValue:
if isinstance(value, list):
return [scanned.get(index, item) for index, item in enumerate(value)]
return scanned.get(None, value)
def _scans_of(key: _AnswerKey, scanned: Sequence[tuple[_AnswerText, str]]) -> Mapping[int | None, str]:
return {text.item: guardrailed for text, guardrailed in scanned if text.key == key}
def _guardrailed_answer(
answer: ExtractionAnswer,
scanned: Sequence[tuple[_AnswerText, str]],
drop_spans: bool,
is_text: bool,
) -> ExtractionAnswer:
if not is_text and any(text.text != guardrailed for text, guardrailed in scanned):
return answer.model_copy(update=_NULLED_ANSWER)
value_scans: Final = _scans_of("value", scanned if is_text else ())
guess_scans: Final = _scans_of("guess", scanned if is_text else ())
update: Final = {
**({"value": _guardrailed_value(answer.value, value_scans)} if value_scans else {}),
**({"guess": _guardrailed_value(answer.guess, guess_scans)} if guess_scans else {}),
**({"span": None, "spans": None} if drop_spans else {}),
}
return answer.model_copy(update=update) if update else answer
async def _scanned_texts(
texts: Sequence[str],
model: object,
guardrail_to_apply: "CustomGuardrail",
request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict
input_type: Literal["request", "response"],
litellm_logging_obj: "LiteLLMLoggingObj | None",
images: Sequence[str] = (),
) -> tuple[str, ...]:
unique_texts: Final = tuple(dict.fromkeys(texts))
inputs: Final = GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None)
if images:
inputs["images"] = list(images)
guardrailed: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict
inputs=inputs,
request_data=request_data,
input_type=input_type,
logging_obj=litellm_logging_obj,
)
returned: Final = guardrailed.get("texts", [])
guardrailed_by_text: Final = MappingProxyType(
{text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)}
)
return tuple(guardrailed_by_text[text] for text in texts)
class ExtractionHandler(BaseTranslation):
async def process_input_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Mapping[str, object]:
state: Final = data.get("state")
states: Final = (state,) if isinstance(state, str) and state else ()
fields: Final = _FIELDS_ADAPTER.validate_python(data.get("fields") or {})
spec_texts: Final = _spec_texts(fields)
images: Final = _image_urls(data.get("images"))
if not states and not spec_texts and not images:
return data
scanned: Final = await _scanned_texts(
(*states, *(text.text for text in spec_texts)),
data.get("model"),
guardrail_to_apply,
data,
"request",
litellm_logging_obj,
images,
)
spec_scans: Final = _scans_by_name(zip(spec_texts, scanned[len(states) :]))
state_rewritten: Final = bool(states) and scanned[0] != states[0]
return {
**data,
**({"state": scanned[0]} if states else {}),
**(
{_metadata_key(data): {**_metadata(data), EXTRACTION_STATE_REWRITTEN_METADATA_KEY: True}}
if state_rewritten
else {}
),
"fields": {name: _guardrailed_spec(spec, spec_scans.get(name, ())) for name, spec in fields.items()},
}
async def process_output_response(
self,
response: object,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict
) -> object:
if not isinstance(response, ExtractionResponse):
return response
texts: Final = tuple(
chain.from_iterable(_answer_texts(name, answer) for name, answer in response.answers.items())
)
state_texts: Final = (response.state_text,) if response.state_text else ()
request: Final = request_data or {}
scanned: Final = (
await _scanned_texts(
(*(text.text for text in texts), *state_texts),
response.model,
guardrail_to_apply,
request,
"response",
litellm_logging_obj,
)
if texts or state_texts
else ()
)
answer_scans: Final = _scans_by_name(zip(texts, scanned))
state_text: Final = next(iter(scanned[len(texts) :]), response.state_text)
specs: Final = _FIELDS_ADAPTER.validate_python(request.get("fields") or {})
return response.model_copy(
update={
"answers": {
name: _guardrailed_answer(
answer,
answer_scans.get(name, ()),
drop_spans=state_text != response.state_text,
is_text=_is_text_answer(answer, specs.get(name)),
)
for name, answer in response.answers.items()
},
**({"state_text": state_text} if state_texts else {}),
}
)

View file

@ -0,0 +1,32 @@
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.cloudflare.common_utils import (
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
unwrap_workers_ai_result,
workers_ai_canonical_model,
workers_ai_default_api_base,
workers_ai_request_model,
workers_ai_run_url,
)
class CloudflareExtractionConfig(BaseExtractionConfig):
api_key_env = ("CLOUDFLARE_API_KEY",)
api_base_env = ("CLOUDFLARE_API_BASE",)
def get_default_api_base(self) -> str | None:
return workers_ai_default_api_base()
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
def canonical_model(self, model: str) -> str:
return workers_ai_canonical_model(model)
def request_model(self, model: str) -> str:
return workers_ai_request_model(model)
def get_complete_url(self, api_base: str, model: str) -> str:
return workers_ai_run_url(api_base, model)
def unwrap_response(self, payload: object) -> object:
return unwrap_workers_ai_result(payload)

View file

@ -76,6 +76,7 @@ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
BaseFileUploadStream,
@ -128,6 +129,7 @@ from litellm.types.containers.main import (
DeleteContainerResult,
)
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
from litellm.types.extraction import ExtractionResponse
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
@ -1700,6 +1702,115 @@ class BaseLLMHTTPHandler:
model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request
)
def _prepare_extraction_request(
self,
model: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
litellm_params: Mapping[str, object],
) -> tuple[str, Mapping[str, str]]:
outbound_headers: Final = provider_config.validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
url: Final = provider_config.get_complete_url(api_base=api_base, model=model)
if logging_obj is not None:
logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # Logging.pre_call is untyped
input=dict(body),
api_key=api_key,
model=model,
additional_args={"api_base": url, "complete_input_dict": dict(body), "headers": outbound_headers},
)
return url, outbound_headers
def extraction(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: HTTPHandler | None = None,
) -> ExtractionResponse:
url, outbound_headers = self._prepare_extraction_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
sync_httpx_client: Final = client if client is not None else get_httpx_client()
try:
response: Final = sync_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
url,
json=dict(body),
headers=dict(outbound_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_extraction_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
)
async def aextraction(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> ExtractionResponse:
url, outbound_headers = self._prepare_extraction_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
async_httpx_client: Final = (
client
if client is not None
else get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider))
)
try:
response: Final = await async_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
url,
json=dict(body),
headers=dict(outbound_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_extraction_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
)
def _prepare_audio_transcription_request(
self,
model: str,
@ -6462,6 +6573,7 @@ class BaseLLMHTTPHandler:
provider_config: Union[
BaseConfig,
BaseDecisionsConfig,
BaseExtractionConfig,
BaseRerankConfig,
BaseResponsesAPIConfig,
BaseImageEditConfig,

View file

@ -0,0 +1,78 @@
import os
import re
from dataclasses import dataclass
from typing import Final
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from .authenticator import (
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
DEFAULT_GITHUB_API_KEY_URL,
DEFAULT_GITHUB_CLIENT_ID,
DEFAULT_GITHUB_DEVICE_CODE_URL,
)
_GHE_SUBDOMAIN_PATTERN: Final = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?")
@dataclass(frozen=True, slots=True)
class GithubHostEndpoints:
host: str
device_code_url: str
access_token_url: str
api_key_url: str
user_url: str
client_id: str
copilot_api_hosts: frozenset[str]
copilot_api_host_suffix: str | None
default_copilot_api_base: str
def normalized_github_host(raw: object) -> str | None:
if raw is None or raw == "":
return GITHUB_COPILOT_DEFAULT_HOST
if not isinstance(raw, str):
return None
host: Final = raw.lower()
if host == GITHUB_COPILOT_DEFAULT_HOST:
return host
if not host.endswith(".ghe.com"):
return None
subdomain: Final = host.removesuffix(".ghe.com")
return host if _GHE_SUBDOMAIN_PATTERN.fullmatch(subdomain) is not None else None
def github_host_endpoints(host: str) -> GithubHostEndpoints:
normalized: Final = normalized_github_host(host)
if normalized is None:
raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
if normalized == GITHUB_COPILOT_DEFAULT_HOST:
return GithubHostEndpoints(
host=normalized,
device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
user_url="https://api.github.com/user",
client_id=client_id,
copilot_api_hosts=frozenset({"githubcopilot.com"}),
copilot_api_host_suffix=".githubcopilot.com",
default_copilot_api_base="https://api.githubcopilot.com",
)
subdomain: Final = normalized.removesuffix(".ghe.com")
api_host: Final = f"api.{subdomain}.ghe.com"
copilot_api_host: Final = f"copilot-api.{subdomain}.ghe.com"
return GithubHostEndpoints(
host=normalized,
device_code_url=f"https://{normalized}/login/device/code",
access_token_url=f"https://{normalized}/login/oauth/access_token",
api_key_url=f"https://{api_host}/copilot_internal/v2/token",
user_url=f"https://{api_host}/user",
client_id=client_id,
copilot_api_hosts=frozenset({copilot_api_host}),
copilot_api_host_suffix=None,
default_copilot_api_base=f"https://{copilot_api_host}",
)

View file

@ -5,8 +5,7 @@ touched on the per-user path.
"""
import asyncio
import hashlib
import os
import hmac
import threading
import time
from collections.abc import Callable, Mapping
@ -27,6 +26,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
GITHUB_COPILOT_AUTH_TYPE_KEY,
GITHUB_COPILOT_HOST_KEY,
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
)
@ -38,14 +38,8 @@ from litellm.exceptions import (
ServiceUnavailableError,
)
from .authenticator import (
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
DEFAULT_GITHUB_API_KEY_URL,
DEFAULT_GITHUB_CLIENT_ID,
DEFAULT_GITHUB_DEVICE_CODE_URL,
github_api_headers,
)
from .common_utils import DEFAULT_GITHUB_COPILOT_API_BASE
from .authenticator import github_api_headers
from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session"
_LLM_PROVIDER: Final = "github_copilot"
@ -98,7 +92,7 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess
)
_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
_IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking
dict[str, "asyncio.Future[GithubCopilotUserSession]"]
dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]]
] = {}
@ -141,7 +135,7 @@ _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
_GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
_COPILOT_ENDPOINTS: Final = TypeAdapter(_CopilotEndpointsPayload)
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, float])
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, str, float])
@dataclass(frozen=True, slots=True, repr=False)
@ -178,53 +172,72 @@ class GithubCopilotDeviceFlowPoll:
access_token: str | None = field(default=None, repr=False)
def validated_copilot_api_base(endpoints_api: object) -> str:
"""Only a genuine *.githubcopilot.com HTTPS endpoint may receive a Copilot
token; anything else falls back to the default base."""
def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str:
if not isinstance(endpoints_api, str):
return DEFAULT_GITHUB_COPILOT_API_BASE
parsed: Final = urlparse(endpoints_api)
return endpoints.default_copilot_api_base
try:
parsed: Final = urlparse(endpoints_api)
port: Final = parsed.port
except ValueError:
return endpoints.default_copilot_api_base
hostname: Final = (parsed.hostname or "").lower()
if (
parsed.scheme.lower() != "https"
or not (hostname == "githubcopilot.com" or hostname.endswith(".githubcopilot.com"))
or not (
hostname in endpoints.copilot_api_hosts
or (endpoints.copilot_api_host_suffix is not None and hostname.endswith(endpoints.copilot_api_host_suffix))
)
or parsed.username is not None
or parsed.password is not None
or port not in (None, 443)
):
return DEFAULT_GITHUB_COPILOT_API_BASE
try:
port: Final = parsed.port
except ValueError:
return DEFAULT_GITHUB_COPILOT_API_BASE
if port not in (None, 443):
return DEFAULT_GITHUB_COPILOT_API_BASE
return endpoints.default_copilot_api_base
return endpoints_api.rstrip("/")
def _session_cache_key(user_id: str, github_token: str) -> str:
return hashlib.sha256(f"{user_id}\x00{github_token}".encode()).hexdigest()
def github_copilot_host_endpoints(credential_name: str) -> GithubHostEndpoints:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
credential: Final = CredentialAccessor.find_credential(credential_name)
host_raw: Final[object] = (
credential.credential_values.get(GITHUB_COPILOT_HOST_KEY) if credential is not None else None
)
host: Final = normalized_github_host(host_raw)
if host is None:
raise BadRequestError(
message=f"Credential '{credential_name}' has an invalid {GITHUB_COPILOT_HOST_KEY}",
llm_provider=_LLM_PROVIDER,
model="",
)
return github_host_endpoints(host)
def evict_copilot_user_session(user_id: str, github_token: str) -> None:
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, github_token))
def _session_cache_key(user_id: str, credential_name: str, github_host: str) -> str:
return f"{user_id}\x00{credential_name}\x00{github_host}"
def _cached_session(cache_key: str) -> GithubCopilotUserSession | None:
def evict_copilot_user_session(user_id: str, credential_name: str, github_host: str) -> None:
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, github_host))
def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSession | None:
cached: Final[object] = _SESSION_CACHE.get_cache(cache_key)
try:
token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
cached_github_token, token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
except ValidationError:
return None
if not hmac.compare_digest(cached_github_token, github_token):
return None
if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time():
return None
return GithubCopilotUserSession(token=token, api_base=api_base)
def _cache_session(cache_key: str, session: GithubCopilotUserSession, expires_at: float) -> None:
def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None:
ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time()
if ttl <= 0:
return
_SESSION_CACHE.set_cache(cache_key, (session.token, session.api_base, expires_at), ttl=ttl)
_SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl)
def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError:
@ -238,7 +251,13 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr
)
def _session_from_response(response: httpx.Response, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
def _session_from_response(
response: httpx.Response,
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
if response.status_code in (401, 403, 404):
_SESSION_CACHE.delete_cache(cache_key)
raise _reject_connection(credential_name)
@ -263,15 +282,17 @@ def _session_from_response(response: httpx.Response, credential_name: str, cache
model="",
) from e
endpoints: Final[object] = payload.get("endpoints")
payload_endpoints: Final[object] = payload.get("endpoints")
endpoints_api: Final[object] = (
_COPILOT_ENDPOINTS.validate_python(endpoints).get("api") if isinstance(endpoints, Mapping) else None
_COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
if isinstance(payload_endpoints, Mapping)
else None
)
api_base: Final = validated_copilot_api_base(endpoints_api)
api_base: Final = validated_copilot_api_base(endpoints_api, endpoints)
session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
if expires_at > 0:
_cache_session(cache_key, session, expires_at)
_cache_session(cache_key, github_token, session, expires_at)
return session
@ -303,7 +324,12 @@ def _expires_at_seconds(raw: object) -> float:
return 0.0
def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
def _do_exchange(
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
import litellm
try:
@ -311,7 +337,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
_SyncGetClient, litellm.module_level_client
)
response: Final = client.get(
DEFAULT_GITHUB_API_KEY_URL,
endpoints.api_key_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -320,7 +346,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
llm_provider=_LLM_PROVIDER,
model="",
) from e
return _session_from_response(response, credential_name, cache_key)
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
def _new_async_github_client() -> _AsyncGitHubClient:
@ -332,11 +358,16 @@ def _new_async_github_client() -> _AsyncGitHubClient:
)(_LLM_PROVIDER)
async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
async def _do_aexchange(
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
client: Final = _new_async_github_client()
try:
response: Final = await client.get(
DEFAULT_GITHUB_API_KEY_URL,
endpoints.api_key_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -345,43 +376,50 @@ async def _do_aexchange(github_token: str, credential_name: str, cache_key: str)
llm_provider=_LLM_PROVIDER,
model="",
) from e
return _session_from_response(response, credential_name, cache_key)
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
def exchange_github_token(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
cache_key: Final = _session_cache_key(user_id, github_token)
cached: Final = _cached_session(cache_key)
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
cached: Final = _cached_session(cache_key, github_token)
if cached is not None:
return cached
lock: Final = _EXCHANGE_LOCKS[int(cache_key, 16) % _SYNC_LOCK_COUNT]
lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
with lock:
rechecked: Final = _cached_session(cache_key)
rechecked: Final = _cached_session(cache_key, github_token)
if rechecked is not None:
return rechecked
return _do_exchange(github_token, credential_name, cache_key)
return _do_exchange(github_token, credential_name, cache_key, endpoints)
async def aexchange_github_token(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
cache_key: Final = _session_cache_key(user_id, github_token)
cached: Final = _cached_session(cache_key)
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
cached: Final = _cached_session(cache_key, github_token)
if cached is not None:
return cached
loop: Final = asyncio.get_running_loop()
in_flight: Final = _IN_FLIGHT.get(cache_key)
if in_flight is not None and not in_flight.done() and in_flight.get_loop() is loop:
return await asyncio.shield(in_flight)
in_flight_entry: Final = _IN_FLIGHT.get(cache_key)
if in_flight_entry is not None:
in_flight_token, in_flight = in_flight_entry
if not hmac.compare_digest(in_flight_token, github_token):
return await _do_aexchange(github_token, credential_name, cache_key, endpoints)
if not in_flight.done() and in_flight.get_loop() is loop:
return await asyncio.shield(in_flight)
future: Final = loop.create_future()
_IN_FLIGHT[cache_key] = future
entry: Final = (github_token, future)
_IN_FLIGHT[cache_key] = entry
try:
session: Final = await _do_aexchange(github_token, credential_name, cache_key)
session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints)
except BaseException as e:
if not future.done():
future.set_exception(e)
@ -392,19 +430,17 @@ async def aexchange_github_token(
future.set_result(session)
return session
finally:
if _IN_FLIGHT.get(cache_key) is future:
if _IN_FLIGHT.get(cache_key) is entry:
_IN_FLIGHT.pop(cache_key, None)
async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
async def astart_device_flow(endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowStart:
client: Final = _new_async_github_client()
device_code_url: Final = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL)
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
try:
response: Final = await client.post(
device_code_url,
endpoints.device_code_url,
headers=github_api_headers(),
json={"client_id": client_id, "scope": _DEVICE_FLOW_SCOPE},
json={"client_id": endpoints.client_id, "scope": _DEVICE_FLOW_SCOPE},
)
except httpx.HTTPStatusError as e:
raise ServiceUnavailableError(
@ -435,16 +471,14 @@ async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
)
async def apoll_device_flow(device_code: str) -> GithubCopilotDeviceFlowPoll:
async def apoll_device_flow(device_code: str, endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowPoll:
client: Final = _new_async_github_client()
access_token_url: Final = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL)
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
try:
response: Final = await client.post(
access_token_url,
endpoints.access_token_url,
headers=github_api_headers(),
json={
"client_id": client_id,
"client_id": endpoints.client_id,
"device_code": device_code,
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
},
@ -500,11 +534,11 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice
)
async def afetch_github_login(github_token: str) -> str:
async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str:
client: Final = _new_async_github_client()
try:
response: Final = await client.get(
"https://api.github.com/user",
endpoints.user_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -530,11 +564,17 @@ async def afetch_github_login(github_token: str) -> str:
return payload["login"]
async def acheck_copilot_seat(user_id: str, github_token: str, credential_name: str) -> GithubCopilotUserSession:
async def acheck_copilot_seat(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
return await aexchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
@ -567,7 +607,11 @@ def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError:
)
def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tuple[str, str]:
def _caller_connection(
kwargs: Mapping[str, object],
credential_name: str,
endpoints: GithubHostEndpoints,
) -> tuple[str, str]:
secret_fields_raw: Final[object] = kwargs.get("secret_fields")
secret_fields: Final = (
cast(Mapping[object, object], secret_fields_raw) # cast-ok: value is Mapping-checked or dict-shaped at runtime
@ -585,8 +629,21 @@ def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tu
if isinstance(credentials_raw, Mapping)
else None
)
token: Final[object] = credentials.get(credential_name) if credentials is not None else None
if not isinstance(user_id, str) or not user_id or not isinstance(token, str) or not token:
connection_raw: Final[object] = credentials.get(credential_name) if credentials is not None else None
connection: Final = (
cast(Mapping[object, object], connection_raw) # cast-ok: runtime Mapping check validates the stored connection
if isinstance(connection_raw, Mapping)
else None
)
token: Final[object] = connection.get("access_token") if connection is not None else None
github_host: Final = normalized_github_host(connection.get("github_host") if connection is not None else None)
if (
not isinstance(user_id, str)
or not user_id
or not isinstance(token, str)
or not token
or github_host != endpoints.host
):
raise _connect_error(credential_name)
return user_id, token
@ -705,11 +762,13 @@ def attach_github_copilot_user_session(
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
return
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
user_id, github_token = _caller_connection(kwargs, credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = exchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
_strip_caller_authorization(kwargs)
@ -725,11 +784,13 @@ async def aattach_github_copilot_user_session(
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
return
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
user_id, github_token = _caller_connection(kwargs, credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = await aexchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
_strip_caller_authorization(kwargs)

View file

@ -16,7 +16,9 @@ The handler is intentionally a thin orchestration layer:
3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS)
vs self-hosted (any host:port, optional TLS) modes.
4. Stream the audio through Riva's ``streaming_response_generator`` and
aggregate ``is_final`` results into a single transcript.
aggregate ``is_final`` results into a single transcript, or, when
``riva_offline`` is set, send the whole buffer through the unary
``offline_recognize`` RPC for NIMs that only ship offline engines.
5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed
on ``_hidden_params`` so cost calculation works.
@ -26,13 +28,19 @@ without the optional STT extras installed.
import asyncio
import inspect
from collections.abc import Callable, Iterable
from collections.abc import Callable, Iterable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
import httpx
import litellm
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_name,
process_audio_file,
)
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.llms.nvidia_riva.audio_transcription.audio_utils import (
resample_to_riva_pcm,
)
@ -45,6 +53,7 @@ from litellm.llms.nvidia_riva.common_utils import (
NvidiaRivaException,
grpc_error_to_litellm_exception,
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.utils import FileTypes, TranscriptionResponse
from litellm.utils import convert_to_model_response_object
@ -67,10 +76,20 @@ class _RivaAuth(Protocol):
"""Opaque ``riva.client.Auth`` handle."""
class _RecognizeFuture(Protocol):
def result(self, timeout: float | None = None) -> object: ...
def done(self) -> bool: ...
def cancel(self) -> bool: ...
class _AsrService(Protocol):
@property
def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ...
class _EndpointingConfig(Protocol):
"""Opaque ``EndpointingConfig`` protobuf message."""
@ -128,6 +147,38 @@ def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding:
return riva_asr_module.AudioEncoding
def _parse_riva_offline(value: object, drop_params: bool) -> bool:
match value:
case None:
return False
case bool():
return value
case str() if (parsed := str_to_bool(value)) is not None:
return parsed
case _ if drop_params:
return False
case _:
raise NvidiaRivaException(
status_code=400,
message=f"riva_offline must be true or false, got {value!r}. Set drop_params to ignore it.",
)
def _deadline_seconds(timeout: float | str | httpx.Timeout, drop_params: bool) -> float:
try:
limit: Final = CompletionTimeout.as_limit(timeout)
except ValueError:
if drop_params:
return COMPLETION_HTTP_FALLBACK_SECONDS
raise NvidiaRivaException(
status_code=400,
message=f"timeout must be a number of seconds, got {timeout!r}. Set drop_params to ignore it.",
) from None
if not isinstance(limit, httpx.Timeout):
return limit
return COMPLETION_HTTP_FALLBACK_SECONDS if limit.read is None else limit.read
class NvidiaRivaAudioTranscription:
"""Sync + async entry point for Riva ASR."""
@ -138,7 +189,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -183,7 +234,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -213,7 +264,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -259,6 +310,10 @@ class NvidiaRivaAudioTranscription:
response_format: Final = request_payload.get("response_format") or "json"
timestamp_granularities: Final = request_payload.get("timestamp_granularities")
drop_params: Final = (
"drop_params" in litellm_params and litellm_params["drop_params"] is True
) or litellm.drop_params is True
riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params)
riva_module, riva_asr_module = _import_riva()
auth_obj: Final = self._construct_auth(
@ -272,10 +327,6 @@ class NvidiaRivaAudioTranscription:
riva_asr_module=riva_asr_module,
recognition_config_dict=recognition_config_dict,
)
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
logging_obj.pre_call(
input=None,
api_key=api_key,
@ -286,24 +337,31 @@ class NvidiaRivaAudioTranscription:
"recognition_config": recognition_config_dict,
"nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")),
"use_ssl": optional_params.get("use_ssl"),
"riva_offline": riva_offline,
},
},
)
try:
asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes)
stream_kwargs: Final[dict[str, object]] = {
"audio_chunks": audio_chunks,
"streaming_config": streaming_config,
}
# Forward the deadline so the stream cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator):
stream_kwargs["timeout"] = float(timeout)
stream: Final = asr_service.streaming_response_generator(**stream_kwargs)
final_results: Final = self._collect_final_results(stream)
responses: Final = (
self._offline_responses(
asr_service=asr_service,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=None if timeout is None else _deadline_seconds(timeout, drop_params),
)
if riva_offline
else self._streaming_responses(
asr_service=asr_service,
riva_asr_module=riva_asr_module,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=timeout,
drop_params=drop_params,
)
)
final_results: Final = self._collect_final_results(responses, require_final=not riva_offline)
except NvidiaRivaException:
raise
except Exception as e:
@ -409,6 +467,54 @@ class NvidiaRivaAudioTranscription:
return config
def _streaming_responses(
self,
asr_service: _AsrService,
riva_asr_module: _RivaAsrModule,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | str | httpx.Timeout | None,
drop_params: bool,
) -> Iterable[object]:
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
return asr_service.streaming_response_generator(
audio_chunks=self._iter_audio_chunks(pcm_bytes),
streaming_config=streaming_config,
**self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params),
)
def _offline_responses(
self,
asr_service: _AsrService,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | None,
) -> Iterable[object]:
pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True)
try:
return (pending.result(timeout=timeout),)
except Exception as e:
if pending.done():
raise
pending.cancel()
raise NvidiaRivaException(
status_code=408,
message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds",
) from e
@classmethod
def _timeout_kwargs(
cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool
) -> Mapping[str, float]:
# Forward the deadline so the call cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is None or not cls._supports_timeout_kwarg(rpc):
return MappingProxyType({})
return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)})
@staticmethod
def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
try:
@ -429,9 +535,9 @@ class NvidiaRivaAudioTranscription:
yield chunk
@staticmethod
def _collect_final_results(stream) -> list[dict[str, object]]:
def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]:
"""
Walk the gRPC stream, ignore empty / non-final chunks, and return a
Walk the gRPC responses, ignore empty / non-final chunks, and return a
list of normalized final-result dicts. Matching the user's note: the
``id`` blocks with no ``results`` are streaming heartbeats and must
be skipped.
@ -440,7 +546,7 @@ class NvidiaRivaAudioTranscription:
for response in stream:
results = getattr(response, "results", None) or []
for result in results:
if not getattr(result, "is_final", False):
if require_final and not getattr(result, "is_final", False):
continue
alternatives = getattr(result, "alternatives", None) or []
if not alternatives:

View file

@ -41,6 +41,9 @@ _GRPC_STATUS_CODE_TO_HTTP: Final[dict] = {
}
_GRPC_MESSAGE_TOO_LARGE: Final = "message larger than max"
def _extract_grpc_status_name(error: object) -> str | None:
"""
Best-effort extraction of a gRPC StatusCode name from an arbitrary error.
@ -79,8 +82,16 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException:
a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code.
"""
status_name: Final = _extract_grpc_status_name(error)
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
detail: Final = _extract_grpc_details(error) or str(error)
if status_name == "RESOURCE_EXHAUSTED" and _GRPC_MESSAGE_TOO_LARGE in detail:
return NvidiaRivaException(
status_code=413,
message=(
f"NVIDIA Riva gRPC error ({status_name}): {detail}. The audio is too large for one offline "
"request to this Riva server: send shorter audio, or drop riva_offline on models that support streaming."
),
)
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}"
return NvidiaRivaException(status_code=http_status, message=message)

View file

@ -8,17 +8,31 @@ repo's own opencode.json cannot redirect model calls. The token is never in argv
the config references it with `{file:<private_dir>/token}`. XDG dirs point at a persisted
LiteLLM-owned root so the user's opencode config and auth are never read, and the session
DB outlives a session for resume. Verified against opencode 1.14.41.
opencode files each session under the folder it ran in and only streams a resumed session's
events from that folder, so resuming one filed elsewhere (by an older LiteLLM that left PWD
unset, or from another sandbox) hangs. The first turn after a resume runs `opencode export` in
the sandbox, and when the session is filed elsewhere it imports a copy under fresh ids, filed
under the sandbox folder, and continues the copy. Import overwrites the folder of a session id it
already has, so reusing the ids would break the original wherever it was saved. Before 1.15.6,
import files a session under the folder named in its JSON, so the copy names the sandbox folder.
A failed or timed-out export or import fails that turn instead of hanging, and the next turn
tries again. Verified against opencode 1.14.41, 1.15.5, 1.15.6 and 1.19.3.
"""
from __future__ import annotations
import itertools
import json
from collections.abc import Mapping, Sequence
import secrets
import string
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from litellm.harness.errors import CapabilityUnsupported, HarnessError, OptionsMismatch
from litellm.harness.options import OpenCodeOptions
from litellm.harness.types import (
@ -33,6 +47,7 @@ from litellm.harness.types import (
)
from litellm.llms.base_llm.harness.transformation import (
BaseCLIHarnessConfig,
HarnessSessionCopy,
HarnessSessionSetup,
HarnessTurnError,
HarnessTurnRequest,
@ -57,6 +72,10 @@ OPENCODE_PROVIDER_NPM: Final = "@ai-sdk/openai-compatible"
# A fixed title skips opencode's extra title-generation model call on the first turn.
OPENCODE_SESSION_TITLE: Final = "litellm-harness"
TOKEN_FILENAME: Final = "token"
RESUME_COPY_FILENAME: Final = "resume-session.json"
ID_ORDERED_PREFIX_LENGTH: Final = 16
ID_RANDOM_LENGTH: Final = 14
ID_RANDOM_ALPHABET: Final = string.ascii_letters + string.digits
INSTRUCTIONS_FILENAME: Final = "instructions.md"
XDG_DIRNAME: Final = "xdg"
XDG_SUBDIRS: Final = ("config", "data", "state", "cache")
@ -162,6 +181,94 @@ class OpenCodeStreamState:
step_texts: Sequence[str] = ()
class ExportedPart(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
class ExportedMessageInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
class ExportedMessage(BaseModel):
model_config = ConfigDict(frozen=True)
info: ExportedMessageInfo
parts: tuple[ExportedPart, ...]
class ExportedSessionInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
directory: str
class ExportedSession(BaseModel):
"""The ids and folder of an `opencode export`; every other field is copied as is."""
model_config = ConfigDict(frozen=True)
info: ExportedSessionInfo
messages: tuple[ExportedMessage, ...]
def ids(self) -> Iterator[str]:
yield self.info.id
for message in self.messages:
yield message.info.id
yield from (part.id for part in message.parts)
EXPORT_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
def fresh_id(old: str) -> str:
"""Keeps the type prefix and time-ordered part of an opencode id, so the copy sorts like the original."""
suffix: Final = "".join(secrets.choice(ID_RANDOM_ALPHABET) for _ in range(ID_RANDOM_LENGTH))
return f"{old[:ID_ORDERED_PREFIX_LENGTH]}{suffix}"
def with_id_fields_replaced(value: JsonValue, fresh: Mapping[str, str], key: str = "") -> JsonValue:
match value:
case str() if key == "id" or key.endswith("ID"):
return fresh.get(value, value)
case list():
return [with_id_fields_replaced(item, fresh, key) for item in value]
case dict():
return {name: with_id_fields_replaced(item, fresh, name) for name, item in value.items()}
case _:
return value
def filed_under(exported: JsonValue, workdir: str) -> JsonValue:
match exported:
case {"info": dict() as info}:
return {**exported, "info": {**info, "directory": workdir}}
case _:
return exported
def session_copy(exported: str, path: str, workdir: str) -> HarnessSessionCopy | None:
"""A copy of an exported session filed under workdir, or None when the session is already filed there."""
try:
raw: Final = EXPORT_JSON.validate_json(exported)
session: Final = ExportedSession.model_validate(raw)
except ValidationError as exc:
raise HarnessError(f"opencode export did not print a session: {exc.errors()[0]['msg']}") from exc
if session.info.directory == workdir:
return None
fresh: Final = MappingProxyType({old: fresh_id(old) for old in session.ids()})
return HarnessSessionCopy(
session_id=fresh[session.info.id],
path=path,
data=json.dumps(filed_under(with_id_fields_replaced(raw, fresh), workdir)).encode("utf-8"),
import_command=(OPENCODE_BINARY, "--pure", "import", path),
)
def _as_dict(value: object) -> Mapping[str, Any]:
return value if isinstance(value, dict) else MappingProxyType({})
@ -362,6 +469,12 @@ class OpenCodeHarnessConfig(BaseCLIHarnessConfig):
# The prompt goes on stdin; opencode appends non-TTY stdin to the message.
return HarnessTurnRequest(argv=argv, env=setup.env, stdin=turn_prompt(ctx, prompt), cwd=ctx.sandbox.workdir)
def transform_resume_export_command(self, native_session_id: str) -> Sequence[str]:
return (OPENCODE_BINARY, "--pure", "export", native_session_id)
def transform_resume_copy(self, exported: str, private_dir: str, workdir: str) -> HarnessSessionCopy | None:
return session_copy(exported, f"{private_dir}/{RESUME_COPY_FILENAME}", workdir)
def create_stream_state(self) -> OpenCodeStreamState:
return OpenCodeStreamState()

View file

@ -5988,6 +5988,7 @@ def completion(
*OPENAI_WIF_KWARGS_KEYS,
*OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS,
"github_copilot_auth_type",
"github_copilot_host",
"github_copilot_user_session",
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
"fireworks_forward_user_id",

View file

@ -1,18 +1,21 @@
import inspect
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.anthropic.pass_through.messages import handler as main
from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.messages import detectors as route_detectors
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
)
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
@ -46,6 +49,7 @@ def _python_amessages() -> PythonAmessages:
_PYTHON_MESSAGES: Final = _python_messages()
_MESSAGES: Final = signature(_PYTHON_MESSAGES)
_FLAG: Final = AsyncFlag("is_async", sync=_MESSAGES, is_set=bool)
_PYTHON_AMESSAGES: Final = _python_amessages()
_AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
@ -71,45 +75,76 @@ def _public_request(
return native_call(legacy, args, _native_kwargs(fields, kwargs))
def _resolved_provider(request: NativeCall) -> str | None:
try:
return get_llm_provider(
str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))
)[1]
except BadRequestError:
return optional_str(request.resolved.get("custom_llm_provider"))
def _python_serves(request: NativeCall) -> bool:
"""Vertex AI serves only Claude through the Anthropic Messages config, as
``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's
translation through Chat Completions. A request Python would answer with classifier
``safeguard_results`` stays on Python too."""
def _resolved(request: NativeCall) -> tuple[str, str | None]:
model: Final = str(request.resolved["model"])
vertex_without_claude: Final = _resolved_provider(request) == "vertex_ai" and "claude" not in model.lower()
return vertex_without_claude or native_route_classifies(model, request.kwargs.get("safeguards"))
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
try:
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
except BadRequestError:
return model, provider
return resolved_model, resolved_provider
def _context(request: NativeCall) -> RouteContext:
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
requested_model: Final = str(request.resolved["model"])
model_info: Final = request.resolved.get("model_info")
api_base: Final = request.resolved.get("api_base")
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
model, provider, model_info, api_base if isinstance(api_base, str) else None
)
if native_config is not None:
return Target.MESSAGES
compacting: Final = request.resolved.get("compaction") is not None
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
return Target.RESPONSES
return Target.CHAT_COMPLETIONS
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection,
Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry,
}
)
_ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**_FEATURES,
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted,
}
)
def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext:
requested_model: Final = str(request.resolved["model"])
model, provider = _resolved(request)
return RouteContext(
Route.MESSAGES,
provider=_resolved_provider(request),
model=str(request.resolved["model"]),
provider=provider,
model=requested_model,
target=_target(request, model, provider),
features=detectors.detected(features, request),
)
_DISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request),
context=lambda request: _context(request, _FEATURES),
)
_ADISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
context=_context,
bypass=_python_serves,
context=lambda request: _context(request, _ASYNC_FEATURES),
)
@ -118,9 +153,11 @@ def anthropic_messages_handler(
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
python: Final = _PYTHON_MESSAGES
if _FLAG.requested(kwargs):
return anthropic_messages(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_MESSAGES,
native=native_call_hook,
@ -131,7 +168,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
python: Final = _PYTHON_AMESSAGES
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_AMESSAGES,
native=native_call_hook,

View file

@ -16580,6 +16580,24 @@
"source": "https://developers.cloudflare.com/workers-ai/platform/pricing/",
"supports_vision": true
},
"cloudflare/clairic": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cloudflare",
"max_input_tokens": 65536,
"mode": "extraction",
"output_cost_per_token": 0.0,
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
"supports_vision": true
},
"cloudflare/@cf/cloudflare/clairic": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cloudflare",
"max_input_tokens": 65536,
"mode": "extraction",
"output_cost_per_token": 0.0,
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
"supports_vision": true
},
"cloudflare/@cf/meta/llama-2-7b-chat-fp16": {
"input_cost_per_token": 1.923e-06,
"litellm_provider": "cloudflare",
@ -67292,6 +67310,7 @@
"fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": {
"cache_read_input_token_cost": 4.5e-08,
"cache_read_input_token_cost_priority": 5.625e-08,
"deprecation_date": "2026-11-07",
"input_cost_per_token": 2.25e-07,
"input_cost_per_token_priority": 2.8125e-07,
"litellm_provider": "fireworks_ai",
@ -67326,6 +67345,7 @@
"fireworks_ai/glm-5p3-flash-us": {
"cache_read_input_token_cost": 4.5e-08,
"cache_read_input_token_cost_priority": 5.625e-08,
"deprecation_date": "2026-11-07",
"input_cost_per_token": 2.25e-07,
"input_cost_per_token_priority": 2.8125e-07,
"litellm_provider": "fireworks_ai",
@ -75062,12 +75082,10 @@
"gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -75078,7 +75096,6 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -75111,12 +75128,10 @@
"vertex_ai/gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -75127,7 +75142,6 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -76927,5 +76941,25 @@
"supported_output_modalities": [
"text"
]
},
"fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash-1017-us": {
"cache_read_input_token_cost": 9e-09,
"cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
"input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/us-only-serverless",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
}
}

View file

@ -10,6 +10,7 @@ from typing import Literal, TypeAlias
from pydantic import ConfigDict, Field, model_validator
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from litellm.types.llms.base import LiteLLMBaseModel
CredentialSource: TypeAlias = Literal["db", "config"]
@ -101,6 +102,7 @@ class UserConnectionFlowHandle(LiteLLMBaseModel):
user_id: str
credential_name: str
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
device_code: str = Field(repr=False)
interval: int
expires_at: float

View file

@ -9,7 +9,14 @@ from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str, signature
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
native_call,
native_call_hook,
optional_str,
signature,
)
__all__ = ("aocr", "ocr")
@ -39,6 +46,7 @@ def _bind_request(
_OCR: Final = signature(_bind_request)
_FLAG: Final = AsyncFlag("aocr", sync=_OCR)
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
@ -59,7 +67,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.OCR,
request=lambda args, kwargs: _public_request("ocr", args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aocr") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -73,9 +80,11 @@ def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
if _FLAG.requested(kwargs):
return aocr(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=runtime.NO_PYTHON,
binding=NATIVE_OCR,
native=native_call_hook,
@ -85,7 +94,7 @@ def ocr(
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=runtime.NO_PYTHON,
binding=NATIVE_AOCR,
native=native_call_hook,

View file

@ -437,6 +437,7 @@ def llm_passthrough_route(
params: QueryParamTypes | None = None,
cookies: CookieTypes | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
guardrailConfig: object = None,
**kwargs,
) -> (
httpx.Response
@ -547,7 +548,7 @@ def llm_passthrough_route(
_request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None)
headers, signed_json_body = provider_config.sign_request(
headers=headers,
litellm_params=litellm_params_dict,
litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig},
request_data=_request_data,
api_base=str(updated_url),
model=model,

View file

@ -1197,6 +1197,8 @@ def should_strip_caller_authorization(
admission credential — a virtual key, an IdP JWT, or an SSO / OIDC /
session token whose ``api_key`` is ``None`` — and must never reach
the upstream, so it is stripped regardless of the ``api_key`` value.
- **Every other auth-enabled mode**: strip when ``Authorization`` is the
credential admission validated.
"""
if _authorization_is_admission_jwt(raw_headers, user_api_key_auth):
return True
@ -1216,7 +1218,21 @@ def should_strip_caller_authorization(
mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream
)
if not (mcp_server.is_oauth_passthrough or is_delegated_oauth):
return False
from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
master_key,
)
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
admission_ran: Final = not is_no_auth_dev_mode(master_key, proxy_general_settings) or (
user_api_key_auth is not None and user_api_key_auth.authenticated_by_custom_auth
)
return admission_ran and (
_authorization_is_litellm_admission_credential_with_settings(
raw_headers, user_api_key_auth, proxy_general_settings
)
)
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if is_delegated_oauth:
@ -1243,16 +1259,39 @@ def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _authorization_is_litellm_admission_credential(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
)
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
return _authorization_is_litellm_admission_credential_with_settings(
raw_headers, user_api_key_auth, proxy_general_settings
)
def _authorization_is_litellm_admission_credential_with_settings(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
proxy_general_settings: Mapping[str, object],
) -> bool:
"""True when ``Authorization`` carries the LiteLLM key admission validated.
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
That is the case when no usable admission header was sent, or when the client repeated the same key in both headers.
"""
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
custom_header_name: Final[object | None] = proxy_general_settings.get("litellm_key_header_name")
custom_admission_header: Final = (
_raw_header_value(raw_headers, custom_header_name.lower())
if isinstance(custom_header_name, str) and custom_header_name
else None
)
admission_header: Final = custom_admission_header or _raw_header_value(raw_headers, "x-litellm-api-key")
authorization: Final = _raw_header_value(raw_headers, "authorization")
if admission_header and authorization:
return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer")

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -1,4 +1,4 @@
:HL["/litellm-asset-prefix/_next/static/chunks/1kid9zr1--h6y.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/3146e697tym4_.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/0ckv5j_fni3xm.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/06yuul4yvgbd0.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.2bn3s6zvc0dyp.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"tree":{"name":"","param":null,"prefetchHints":4176,"slots":{"children":{"name":"(dashboard)","param":null,"prefetchHints":4192,"slots":{"children":{"name":"__PAGE__","param":null,"prefetchHints":4256,"slots":null}}}}},"staleTime":300,"buildId":"kXnLzJ6ylsRPmgSkCkCKM"}
0:{"tree":{"name":"","param":null,"prefetchHints":4176,"slots":{"children":{"name":"(dashboard)","param":null,"prefetchHints":4192,"slots":{"children":{"name":"__PAGE__","param":null,"prefetchHints":4256,"slots":null}}}}},"staleTime":300,"buildId":"u_MWVLx_WA9i-AM2JdCsp"}

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1 @@
(globalThis.TURBOPACK||(globalThis.TURBOPACK=[])).push(["object"==typeof document?document.currentScript:void 0,516430,t=>{"use strict";var a=t.i(180127);t.s(["ArrowLeftIcon",()=>a.default])},339402,t=>{"use strict";let a=(0,t.i(475254).default)("code",[["path",{d:"m16 18 6-6-6-6",key:"eg8j8"}],["path",{d:"m8 6-6 6 6 6",key:"ppft3o"}]]);t.s(["default",0,a])},837007,t=>{"use strict";var a=t.i(603908);t.s(["PlusIcon",()=>a.default])},440987,t=>{"use strict";var a=t.i(903446);t.s(["SettingsIcon",()=>a.default])},569074,t=>{"use strict";let a=(0,t.i(475254).default)("upload",[["path",{d:"M12 3v12",key:"1x0j5s"}],["path",{d:"m17 8-5-5-5 5",key:"7q97r8"}],["path",{d:"M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4",key:"ih7n3h"}]]);t.s(["Upload",0,a],569074)},868499,t=>{"use strict";var a=t.i(843476);t.s([],558762),t.i(558762);var e=t.i(366250),o=t.i(402820),r=t.i(156736),l=t.i(209793),i=t.i(784324),s=t.i(264951),n=t.i(77173);let d=t.i(313488).DialogTrigger;var u=t.i(974217),g=t.i(325326),c=t.i(301807);let p={modal:!0,disablePointerDismissal:!0,role:"alertdialog"};class f extends g.DialogHandle{constructor(t){super(t??new c.DialogStore(p)),t&&this.store.update(p)}}t.s(["Backdrop",()=>o.DialogBackdrop,"Close",()=>r.DialogClose,"Description",()=>l.DialogDescription,"Handle",0,f,"Popup",()=>i.DialogPopup,"Portal",()=>s.DialogPortal,"Root",0,function(t){return(0,e.useRenderDialogRoot)(t,"alert-dialog")},"Title",()=>n.DialogTitle,"Trigger",0,d,"Viewport",()=>u.DialogViewport,"createHandle",0,function(){return new f}],734604);var m=t.i(734604),m=m,x=t.i(196631),D=t.i(519455);function j({...t}){return(0,a.jsx)(m.Portal,{"data-slot":"alert-dialog-portal",...t})}function v({className:t,...e}){return(0,a.jsx)(m.Backdrop,{"data-slot":"alert-dialog-overlay",className:(0,x.cn)("fixed inset-0 isolate z-popup bg-black/10 duration-100 supports-backdrop-filter:backdrop-blur-xs data-open:animate-in data-open:fade-in-0 data-closed:animate-out data-closed:fade-out-0",t),...e})}t.s(["AlertDialog",0,function({...t}){return(0,a.jsx)(m.Root,{"data-slot":"alert-dialog",...t})},"AlertDialogAction",0,function({className:t,variant:e="default",size:o="default",...r}){return(0,a.jsx)(m.Close,{"data-slot":"alert-dialog-action",className:(0,x.cn)(t),render:(0,a.jsx)(D.Button,{variant:e,size:o}),...r})},"AlertDialogCancel",0,function({className:t,variant:e="outline",size:o="default",...r}){return(0,a.jsx)(m.Close,{"data-slot":"alert-dialog-cancel",className:(0,x.cn)(t),render:(0,a.jsx)(D.Button,{variant:e,size:o}),...r})},"AlertDialogContent",0,function({className:t,size:e="default",...o}){return(0,a.jsxs)(j,{children:[(0,a.jsx)(v,{}),(0,a.jsx)(m.Popup,{"data-slot":"alert-dialog-content","data-size":e,className:(0,x.cn)("group/alert-dialog-content fixed top-1/2 left-1/2 z-popup grid w-full -translate-x-1/2 -translate-y-1/2 gap-6 rounded-xl bg-popover p-6 text-popover-foreground ring-1 ring-foreground/10 duration-100 outline-none data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-lg data-open:animate-in data-open:fade-in-0 data-open:zoom-in-95 data-closed:animate-out data-closed:fade-out-0 data-closed:zoom-out-95",t),...o})]})},"AlertDialogDescription",0,function({className:t,...e}){return(0,a.jsx)(m.Description,{"data-slot":"alert-dialog-description",className:(0,x.cn)("text-sm text-balance text-muted-foreground md:text-pretty *:[a]:underline *:[a]:underline-offset-3 *:[a]:hover:text-foreground",t),...e})},"AlertDialogFooter",0,function({className:t,...e}){return(0,a.jsx)("div",{"data-slot":"alert-dialog-footer",className:(0,x.cn)("flex flex-col-reverse gap-2 group-data-[size=sm]/alert-dialog-content:grid group-data-[size=sm]/alert-dialog-content:grid-cols-2 sm:flex-row sm:justify-end",t),...e})},"AlertDialogHeader",0,function({className:t,...e}){return(0,a.jsx)("div",{"data-slot":"alert-dialog-header",className:(0,x.cn)("grid grid-rows-[auto_1fr] place-items-center gap-1.5 text-center has-data-[slot=alert-dialog-media]:grid-rows-[auto_auto_1fr] has-data-[slot=alert-dialog-media]:gap-x-6 sm:group-data-[size=default]/alert-dialog-content:place-items-start sm:group-data-[size=default]/alert-dialog-content:text-left sm:group-data-[size=default]/alert-dialog-content:has-data-[slot=alert-dialog-media]:grid-rows-[auto_1fr]",t),...e})},"AlertDialogTitle",0,function({className:t,...e}){return(0,a.jsx)(m.Title,{"data-slot":"alert-dialog-title",className:(0,x.cn)("text-lg font-medium sm:group-data-[size=default]/alert-dialog-content:group-has-data-[slot=alert-dialog-media]/alert-dialog-content:col-start-2",t),...e})},"AlertDialogTrigger",0,function({...t}){return(0,a.jsx)(m.Trigger,{"data-slot":"alert-dialog-trigger",...t})}],868499)},841238,t=>{"use strict";let a=t=>"object"==typeof t&&null!==t;t.s(["streamErrorMessage",0,t=>{if(!a(t)||!t.error)return null;let{error:e}=t;return"string"==typeof e?e:a(e)&&"string"==typeof e.message?e.message:JSON.stringify(e)}])}]);

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1 @@
(globalThis.TURBOPACK||(globalThis.TURBOPACK=[])).push(["object"==typeof document?document.currentScript:void 0,180127,e=>{"use strict";let t=(0,e.i(475254).default)("arrow-left",[["path",{d:"m12 19-7-7 7-7",key:"1l729n"}],["path",{d:"M19 12H5",key:"x3x0zl"}]]);e.s(["default",0,t])},871689,e=>{"use strict";var t=e.i(180127);e.s(["ArrowLeft",()=>t.default])},595468,e=>{"use strict";var t=e.i(123287);e.s(["CheckCircle2",()=>t.default])},541071,e=>{"use strict";var t=e.i(373488);e.s(["MoreHorizontal",()=>t.default])},332102,e=>{"use strict";let t=(0,e.i(475254).default)("inbox",[["polyline",{points:"22 12 16 12 14 15 10 15 8 12 2 12",key:"o97t9d"}],["path",{d:"M5.45 5.11 2 12v6a2 2 0 0 0 2 2h16a2 2 0 0 0 2-2v-6l-3.45-6.89A2 2 0 0 0 16.76 4H7.24a2 2 0 0 0-1.79 1.11z",key:"oot6mr"}]]);e.s(["Inbox",0,t],332102)},878894,e=>{"use strict";var t=e.i(582458);e.s(["AlertTriangle",()=>t.default])},512154,e=>{e.q("/litellm-asset-prefix/_next/static/media/bing.3b9zkaag7urkm.png")},764453,e=>{e.q("/litellm-asset-prefix/_next/static/media/dataforseo.1g2jptyl8rcb1.png")},341367,e=>{e.q("/litellm-asset-prefix/_next/static/media/exa_ai.36h3hrkelbgj-.png")},732731,e=>{e.q("/litellm-asset-prefix/_next/static/media/google_pse.3hii8gkiytuod.png")},601739,e=>{e.q("/litellm-asset-prefix/_next/static/media/nimble.0ors74qocyffr.png")},911676,e=>{e.q("/litellm-asset-prefix/_next/static/media/parallel_ai.0jx5g5pf0u355.png")},692745,e=>{e.q("/litellm-asset-prefix/_next/static/media/perplexity.2zhky1a8ufk3x.png")},380084,e=>{e.q("/litellm-asset-prefix/_next/static/media/tavily.15dorlkyzxydf.png")},450240,e=>{"use strict";var t=e.i(843476),i=e.i(286536),s=e.i(77705),a=e.i(271645),l=e.i(950594);let r=a.forwardRef(({className:e,groupClassName:r,disabled:n,...p},d)=>{let[o,u]=a.useState(!1);return(0,t.jsxs)(l.InputGroup,{className:r,children:[(0,t.jsx)(l.InputGroupInput,{...p,ref:d,type:o?"text":"password",disabled:n,className:e}),(0,t.jsx)(l.InputGroupAddon,{align:"inline-end",children:(0,t.jsx)(l.InputGroupButton,{size:"icon-xs",disabled:n,"aria-label":o?"Hide password":"Show password",onClick:()=>u(e=>!e),children:o?(0,t.jsx)(s.EyeOff,{}):(0,t.jsx)(i.Eye,{})})})]})});r.displayName="PasswordInput",e.s(["PasswordInput",0,r])}]);

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1,35 @@
(globalThis.TURBOPACK||(globalThis.TURBOPACK=[])).push(["object"==typeof document?document.currentScript:void 0,972520,e=>{"use strict";let t=(0,e.i(475254).default)("arrow-right",[["path",{d:"M5 12h14",key:"1ays0h"}],["path",{d:"m12 5 7 7-7 7",key:"xquz4c"}]]);e.s(["ArrowRight",0,t],972520)},411929,e=>{"use strict";var t=e.i(843476),r=e.i(271645),s=e.i(972520),a=e.i(174886),o=e.i(519455),n=e.i(515288),i=e.i(624687),l=e.i(571303),d=e.i(602869),c=e.i(417385);let u=({accessToken:e})=>{let[u,m]=(0,r.useState)(`{
"model": "openai/gpt-4o",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "Explain quantum computing in simple terms"
}
],
"temperature": 0.7,
"max_tokens": 500,
"stream": true
}`),[p,h]=(0,r.useState)(""),[x,f]=(0,r.useState)(!1),g=async()=>{f(!0);try{let a;try{a=JSON.parse(u)}catch(e){c.toast.fromError("Invalid JSON in request body"),f(!1);return}let o={call_type:"completion",request_body:a};if(!e){c.toast.fromError("No access token found"),f(!1);return}let n=await (0,d.transformRequestCall)(e,o);if(n.raw_request_api_base&&n.raw_request_body){var t,r,s;let e,a,o=(t=n.raw_request_api_base,r=n.raw_request_body,s=n.raw_request_headers||{},e=JSON.stringify(r,null,2).split("\n").map(e=>` ${e}`).join("\n"),a=Object.entries(s).map(([e,t])=>`-H '${e}: ${t}'`).join(" \\\n "),`curl -X POST \\
${t} \\
${a?`${a} \\
`:""}-H 'Content-Type: application/json' \\
-d '{
${e}
}'`);h(o),c.toast.success("Request transformed successfully")}else{let e="string"==typeof n?n:JSON.stringify(n);h(e),c.toast.info("Transformed request received in unexpected format")}}catch(e){console.error("Error transforming request:",e),c.toast.fromError("Failed to transform request")}finally{f(!1)}};return(0,t.jsxs)("div",{className:"p-2",children:[(0,t.jsx)("h1",{className:"text-lg font-medium text-foreground",children:"Playground"}),(0,t.jsx)("p",{className:"text-sm text-muted-foreground",children:"See how LiteLLM transforms your request for the specified provider."}),(0,t.jsxs)("div",{className:"mt-4 grid grid-cols-1 gap-4 lg:grid-cols-2",children:[(0,t.jsxs)(n.Card,{children:[(0,t.jsxs)(n.CardHeader,{children:[(0,t.jsx)(n.CardTitle,{className:"text-2xl font-bold",children:"Original Request"}),(0,t.jsx)(n.CardDescription,{children:"The request you would send to LiteLLM /chat/completions endpoint."})]}),(0,t.jsx)(n.CardContent,{children:(0,t.jsx)(i.Textarea,{className:"h-72 resize-none p-4 font-mono text-sm field-sizing-fixed",value:u,onChange:e=>m(e.target.value),onKeyDown:e=>{(e.metaKey||e.ctrlKey)&&"Enter"===e.key&&(e.preventDefault(),g())},placeholder:"Press Cmd/Ctrl + Enter to transform"})}),(0,t.jsx)(n.CardFooter,{className:"justify-end",children:(0,t.jsxs)(o.Button,{onClick:g,disabled:x,children:[(0,t.jsx)("span",{children:"Transform"}),x?(0,t.jsx)(l.UiLoadingSpinner,{className:"size-4"}):(0,t.jsx)(s.ArrowRight,{})]})})]}),(0,t.jsxs)(n.Card,{children:[(0,t.jsxs)(n.CardHeader,{children:[(0,t.jsx)(n.CardTitle,{className:"text-2xl font-bold",children:"Transformed Request"}),(0,t.jsx)(n.CardDescription,{children:"How LiteLLM transforms your request for the specified provider."}),(0,t.jsx)("p",{className:"mt-2 text-xs text-muted-foreground",children:"Note: Sensitive headers are not shown."})]}),(0,t.jsx)(n.CardContent,{children:(0,t.jsxs)("div",{className:"relative rounded-md bg-muted",children:[(0,t.jsx)("pre",{className:"h-72 overflow-auto p-4 font-mono text-sm",children:p||`curl -X POST \\
https://api.openai.com/v1/chat/completions \\
-H 'Authorization: Bearer sk-xxx' \\
-H 'Content-Type: application/json' \\
-d '{
"model": "gpt-4",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
}
],
"temperature": 0.7
}'`}),(0,t.jsx)(o.Button,{variant:"ghost",size:"icon-sm","aria-label":"Copy to clipboard",className:"absolute top-2 right-2",onClick:()=>{navigator.clipboard.writeText(p||""),c.toast.success("Copied to clipboard")},children:(0,t.jsx)(a.Copy,{})})]})})]})]}),(0,t.jsx)("div",{className:"mt-4 text-right",children:(0,t.jsxs)("p",{className:"text-sm text-muted-foreground",children:["Found an error? File an issue"," ",(0,t.jsx)("a",{className:"underline underline-offset-4",href:"https://github.com/BerriAI/litellm/issues",target:"_blank",rel:"noopener noreferrer",children:"here"}),"."]})})]})};var m=e.i(135214);e.s(["default",0,function(){let{accessToken:e}=(0,m.default)();return(0,t.jsx)(u,{accessToken:e})}],411929)}]);

View file

@ -0,0 +1 @@
(globalThis.TURBOPACK||(globalThis.TURBOPACK=[])).push(["object"==typeof document?document.currentScript:void 0,75533,t=>{"use strict";var e=t.i(843476),u=t.i(135214),i=t.i(936578),l=t.i(601964);t.s(["default",0,function(){let{isLoading:t,isAuthorized:o}=(0,u.default)();return t||!o?(0,e.jsx)(i.default,{}):(0,e.jsx)(l.default,{})}])}]);

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