mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge remote-tracking branch 'berri/litellm_internal_staging' into litellm_mistral_ocr_batches
This commit is contained in:
commit
1d911c766f
329 changed files with 24922 additions and 2730 deletions
4
.github/workflows/test-linting.yml
vendored
4
.github/workflows/test-linting.yml
vendored
|
|
@ -113,7 +113,7 @@ jobs:
|
|||
- name: Check ruff format
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- ':(glob)litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
|
||||
echo "No changed litellm Python files to check with ruff format."
|
||||
exit 0
|
||||
|
|
@ -172,7 +172,7 @@ jobs:
|
|||
- name: Check tests/e2e basedpyright (zero errors)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
|
||||
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' | grep -q .; then
|
||||
uv run --no-sync basedpyright tests/e2e
|
||||
else
|
||||
echo "No changed tests/e2e Python files; skipping."
|
||||
|
|
|
|||
4
Makefile
4
Makefile
|
|
@ -150,8 +150,8 @@ lint-install:
|
|||
# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
|
||||
# only the litellm Python files changed vs the base are checked, so a pre-existing
|
||||
# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
|
||||
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
|
||||
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
|
||||
# recursively, so 'litellm/*.py' covers top-level files and nested modules alike,
|
||||
# the same set CI's ':(glob)litellm/**/*.py' selects.
|
||||
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
||||
@base_ref=$$($(RESOLVE_BASE)) && \
|
||||
changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \
|
||||
|
|
|
|||
|
|
@ -0,0 +1,15 @@
|
|||
-- DropForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
-- AddForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
|
@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
|
|||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String?
|
||||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
|
|
|
|||
|
|
@ -45,9 +45,11 @@ from typing import (
|
|||
TYPE_CHECKING,
|
||||
Union,
|
||||
)
|
||||
from collections.abc import Mapping
|
||||
from litellm.types.integrations.datadog import DatadogInitParams
|
||||
from litellm.types.integrations.newrelic import NewRelicInitParams
|
||||
from litellm.litellm_core_utils.core_helpers import drop_params_env_flag
|
||||
from litellm.types.integrations.pointfive import PointFiveInitParams
|
||||
from litellm._logging import (
|
||||
set_verbose,
|
||||
_turn_on_debug,
|
||||
|
|
@ -154,6 +156,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"smtp_email",
|
||||
"deepeval",
|
||||
"s3_v2",
|
||||
"pointfive",
|
||||
"aws_sqs",
|
||||
"vector_store_pre_call_hook",
|
||||
"dotprompt",
|
||||
|
|
@ -439,6 +442,7 @@ s3_audit_callback_params: Optional[Dict] = None
|
|||
datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None
|
||||
datadog_params: Optional[Union[DatadogInitParams, Dict]] = None
|
||||
newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None
|
||||
pointfive_params: Optional[Union[PointFiveInitParams, Mapping[str, object]]] = None
|
||||
aws_sqs_callback_params: Optional[Dict] = None
|
||||
generic_logger_headers: Optional[Dict] = None
|
||||
default_key_generate_params: Optional[Dict] = None
|
||||
|
|
@ -475,6 +479,7 @@ prometheus_metrics_config: Optional[List] = None
|
|||
prometheus_exclude_metrics: Optional[List[str]] = None
|
||||
prometheus_exclude_labels: Optional[List[str]] = None
|
||||
prometheus_emit_stream_label: bool = False
|
||||
prometheus_emit_input_sequence_length_label: bool = False
|
||||
prometheus_deployment_and_latency_caller_identity: Literal[
|
||||
"api_key_alias",
|
||||
"user_email",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
BedrockAgentCoreA2ATransformation,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
|
|
@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
Returns:
|
||||
A2A JSON-RPC response dict from the AgentCore agent
|
||||
"""
|
||||
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
url, headers, body = await run_aws_signing(
|
||||
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
Yields:
|
||||
A2A streaming response events from the AgentCore agent
|
||||
"""
|
||||
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
url, headers, body = await run_aws_signing(
|
||||
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ import litellm
|
|||
from litellm import ModelResponse
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
responses_reasoning_item_from_thinking_blocks,
|
||||
responses_reasoning_items_from_thinking_blocks,
|
||||
)
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.bridges.completion_transformation import (
|
||||
|
|
@ -129,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]:
|
|||
return stored
|
||||
raw_blocks: Final = msg.get("thinking_blocks") or ()
|
||||
blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json
|
||||
from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks)
|
||||
return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload
|
||||
replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks)
|
||||
return [dict(item) for item in replayed] # mutable-ok: API message payload
|
||||
|
||||
|
||||
def _build_reasoning_item(
|
||||
|
|
@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
|
|||
provider_specific_fields: Mapping[str, object]
|
||||
|
||||
|
||||
def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
|
||||
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
|
||||
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
|
||||
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
|
||||
string payload in ``input`` rather than ``arguments``; both map to
|
||||
|
|
@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
# Tool calls accumulate into the single trailing tool_calls choice
|
||||
# like the typed branches above; a choice per call would hide every
|
||||
# call after choices[0] from chat clients
|
||||
accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
|
||||
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
|
||||
tool_call_index += 1
|
||||
elif handle_raw_dict_callback is not None:
|
||||
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
|
||||
|
|
@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
# New output item added
|
||||
output_item = parsed_chunk.get("item", {})
|
||||
if output_item.get("type") in ("function_call", "custom_tool_call"):
|
||||
converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
|
||||
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
|
||||
provider_specific_fields: Final = converted.get("provider_specific_fields")
|
||||
|
||||
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
|
||||
|
|
@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
index=0,
|
||||
delta=Delta(
|
||||
tool_calls=(
|
||||
_tool_call_dict_from_output_item(
|
||||
tool_call_dict_from_output_item(
|
||||
output_item, parsed_chunk.get("output_index", 0)
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -398,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
|
|||
minimum=1,
|
||||
maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
|
||||
)
|
||||
TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range(
|
||||
"TOKEN_COUNTER_MAX_EXACT_CHARS",
|
||||
default=4_000_000,
|
||||
minimum=1,
|
||||
maximum=1_000_000_000,
|
||||
)
|
||||
TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range(
|
||||
"TOKEN_COUNTER_MAX_CONCURRENT_COUNTS",
|
||||
default=4,
|
||||
minimum=1,
|
||||
maximum=256,
|
||||
)
|
||||
MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
|
||||
MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
|
||||
OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
|
||||
|
|
@ -570,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
|
|||
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
|
||||
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
|
||||
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
|
||||
AWS_SIGNING_MAX_THREADS: Final = 16
|
||||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
|
||||
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
|
||||
)
|
||||
|
|
@ -1669,6 +1682,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT
|
|||
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
|
||||
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
|
||||
PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25")))
|
||||
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
|
||||
RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
|
||||
|
|
|
|||
|
|
@ -54,18 +54,22 @@ def missing_streamable_http_client_error() -> ImportError:
|
|||
)
|
||||
|
||||
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import (
|
||||
METHOD_NOT_FOUND,
|
||||
ClientResult,
|
||||
GetPromptRequestParams,
|
||||
GetPromptResult,
|
||||
ListPromptsResult,
|
||||
ListResourcesResult,
|
||||
ListResourceTemplatesResult,
|
||||
Prompt,
|
||||
ResourceTemplate,
|
||||
ServerNotification,
|
||||
ServerRequest,
|
||||
TextContent,
|
||||
)
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
||||
|
|
@ -777,8 +781,19 @@ class MCPClient:
|
|||
"""List available prompts from the server."""
|
||||
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_prompts_operation(session: ClientSession):
|
||||
return await session.list_prompts()
|
||||
async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult:
|
||||
capabilities: Final = session.get_server_capabilities()
|
||||
if capabilities is not None and capabilities.prompts is None:
|
||||
return ListPromptsResult(prompts=[])
|
||||
try:
|
||||
return await session.list_prompts()
|
||||
except McpError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
verbose_logger.debug(
|
||||
"MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error
|
||||
)
|
||||
return ListPromptsResult(prompts=[])
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_prompts_operation)
|
||||
|
|
@ -854,8 +869,19 @@ class MCPClient:
|
|||
"""List available resources from the server."""
|
||||
verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_resources_operation(session: ClientSession):
|
||||
return await session.list_resources()
|
||||
async def _list_resources_operation(session: ClientSession) -> ListResourcesResult:
|
||||
capabilities: Final = session.get_server_capabilities()
|
||||
if capabilities is not None and capabilities.resources is None:
|
||||
return ListResourcesResult(resources=[])
|
||||
try:
|
||||
return await session.list_resources()
|
||||
except McpError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
verbose_logger.debug(
|
||||
"MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error
|
||||
)
|
||||
return ListResourcesResult(resources=[])
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_resources_operation)
|
||||
|
|
@ -890,8 +916,19 @@ class MCPClient:
|
|||
"""List available resource templates from the server."""
|
||||
verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_resource_templates_operation(session: ClientSession):
|
||||
return await session.list_resource_templates()
|
||||
async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult:
|
||||
capabilities: Final = session.get_server_capabilities()
|
||||
if capabilities is not None and capabilities.resources is None:
|
||||
return ListResourceTemplatesResult(resourceTemplates=[])
|
||||
try:
|
||||
return await session.list_resource_templates()
|
||||
except McpError as error:
|
||||
if error.error.code != METHOD_NOT_FOUND:
|
||||
raise
|
||||
verbose_logger.debug(
|
||||
"MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error
|
||||
)
|
||||
return ListResourceTemplatesResult(resourceTemplates=[])
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_resource_templates_operation)
|
||||
|
|
|
|||
|
|
@ -378,6 +378,27 @@
|
|||
},
|
||||
"description": "OpenTelemetry Logging Integration"
|
||||
},
|
||||
{
|
||||
"id": "pointfive",
|
||||
"displayName": "PointFive",
|
||||
"logo": "pointfive.png",
|
||||
"supports_key_team_logging": false,
|
||||
"dynamic_params": {
|
||||
"POINTFIVE_API_KEY": {
|
||||
"type": "password",
|
||||
"ui_name": "API Key",
|
||||
"description": "PointFive API key, used to request an upload url for each batch of logs",
|
||||
"required": true
|
||||
},
|
||||
"POINTFIVE_API_URL": {
|
||||
"type": "text",
|
||||
"ui_name": "API URL",
|
||||
"description": "PointFive API endpoint. Leave blank to use https://api.pointfive.co/api/v1/ingestion",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "PointFive Logging Integration"
|
||||
},
|
||||
{
|
||||
"id": "s3",
|
||||
"displayName": "S3",
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
# On success, logs events to Langfuse
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from datetime import datetime
|
||||
|
|
@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None:
|
|||
return value if isinstance(value, dict) else None
|
||||
|
||||
|
||||
def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]:
|
||||
"""Header pairs with the key type widened back to what a caller-supplied dict can actually hold."""
|
||||
return mapping.items()
|
||||
|
||||
|
||||
def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool:
|
||||
if not isinstance(trace_id, str) or not isinstance(session_id, str):
|
||||
return False
|
||||
request: Final = _object_mapping(proxy_server_request)
|
||||
raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None
|
||||
if raw_headers is None:
|
||||
return False
|
||||
headers: Final = MappingProxyType(
|
||||
{key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)}
|
||||
)
|
||||
if headers.get("x-litellm-trace-id"):
|
||||
return False
|
||||
if headers.get("langfuse_trace_id") is not None:
|
||||
return False
|
||||
if trace_id != session_id and headers.get("langfuse_session_id") != session_id:
|
||||
return False
|
||||
if headers.get("x-litellm-session-id") == trace_id:
|
||||
return True
|
||||
if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None:
|
||||
return False
|
||||
user_agent: Final = headers.get("user-agent")
|
||||
codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None
|
||||
return any(
|
||||
value == trace_id
|
||||
and (
|
||||
key == "x-session-id"
|
||||
or re.fullmatch(r"x-.+-session-id", key) is not None
|
||||
or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id"))
|
||||
)
|
||||
for key, value in headers.items()
|
||||
)
|
||||
|
||||
|
||||
class _UsageObject(Protocol):
|
||||
"""Token-count surface the Langfuse logger reads off a response usage payload."""
|
||||
|
||||
|
|
@ -609,6 +648,18 @@ class LangFuseLogger:
|
|||
# This allows continuing an existing trace while still returning the correct trace_id
|
||||
if existing_trace_id is not None:
|
||||
trace_id = existing_trace_id
|
||||
resolved_trace_id: Final = (
|
||||
litellm_call_id or trace_id
|
||||
if existing_trace_id is None
|
||||
and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request"))
|
||||
else trace_id
|
||||
)
|
||||
if resolved_trace_id != trace_id:
|
||||
verbose_logger.debug(
|
||||
"Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace",
|
||||
trace_id,
|
||||
resolved_trace_id,
|
||||
)
|
||||
requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
|
||||
update_trace_keys: Final = (
|
||||
requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else ()
|
||||
|
|
@ -663,7 +714,7 @@ class LangFuseLogger:
|
|||
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
|
||||
else: # don't overwrite an existing trace
|
||||
trace_params = {
|
||||
"id": trace_id,
|
||||
"id": resolved_trace_id,
|
||||
"name": trace_name,
|
||||
"session_id": session_id,
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
|
|
@ -845,13 +896,13 @@ class LangFuseLogger:
|
|||
# Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value
|
||||
# to match expected test behavior
|
||||
if hasattr(generation_client, "trace_id") and generation_client.trace_id:
|
||||
if generation_client.trace_id != trace_id:
|
||||
if generation_client.trace_id != resolved_trace_id:
|
||||
verbose_logger.warning(
|
||||
"Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.",
|
||||
trace_id,
|
||||
resolved_trace_id,
|
||||
generation_client.trace_id,
|
||||
)
|
||||
return trace_id, generation_id
|
||||
return resolved_trace_id, generation_id
|
||||
except Exception:
|
||||
verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc())
|
||||
return None, None
|
||||
|
|
|
|||
5
litellm/integrations/pointfive/__init__.py
Normal file
5
litellm/integrations/pointfive/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""PointFive logging integration for LiteLLM."""
|
||||
|
||||
from litellm.integrations.pointfive.logger import PointFiveLogger
|
||||
|
||||
__all__ = ("PointFiveLogger",)
|
||||
304
litellm/integrations/pointfive/logger.py
Normal file
304
litellm/integrations/pointfive/logger.py
Normal file
|
|
@ -0,0 +1,304 @@
|
|||
"""
|
||||
PointFive logging integration.
|
||||
|
||||
Buffers ``StandardLoggingPayload`` records and ships each flush as one gzipped
|
||||
newline-delimited JSON object, rather than one object per request. Uploads go through a
|
||||
presigned URL issued by the PointFive API, so the proxy needs no cloud credentials and
|
||||
runs unchanged wherever it is hosted.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.integrations.pointfive.payload import chunk_lines, encode_lines, serialize_records
|
||||
from litellm.integrations.pointfive.upload_client import PointFiveUploadClient, PointFiveUploadError
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
redacted_standard_logging_payload,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, httpxSpecialProvider
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.integrations.pointfive import DEFAULT_API_URL, PointFiveInitParams, PointFiveUploadFailure
|
||||
|
||||
_ENV_REFERENCE_PREFIX: Final = "os.environ/"
|
||||
|
||||
|
||||
def _resolved_secret(value: str | None) -> str | None:
|
||||
"""
|
||||
Resolve a config value that may name a secret, in any shape the secret manager accepts.
|
||||
|
||||
A reference that resolves to nothing stays unresolved rather than falling back to its own
|
||||
text, so an unset ``os.environ/NAME`` reports a missing key instead of being sent as one.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
resolved: Final = get_secret_str(value)
|
||||
if resolved:
|
||||
return resolved
|
||||
return None if value.startswith(_ENV_REFERENCE_PREFIX) else value
|
||||
|
||||
|
||||
def _configured_params() -> PointFiveInitParams:
|
||||
"""Read ``litellm.pointfive_params``, validating a raw config dict on the way through."""
|
||||
configured: Final = litellm.pointfive_params
|
||||
if isinstance(configured, PointFiveInitParams):
|
||||
return configured
|
||||
if isinstance(configured, Mapping):
|
||||
return PointFiveInitParams.model_validate(configured)
|
||||
return PointFiveInitParams()
|
||||
|
||||
|
||||
def _resolved_api_key(params: PointFiveInitParams) -> str | None:
|
||||
"""Prefer the configured key, falling back to the environment the proxy UI writes."""
|
||||
return _resolved_secret(params.api_key) or get_secret_str("POINTFIVE_API_KEY")
|
||||
|
||||
|
||||
def _resolved_api_url(params: PointFiveInitParams) -> str:
|
||||
"""Prefer the configured url, then the environment, then the public endpoint."""
|
||||
return _resolved_secret(params.api_url) or get_secret_str("POINTFIVE_API_URL") or DEFAULT_API_URL
|
||||
|
||||
|
||||
def _upload_client_for(params: PointFiveInitParams) -> PointFiveUploadClient:
|
||||
"""
|
||||
Build an upload client for the key and url configured right now.
|
||||
|
||||
Resolved per call rather than kept: the proxy ui writes new values into the
|
||||
environment of a running proxy, and reading them once would need a restart to take
|
||||
effect. ``get_async_httpx_client`` is cached, so this reuses the same connections.
|
||||
"""
|
||||
api_key: Final = _resolved_api_key(params)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"pointfive logging requires an api key. Set POINTFIVE_API_KEY, or "
|
||||
"litellm_settings.pointfive_params.api_key in config.yaml"
|
||||
)
|
||||
return PointFiveUploadClient(
|
||||
api_key=api_key,
|
||||
api_url=_resolved_api_url(params),
|
||||
http_client=get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback),
|
||||
max_retries=params.max_upload_retries,
|
||||
)
|
||||
|
||||
|
||||
class PointFiveLogger(CustomBatchLogger):
|
||||
"""Batching callback that ships LiteLLM request logs to PointFive."""
|
||||
|
||||
preserve_events_added_during_flush = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params: PointFiveInitParams | None = None,
|
||||
upload_client: PointFiveUploadClient | None = None,
|
||||
start_periodic_flush: bool = True,
|
||||
) -> None:
|
||||
resolved: Final = params if params is not None else _configured_params()
|
||||
self.max_batch_bytes: Final = resolved.max_batch_bytes
|
||||
self.params: Final = resolved
|
||||
self.given_upload_client: Final = upload_client
|
||||
if upload_client is None:
|
||||
_upload_client_for(resolved) # refuse to start without a key, rather than at the first flush
|
||||
super().__init__(
|
||||
flush_lock=asyncio.Lock(),
|
||||
batch_size=resolved.batch_size,
|
||||
flush_interval=resolved.flush_interval,
|
||||
turn_off_message_logging=bool(resolved.turn_off_message_logging),
|
||||
)
|
||||
self._flushing: bool = False
|
||||
self._batch_flush_task: asyncio.Task[None] | None = None
|
||||
self._periodic_flush_task: asyncio.Task[None] | None = (
|
||||
self._start_periodic_flush_task() if start_periodic_flush else None
|
||||
)
|
||||
|
||||
@property
|
||||
def upload_client(self) -> PointFiveUploadClient:
|
||||
"""The client for the currently configured key and url, so a ui edit needs no restart."""
|
||||
if self.given_upload_client is not None:
|
||||
return self.given_upload_client
|
||||
return _upload_client_for(self.params)
|
||||
|
||||
def _start_periodic_flush_task(self) -> asyncio.Task[None] | None:
|
||||
"""Start the periodic flush only once an event loop is actually running."""
|
||||
try:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
return loop.create_task(self.periodic_flush())
|
||||
|
||||
def _start_batch_flush_task(self) -> None:
|
||||
"""
|
||||
Upload a full batch in the background, so no request waits on PointFive.
|
||||
|
||||
Awaiting it here put the upload, its retries and their backoff on the caller's
|
||||
path, and a hung api held a response open for as long as the attempts took.
|
||||
"""
|
||||
if self._batch_flush_task is not None and not self._batch_flush_task.done():
|
||||
return
|
||||
try:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return
|
||||
self._batch_flush_task = loop.create_task(self.flush_queue(skip_if_flushing=True))
|
||||
|
||||
def _flush_task_is_alive(self) -> bool:
|
||||
"""A task whose loop has been closed never runs again, yet never reports itself done."""
|
||||
task: Final = self._periodic_flush_task
|
||||
return task is not None and not task.done() and not task.get_loop().is_closed()
|
||||
|
||||
async def periodic_flush(self) -> None:
|
||||
"""
|
||||
Report in straight away, then flush on the interval as usual.
|
||||
|
||||
The inherited loop sleeps first, so a proxy that has just loaded the callback says
|
||||
nothing for a whole interval, five minutes by default. PointFive shows the integration
|
||||
as still waiting for its first call for all that time, which reads as a broken setup
|
||||
rather than an idle one. An empty queue makes this first cycle a ping, so a proxy with
|
||||
no traffic yet announces itself without uploading an object that holds no records.
|
||||
"""
|
||||
await self.flush_queue(skip_if_flushing=True)
|
||||
await super().periodic_flush()
|
||||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
response_obj: object,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
await self._enqueue(kwargs)
|
||||
|
||||
async def async_log_failure_event(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
response_obj: object,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
await self._enqueue(kwargs)
|
||||
|
||||
async def _enqueue(self, kwargs: Mapping[str, object]) -> None:
|
||||
"""Buffer one record, flushing early once the batch threshold is reached."""
|
||||
try:
|
||||
if not self._flush_task_is_alive():
|
||||
self._periodic_flush_task = self._start_periodic_flush_task()
|
||||
|
||||
record: Final = self._record_for(kwargs)
|
||||
if record is None:
|
||||
verbose_logger.debug("pointfive: event carried no standard_logging_object, skipping")
|
||||
return
|
||||
|
||||
self.log_queue.append(record)
|
||||
self._drop_overflow()
|
||||
if len(self.log_queue) >= self.batch_size:
|
||||
self._start_batch_flush_task()
|
||||
except Exception: # noqa: BLE001 # logging must never break the request path
|
||||
verbose_logger.exception("pointfive: failed to queue an event")
|
||||
|
||||
def _record_for(self, kwargs: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
"""
|
||||
The record to buffer, redacted the way the framework would have redacted it.
|
||||
|
||||
A success reaches a callback already redacted, an async failure does not, so both
|
||||
the excluded-field list and this callback's own setting are applied here, then the
|
||||
global, per-request and header settings that only the framework's predicate knows.
|
||||
"""
|
||||
details: Final = self.redact_standard_logging_payload_from_model_call_details(
|
||||
dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict
|
||||
)
|
||||
payload: Final = details.get("standard_logging_object")
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
if should_redact_message_logging(details):
|
||||
return redacted_standard_logging_payload(payload)
|
||||
return payload
|
||||
|
||||
def _drop_overflow(self) -> None:
|
||||
"""
|
||||
Hold the queue to its cap as records arrive, not only after a flush has failed.
|
||||
|
||||
Never while a flush is running: it holds a snapshot taken by length, and trimming
|
||||
the front underneath it would make the post-flush drain remove records that arrived
|
||||
during the upload and were never sent. The next arrival after the flush trims.
|
||||
"""
|
||||
if self._flushing:
|
||||
return
|
||||
overflow: Final = len(self.log_queue) - self.max_queue_size
|
||||
if overflow <= 0:
|
||||
return
|
||||
del self.log_queue[:overflow]
|
||||
verbose_logger.warning("pointfive: queue over %s records, dropped %s oldest", self.max_queue_size, overflow)
|
||||
|
||||
async def flush_queue(self, skip_if_flushing: bool = False) -> None:
|
||||
"""
|
||||
Flush as usual, or report liveness when there is nothing to send.
|
||||
|
||||
``CustomBatchLogger`` skips an empty queue entirely, so without this an idle proxy
|
||||
would look identical to a dead one.
|
||||
|
||||
``skip_if_flushing`` is what a full batch, and the loop's opening cycle, pass. Uploading one takes seconds, and
|
||||
every event arriving meanwhile crosses the threshold too, so each would queue on the
|
||||
flush lock and then ship the handful of records left behind it. That turns one burst
|
||||
into a stream of tiny objects, which is what batching exists to avoid. The running
|
||||
flush already carries what is queued, and the interval catches whatever it missed.
|
||||
"""
|
||||
if not self.log_queue:
|
||||
await self._ping()
|
||||
return
|
||||
if skip_if_flushing and self._flushing:
|
||||
return
|
||||
|
||||
self._flushing = True
|
||||
try:
|
||||
await super().flush_queue()
|
||||
finally:
|
||||
self._flushing = False
|
||||
|
||||
async def async_health_check(self) -> IntegrationHealthCheckStatus:
|
||||
"""Answer the proxy ui test button by asking the api whether it accepts this key."""
|
||||
try:
|
||||
failure: Final = await self.upload_client.ping()
|
||||
except ValueError as missing_key:
|
||||
return IntegrationHealthCheckStatus(status="unhealthy", error_message=str(missing_key))
|
||||
if failure is not None:
|
||||
return IntegrationHealthCheckStatus(status="unhealthy", error_message=failure.detail)
|
||||
return IntegrationHealthCheckStatus(status="healthy", error_message=None)
|
||||
|
||||
async def _ping(self) -> None:
|
||||
"""Report liveness, never failing the flush over it."""
|
||||
try:
|
||||
failure: Final = await self.upload_client.ping()
|
||||
except ValueError as missing_key:
|
||||
verbose_logger.warning("pointfive: liveness ping skipped, %s", missing_key)
|
||||
return
|
||||
if failure is not None:
|
||||
verbose_logger.warning("pointfive: liveness ping failed, %s", failure.detail)
|
||||
|
||||
async def async_send_batch(self) -> None:
|
||||
"""
|
||||
Upload everything queued, split into objects of at most ``max_batch_bytes``.
|
||||
|
||||
A retryable failure propagates so ``CustomBatchLogger`` keeps the rest of the batch
|
||||
for the next flush; the records already shipped or already refused leave the queue
|
||||
first, so a retry re-sends at most the object that failed. A rejection the server
|
||||
will refuse again drops that object, since holding it would block every record
|
||||
queued behind it.
|
||||
"""
|
||||
pending: Final = tuple(self.log_queue)
|
||||
if not pending:
|
||||
return
|
||||
|
||||
client: Final = self.upload_client
|
||||
chunks: Final = chunk_lines(serialize_records(pending), self.max_batch_bytes)
|
||||
for index, chunk in enumerate(chunks):
|
||||
outcome = await client.upload(await encode_lines(chunk))
|
||||
if not isinstance(outcome, PointFiveUploadFailure):
|
||||
continue
|
||||
if outcome.retryable:
|
||||
del self.log_queue[: sum(len(shipped) for shipped in chunks[:index])]
|
||||
raise PointFiveUploadError(outcome.detail)
|
||||
verbose_logger.error("pointfive: dropping %s records, %s", len(chunk), outcome.detail)
|
||||
53
litellm/integrations/pointfive/payload.py
Normal file
53
litellm/integrations/pointfive/payload.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
"""Turns buffered log records into the gzipped NDJSON objects that get uploaded."""
|
||||
|
||||
import gzip
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from itertools import accumulate, groupby, islice
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
_NEWLINE_BYTES: Final = 1
|
||||
|
||||
|
||||
def serialize_records(records: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
"""Serialize each record to one JSON line."""
|
||||
return tuple(safe_dumps(record) for record in records)
|
||||
|
||||
|
||||
def _encoded_size(line: str) -> int:
|
||||
return len(line.encode("utf-8")) + _NEWLINE_BYTES
|
||||
|
||||
|
||||
def _object_indices(sizes: Sequence[int], max_bytes: int) -> Iterator[int]:
|
||||
"""Number each line with the object it belongs to, opening a new one on overflow."""
|
||||
|
||||
def advance(state: tuple[int, int], size: int) -> tuple[int, int]:
|
||||
index, used = state
|
||||
return (index + 1, size) if used and used + size > max_bytes else (index, used + size)
|
||||
|
||||
return (index for index, _ in islice(accumulate(sizes, advance, initial=(0, 0)), 1, None))
|
||||
|
||||
|
||||
def chunk_lines(lines: Sequence[str], max_bytes: int) -> tuple[tuple[str, ...], ...]:
|
||||
"""
|
||||
Group serialized lines into objects of at most ``max_bytes`` uncompressed.
|
||||
|
||||
A line above the bound on its own still becomes its own object. A record cannot be
|
||||
split, and holding it back would stall every record queued behind it.
|
||||
"""
|
||||
sizes: Final = tuple(_encoded_size(line) for line in lines)
|
||||
numbered: Final = zip(_object_indices(sizes, max_bytes), lines, strict=True)
|
||||
return tuple(tuple(line for _, line in group) for _, group in groupby(numbered, lambda pair: pair[0]))
|
||||
|
||||
|
||||
async def encode_lines(lines: Sequence[str]) -> bytes:
|
||||
"""
|
||||
Join lines as NDJSON and gzip them off the event loop.
|
||||
|
||||
An object can be several megabytes, and compressing that inline would block the
|
||||
proxy for as long as it takes.
|
||||
"""
|
||||
compress: Final = asyncify(gzip.compress)
|
||||
return await compress("\n".join(lines).encode("utf-8"))
|
||||
194
litellm/integrations/pointfive/upload_client.py
Normal file
194
litellm/integrations/pointfive/upload_client.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
"""
|
||||
Uploads one batch to PointFive through a presigned URL.
|
||||
|
||||
The proxy holds no cloud credentials. For every batch it asks the PointFive API for a
|
||||
single-use presigned URL and PUTs the bytes there, so the same plugin runs unchanged on
|
||||
AWS, GCP, Azure or on-prem. The server picks the object key, so the proxy never chooses
|
||||
where its data lands.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.integrations.pointfive import (
|
||||
RETRYABLE_UPLOAD_STATUS_CODES,
|
||||
PointFiveUploadFailure,
|
||||
PointFiveUploadTarget,
|
||||
)
|
||||
|
||||
UPLOAD_KIND: Final = "LITELLM"
|
||||
UPLOAD_URL_PATH: Final = "/upload-url"
|
||||
PING_PATH: Final = "/ping"
|
||||
PUT_HEADERS: Final = MappingProxyType({"Content-Type": "application/x-ndjson", "Content-Encoding": "gzip"})
|
||||
|
||||
|
||||
class _PresignRequest(BaseModel):
|
||||
kind: str = UPLOAD_KIND
|
||||
byte_count: int = Field(serialization_alias="byteCount")
|
||||
|
||||
|
||||
class _PingRequest(BaseModel):
|
||||
kind: str = UPLOAD_KIND
|
||||
|
||||
|
||||
class _TargetPayload(BaseModel):
|
||||
upload_url: str = Field(alias="uploadUrl")
|
||||
object_key: str = Field(alias="objectKey")
|
||||
|
||||
|
||||
class _ErrorPayload(BaseModel):
|
||||
error: str = ""
|
||||
|
||||
|
||||
class PointFiveUploadError(Exception):
|
||||
"""A batch could not be uploaded and the failure is worth retrying."""
|
||||
|
||||
|
||||
def _failure_for(response: httpx.Response, what: str) -> PointFiveUploadFailure:
|
||||
detail: Final = f"{what} returned {response.status_code}"
|
||||
reason: Final = _refusal_reason(response.text)
|
||||
return PointFiveUploadFailure(
|
||||
f"{detail}, {reason}" if reason else detail,
|
||||
retryable=response.status_code in RETRYABLE_UPLOAD_STATUS_CODES,
|
||||
)
|
||||
|
||||
|
||||
def _refusal_reason(body: str) -> str:
|
||||
try:
|
||||
return _ErrorPayload.model_validate_json(body).error
|
||||
except ValidationError:
|
||||
return ""
|
||||
|
||||
|
||||
def _parse_target(body: str) -> PointFiveUploadTarget | PointFiveUploadFailure:
|
||||
try:
|
||||
target: Final = _TargetPayload.model_validate_json(body)
|
||||
except ValidationError:
|
||||
return PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False)
|
||||
return PointFiveUploadTarget(upload_url=target.upload_url, object_key=target.object_key)
|
||||
|
||||
|
||||
class PointFiveUploadClient:
|
||||
"""Presigns and uploads one batch at a time."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
api_url: str,
|
||||
http_client: AsyncHTTPHandler,
|
||||
max_retries: int,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
validate_upload_url: Callable[[str], tuple[str, str]] = validate_url,
|
||||
) -> None:
|
||||
self.api_key: Final = api_key
|
||||
self.api_url: Final = api_url.rstrip("/")
|
||||
self.http_client: Final = http_client
|
||||
self.max_retries: Final = max_retries
|
||||
self.sleep: Final = sleep
|
||||
self.validate_upload_url: Final = validate_upload_url
|
||||
|
||||
async def upload(self, body: bytes) -> str | PointFiveUploadFailure:
|
||||
"""
|
||||
Upload one gzipped batch, returning the object key it landed at.
|
||||
|
||||
Every attempt presigns again, so a retry never reuses a URL that has expired or
|
||||
has already been consumed.
|
||||
"""
|
||||
for attempt in range(self.max_retries):
|
||||
match await self._upload_once(body):
|
||||
case PointFiveUploadFailure(retryable=True) as failure:
|
||||
if attempt + 1 >= self.max_retries:
|
||||
return PointFiveUploadFailure(
|
||||
f"{failure.detail}, gave up after {self.max_retries} attempts", retryable=True
|
||||
)
|
||||
await self.sleep(float(1 << attempt))
|
||||
case outcome:
|
||||
return outcome
|
||||
return PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False)
|
||||
|
||||
async def _upload_once(self, body: bytes) -> str | PointFiveUploadFailure:
|
||||
target: Final = await self._presign(len(body))
|
||||
if isinstance(target, PointFiveUploadFailure):
|
||||
return target
|
||||
|
||||
rejection: Final = await self._put(target, body)
|
||||
if rejection is not None:
|
||||
return rejection
|
||||
|
||||
verbose_logger.debug("pointfive: uploaded %s gzipped bytes to %s", len(body), target.object_key)
|
||||
return target.object_key
|
||||
|
||||
async def ping(self) -> PointFiveUploadFailure | None:
|
||||
"""Report that the proxy is alive when it has nothing to upload."""
|
||||
body: Final = await self._post(PING_PATH, _PingRequest())
|
||||
if isinstance(body, PointFiveUploadFailure):
|
||||
return body
|
||||
return None
|
||||
|
||||
async def _presign(self, byte_count: int) -> PointFiveUploadTarget | PointFiveUploadFailure:
|
||||
"""Ask the PointFive API for a presigned URL sized to this batch."""
|
||||
body: Final = await self._post(UPLOAD_URL_PATH, _PresignRequest(byte_count=byte_count))
|
||||
if isinstance(body, PointFiveUploadFailure):
|
||||
return body
|
||||
return _parse_target(body)
|
||||
|
||||
async def _post(self, path: str, request: BaseModel) -> str | PointFiveUploadFailure:
|
||||
"""POST one JSON request to the PointFive ingestion API and return its raw body."""
|
||||
try:
|
||||
response: Final = await self.http_client.post(
|
||||
self.api_url + path,
|
||||
json=request.model_dump(by_alias=True),
|
||||
headers={ # mutable-ok: AsyncHTTPHandler.post types headers as dict
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
return _failure_for(e.response, "pointfive api")
|
||||
except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt
|
||||
return PointFiveUploadFailure(f"pointfive api unreachable: {type(e).__name__}", retryable=True)
|
||||
return response.text
|
||||
|
||||
async def _put(self, target: PointFiveUploadTarget, body: bytes) -> PointFiveUploadFailure | None:
|
||||
"""
|
||||
PUT the batch to the presigned URL, which carries its own authorization.
|
||||
|
||||
The server chose that URL, so it is treated like any other externally supplied
|
||||
destination: the host is checked against blocked networks before connecting, and
|
||||
a redirect is refused rather than followed. A presigned URL never legitimately
|
||||
redirects, and following one would let a compromised endpoint point the proxy at
|
||||
an internal service.
|
||||
"""
|
||||
destination: Final = self._destination(target.upload_url)
|
||||
if isinstance(destination, PointFiveUploadFailure):
|
||||
return destination
|
||||
url, host = destination
|
||||
headers: Final = dict(PUT_HEADERS, Host=host) if host else dict(PUT_HEADERS) # mutable-ok: put wants dict
|
||||
try:
|
||||
await self.http_client.put(url, data=body, headers=headers, follow_redirects=False)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.is_redirect:
|
||||
return PointFiveUploadFailure(
|
||||
f"presigned upload redirected with {e.response.status_code}, refusing to follow", retryable=False
|
||||
)
|
||||
return _failure_for(e.response, "presigned upload")
|
||||
except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt
|
||||
return PointFiveUploadFailure(f"presigned upload unreachable: {type(e).__name__}", retryable=True)
|
||||
return None
|
||||
|
||||
def _destination(self, upload_url: str) -> tuple[str, str | None] | PointFiveUploadFailure:
|
||||
if not getattr(litellm, "user_url_validation", True):
|
||||
return upload_url, None
|
||||
try:
|
||||
return self.validate_upload_url(upload_url)
|
||||
except SSRFError as e:
|
||||
return PointFiveUploadFailure(f"presigned upload url refused: {e}", retryable=False)
|
||||
|
|
@ -246,6 +246,7 @@ class PrometheusLogger(CustomLogger):
|
|||
# logger so toggling these flags only takes effect after a
|
||||
# restart, keeping init-time and runtime label sets in sync.
|
||||
self._cached_metric_labels: dict[str, list[str]] = {}
|
||||
self._emit_input_sequence_length_label = litellm.prometheus_emit_input_sequence_length_label is True
|
||||
|
||||
_custom_buckets: Final = litellm.prometheus_latency_buckets
|
||||
self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS
|
||||
|
|
@ -1522,6 +1523,11 @@ class PrometheusLogger(CustomLogger):
|
|||
# 2. Pyright does not allow us to run isinstance(standard_logging_payload, StandardLoggingPayload) <- this would be ideal
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
input_sequence_length=(
|
||||
self._get_input_sequence_length(standard_logging_payload, kwargs, response_obj)
|
||||
if self._emit_input_sequence_length_label
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# set x-ratelimit headers
|
||||
|
|
@ -2192,6 +2198,36 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens)
|
||||
|
||||
@staticmethod
|
||||
def _get_input_sequence_length(
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
kwargs: Mapping[str, object],
|
||||
response_obj: object,
|
||||
) -> str:
|
||||
prompt_tokens: Final = standard_logging_payload.get("prompt_tokens")
|
||||
if prompt_tokens:
|
||||
return get_input_sequence_length_bucket(prompt_tokens)
|
||||
combined_usage: Final = kwargs.get("combined_usage_object")
|
||||
if (
|
||||
combined_usage is not None
|
||||
and getattr(kwargs.get("_litellm_upstream_reported_usage"), "total_tokens", None) is not None
|
||||
):
|
||||
return get_input_sequence_length_bucket(None)
|
||||
reported_usage: Final = (
|
||||
response_obj.get("usage") if isinstance(response_obj, dict) else getattr(response_obj, "usage", None)
|
||||
)
|
||||
if reported_usage is None and combined_usage is None:
|
||||
return get_input_sequence_length_bucket(None)
|
||||
usage_metadata: Final = standard_logging_payload["metadata"].get("usage_object")
|
||||
if isinstance(usage_metadata, Mapping):
|
||||
return get_input_sequence_length_bucket(usage_metadata.get("prompt_tokens"))
|
||||
if combined_usage is None and isinstance(response_obj, dict):
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
normalized_usage: Final[Mapping[str, object]] = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
|
||||
return get_input_sequence_length_bucket(normalized_usage.get("prompt_tokens"))
|
||||
return get_input_sequence_length_bucket(prompt_tokens)
|
||||
|
||||
def _set_latency_metrics(
|
||||
self,
|
||||
kwargs: dict,
|
||||
|
|
@ -2202,7 +2238,16 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_team_alias: str | None,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
label_context: PrometheusLabelFactoryContext | None = None,
|
||||
input_sequence_length: str | None = None,
|
||||
):
|
||||
latency_enum_values: Final = (
|
||||
replace(enum_values, input_sequence_length=input_sequence_length)
|
||||
if input_sequence_length is not None
|
||||
else enum_values
|
||||
)
|
||||
latency_label_context: Final = (
|
||||
PrometheusLabelFactoryContext(latency_enum_values) if input_sequence_length is not None else label_context
|
||||
)
|
||||
# latency metrics
|
||||
end_time: Final[datetime] = kwargs.get("end_time") or datetime.now()
|
||||
start_time: Final[datetime | None] = kwargs.get("start_time")
|
||||
|
|
@ -2220,8 +2265,8 @@ class PrometheusLogger(CustomLogger):
|
|||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_llm_api_time_to_first_token_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
enum_values=latency_enum_values,
|
||||
label_context=latency_label_context,
|
||||
)
|
||||
self.litellm_llm_api_time_to_first_token_metric.labels(**_ttft_labels).observe(time_to_first_token_seconds)
|
||||
self._track_end_user_metric_series(
|
||||
|
|
@ -2241,8 +2286,8 @@ class PrometheusLogger(CustomLogger):
|
|||
if api_call_total_time_seconds is not None:
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_llm_api_latency_metric"),
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
enum_values=latency_enum_values,
|
||||
label_context=latency_label_context,
|
||||
)
|
||||
self.litellm_llm_api_latency_metric.labels(**_labels).observe(api_call_total_time_seconds)
|
||||
self._track_end_user_metric_series(
|
||||
|
|
@ -2272,8 +2317,8 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_total_latency_metric"),
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
enum_values=latency_enum_values,
|
||||
label_context=latency_label_context,
|
||||
)
|
||||
self.litellm_request_total_latency_metric.labels(**_labels).observe(_observed_total_time_seconds)
|
||||
self._track_end_user_metric_series(
|
||||
|
|
|
|||
|
|
@ -10,9 +10,11 @@ import asyncio
|
|||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final, cast
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
|
||||
|
|
@ -24,7 +26,7 @@ from litellm.integrations.s3 import (
|
|||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -35,6 +37,9 @@ from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
|
|||
|
||||
from .custom_batch_logger import CustomBatchLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
|
||||
class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
||||
def __init__(
|
||||
|
|
@ -232,6 +237,26 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}"
|
||||
)
|
||||
|
||||
def _sign_put(
|
||||
self, credentials: "Credentials", url: str, json_string: str, headers: Mapping[str, str]
|
||||
) -> dict[str, str]: # mutable-ok: [LIT001] AsyncHTTPHandler.put/HTTPHandler.put only accept dict headers
|
||||
"""
|
||||
``RefreshableCredentials`` (IMDS roles) may refresh between the access key, secret and token
|
||||
reads SigV4 performs, producing a mixed-generation signature that S3 rejects with 403.
|
||||
Freezing first makes the three values one atomic snapshot.
|
||||
"""
|
||||
from botocore.auth import S3SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
|
||||
frozen: Final = (
|
||||
credentials.get_frozen_credentials() if isinstance(credentials, RefreshableCredentials) else credentials
|
||||
)
|
||||
aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=dict(headers))
|
||||
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
|
||||
S3SigV4Auth(frozen, "s3", aws_region_name).add_auth(aws_request)
|
||||
return dict(aws_request.headers.items())
|
||||
|
||||
def _sse_headers(self) -> Mapping[str, str]:
|
||||
candidates: Final = {
|
||||
"x-amz-server-side-encryption": self.s3_server_side_encryption,
|
||||
|
|
@ -317,26 +342,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
try:
|
||||
import base64
|
||||
import hashlib
|
||||
|
||||
from botocore.auth import S3SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
try:
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
|
||||
asyncified_get_credentials: Final = asyncify(self.get_credentials)
|
||||
credentials: Final = await asyncified_get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
aws_session_token=self.s3_aws_session_token,
|
||||
aws_region_name=self.s3_region_name,
|
||||
aws_session_name=self.s3_aws_session_name,
|
||||
aws_profile_name=self.s3_aws_profile_name,
|
||||
aws_role_name=self.s3_aws_role_name,
|
||||
aws_web_identity_token=self.s3_aws_web_identity_token,
|
||||
aws_sts_endpoint=self.s3_aws_sts_endpoint,
|
||||
)
|
||||
|
||||
verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key)
|
||||
verbose_logger.debug("s3_v2 logger - s3_verify setting: %s", self.s3_verify)
|
||||
|
|
@ -363,19 +374,28 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
**self._sse_headers(),
|
||||
}
|
||||
|
||||
# Sign the request
|
||||
aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers)
|
||||
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
|
||||
S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
||||
async def signed_put() -> httpx.Response:
|
||||
credentials: Final = await asyncified_get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
aws_session_token=self.s3_aws_session_token,
|
||||
aws_region_name=self.s3_region_name,
|
||||
aws_session_name=self.s3_aws_session_name,
|
||||
aws_profile_name=self.s3_aws_profile_name,
|
||||
aws_role_name=self.s3_aws_role_name,
|
||||
aws_web_identity_token=self.s3_aws_web_identity_token,
|
||||
aws_sts_endpoint=self.s3_aws_sts_endpoint,
|
||||
)
|
||||
signed_headers: Final = await run_aws_signing(self._sign_put, credentials, url, json_string, headers)
|
||||
try:
|
||||
return await self.async_httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
except httpx.HTTPStatusError as error:
|
||||
return error.response
|
||||
|
||||
# Prepare the signed headers
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
||||
# Make the request with retry for transient S3 errors (500/503)
|
||||
max_retries: Final = 3
|
||||
for attempt in range(max_retries):
|
||||
response = await self.async_httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
||||
response = await signed_put()
|
||||
if response.status_code in (403, 500, 503) and attempt < max_retries - 1:
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
"S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
|
||||
|
|
@ -479,20 +499,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
try:
|
||||
import base64
|
||||
import hashlib
|
||||
|
||||
from botocore.auth import S3SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
try:
|
||||
verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key)
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
aws_session_token=self.s3_aws_session_token,
|
||||
aws_region_name=self.s3_region_name,
|
||||
)
|
||||
|
||||
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
|
||||
|
||||
|
|
@ -516,22 +526,24 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
**self._sse_headers(),
|
||||
}
|
||||
|
||||
# Sign the request
|
||||
aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers)
|
||||
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
|
||||
S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
||||
|
||||
# Prepare the signed headers
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
||||
httpx_client: Final = _get_httpx_client(
|
||||
params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None)
|
||||
)
|
||||
# Make the request with retry for transient S3 errors (500/503)
|
||||
|
||||
def signed_put() -> httpx.Response:
|
||||
credentials: Final = self.get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
aws_session_token=self.s3_aws_session_token,
|
||||
aws_region_name=self.s3_region_name,
|
||||
)
|
||||
signed_headers: Final = self._sign_put(credentials, url, json_string, headers)
|
||||
return httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
|
||||
max_retries: Final = 3
|
||||
for attempt in range(max_retries):
|
||||
response = httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
||||
response = signed_put()
|
||||
if response.status_code in (403, 500, 503) and attempt < max_retries - 1:
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
"S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
|
||||
|
|
@ -597,7 +609,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
|
||||
# Sign the request
|
||||
aws_request: Final = AWSRequest(method="GET", url=url, headers=headers)
|
||||
S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
|
||||
await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request)
|
||||
|
||||
# Prepare the signed headers
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm.constants import (
|
|||
SQS_SEND_MESSAGE_ACTION,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM):
|
|||
data=prepped.body,
|
||||
headers=prepped.headers,
|
||||
)
|
||||
SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request)
|
||||
await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request)
|
||||
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm.integrations.newrelic import NewRelicLogger
|
|||
from litellm.integrations.openmeter import OpenMeterLogger
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
from litellm.integrations.opik.opik import OpikLogger
|
||||
from litellm.integrations.pointfive import PointFiveLogger
|
||||
from litellm.integrations.posthog import PostHogLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
|
|
@ -95,6 +96,7 @@ class CustomLoggerRegistry:
|
|||
"agentops": AgentOps,
|
||||
"deepeval": DeepEvalLogger,
|
||||
"s3_v2": S3Logger,
|
||||
"pointfive": PointFiveLogger,
|
||||
"aws_sqs": SQSLogger,
|
||||
"dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler,
|
||||
"dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
|
|||
"aws_web_identity_token",
|
||||
"aws_sts_endpoint",
|
||||
"aws_external_id",
|
||||
"aws_session_tags",
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"aws_bedrock_project_id",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -120,7 +120,6 @@ from litellm.types.utils import (
|
|||
CachingDetails,
|
||||
CallTypes,
|
||||
CostBreakdown,
|
||||
CostResponseTypes,
|
||||
CustomPricingLiteLLMParams,
|
||||
DynamicPromptManagementParamLiteral,
|
||||
EmbeddingResponse,
|
||||
|
|
@ -183,6 +182,7 @@ from ..integrations.lunary import LunaryLogger
|
|||
from ..integrations.newrelic import NewRelicLogger
|
||||
from ..integrations.openmeter import OpenMeterLogger
|
||||
from ..integrations.opik.opik import OpikLogger
|
||||
from ..integrations.pointfive import PointFiveLogger
|
||||
from ..integrations.posthog import PostHogLogger
|
||||
from ..integrations.prompt_layer import PromptLayerLogger
|
||||
from ..integrations.s3 import S3Logger
|
||||
|
|
@ -204,7 +204,7 @@ if TYPE_CHECKING:
|
|||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse
|
||||
try:
|
||||
from litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
EnterpriseCallbackControls,
|
||||
|
|
@ -2381,7 +2381,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self,
|
||||
raw_bytes: list[bytes],
|
||||
provider_config: "BasePassthroughConfig",
|
||||
) -> Optional["CostResponseTypes"]:
|
||||
) -> Optional["LoggedRelayResponse"]:
|
||||
all_chunks: Final = provider_config._convert_raw_bytes_to_str_lines(raw_bytes)
|
||||
complete_streaming_response: Final = provider_config.handle_logging_collected_chunks(
|
||||
all_chunks=all_chunks,
|
||||
|
|
@ -4377,6 +4377,14 @@ def _init_custom_logger_compatible_class(
|
|||
_s3_v2_logger: Final = S3V2Logger()
|
||||
_in_memory_loggers.append(_s3_v2_logger)
|
||||
return _s3_v2_logger
|
||||
elif logging_integration == "pointfive":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, PointFiveLogger):
|
||||
return callback
|
||||
|
||||
_pointfive_logger: Final = PointFiveLogger()
|
||||
_in_memory_loggers.append(_pointfive_logger)
|
||||
return _pointfive_logger
|
||||
elif logging_integration == "aws_sqs":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, SQSLogger):
|
||||
|
|
@ -5065,6 +5073,10 @@ def get_custom_logger_compatible_class(
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, S3V2Logger):
|
||||
return callback
|
||||
elif logging_integration == "pointfive":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, PointFiveLogger):
|
||||
return callback
|
||||
elif logging_integration == "aws_sqs":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, SQSLogger):
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import json
|
|||
import re
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Iterable, Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
import litellm
|
||||
|
|
@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
|
|||
del choice.enhancements
|
||||
|
||||
|
||||
def _invalid_choices_message(response_object: Mapping[str, object]) -> str:
|
||||
raw_keys: Final = list(response_object.keys())
|
||||
if "choices" not in response_object:
|
||||
return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}"
|
||||
return (
|
||||
f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). "
|
||||
f"Raw keys: {raw_keys}"
|
||||
)
|
||||
|
||||
|
||||
async def convert_to_streaming_response_async(
|
||||
response_object: dict | None = None,
|
||||
):
|
||||
|
|
@ -179,14 +189,12 @@ async def convert_to_streaming_response_async(
|
|||
|
||||
choice_list: Final[list[StreamingChoices]] = []
|
||||
|
||||
if not response_object.get("choices"):
|
||||
if not isinstance(response_object.get("choices"), list):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
message=_invalid_choices_message(response_object),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
|
@ -287,14 +295,12 @@ def convert_to_streaming_response(
|
|||
model_response_object: Final = ModelResponseStream()
|
||||
choice_list: Final[list[StreamingChoices]] = []
|
||||
|
||||
if not response_object.get("choices"):
|
||||
if not isinstance(response_object.get("choices"), list):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
message=_invalid_choices_message(response_object),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
|
@ -623,15 +629,12 @@ def convert_to_model_response_object(
|
|||
return convert_to_streaming_response(response_object=response_object)
|
||||
choice_list: Final[list[Choices]] = []
|
||||
|
||||
if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable):
|
||||
if not isinstance(response_object.get("choices"), list):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"LiteLLM: provider returned a response with no 'choices'. "
|
||||
f"Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
message=_invalid_choices_message(response_object),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -441,7 +441,15 @@ class LoggingCallbackManager:
|
|||
|
||||
return result
|
||||
|
||||
def get_callback_objects(self) -> tuple[tuple[str, CustomLogger | Callable], ...]:
|
||||
return tuple(
|
||||
(self._get_callback_string(callback), callback)
|
||||
for callback in self._get_all_callbacks()
|
||||
if not isinstance(callback, str)
|
||||
)
|
||||
|
||||
def _get_callback_string(self, callback: CustomLogger | Callable | str) -> str:
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
from litellm.litellm_core_utils.custom_logger_registry import (
|
||||
CustomLoggerRegistry,
|
||||
)
|
||||
|
|
@ -449,6 +457,8 @@ class LoggingCallbackManager:
|
|||
"""Convert a callback to its string representation"""
|
||||
if isinstance(callback, str):
|
||||
return callback
|
||||
elif isinstance(callback, OpenTelemetry) and callback.callback_name is not None:
|
||||
return callback.callback_name
|
||||
elif isinstance(callback, CustomLogger):
|
||||
# Try to get the string representation from the registry
|
||||
callback_str: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback))
|
||||
|
|
|
|||
|
|
@ -6,8 +6,8 @@ import io
|
|||
import json
|
||||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from itertools import groupby
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from itertools import groupby, islice
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
|
|
@ -1320,17 +1320,128 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin
|
|||
return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo
|
||||
|
||||
|
||||
def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
_SUBSCHEMA_KEYWORDS: Final = frozenset(
|
||||
{
|
||||
"additionalItems",
|
||||
"additionalProperties",
|
||||
"contains",
|
||||
"else",
|
||||
"if",
|
||||
"items",
|
||||
"not",
|
||||
"propertyNames",
|
||||
"then",
|
||||
"unevaluatedItems",
|
||||
"unevaluatedProperties",
|
||||
}
|
||||
)
|
||||
_SUBSCHEMA_LIST_KEYWORDS: Final = frozenset({"allOf", "anyOf", "items", "oneOf", "prefixItems"})
|
||||
_SUBSCHEMA_MAP_KEYWORDS: Final = frozenset(
|
||||
{"$defs", "definitions", "dependentSchemas", "patternProperties", "properties"}
|
||||
)
|
||||
|
||||
_MAX_SCHEMA_NESTING: Final = 1024
|
||||
|
||||
|
||||
def drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Drop every regex in a schema position that Python's ``re`` cannot compile.
|
||||
|
||||
OpenAI validates tool ``parameters`` against the 2020-12 metaschema with
|
||||
``jsonschema``'s format checker, which hands each ``pattern`` value and each
|
||||
``patternProperties`` key to ``re.compile``, so a regex written for an
|
||||
ECMA-262 engine (Unicode property escapes such as ``\\p{Cc}``, as in Claude
|
||||
Code's ``Artifact`` tool) is refused with "'...' is not a 'regex'" by every
|
||||
model family on both the chat and Responses wires. Only schema positions are
|
||||
walked (properties, items, combinators, ``$defs`` and the other applicators),
|
||||
so a ``pattern`` key inside ``default``, ``examples``, ``const`` or vendor
|
||||
extensions is data and stays. Outside strict mode the keyword is only a
|
||||
hint, so dropping it costs the model a constraint and the caller nothing.
|
||||
Compilable regexes and everything else pass through, the input is never
|
||||
mutated, and the same object comes back when nothing was dropped. The walk
|
||||
is level-order rather than recursive, rebuilt deepest level first, and stops
|
||||
at more schema levels than a JSON parser admits, so a cyclic schema built in
|
||||
code cannot spin it.
|
||||
"""
|
||||
rebuilt: dict[int, Mapping[str, object]] = {} # mutable-ok: per-call memo of rewritten nodes, deepest level first
|
||||
for level in reversed(tuple(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING))):
|
||||
rebuilt.update(
|
||||
(id(node), rewritten)
|
||||
for node in level
|
||||
if (rewritten := _node_without_non_python_regex(node, rebuilt)) is not node
|
||||
)
|
||||
return rebuilt.get(id(schema), schema)
|
||||
|
||||
|
||||
def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, object], ...]]:
|
||||
frontier: tuple[Mapping[str, object], ...] = (schema,) # rebind-ok: level-order cursor, one level a round
|
||||
while frontier:
|
||||
yield frontier
|
||||
frontier = tuple(child for node in frontier for child in _subschemas(node))
|
||||
|
||||
|
||||
def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]:
|
||||
for key, value in node.items():
|
||||
if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict):
|
||||
yield from (sub for sub in value.values() if isinstance(sub, dict))
|
||||
elif key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list):
|
||||
yield from (sub for sub in value if isinstance(sub, dict))
|
||||
elif key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict):
|
||||
yield value
|
||||
|
||||
|
||||
def _node_without_non_python_regex(
|
||||
node: Mapping[str, object], rebuilt: Mapping[int, Mapping[str, object]]
|
||||
) -> Mapping[str, object]:
|
||||
kept: Final = { # mutable-ok: tool parameters are JSON dicts
|
||||
key: _keyword_value_rebuilt(key, value, rebuilt)
|
||||
for key, value in node.items()
|
||||
if key != "pattern" or not isinstance(value, str) or _is_python_regex(value)
|
||||
}
|
||||
return node if len(kept) == len(node) and all(kept[key] is node[key] for key in kept) else kept
|
||||
|
||||
|
||||
def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mapping[str, object]]) -> object:
|
||||
if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict):
|
||||
kept: Final = { # mutable-ok: tool parameters are JSON dicts
|
||||
name: rebuilt.get(id(sub), sub)
|
||||
for name, sub in value.items()
|
||||
if key != "patternProperties" or not isinstance(name, str) or _is_python_regex(name)
|
||||
}
|
||||
return value if len(kept) == len(value) and all(kept[name] is value[name] for name in kept) else kept
|
||||
if key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list):
|
||||
items: Final = [rebuilt.get(id(sub), sub) for sub in value] # mutable-ok: tool parameters are JSON lists
|
||||
return value if all(new is old for new, old in zip(items, value, strict=True)) else items
|
||||
if key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict):
|
||||
return rebuilt.get(id(value), value)
|
||||
return value
|
||||
|
||||
|
||||
def _is_python_regex(pattern: str) -> bool:
|
||||
try:
|
||||
re.compile(pattern)
|
||||
except (re.error, RecursionError):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def flatten_combinators_and_drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return flatten_top_level_schema_combinators(drop_non_python_regex_patterns(schema))
|
||||
|
||||
|
||||
def tool_with_sanitized_parameters(
|
||||
tool: Mapping[str, object],
|
||||
sanitize: Callable[[Mapping[str, object]], Mapping[str, object]],
|
||||
) -> Mapping[str, object]:
|
||||
function: Final = tool.get("function")
|
||||
if not isinstance(function, dict):
|
||||
return tool
|
||||
parameters: Final = function.get("parameters")
|
||||
if not isinstance(parameters, dict):
|
||||
return tool
|
||||
flattened: Final = flatten_top_level_schema_combinators(parameters)
|
||||
if flattened is parameters:
|
||||
sanitized: Final = sanitize(parameters)
|
||||
if sanitized is parameters:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts
|
||||
return {**tool, "function": {**function, "parameters": sanitized}} # mutable-ok: request tools are JSON dicts
|
||||
|
||||
|
||||
def _get_image_mime_type_from_url(url: str) -> str | None:
|
||||
|
|
@ -1823,14 +1934,11 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]:
|
|||
return None, message_content
|
||||
|
||||
|
||||
def _readable_thinking_text(
|
||||
block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
|
||||
) -> str:
|
||||
def _readable_thinking_text(block: Mapping[str, object]) -> str:
|
||||
"""The text a chat model can read back, empty for redacted blocks and malformed ones."""
|
||||
if block.get("type") != "thinking":
|
||||
return ""
|
||||
thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag
|
||||
return str(thinking or "")
|
||||
return str(block.get("thinking") or "")
|
||||
|
||||
|
||||
def reasoning_content_from_thinking_blocks(
|
||||
|
|
@ -1843,24 +1951,125 @@ def reasoning_content_from_thinking_blocks(
|
|||
return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block)))
|
||||
|
||||
|
||||
def responses_reasoning_item_from_thinking_blocks(
|
||||
thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
|
||||
) -> ChatCompletionReasoningItem | None:
|
||||
"""Build a Responses API `reasoning` input item from Anthropic thinking blocks.
|
||||
ENCRYPTED_REASONING_SIGNATURE_PREFIX: Final = "litellm_encrypted_reasoning:"
|
||||
|
||||
The item carries no `id`: the Responses API rejects an empty one and 404s on any id it
|
||||
did not mint itself, while an item without an id is always accepted.
|
||||
|
||||
def encrypted_reasoning_signature(encrypted_content: str) -> str:
|
||||
"""The opaque value a Responses API reasoning item's `encrypted_content` travels in.
|
||||
|
||||
Anthropic clients echo a thinking block's `signature` and a redacted block's `data`
|
||||
back verbatim, so either field can carry the encrypted reasoning across turns; the
|
||||
prefix tells the two apart from a signature Anthropic minted.
|
||||
"""
|
||||
return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}"
|
||||
|
||||
|
||||
def _carries_encrypted_reasoning(signature: object) -> bool:
|
||||
return isinstance(signature, str) and signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX)
|
||||
|
||||
|
||||
def encrypted_content_from_signature(signature: object) -> str | None:
|
||||
if not isinstance(signature, str) or not _carries_encrypted_reasoning(signature):
|
||||
return None
|
||||
return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None
|
||||
|
||||
|
||||
def _encrypted_reasoning_field(block: Mapping[str, object]) -> object:
|
||||
match block.get("type"):
|
||||
case "thinking":
|
||||
return block.get("signature")
|
||||
case "redacted_thinking":
|
||||
return block.get("data")
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
|
||||
return encrypted_content_from_signature(_encrypted_reasoning_field(block))
|
||||
|
||||
|
||||
def is_encrypted_reasoning_block(block: object) -> bool:
|
||||
"""A thinking or redacted_thinking block carrying Responses API encrypted reasoning.
|
||||
|
||||
Only the Responses API that minted the content can read it back, so an Anthropic
|
||||
backend has to drop such a block rather than fail signature verification on it.
|
||||
"""
|
||||
if not isinstance(block, Mapping):
|
||||
return False
|
||||
mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance
|
||||
return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping))
|
||||
|
||||
|
||||
def strip_encrypted_reasoning_from_messages(messages: object) -> None:
|
||||
"""Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from
|
||||
Anthropic-shaped history.
|
||||
|
||||
The whole block goes, the way #40280 drops undecryptable Responses ``input`` items: a
|
||||
provider that did not mint the block rejects it signed (a foreign signature) and unsigned
|
||||
(a missing signature) alike, so keeping its text as an unsigned thinking block only moves
|
||||
the 400 from the router to the provider.
|
||||
|
||||
Mutates the content lists in place: the router's fallback snapshot shares these
|
||||
message objects, so a rebound list would replay the stripped blocks on the fallback hop.
|
||||
"""
|
||||
if not isinstance(messages, list):
|
||||
return
|
||||
for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json
|
||||
_strip_encrypted_reasoning_from_blocks(content)
|
||||
|
||||
|
||||
def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
|
||||
return (
|
||||
cast(list[object], content) # cast-ok: narrowed by isinstance
|
||||
for message in messages
|
||||
if isinstance(message, Mapping)
|
||||
for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance
|
||||
if isinstance(content, list)
|
||||
)
|
||||
|
||||
|
||||
def _strip_encrypted_reasoning_from_blocks(content: object) -> None:
|
||||
blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance
|
||||
kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block))
|
||||
blocks[:] = kept # rebind-ok: shared with fallback snapshot
|
||||
|
||||
|
||||
def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str:
|
||||
index, block = indexed_block
|
||||
return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary"
|
||||
|
||||
|
||||
def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> ChatCompletionReasoningItem | None:
|
||||
summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload
|
||||
ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text)
|
||||
for block in thinking_blocks
|
||||
for block in group
|
||||
if (text := _readable_thinking_text(block))
|
||||
]
|
||||
encrypted_content: Final = encrypted_content_of_block(group[0])
|
||||
if encrypted_content is not None:
|
||||
return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content)
|
||||
if not summary:
|
||||
return None
|
||||
return ChatCompletionReasoningItem(type="reasoning", summary=summary)
|
||||
|
||||
|
||||
def responses_reasoning_items_from_thinking_blocks(
|
||||
thinking_blocks: Iterable[Mapping[str, object]],
|
||||
) -> tuple[ChatCompletionReasoningItem, ...]:
|
||||
"""Build Responses API `reasoning` input items from Anthropic thinking blocks.
|
||||
|
||||
A block carrying encrypted reasoning replays the item it came from byte for byte;
|
||||
a run of plain thinking blocks collapses into one summary-only item. No item carries
|
||||
an `id`: the Responses API 404s on any id it did not mint itself and rejects an empty
|
||||
one, while an item without an id is always accepted.
|
||||
"""
|
||||
return tuple(
|
||||
item
|
||||
for _, group in groupby(enumerate(thinking_blocks), key=_reasoning_replay_group_key)
|
||||
if (item := _reasoning_item_from_block_group(tuple(block for _, block in group))) is not None
|
||||
)
|
||||
|
||||
|
||||
def _parse_content_for_reasoning(
|
||||
message_text: str | None,
|
||||
) -> tuple[str | None, str | None]:
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from litellm.types.utils import GenericImageParsingChunk
|
|||
from .common_utils import (
|
||||
convert_content_list_to_str,
|
||||
infer_content_type_from_url_and_content,
|
||||
is_encrypted_reasoning_block,
|
||||
is_non_content_values_set,
|
||||
parse_tool_call_arguments,
|
||||
)
|
||||
|
|
@ -2299,13 +2300,16 @@ def sanitize_messages_for_tool_calling(
|
|||
|
||||
|
||||
def _is_unsignable_thinking_block(block: object) -> bool:
|
||||
"""A `thinking` block that Anthropic cannot accept on input.
|
||||
"""A thinking block that Anthropic cannot accept on input.
|
||||
|
||||
Anthropic verifies the thinking signature cryptographically, so a block whose
|
||||
signature is null, empty, or missing (e.g. from an open-source reasoning model)
|
||||
is rejected with a 400 and must be dropped rather than blanked or repaired.
|
||||
`redacted_thinking` blocks carry no signature and are always kept.
|
||||
is rejected with a 400 and must be dropped rather than blanked or repaired, and
|
||||
so is a block whose signature or data carries another provider's encrypted
|
||||
reasoning. A `redacted_thinking` block Anthropic minted is always kept.
|
||||
"""
|
||||
if is_encrypted_reasoning_block(block):
|
||||
return True
|
||||
if not isinstance(block, dict) or block.get("type") != "thinking":
|
||||
return False
|
||||
signature: Final = block.get("signature")
|
||||
|
|
|
|||
|
|
@ -162,6 +162,19 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
|||
output_item["arguments"] = redacted_str
|
||||
|
||||
|
||||
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""
|
||||
Return a copy of a ``StandardLoggingPayload`` with its messages and response redacted.
|
||||
|
||||
The success path redacts through ``perform_redaction`` before a callback ever sees the
|
||||
payload, but the failure path does not, so a callback that batches both has to redact
|
||||
the ones it is handed.
|
||||
"""
|
||||
redacted: Final = copy.deepcopy(dict(payload)) # mutable-ok: redacted in place below
|
||||
_redact_standard_logging_object({"standard_logging_object": redacted}) # mutable-ok: the callee's shape
|
||||
return redacted
|
||||
|
||||
|
||||
def _redact_standard_logging_object(model_call_details: dict):
|
||||
"""Redact messages and response inside standard_logging_object if present."""
|
||||
standard_logging_object: Final = model_call_details.get("standard_logging_object")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import base64
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast
|
||||
|
|
@ -210,7 +210,7 @@ def apply_grounding_request_counts(
|
|||
|
||||
|
||||
class ChunkProcessor:
|
||||
def __init__(self, chunks: list, messages: list | None = None):
|
||||
def __init__(self, chunks: list, messages: Sequence | None = None):
|
||||
self.chunks = self._sort_chunks(chunks)
|
||||
self.messages = messages
|
||||
self.first_chunk = chunks[0]
|
||||
|
|
@ -1004,8 +1004,9 @@ class ChunkProcessor:
|
|||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
model: str,
|
||||
completion_output: str,
|
||||
messages: list | None = None,
|
||||
messages: Sequence | None = None,
|
||||
reasoning_tokens: int | None = None,
|
||||
count_prompt_tokens: Callable[[], int] | None = None,
|
||||
) -> Usage:
|
||||
"""
|
||||
Calculate usage for the given chunks.
|
||||
|
|
@ -1030,7 +1031,9 @@ class ChunkProcessor:
|
|||
cost: Final[float | None] = calculated_usage_per_chunk["cost"]
|
||||
|
||||
try:
|
||||
returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages)
|
||||
returned_usage.prompt_tokens = prompt_tokens or (
|
||||
count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages)
|
||||
)
|
||||
except Exception: # don't allow this failing to block a complete streaming response from being returned
|
||||
print_verbose("token_counter failed, assuming prompt tokens is 0")
|
||||
returned_usage.prompt_tokens = 0
|
||||
|
|
|
|||
|
|
@ -1473,17 +1473,14 @@ class CustomStreamWrapper:
|
|||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "cached_response":
|
||||
cached_chunk: Final = cast(ModelResponseStream, chunk)
|
||||
chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason
|
||||
cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None
|
||||
chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None
|
||||
response_obj = {
|
||||
"text": cached_chunk.choices[0].delta.content,
|
||||
"text": cached_choice.delta.content if cached_choice is not None else None,
|
||||
"is_finished": chunk_finish_reason is not None,
|
||||
"finish_reason": chunk_finish_reason,
|
||||
"original_chunk": cached_chunk,
|
||||
"tool_calls": (
|
||||
cached_chunk.choices[0].delta.tool_calls
|
||||
if hasattr(cached_chunk.choices[0].delta, "tool_calls")
|
||||
else None
|
||||
),
|
||||
"tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None),
|
||||
}
|
||||
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
|
|
|
|||
|
|
@ -3,11 +3,15 @@
|
|||
import base64
|
||||
import io
|
||||
import struct
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
import httpx
|
||||
import tiktoken
|
||||
from tokenizers import Tokenizer
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -21,7 +25,10 @@ from litellm.constants import (
|
|||
MAX_TILE_HEIGHT,
|
||||
MAX_TILE_WIDTH,
|
||||
TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
|
||||
TOKEN_COUNTER_MAX_CONCURRENT_COUNTS,
|
||||
TOKEN_COUNTER_MAX_EXACT_CHARS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
|
||||
from litellm.litellm_core_utils.url_utils import safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
|
|
@ -172,6 +179,13 @@ def calculate_tiles_needed(
|
|||
return total_tiles
|
||||
|
||||
|
||||
def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int:
|
||||
largest_tile_count: Final = calculate_tiles_needed(
|
||||
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES
|
||||
)
|
||||
return base_tokens + (base_tokens * 2) * largest_tile_count
|
||||
|
||||
|
||||
def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]:
|
||||
return struct.unpack(fmt, buffer)
|
||||
|
||||
|
|
@ -317,6 +331,32 @@ TokenCounterFunction = Callable[[str], int]
|
|||
Type for a function that counts tokens in a string.
|
||||
"""
|
||||
|
||||
EXTRAPOLATION_SAMPLES: Final = 16
|
||||
T_ParamSpec: Final = ParamSpec("T_ParamSpec")
|
||||
T_Retval = TypeVar("T_Retval")
|
||||
_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter")
|
||||
|
||||
|
||||
def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter:
|
||||
existing: Final = _COUNT_OFFLOAD_LIMITER.get(None)
|
||||
if existing is not None:
|
||||
return existing
|
||||
created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS)
|
||||
_COUNT_OFFLOAD_LIMITER.set(created)
|
||||
return created
|
||||
|
||||
|
||||
def offload_token_count(
|
||||
function: Callable[T_ParamSpec, T_Retval],
|
||||
) -> Callable[T_ParamSpec, Awaitable[T_Retval]]:
|
||||
async def offloaded(
|
||||
*args: T_ParamSpec.args,
|
||||
**kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract
|
||||
) -> T_Retval:
|
||||
return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs)
|
||||
|
||||
return offloaded
|
||||
|
||||
|
||||
def _get_tiktoken_count_function(
|
||||
encode_length: Callable[[str], int],
|
||||
|
|
@ -538,9 +578,40 @@ def _count_extra(
|
|||
return num_tokens
|
||||
|
||||
|
||||
def _get_extrapolating_count_function(
|
||||
count_exactly: TokenCounterFunction,
|
||||
max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS,
|
||||
) -> TokenCounterFunction:
|
||||
def count_tokens(text: str) -> int:
|
||||
if len(text) <= max_exact_chars:
|
||||
return count_exactly(text)
|
||||
samples: Final = _evenly_spaced_samples(text, max_exact_chars)
|
||||
sampled_chars: Final = sum(len(sample) for sample in samples)
|
||||
return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars)
|
||||
|
||||
return count_tokens
|
||||
|
||||
|
||||
def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]:
|
||||
sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars)
|
||||
sample_chars: Final = total_chars // sample_count
|
||||
last_start: Final = len(text) - sample_chars
|
||||
return tuple(
|
||||
text[start : start + sample_chars]
|
||||
for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count))
|
||||
)
|
||||
|
||||
|
||||
def _get_count_function(
|
||||
model: str | None,
|
||||
custom_tokenizer: dict | SelectTokenizerResponse | None = None,
|
||||
) -> TokenCounterFunction:
|
||||
return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer))
|
||||
|
||||
|
||||
def _get_exact_count_function(
|
||||
model: str | None,
|
||||
custom_tokenizer: dict | SelectTokenizerResponse | None = None,
|
||||
) -> TokenCounterFunction:
|
||||
"""
|
||||
Get the function to count tokens based on the model and custom tokenizer."""
|
||||
|
|
@ -549,10 +620,10 @@ def _get_count_function(
|
|||
if model is not None or custom_tokenizer is not None:
|
||||
tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
|
||||
if tokenizer_json["type"] == "huggingface_tokenizer":
|
||||
tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"]
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
enc: Final = tokenizer_json["tokenizer"].encode(text)
|
||||
return len(enc.ids)
|
||||
return len(tokenizer.encode_batch_fast([text])[0])
|
||||
|
||||
return count_tokens
|
||||
elif tokenizer_json["type"] == "openai_tokenizer":
|
||||
|
|
|
|||
|
|
@ -13,10 +13,11 @@ Pattern Overview:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from collections.abc import Mapping, MutableSequence, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
|
@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
merge_guardrailed_scoped_messages,
|
||||
merge_returned_tools_into_request_tools,
|
||||
scoped_structured_message_indices,
|
||||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
|
|
@ -153,6 +155,46 @@ class ExtractedInput:
|
|||
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ToolCallShape:
|
||||
name: str | None
|
||||
arguments: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SSEFieldRewrite:
|
||||
"""One field of one nested section of a buffered SSE event, rewritten."""
|
||||
|
||||
section: str
|
||||
field: str
|
||||
value: object
|
||||
|
||||
|
||||
class _SSEEventRewriter(Protocol):
|
||||
def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ...
|
||||
|
||||
|
||||
def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]:
|
||||
rewrite: Final = rewrite_event(event)
|
||||
section: Final = None if rewrite is None else event.get(rewrite.section)
|
||||
if rewrite is None or not isinstance(section, Mapping):
|
||||
return event
|
||||
return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict
|
||||
|
||||
|
||||
def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
|
||||
"""The guardrail-visible shape of each tool call, whether the guardrail handed
|
||||
back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
|
||||
functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls)
|
||||
return tuple(
|
||||
_ToolCallShape(
|
||||
name=name if isinstance(name := stream_item_field(function, "name"), str) else None,
|
||||
arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "",
|
||||
)
|
||||
for function in functions
|
||||
)
|
||||
|
||||
|
||||
class _AnthropicSSEDelta(TypedDict, total=False):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
|
|
@ -170,12 +212,18 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
them through guardrail rewrites; downstream provider handling is out of scope.
|
||||
"""
|
||||
|
||||
delivers_ended_stream_text_rewrites = True
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
|
||||
def post_call_hook_response(self, response: object) -> object:
|
||||
if not isinstance(response, ModelResponse):
|
||||
return response
|
||||
return self.adapter.translate_openai_response_to_anthropic(response)
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_usage_response(
|
||||
responses_so_far: Sequence[object],
|
||||
|
|
@ -1050,6 +1098,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
first_choice.message.tool_calls,
|
||||
)
|
||||
string_so_far = first_choice.message.content
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
|
||||
guardrail_inputs: Final = GenericGuardrailAPIInputs()
|
||||
if string_so_far:
|
||||
guardrail_inputs["texts"] = [string_so_far]
|
||||
|
|
@ -1084,6 +1133,19 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
and guardrailed_texts[0] != string_so_far
|
||||
):
|
||||
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
|
||||
if deliver_ended_stream_rewrites:
|
||||
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
|
||||
self._write_ended_stream_tool_call_rewrites(
|
||||
responses_so_far,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list)
|
||||
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tool_calls_list or ()
|
||||
),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
|
||||
return responses_so_far
|
||||
|
|
@ -1206,44 +1268,124 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
@staticmethod
|
||||
def _write_ended_stream_text_rewrite(
|
||||
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
rewritten_text: str,
|
||||
) -> None:
|
||||
"""Deliver an ended-stream guardrail text rewrite by rewriting the
|
||||
buffered chunks in place: the first ``text_delta`` carries the full
|
||||
rewritten text and every later one is blanked, leaving the surrounding
|
||||
message and content-block framing untouched. Handles both chunk formats
|
||||
this stream carries (parsed event dicts and raw SSE bytes)."""
|
||||
message and content-block framing untouched."""
|
||||
replacements: Final = chain((rewritten_text,), repeat(""))
|
||||
for idx, item in enumerate(responses_so_far):
|
||||
if isinstance(item, dict):
|
||||
delta = item.get("delta")
|
||||
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
|
||||
if delta.get("type") == "text_delta":
|
||||
delta["text"] = next(replacements)
|
||||
elif isinstance(item, (bytes, bytearray)):
|
||||
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
|
||||
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
|
||||
)
|
||||
|
||||
def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
delta: Final = event.get("delta")
|
||||
if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
|
||||
return None
|
||||
if delta.get("type") != "text_delta":
|
||||
return None
|
||||
return _SSEFieldRewrite("delta", "text", next(replacements))
|
||||
|
||||
AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
|
||||
|
||||
@classmethod
|
||||
def _write_ended_stream_tool_call_rewrites(
|
||||
cls,
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
*,
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
guardrail_name: str,
|
||||
) -> None:
|
||||
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
|
||||
buffered chunks in place: the rebuilt response lists tool calls in the
|
||||
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
|
||||
on the nth block, its first ``input_json_delta`` carrying the full rewritten
|
||||
arguments, every later one blanked, and ``content_block_start`` carrying the
|
||||
rewritten name. Blocks that do not line up with the rebuilt tool calls make
|
||||
the rewrite undeliverable, so the pipeline executor discards it and releases
|
||||
the original chunks."""
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
block_indices: Final = tuple(
|
||||
index
|
||||
for item in responses_so_far
|
||||
for event in cls._iter_sse_events(item)
|
||||
if event.get("type") == "content_block_start"
|
||||
and isinstance(block := event.get("content_block"), Mapping)
|
||||
and block.get("type") == "tool_use"
|
||||
and isinstance(index := event.get("index"), int)
|
||||
)
|
||||
if len(block_indices) != len(post_guardrail_tool_calls):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(guardrail_name)
|
||||
rewrites_by_block: Final = MappingProxyType(
|
||||
{
|
||||
index: after
|
||||
for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if after != before
|
||||
}
|
||||
)
|
||||
argument_replacements: Final = MappingProxyType(
|
||||
{index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()}
|
||||
)
|
||||
|
||||
def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
index: Final = event.get("index")
|
||||
if not isinstance(index, int) or index not in rewrites_by_block:
|
||||
return None
|
||||
match event.get("type"):
|
||||
case "content_block_start":
|
||||
name: Final = rewrites_by_block[index].name
|
||||
if name is None:
|
||||
return None
|
||||
return _SSEFieldRewrite("content_block", "name", name)
|
||||
case "content_block_delta":
|
||||
delta: Final = event.get("delta")
|
||||
if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta":
|
||||
return None
|
||||
return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index]))
|
||||
case _:
|
||||
return None
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
|
||||
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
|
||||
replacement text, leaving all other events and framing byte-identical."""
|
||||
def _rewrite_ended_stream_events(
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
rewrite_event: _SSEEventRewriter,
|
||||
) -> None:
|
||||
"""Replace every buffered event ``rewrite_event`` returns a rewrite for, in
|
||||
both chunk formats this stream carries (parsed event dicts and raw SSE
|
||||
bytes), leaving every other event and the framing untouched."""
|
||||
rewritten_items: Final = tuple(
|
||||
AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far
|
||||
)
|
||||
responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object:
|
||||
if isinstance(item, dict):
|
||||
return _rewritten_event(_as_str_mapping(item), rewrite_event)
|
||||
if isinstance(item, (bytes, bytearray)):
|
||||
return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event)
|
||||
return item
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes:
|
||||
"""Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites,
|
||||
leaving all other events and framing byte-identical."""
|
||||
try:
|
||||
decoded: Final = sse_bytes.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return sse_bytes
|
||||
return "\n\n".join(
|
||||
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
|
||||
"\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n"))
|
||||
for block in decoded.split("\n\n")
|
||||
).encode("utf-8")
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
|
||||
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
|
||||
def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str:
|
||||
if not line.startswith("data:"):
|
||||
return line
|
||||
try:
|
||||
|
|
@ -1252,14 +1394,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
except json.JSONDecodeError:
|
||||
return line
|
||||
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
|
||||
if not isinstance(data, dict):
|
||||
return line
|
||||
delta: Final = data.get("delta")
|
||||
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
|
||||
return line
|
||||
return "data: " + json.dumps(
|
||||
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
|
||||
)
|
||||
rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event)
|
||||
return line if rewritten is data else "data: " + json.dumps(rewritten)
|
||||
|
||||
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
|
||||
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_file_ids_from_messages,
|
||||
is_encrypted_reasoning_block,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
THOUGHT_SIGNATURE_SEPARATOR,
|
||||
|
|
@ -72,8 +73,13 @@ _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
|||
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
|
||||
|
||||
|
||||
def is_claude_code_user_agent(user_agent: str) -> bool:
|
||||
return user_agent.startswith("claude-cli/")
|
||||
"""Claude Code sends its API calls through the Anthropic SDK as `claude-cli/<version>` and its own
|
||||
fetches, such as gateway model discovery, as `claude-code/<version>`"""
|
||||
return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES)
|
||||
|
||||
|
||||
def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
|
||||
|
|
@ -1201,6 +1207,32 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A
|
|||
return out
|
||||
|
||||
|
||||
def _without_encrypted_reasoning_blocks(message: dict) -> dict | None: # mutable-ok: Anthropic message payload shape
|
||||
if not isinstance(message, Mapping):
|
||||
return message
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return message
|
||||
kept: Final = [b for b in content if not is_encrypted_reasoning_block(b)] # mutable-ok: API message payload
|
||||
if len(kept) == len(content):
|
||||
return message
|
||||
if not kept:
|
||||
return None
|
||||
return {**message, "content": kept} # mutable-ok: API message payload
|
||||
|
||||
|
||||
def strip_encrypted_reasoning_blocks_from_anthropic_messages(
|
||||
messages: Sequence[dict], # mutable-ok: Anthropic message payload shape
|
||||
) -> list[dict]: # mutable-ok: AnthropicMessagesRequest.messages is typed list[dict]
|
||||
"""
|
||||
Drop thinking / redacted_thinking blocks that carry another provider's encrypted
|
||||
reasoning (a turn the Responses API bridge served) before the request reaches
|
||||
Anthropic, which cannot verify them. Anthropic's own signed blocks are kept.
|
||||
"""
|
||||
stripped: Final = (_without_encrypted_reasoning_blocks(m) for m in messages)
|
||||
return [m for m in stripped if m is not None] # mutable-ok: API message payload
|
||||
|
||||
|
||||
def strip_thinking_blocks_from_anthropic_messages_request_dict(
|
||||
data: dict[str, Any],
|
||||
) -> None:
|
||||
|
|
@ -1629,11 +1661,16 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
|
|||
|
||||
|
||||
def _anthropic_model_entry(
|
||||
model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str]
|
||||
model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str], listed_ids: Mapping[str, str]
|
||||
) -> Mapping[str, object]:
|
||||
listed_id: Final = listed_ids.get(model["id"])
|
||||
source: Final[Mapping[str, object]] = (
|
||||
MappingProxyType({"source_model": model["id"]}) if listed_id is not None else MappingProxyType({})
|
||||
)
|
||||
return { # mutable-ok: JSON response body, serialized by the route and never mutated
|
||||
"type": "model",
|
||||
"id": model["id"],
|
||||
"id": listed_id or model["id"],
|
||||
**source,
|
||||
"display_name": display_names.get(model["id"], model["id"]),
|
||||
"created_at": created_at,
|
||||
"max_input_tokens": model.get("max_input_tokens"),
|
||||
|
|
@ -1644,6 +1681,7 @@ def _anthropic_model_entry(
|
|||
def create_anthropic_model_list_response(
|
||||
models: Sequence[ModelInfoResponse],
|
||||
display_names: Mapping[str, str] = MappingProxyType({}),
|
||||
listed_ids: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> Mapping[str, object]:
|
||||
"""Build the Anthropic-native /v1/models envelope.
|
||||
|
||||
|
|
@ -1653,17 +1691,19 @@ def create_anthropic_model_list_response(
|
|||
over from the OpenAI-shaped listing, named as the Messages API names them, and
|
||||
are always present because the vendor shape declares them nullable, not optional.
|
||||
display_names maps a listed model id to a configured human-readable name; ids
|
||||
without an entry fall back to the id itself, matching the vendor behavior
|
||||
without an entry fall back to the id itself, matching the vendor behavior.
|
||||
listed_ids maps a model id to the id the caller should see it under (the Claude
|
||||
Code view); ids without an entry are listed as they are
|
||||
"""
|
||||
created_at: Final = (
|
||||
datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
)
|
||||
data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated
|
||||
_anthropic_model_entry(model, created_at, display_names) for model in models
|
||||
_anthropic_model_entry(model, created_at, display_names, listed_ids) for model in models
|
||||
]
|
||||
return { # mutable-ok: JSON response body, serialized by the route and never mutated
|
||||
"data": data,
|
||||
"has_more": False,
|
||||
"first_id": models[0]["id"] if models else None,
|
||||
"last_id": models[-1]["id"] if models else None,
|
||||
"first_id": data[0]["id"] if data else None,
|
||||
"last_id": data[-1]["id"] if data else None,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -113,6 +113,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import (
|
|||
from litellm.llms.anthropic.common_utils import (
|
||||
is_empty_unsigned_thinking_block,
|
||||
normalize_anthropic_tool_use_id,
|
||||
strip_encrypted_reasoning_blocks_from_anthropic_messages,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management import (
|
||||
PolyfillResult,
|
||||
|
|
@ -417,7 +418,8 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
model: str | None = None,
|
||||
) -> list:
|
||||
new_messages: Final[list[AllMessageValues]] = []
|
||||
for m in messages:
|
||||
replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
|
||||
for m in replayable_messages:
|
||||
user_message: ChatCompletionUserMessage | None = None
|
||||
tool_message_list: list[ChatCompletionToolMessage] = []
|
||||
new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = []
|
||||
|
|
@ -1487,8 +1489,9 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
anthropic_content.insert(0, polyfill_result.compaction_block)
|
||||
|
||||
## extract finish reason
|
||||
openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop"
|
||||
translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic(
|
||||
openai_finish_reason=response.choices[0].finish_reason
|
||||
openai_finish_reason=openai_finish_reason
|
||||
)
|
||||
anthropic_finish_reason: Final = (
|
||||
"refusal"
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from ...common_utils import (
|
|||
AnthropicModelInfo,
|
||||
optionally_handle_anthropic_oauth,
|
||||
strip_advisor_blocks_from_messages,
|
||||
strip_encrypted_reasoning_blocks_from_anthropic_messages,
|
||||
)
|
||||
|
||||
DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
|
||||
|
|
@ -613,7 +614,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
messages = strip_advisor_blocks_from_messages(messages)
|
||||
|
||||
anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest(
|
||||
messages=messages,
|
||||
messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(messages),
|
||||
max_tokens=max_tokens,
|
||||
model=model,
|
||||
**anthropic_messages_optional_request_params,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
|
|||
AnthropicMessagesResponse,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from ..utils import litellm_logging_obj_from_kwargs, local_model_name
|
||||
from .streaming_iterator import AnthropicResponsesStreamWrapper
|
||||
|
|
@ -34,6 +35,15 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str,
|
|||
return extra_kwargs or {}
|
||||
|
||||
|
||||
def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: object) -> bool:
|
||||
provider: Final = (
|
||||
custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1]
|
||||
)
|
||||
provider_model: Final = local_model_name(model, provider)
|
||||
responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(provider, provider_model)
|
||||
return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model)
|
||||
|
||||
|
||||
def _build_responses_kwargs(
|
||||
*,
|
||||
max_tokens: int,
|
||||
|
|
@ -85,8 +95,13 @@ def _build_responses_kwargs(
|
|||
request_data["output_format"] = output_format
|
||||
|
||||
anthropic_request: Final = AnthropicMessagesRequest(**request_data)
|
||||
responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request)
|
||||
forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs)
|
||||
responses_kwargs: Final = _ADAPTER.translate_request(
|
||||
anthropic_request,
|
||||
include_encrypted_reasoning=_provider_returns_encrypted_reasoning(
|
||||
model, forwarded_kwargs.get("custom_llm_provider")
|
||||
),
|
||||
)
|
||||
|
||||
# Normalize reasoning effort based on model capabilities
|
||||
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
|
||||
|
|
@ -111,7 +126,7 @@ def _build_responses_kwargs(
|
|||
responses_kwargs["stream"] = True
|
||||
|
||||
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
|
||||
excluded: Final = {"anthropic_messages"}
|
||||
excluded: Final = frozenset(("anthropic_messages",))
|
||||
for key, value in forwarded_kwargs.items():
|
||||
if key == "litellm_logging_obj" and value is not None:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -132,6 +147,14 @@ def _build_responses_kwargs(
|
|||
if explicit_prompt_cache_key is not None:
|
||||
responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key
|
||||
|
||||
deployment_include: Final = forwarded_kwargs.get("include")
|
||||
bridge_include: Final = responses_kwargs.get("include")
|
||||
if isinstance(deployment_include, list) and isinstance(bridge_include, list):
|
||||
responses_kwargs["include"] = [
|
||||
*bridge_include,
|
||||
*(item for item in deployment_include if item not in bridge_include),
|
||||
]
|
||||
|
||||
return responses_kwargs
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9,13 +9,19 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
refusal_stop_details,
|
||||
responses_output_refusal_text,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
|
||||
|
||||
from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
|
||||
from .transformation import (
|
||||
REASONING_SUMMARY_PART_SEPARATOR,
|
||||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
|
|
@ -29,9 +35,10 @@ class AnthropicResponsesStreamWrapper:
|
|||
response.created -> message_start
|
||||
response.output_item.added -> content_block_start (if message/function_call)
|
||||
response.output_text.delta -> content_block_delta (text_delta)
|
||||
response.reasoning_summary_part.added -> content_block_delta (thinking_delta separator)
|
||||
response.reasoning_summary_text.delta -> content_block_delta (thinking_delta)
|
||||
response.function_call_arguments.delta -> content_block_delta (input_json_delta)
|
||||
response.output_item.done -> content_block_stop
|
||||
response.output_item.done -> content_block_delta (signature_delta) + content_block_stop
|
||||
response.completed -> message_delta + message_stop
|
||||
"""
|
||||
|
||||
|
|
@ -94,6 +101,38 @@ class AnthropicResponsesStreamWrapper:
|
|||
)
|
||||
return block_idx
|
||||
|
||||
@staticmethod
|
||||
def _field(source: object, name: str) -> object:
|
||||
return source.get(name) if isinstance(source, dict) else getattr(source, name, None)
|
||||
|
||||
def _close_reasoning_item(self, item: object, item_id: str | None) -> None:
|
||||
block_idx: Final = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
|
||||
encrypted_content: Final = self._field(item, "encrypted_content")
|
||||
signature: Final = (
|
||||
encrypted_reasoning_signature(encrypted_content)
|
||||
if isinstance(encrypted_content, str) and encrypted_content
|
||||
else None
|
||||
)
|
||||
if block_idx < 0 and signature is None:
|
||||
return
|
||||
if block_idx < 0:
|
||||
redacted_idx: Final = self._open_block(
|
||||
item_id,
|
||||
{"type": "redacted_thinking", "data": signature}, # mutable-ok: API message payload
|
||||
)
|
||||
stop: Final = {"type": "content_block_stop", "index": redacted_idx} # mutable-ok: API message payload
|
||||
self._chunk_queue.append(stop)
|
||||
return
|
||||
if signature is not None:
|
||||
self._chunk_queue.append(
|
||||
{ # mutable-ok: API message payload
|
||||
"type": "content_block_delta",
|
||||
"index": block_idx,
|
||||
"delta": {"type": "signature_delta", "signature": signature}, # mutable-ok: API message payload
|
||||
}
|
||||
)
|
||||
self._chunk_queue.append({"type": "content_block_stop", "index": block_idx}) # mutable-ok: API message payload
|
||||
|
||||
def _process_event(self, event: object) -> None:
|
||||
"""Convert one Responses API event into zero or more Anthropic chunks queued for emission."""
|
||||
event_type = getattr(event, "type", None)
|
||||
|
|
@ -175,6 +214,26 @@ class AnthropicResponsesStreamWrapper:
|
|||
)
|
||||
return
|
||||
|
||||
if event_type == "response.reasoning_summary_part.added":
|
||||
part_item_id: Final = self._field(event, "item_id")
|
||||
summary_index: Final = self._field(event, "summary_index")
|
||||
part_block_idx: Final = (
|
||||
self._item_id_to_block_index.get(part_item_id, -1) if isinstance(part_item_id, str) else -1
|
||||
)
|
||||
if part_block_idx < 0 or not isinstance(summary_index, int) or summary_index == 0:
|
||||
return
|
||||
self._chunk_queue.append(
|
||||
{ # mutable-ok: API message payload
|
||||
"type": "content_block_delta",
|
||||
"index": part_block_idx,
|
||||
"delta": { # mutable-ok: API message payload
|
||||
"type": "thinking_delta",
|
||||
"thinking": REASONING_SUMMARY_PART_SEPARATOR,
|
||||
},
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
# ---- reasoning summary text delta ----
|
||||
if event_type == "response.reasoning_summary_text.delta":
|
||||
item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None)
|
||||
|
|
@ -220,6 +279,9 @@ class AnthropicResponsesStreamWrapper:
|
|||
item_id = (
|
||||
getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None
|
||||
)
|
||||
if self._field(item, "type") == "reasoning":
|
||||
self._close_reasoning_item(item, item_id)
|
||||
return
|
||||
block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
|
||||
if block_idx < 0:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -13,7 +13,8 @@ from typing import Any, Final, cast
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
TOOL_RESULT_IMAGE_BOUNDARY,
|
||||
TOOL_RESULT_IMAGE_PLACEHOLDER,
|
||||
responses_reasoning_item_from_thinking_blocks,
|
||||
encrypted_reasoning_signature,
|
||||
responses_reasoning_items_from_thinking_blocks,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.litellm_core_utils.reasoning_effort_utils import (
|
||||
|
|
@ -33,6 +34,7 @@ from litellm.types.llms.anthropic import (
|
|||
AnthropicFinishReason,
|
||||
AnthropicMessagesRequest,
|
||||
AnthropicMessagesToolChoice,
|
||||
AnthropicResponseContentBlockRedactedThinking,
|
||||
AnthropicResponseContentBlockText,
|
||||
AnthropicResponseContentBlockThinking,
|
||||
AnthropicResponseContentBlockToolUse,
|
||||
|
|
@ -43,11 +45,13 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
|
|||
AnthropicUsage,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionThinkingBlock,
|
||||
ResponseAPIUsage,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
|
||||
REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n"
|
||||
RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content"
|
||||
|
||||
|
||||
class LiteLLMAnthropicToResponsesAPIAdapter:
|
||||
"""
|
||||
|
|
@ -163,49 +167,55 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
return str(getattr(part, "text", None) or "")
|
||||
|
||||
@classmethod
|
||||
def _thinking_blocks_from_reasoning_item(
|
||||
def _thinking_block_from_reasoning_item(
|
||||
cls,
|
||||
summary: Iterable[object],
|
||||
) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload
|
||||
"""Anthropic thinking blocks for one Responses reasoning item.
|
||||
encrypted_content: object,
|
||||
) -> dict[str, Any] | None: # mutable-ok: API message payload
|
||||
"""The one Anthropic block for a Responses reasoning item.
|
||||
|
||||
The signature stays empty: only Anthropic can sign a thinking block, and a stand-in
|
||||
value would be replayed as a real one and rejected by every backend that verifies it.
|
||||
The item's encrypted reasoning rides the block's opaque field (`signature`, or
|
||||
`data` when there is no summary text) so the client echoes it back and the next
|
||||
turn replays the very item OpenAI produced; without it the signature stays empty,
|
||||
since only Anthropic can sign a thinking block.
|
||||
"""
|
||||
return tuple(
|
||||
AnthropicResponseContentBlockThinking(
|
||||
type="thinking",
|
||||
thinking=text,
|
||||
signature=None,
|
||||
).model_dump()
|
||||
for part in summary
|
||||
if (text := cls._summary_part_text(part))
|
||||
text: Final = REASONING_SUMMARY_PART_SEPARATOR.join(
|
||||
part_text for part in summary if (part_text := cls._summary_part_text(part))
|
||||
)
|
||||
if not isinstance(encrypted_content, str) or not encrypted_content:
|
||||
if not text:
|
||||
return None
|
||||
return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=None).model_dump()
|
||||
signature: Final = encrypted_reasoning_signature(encrypted_content)
|
||||
if not text:
|
||||
return AnthropicResponseContentBlockRedactedThinking(type="redacted_thinking", data=signature).model_dump()
|
||||
return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=signature).model_dump()
|
||||
|
||||
@staticmethod
|
||||
def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str:
|
||||
"""Group a run of consecutive thinking blocks together; keep every other block alone."""
|
||||
index, block = indexed_block
|
||||
return "thinking" if block.get("type") == "thinking" else f"block:{index}"
|
||||
return "thinking" if block.get("type") in ("thinking", "redacted_thinking") else f"block:{index}"
|
||||
|
||||
@classmethod
|
||||
def _assistant_group_to_input_item(
|
||||
def _assistant_group_to_input_items(
|
||||
cls, group: tuple[Mapping[str, object], ...]
|
||||
) -> dict[str, Any] | None: # mutable-ok: API message payload
|
||||
) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload
|
||||
first: Final = group[0]
|
||||
btype: Final = first.get("type")
|
||||
if btype == "thinking":
|
||||
blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload
|
||||
reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks)
|
||||
return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload
|
||||
if btype in ("thinking", "redacted_thinking"):
|
||||
replayed: Final = responses_reasoning_items_from_thinking_blocks(group)
|
||||
return tuple(dict(item) for item in replayed) # mutable-ok: API message payload
|
||||
if btype == "tool_use":
|
||||
return { # mutable-ok: API message payload
|
||||
"type": "function_call",
|
||||
"call_id": first.get("id", ""),
|
||||
"name": first.get("name", ""),
|
||||
"arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload
|
||||
}
|
||||
return None
|
||||
return (
|
||||
{ # mutable-ok: API message payload
|
||||
"type": "function_call",
|
||||
"call_id": first.get("id", ""),
|
||||
"name": first.get("name", ""),
|
||||
"arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload
|
||||
},
|
||||
)
|
||||
return ()
|
||||
|
||||
def translate_messages_to_responses_input(
|
||||
self,
|
||||
|
|
@ -362,7 +372,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
input_items.extend(
|
||||
item
|
||||
for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key)
|
||||
if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None
|
||||
for item in self._assistant_group_to_input_items(tuple(block for _, block in group))
|
||||
)
|
||||
asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload
|
||||
{"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload
|
||||
|
|
@ -495,10 +505,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
def translate_request(
|
||||
self,
|
||||
anthropic_request: AnthropicMessagesRequest,
|
||||
include_encrypted_reasoning: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Translate a full Anthropic /v1/messages request dict to
|
||||
litellm.responses() / litellm.aresponses() kwargs.
|
||||
|
||||
``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content``
|
||||
on every call, so a reasoning model's items can be replayed intact next turn even
|
||||
when the client sent no ``thinking`` block; pass False for a provider whose
|
||||
Responses API rejects ``include``.
|
||||
"""
|
||||
model: Final[str] = anthropic_request["model"]
|
||||
messages_list: Final = cast(
|
||||
|
|
@ -528,6 +544,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
"model": model,
|
||||
"input": input_items,
|
||||
}
|
||||
if include_encrypted_reasoning:
|
||||
responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: API request payload
|
||||
|
||||
if system and not developer_parts:
|
||||
if isinstance(system, str):
|
||||
|
|
@ -634,7 +652,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
|
||||
for item in response.output:
|
||||
if isinstance(item, ResponseReasoningItem):
|
||||
content.extend(self._thinking_blocks_from_reasoning_item(item.summary))
|
||||
reasoning_block = self._thinking_block_from_reasoning_item(item.summary, item.encrypted_content)
|
||||
if reasoning_block is not None:
|
||||
content.append(reasoning_block)
|
||||
|
||||
elif isinstance(item, ResponseOutputMessage):
|
||||
for part in item.content:
|
||||
|
|
@ -684,11 +704,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
).model_dump()
|
||||
)
|
||||
elif item_type == "reasoning":
|
||||
content.extend(
|
||||
self._thinking_blocks_from_reasoning_item(
|
||||
cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json
|
||||
)
|
||||
reasoning_block = self._thinking_block_from_reasoning_item(
|
||||
cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json
|
||||
item.get("encrypted_content"),
|
||||
)
|
||||
if reasoning_block is not None:
|
||||
content.append(reasoning_block)
|
||||
elif item_type == "function_call":
|
||||
try:
|
||||
input_data = json.loads(item.get("arguments", "{}"))
|
||||
|
|
|
|||
|
|
@ -7,8 +7,9 @@ from httpx._models import Headers, Response
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
flatten_combinators_and_drop_non_python_regex_patterns,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
tool_with_sanitized_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_azure_openai_messages,
|
||||
|
|
@ -39,14 +40,17 @@ else:
|
|||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
def sanitized_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
tools: Final = optional_params.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
sanitized: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns)
|
||||
if isinstance(tool, dict)
|
||||
else tool
|
||||
for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
return MappingProxyType({"tools": sanitized})
|
||||
|
||||
|
||||
class AzureOpenAIConfig(BaseConfig):
|
||||
|
|
@ -278,7 +282,7 @@ class AzureOpenAIConfig(BaseConfig):
|
|||
"model": model,
|
||||
"messages": azure_messages,
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
**sanitized_tools_update(optional_params),
|
||||
}
|
||||
|
||||
def transform_response(
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.utils import get_model_info, supports_reasoning
|
||||
|
||||
from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig
|
||||
from .gpt_transformation import flattened_tools_update
|
||||
from .gpt_transformation import sanitized_tools_update
|
||||
|
||||
|
||||
class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
||||
|
|
@ -111,6 +111,6 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
|||
model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name
|
||||
flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
**sanitized_tools_update(optional_params),
|
||||
}
|
||||
return super().transform_request(model, messages, flattened_params, litellm_params, headers)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import DEFAULT_MAX_RETRIES
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.openai.common_utils import BaseOpenAILLM
|
||||
from litellm.secret_managers.get_azure_ad_token_provider import (
|
||||
|
|
@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM):
|
|||
if scope is None:
|
||||
scope = "https://cognitiveservices.azure.com/.default"
|
||||
|
||||
max_retries: Final = litellm_params.get("max_retries")
|
||||
configured_max_retries: Final = litellm_params.get("max_retries")
|
||||
max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries
|
||||
timeout: Final = litellm_params.get("timeout")
|
||||
if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret:
|
||||
verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth")
|
||||
|
|
@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM):
|
|||
else:
|
||||
azure_client_params["http_client"] = self._get_sync_http_client()
|
||||
|
||||
if max_retries is not None:
|
||||
azure_client_params["max_retries"] = max_retries
|
||||
azure_client_params["max_retries"] = max_retries
|
||||
if timeout is not None:
|
||||
azure_client_params["timeout"] = timeout
|
||||
|
||||
|
|
|
|||
|
|
@ -1,24 +1,102 @@
|
|||
import re
|
||||
from collections.abc import Callable, Collection, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
|
||||
import httpx
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.azure.common_utils import BaseAzureLLM
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
|
||||
from litellm.llms.base_llm.passthrough.transformation import (
|
||||
BasePassthroughConfig,
|
||||
RelayShape,
|
||||
logged_relay_shape,
|
||||
replace_path_segment,
|
||||
strip_leading_model_segment,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from httpx import URL
|
||||
|
||||
from litellm.types.utils import CostResponseTypes
|
||||
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
|
||||
|
||||
|
||||
class RelayedChatRequest(BaseModel):
|
||||
messages: Sequence[Mapping[str, object]] | None = None
|
||||
|
||||
|
||||
class RelayedCallDetails(BaseModel):
|
||||
request_data: RelayedChatRequest | None = None
|
||||
|
||||
|
||||
def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None:
|
||||
try:
|
||||
details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details)
|
||||
except ValidationError:
|
||||
return None
|
||||
return details.request_data.messages if details.request_data else None
|
||||
|
||||
|
||||
RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate)
|
||||
|
||||
OPENAI_RELAY_SHAPES: Final = (
|
||||
RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate),
|
||||
RESPONSES_RELAY_SHAPE,
|
||||
RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate),
|
||||
)
|
||||
|
||||
|
||||
def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None:
|
||||
"""A streaming logging object assembles the logged response from the terminal event, not from its body."""
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
|
||||
terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks)
|
||||
if terminal_event is None:
|
||||
return None
|
||||
logging_obj.call_type = (
|
||||
RESPONSES_RELAY_SHAPE.call_type.value
|
||||
) # rebind-ok: routes cost calculation to the relayed shape's pricing path
|
||||
return terminal_event
|
||||
|
||||
|
||||
AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(?<![^/])openai/deployments/([^/]+)")
|
||||
|
||||
|
||||
def azure_router_model_in_endpoint(endpoint: str, router_models: Collection[str]) -> str | None:
|
||||
parts: Final = endpoint.split("/")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
return next((part for part in parts if part in router_models), None)
|
||||
|
||||
|
||||
def foreign_azure_deployment(
|
||||
endpoint: str, model_group: str, served_models: Callable[[], Collection[str]]
|
||||
) -> str | None:
|
||||
match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint)
|
||||
if match is None:
|
||||
return None
|
||||
deployment: Final = match.group(1)
|
||||
if deployment == model_group:
|
||||
return None
|
||||
served: Final = frozenset(name.casefold() for name in served_models())
|
||||
return None if deployment.casefold() in served else deployment
|
||||
|
||||
|
||||
def without_api_version(api_base: str) -> str:
|
||||
url: Final = httpx.URL(api_base)
|
||||
kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version")
|
||||
return str(url.copy_with(params=httpx.QueryParams(kept_params)))
|
||||
|
||||
|
||||
class AzurePassthroughConfig(BasePassthroughConfig):
|
||||
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
|
||||
return "stream" in request_data
|
||||
return bool(request_data.get("stream"))
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
@ -36,14 +114,17 @@ class AzurePassthroughConfig(BasePassthroughConfig):
|
|||
|
||||
litellm_metadata: Final = litellm_params.get("litellm_metadata") or {}
|
||||
model_group: Final = litellm_metadata.get("model_group")
|
||||
if model_group and model_group in endpoint:
|
||||
endpoint = endpoint.replace(model_group, model)
|
||||
routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint
|
||||
native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,))
|
||||
|
||||
caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None
|
||||
relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url
|
||||
complete_url: Final = BaseAzureLLM._get_base_azure_url(
|
||||
api_base=base_target_url,
|
||||
litellm_params=litellm_params,
|
||||
route=endpoint,
|
||||
default_api_version=litellm_params.get("api_version"),
|
||||
api_base=relay_base,
|
||||
litellm_params=MappingProxyType(
|
||||
{**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")}
|
||||
),
|
||||
route=native_endpoint,
|
||||
)
|
||||
return (
|
||||
httpx.URL(complete_url),
|
||||
|
|
@ -92,13 +173,13 @@ class AzurePassthroughConfig(BasePassthroughConfig):
|
|||
request_data: dict,
|
||||
logging_obj: Logging,
|
||||
endpoint: str,
|
||||
) -> Optional["CostResponseTypes"]:
|
||||
) -> Optional["LoggedRelayResponse"]:
|
||||
from litellm import encoding
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
if "chat/completions" not in endpoint:
|
||||
return None
|
||||
return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
|
||||
|
||||
openai_chat_config: Final = OpenAIGPTConfig()
|
||||
|
||||
|
|
@ -116,3 +197,27 @@ class AzurePassthroughConfig(BasePassthroughConfig):
|
|||
)
|
||||
|
||||
return litellm_model_response
|
||||
|
||||
def handle_logging_collected_chunks(
|
||||
self,
|
||||
all_chunks: Sequence[str],
|
||||
litellm_logging_obj: Logging,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
endpoint: str,
|
||||
) -> Optional["LoggedRelayResponse"]:
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
|
||||
OpenAIPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix):
|
||||
return logged_responses_stream(all_chunks, litellm_logging_obj)
|
||||
if "chat/completions" not in endpoint:
|
||||
return None
|
||||
|
||||
return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser
|
||||
all_chunks=all_chunks,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
model=model,
|
||||
messages=_relayed_messages(litellm_logging_obj),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ import copy
|
|||
import enum
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from httpx import Response
|
||||
|
|
@ -15,7 +14,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
filter_value_from_dict,
|
||||
)
|
||||
from litellm.llms.azure.common_utils import BaseAzureLLM
|
||||
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
api_key_header_for_base,
|
||||
is_foundry_model_inference_base,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
|
||||
from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
|
||||
|
|
@ -146,11 +148,7 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
"""
|
||||
Returns True if the request should use `api-key` header for authentication.
|
||||
"""
|
||||
parsed_url: Final = urlparse(api_base)
|
||||
host: Final = parsed_url.hostname
|
||||
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
|
||||
return True
|
||||
return False
|
||||
return api_key_header_for_base(api_base) == "api-key"
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
|
|||
return "/openai/deployments" not in parsed.path
|
||||
|
||||
|
||||
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
|
||||
host: Final = urlparse(api_base).hostname if api_base else None
|
||||
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
|
||||
return "api-key"
|
||||
return "Authorization"
|
||||
|
||||
|
||||
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
|
||||
"""
|
||||
Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment.
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
|
|||
"""
|
||||
Get the appropriate image edit config for an Azure AI model.
|
||||
|
||||
- MAI models use /mai/v1/images/edits with multipart form data and size
|
||||
- MAI models use /mai/v1/images/edits with multipart form data
|
||||
- FLUX 2 models use JSON with base64 image
|
||||
- FLUX 1 models use multipart/form-data
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from httpx._types import RequestFiles
|
||||
|
|
@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import (
|
|||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
from litellm.types.llms.openai import FileTypes
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
|
@ -26,65 +25,8 @@ if TYPE_CHECKING:
|
|||
class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig):
|
||||
"""Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5)."""
|
||||
|
||||
DEFAULT_SIZE = "1024x1024"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
return ["prompt", "image", "model", "n", "size"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
image_edit_optional_params: ImageEditOptionalRequestParams,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
optional_params: Final[dict[str, Any]] = {}
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
|
||||
for key, value in dict(image_edit_optional_params).items():
|
||||
if value is None or key in optional_params:
|
||||
continue
|
||||
|
||||
if key in supported_params:
|
||||
if key == "size" and value:
|
||||
size_param = cast(str, value)
|
||||
self._validate_size_param(size_param)
|
||||
optional_params[key] = size_param
|
||||
else:
|
||||
optional_params[key] = value
|
||||
elif not drop_params:
|
||||
raise ValueError(
|
||||
f"Parameter {key} is not supported for model {model}. "
|
||||
f"Supported parameters are {supported_params}. "
|
||||
f"Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
if "size" not in optional_params:
|
||||
optional_params["size"] = self.DEFAULT_SIZE
|
||||
|
||||
return optional_params
|
||||
|
||||
def _validate_size_param(self, size: str) -> None:
|
||||
known_sizes: Final = {
|
||||
"1024x1024",
|
||||
"1792x1024",
|
||||
"1024x1792",
|
||||
"512x512",
|
||||
"256x256",
|
||||
}
|
||||
|
||||
if size in known_sizes:
|
||||
return
|
||||
|
||||
if "x" in size:
|
||||
try:
|
||||
tuple(map(int, size.lower().split("x", 1)))
|
||||
return
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
|
||||
)
|
||||
return ["prompt", "image", "model", "n"]
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
|
@ -21,6 +22,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
DEFAULT_WIDTH = 1024
|
||||
DEFAULT_HEIGHT = 1024
|
||||
|
||||
MAX_IMAGES_PER_REQUEST: Final = 1
|
||||
MIN_DIMENSION_PX: Final = 768
|
||||
MAX_TOTAL_PX: Final = 1_056_768
|
||||
|
||||
@staticmethod
|
||||
def get_mai_image_generation_url(
|
||||
api_base: str | None,
|
||||
|
|
@ -145,16 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
|
||||
if k in supported_params:
|
||||
if k == "size" and v:
|
||||
self._map_size_param(v, optional_params)
|
||||
self._map_size_param(v, optional_params, model)
|
||||
elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST:
|
||||
if not drop_params:
|
||||
raise self._unsupported(
|
||||
model,
|
||||
f"n={v} is not supported for model {model}. The Azure AI MAI image "
|
||||
f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per "
|
||||
"request and ignores any count, so a larger value would silently "
|
||||
"return fewer images than requested. Send one request per image, or "
|
||||
"set drop_params=True to drop n.",
|
||||
)
|
||||
else:
|
||||
optional_params[k] = v
|
||||
elif k in ("width", "height"):
|
||||
optional_params[k] = v
|
||||
elif not drop_params:
|
||||
raise ValueError(
|
||||
raise self._unsupported(
|
||||
model,
|
||||
f"Parameter {k} is not supported for model {model}. "
|
||||
f"Supported parameters are {supported_params} and width/height. "
|
||||
f"Set drop_params=True to drop unsupported parameters."
|
||||
f"Set drop_params=True to drop unsupported parameters.",
|
||||
)
|
||||
|
||||
if "width" not in optional_params:
|
||||
|
|
@ -165,7 +181,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
optional_params.pop("size", None)
|
||||
return optional_params
|
||||
|
||||
def _map_size_param(self, size: str, optional_params: dict) -> None:
|
||||
@staticmethod
|
||||
def _unsupported(model: str, message: str) -> UnsupportedParamsError:
|
||||
return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model)
|
||||
|
||||
def _image_count(self, n: object, model: str) -> int:
|
||||
if isinstance(n, int):
|
||||
return n
|
||||
try:
|
||||
return int(str(n))
|
||||
except ValueError:
|
||||
raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.")
|
||||
|
||||
def _map_size_param(self, size: str, optional_params: dict, model: str) -> None:
|
||||
size_mapping: Final = {
|
||||
"1024x1024": (1024, 1024),
|
||||
"1792x1024": (1792, 1024),
|
||||
|
|
@ -176,19 +204,36 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
|
||||
if size in size_mapping:
|
||||
width, height = size_mapping[size]
|
||||
optional_params["width"] = width
|
||||
optional_params["height"] = height
|
||||
elif "x" in size:
|
||||
try:
|
||||
width, height = map(int, size.lower().split("x"))
|
||||
optional_params["width"] = width
|
||||
optional_params["height"] = height
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').")
|
||||
raise self._unsupported(
|
||||
model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
raise self._unsupported(
|
||||
model,
|
||||
f"Unsupported size value: '{size}'. "
|
||||
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
|
||||
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.",
|
||||
)
|
||||
|
||||
self._validate_dimensions(model=model, size=size, width=width, height=height)
|
||||
optional_params["width"] = width
|
||||
optional_params["height"] = height
|
||||
|
||||
def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None:
|
||||
if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX:
|
||||
raise self._unsupported(
|
||||
model,
|
||||
f"Unsupported size value: '{size}'. Azure AI MAI image models require width and "
|
||||
f"height of at least {self.MIN_DIMENSION_PX} pixels.",
|
||||
)
|
||||
if width * height > self.MAX_TOTAL_PX:
|
||||
raise self._unsupported(
|
||||
model,
|
||||
f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most "
|
||||
f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).",
|
||||
)
|
||||
|
||||
def transform_image_generation_response(
|
||||
|
|
|
|||
232
litellm/llms/azure_ai/passthrough/transformation.py
Normal file
232
litellm/llms/azure_ai/passthrough/transformation.py
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
api_key_header_for_base,
|
||||
get_azure_ai_auth_headers,
|
||||
)
|
||||
from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config
|
||||
from litellm.llms.base_llm.passthrough.transformation import (
|
||||
BasePassthroughConfig,
|
||||
RelayShape,
|
||||
logged_relay_shape,
|
||||
strip_leading_model_segment,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from httpx import URL, Response
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
|
||||
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
|
||||
|
||||
|
||||
EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
class PassthroughMetadata(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
model_group: str = ""
|
||||
|
||||
|
||||
def model_group_from(litellm_params: Mapping[str, object]) -> str:
|
||||
try:
|
||||
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
|
||||
except ValidationError:
|
||||
return ""
|
||||
|
||||
|
||||
def api_version_from(litellm_params: Mapping[str, object]) -> str | None:
|
||||
try:
|
||||
return TypeAdapter(str | None).validate_python(litellm_params.get("api_version"))
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def foundry_root(api_base: str) -> str:
|
||||
url: Final = httpx.URL(api_base)
|
||||
segments: Final = tuple(segment for segment in url.path.split("/") if segment)
|
||||
root_segments: Final = segments[: segments.index("models")] if "models" in segments else segments
|
||||
return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/")
|
||||
|
||||
|
||||
def is_repeated_native_prefix(native_segments: tuple[str, ...], overlap: int) -> bool:
|
||||
return overlap == len(native_segments) or native_segments[0] == "openai"
|
||||
|
||||
|
||||
def without_repeated_native_prefix(root: str, native_endpoint: str) -> str:
|
||||
url: Final = httpx.URL(root)
|
||||
root_segments: Final = tuple(segment for segment in url.path.split("/") if segment)
|
||||
native_segments: Final = tuple(segment.casefold() for segment in native_endpoint.split("/") if segment)
|
||||
overlap: Final = next(
|
||||
(
|
||||
length
|
||||
for length in range(min(len(root_segments), len(native_segments)), 0, -1)
|
||||
if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length]
|
||||
and is_repeated_native_prefix(native_segments, length)
|
||||
),
|
||||
0,
|
||||
)
|
||||
kept_segments: Final = root_segments[: len(root_segments) - overlap]
|
||||
return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/")
|
||||
|
||||
|
||||
def relay_query_params(
|
||||
request_query_params: Mapping[str, object] | None,
|
||||
deployment_api_version: str | None,
|
||||
api_base: str,
|
||||
) -> Mapping[str, object] | None:
|
||||
if request_query_params and "api-version" in request_query_params:
|
||||
return request_query_params
|
||||
api_version: Final = deployment_api_version or httpx.URL(api_base).params.get("api-version")
|
||||
if api_version is None:
|
||||
return request_query_params
|
||||
return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version})
|
||||
|
||||
|
||||
def relayed_body(httpx_response: Response) -> str | dict:
|
||||
try:
|
||||
body: Final[object] = httpx_response.json()
|
||||
except ValueError:
|
||||
return httpx_response.text
|
||||
return body if isinstance(body, dict) else httpx_response.text
|
||||
|
||||
|
||||
FOUNDRY_RELAY_SHAPES: Final = (
|
||||
RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate),
|
||||
RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate),
|
||||
)
|
||||
|
||||
|
||||
class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig):
|
||||
def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None:
|
||||
super().__init__()
|
||||
self.ocr_config_for: Final = ocr_config_for
|
||||
|
||||
def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool:
|
||||
return bool(request_data.get("stream"))
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
request_query_params: Mapping[str, object] | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[URL, str]:
|
||||
base_target_url: Final = self.get_api_base(api_base)
|
||||
if base_target_url is None:
|
||||
raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE")
|
||||
|
||||
native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params)))
|
||||
root: Final = without_repeated_native_prefix(foundry_root(base_target_url), native_endpoint)
|
||||
query_params: Final = relay_query_params(
|
||||
request_query_params, api_version_from(litellm_params), base_target_url
|
||||
)
|
||||
return (self.format_url(native_endpoint, root, query_params), root)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx
|
||||
auth_headers: Final = get_azure_ai_auth_headers(
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
api_key_header=api_key_header_for_base(api_base),
|
||||
)
|
||||
return {**headers, **auth_headers} # mutable-ok: base class contract returns dict for httpx
|
||||
|
||||
def logging_non_streaming_response(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
httpx_response: Response,
|
||||
request_data: Mapping[str, object],
|
||||
logging_obj: Logging,
|
||||
endpoint: str,
|
||||
) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None:
|
||||
from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig
|
||||
|
||||
chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
httpx_response=httpx_response,
|
||||
request_data=dict(request_data), # mutable-ok: AzurePassthroughConfig wants a dict
|
||||
logging_obj=logging_obj,
|
||||
endpoint=endpoint,
|
||||
)
|
||||
if chat_result is not None:
|
||||
return chat_result
|
||||
ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint)
|
||||
if ocr_result is not None:
|
||||
return ocr_result
|
||||
foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
|
||||
if foundry_result is not None:
|
||||
return foundry_result
|
||||
return StandardPassThroughResponseObject(response=relayed_body(httpx_response))
|
||||
|
||||
def logged_ocr_response(
|
||||
self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str
|
||||
) -> OCRResponse | None:
|
||||
ocr_config: Final = self.ocr_config_for(model)
|
||||
if ocr_config is None or httpx_response.status_code != 200:
|
||||
return None
|
||||
relayed_url: Final = httpx_response.request.url
|
||||
relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/")
|
||||
ocr_url: Final = httpx.URL(
|
||||
ocr_config.get_complete_url(
|
||||
api_base=relayed_origin,
|
||||
model=model,
|
||||
optional_params={}, # mutable-ok: BaseOCRConfig wants a dict
|
||||
)
|
||||
)
|
||||
known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params))
|
||||
native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes)
|
||||
if f"/{native_endpoint.strip('/')}" != ocr_url.path:
|
||||
return None
|
||||
try:
|
||||
ocr_response: Final = ocr_config.transform_ocr_response(
|
||||
model=model, raw_response=httpx_response, logging_obj=logging_obj
|
||||
)
|
||||
except (ValueError, AttributeError) as error:
|
||||
verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error)
|
||||
return None
|
||||
logging_obj.call_type = CallTypes.aocr.value # rebind-ok: routes cost calculation to the per-page OCR path
|
||||
return ocr_response
|
||||
|
||||
def handle_logging_collected_chunks(
|
||||
self,
|
||||
all_chunks: Sequence[str],
|
||||
litellm_logging_obj: Logging,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
endpoint: str,
|
||||
) -> LoggedRelayResponse | None:
|
||||
from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig
|
||||
|
||||
return AzurePassthroughConfig().handle_logging_collected_chunks(
|
||||
all_chunks=all_chunks,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
endpoint=endpoint,
|
||||
)
|
||||
|
|
@ -52,13 +52,28 @@ class StreamingScanKey:
|
|||
|
||||
|
||||
class BaseTranslation(ABC):
|
||||
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
|
||||
delivers_ended_stream_rewrites: ClassVar[bool] = False
|
||||
"""Whether ``process_output_streaming_response`` accepts
|
||||
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
|
||||
stream, writes guardrail text rewrites back across ``responses_so_far`` so
|
||||
a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
|
||||
text rewrites on every other translation, are undeliverable: the pipeline
|
||||
executor discards them and releases the original chunks."""
|
||||
stream, writes guardrail text and tool-call rewrites back across
|
||||
``responses_so_far`` so a buffered pipeline can release rewritten chunks,
|
||||
raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites
|
||||
on every other translation are undeliverable: the pipeline executor
|
||||
discards them and releases the original chunks."""
|
||||
|
||||
assembles_streamed_response: ClassVar[bool] = False
|
||||
"""Whether ``process_output_streaming_response`` stores the assembled response of an
|
||||
ended stream under ``request_data["response"]`` before scanning it, the way the chat,
|
||||
Responses, and Messages translations do. A streaming pipeline runs a guardrail that only
|
||||
has the legacy post-call hook against that response, so on a translation without it such
|
||||
a guardrail keeps running on its own."""
|
||||
|
||||
def post_call_hook_response(self, response: object) -> object:
|
||||
"""The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
|
||||
the object the translation stores under ``request_data["response"]`` while scanning an
|
||||
ended stream. Chat and Responses scan that shape already; a translation that scans a
|
||||
different one (Messages scans an OpenAI-shaped ModelResponse) overrides this."""
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def transform_user_api_key_dict_to_metadata(
|
||||
|
|
@ -175,9 +190,9 @@ class BaseTranslation(ABC):
|
|||
transformations (see ``StreamTransformSink``); base handlers ignore it.
|
||||
``deliver_ended_stream_rewrites`` is passed True only when the caller
|
||||
holds the whole buffered stream and the subclass declares
|
||||
``delivers_ended_stream_text_rewrites``: the handler then writes
|
||||
guardrail text rewrites back across ``responses_so_far`` instead of
|
||||
discarding them.
|
||||
``delivers_ended_stream_rewrites``: the handler then writes
|
||||
guardrail text and tool-call rewrites back across ``responses_so_far``
|
||||
instead of discarding them.
|
||||
"""
|
||||
return responses_so_far
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,14 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from abc import abstractmethod
|
||||
from typing import TYPE_CHECKING, Final, Optional, Union
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
from ..base_utils import BaseLLMModelInfo
|
||||
|
||||
|
|
@ -7,9 +16,68 @@ if TYPE_CHECKING:
|
|||
from httpx import URL, Headers, Response
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.utils import CostResponseTypes
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject
|
||||
|
||||
from ..chat.transformation import BaseLLMException
|
||||
from ..ocr.transformation import OCRResponse
|
||||
|
||||
LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent
|
||||
|
||||
|
||||
RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str:
|
||||
path: Final = endpoint.lstrip("/")
|
||||
for model_name in model_names:
|
||||
if not model_name:
|
||||
continue
|
||||
if path == model_name:
|
||||
return ""
|
||||
if path.startswith(f"{model_name}/"):
|
||||
return path[len(model_name) + 1 :]
|
||||
return path
|
||||
|
||||
|
||||
def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str:
|
||||
bounded_segment: Final = re.compile(rf"(?<![^/]){re.escape(segment)}(?![^/:])")
|
||||
return bounded_segment.sub(lambda _: replacement, endpoint)
|
||||
|
||||
|
||||
def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None:
|
||||
if httpx_response.status_code != 200:
|
||||
return None
|
||||
try:
|
||||
return RELAYED_JSON_OBJECT.validate_python(httpx_response.json())
|
||||
except (ValueError, ValidationError):
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RelayShape:
|
||||
path_suffix: str
|
||||
call_type: CallTypes
|
||||
parse: Callable[[Mapping[str, object]], LoggedRelayResponse]
|
||||
|
||||
|
||||
def logged_relay_shape(
|
||||
shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str
|
||||
) -> LoggedRelayResponse | None:
|
||||
relayed_path: Final = f"/{endpoint.strip('/')}"
|
||||
shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None)
|
||||
body: Final = relayed_json_object(httpx_response) if shape else None
|
||||
if shape is None or body is None:
|
||||
return None
|
||||
try:
|
||||
parsed: Final = shape.parse(body)
|
||||
except ValidationError:
|
||||
return None
|
||||
logging_obj.call_type = (
|
||||
shape.call_type.value
|
||||
) # rebind-ok: routes cost calculation to the relayed shape's pricing path
|
||||
return parsed
|
||||
|
||||
|
||||
class BasePassthroughConfig(BaseLLMModelInfo):
|
||||
|
|
@ -23,8 +91,8 @@ class BasePassthroughConfig(BaseLLMModelInfo):
|
|||
self,
|
||||
endpoint: str,
|
||||
base_target_url: str,
|
||||
request_query_params: dict | None,
|
||||
) -> "URL":
|
||||
request_query_params: Mapping[str, object] | None,
|
||||
) -> URL:
|
||||
"""
|
||||
Helper function to add query params to the url
|
||||
Args:
|
||||
|
|
@ -58,7 +126,7 @@ class BasePassthroughConfig(BaseLLMModelInfo):
|
|||
endpoint: str,
|
||||
request_query_params: dict | None,
|
||||
litellm_params: dict,
|
||||
) -> tuple["URL", str]:
|
||||
) -> tuple[URL, str]:
|
||||
"""
|
||||
Get the complete url for the request
|
||||
Returns:
|
||||
|
|
@ -88,9 +156,7 @@ class BasePassthroughConfig(BaseLLMModelInfo):
|
|||
"""
|
||||
return headers, None
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, "Headers"]
|
||||
) -> "BaseLLMException":
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException:
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
|
||||
|
|
@ -99,21 +165,21 @@ class BasePassthroughConfig(BaseLLMModelInfo):
|
|||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
httpx_response: "Response",
|
||||
httpx_response: Response,
|
||||
request_data: dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
endpoint: str,
|
||||
) -> Optional["CostResponseTypes"]:
|
||||
) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None:
|
||||
pass
|
||||
|
||||
def handle_logging_collected_chunks(
|
||||
self,
|
||||
all_chunks: list[str],
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
endpoint: str,
|
||||
) -> Optional["CostResponseTypes"]:
|
||||
) -> LoggedRelayResponse | None:
|
||||
return None
|
||||
|
||||
def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]:
|
||||
|
|
|
|||
|
|
@ -1,21 +1,27 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import contextvars
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import urllib.parse
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime
|
||||
from functools import partial
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import (
|
||||
AWS_SIGNING_MAX_THREADS,
|
||||
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
|
||||
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES,
|
||||
BEDROCK_IAM_CACHE_MAX_ENTRIES,
|
||||
|
|
@ -26,6 +32,7 @@ from litellm.constants import (
|
|||
from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.secret_managers.main import get_secret, get_secret_str
|
||||
from litellm.types.llms.bedrock import AwsSessionTag
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.awsrequest import AWSPreparedRequest
|
||||
|
|
@ -47,6 +54,47 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile(
|
|||
|
||||
SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"})
|
||||
|
||||
_AWS_SESSION_TAGS_ADAPTER: Final[TypeAdapter[tuple[AwsSessionTag, ...]]] = TypeAdapter(tuple[AwsSessionTag, ...])
|
||||
|
||||
|
||||
def _canonical_aws_session_tags(raw_tags: object) -> tuple[AwsSessionTag, ...] | None:
|
||||
if raw_tags is None:
|
||||
return None
|
||||
try:
|
||||
validated: Final = _AWS_SESSION_TAGS_ADAPTER.validate_python(raw_tags)
|
||||
except ValidationError as e:
|
||||
raise ValueError(
|
||||
"Invalid 'aws_session_tags' value. Expected a list of {'Key': <str>, 'Value': <str>} dicts, "
|
||||
f"e.g. [{{'Key': 'team', 'Value': 'genai'}}]. Got: {raw_tags!r}"
|
||||
) from e
|
||||
return tuple(sorted(validated, key=lambda tag: tag["Key"]))
|
||||
|
||||
|
||||
class _AssumeRoleParams(TypedDict):
|
||||
RoleArn: ReadOnly[str]
|
||||
RoleSessionName: ReadOnly[str]
|
||||
ExternalId: ReadOnly[NotRequired[str]]
|
||||
Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]]
|
||||
|
||||
|
||||
def _assume_role_params(
|
||||
aws_role_name: str,
|
||||
aws_session_name: str,
|
||||
aws_external_id: str | None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None,
|
||||
) -> _AssumeRoleParams:
|
||||
match (aws_external_id, tuple(aws_session_tags or ())):
|
||||
case (None, ()):
|
||||
return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name)
|
||||
case (None, tags):
|
||||
return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, Tags=tags)
|
||||
case (external_id, ()):
|
||||
return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id)
|
||||
case (external_id, tags):
|
||||
return _AssumeRoleParams(
|
||||
RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id, Tags=tags
|
||||
)
|
||||
|
||||
|
||||
class BedrockRequestTarget(BaseModel):
|
||||
aws_region_name: str
|
||||
|
|
@ -80,7 +128,11 @@ class AwsAuthError(Exception):
|
|||
super().__init__(self.message) # Call the base class constructor with the parameters it needs
|
||||
|
||||
|
||||
class BaseAWSLLM:
|
||||
class SignsRequestsWithAWS:
|
||||
pass
|
||||
|
||||
|
||||
class BaseAWSLLM(SignsRequestsWithAWS):
|
||||
# Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
|
||||
# Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static
|
||||
# access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient
|
||||
|
|
@ -120,6 +172,7 @@ class BaseAWSLLM:
|
|||
"aws_sts_endpoint",
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"aws_external_id",
|
||||
"aws_session_tags",
|
||||
]
|
||||
|
||||
def _get_ssl_verify(self, ssl_verify: bool | str | None = None):
|
||||
|
|
@ -137,7 +190,7 @@ class BaseAWSLLM:
|
|||
|
||||
return get_ssl_verify(ssl_verify=ssl_verify)
|
||||
|
||||
def get_cache_key(self, credential_args: Mapping[str, str | bool | None]) -> str:
|
||||
def get_cache_key(self, credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None]) -> str:
|
||||
"""
|
||||
Generate a unique cache key based on the credential arguments.
|
||||
"""
|
||||
|
|
@ -147,7 +200,7 @@ class BaseAWSLLM:
|
|||
|
||||
def _get_or_set_cached_credentials(
|
||||
self,
|
||||
credential_args: Mapping[str, str | bool | None],
|
||||
credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None],
|
||||
credential_fetcher: Callable[[], tuple[Credentials, int | None]],
|
||||
) -> Any:
|
||||
"""
|
||||
|
|
@ -222,6 +275,7 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token: str | None = None,
|
||||
aws_sts_endpoint: str | None = None,
|
||||
aws_external_id: str | None = None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -258,6 +312,7 @@ class BaseAWSLLM:
|
|||
(aws_external_id, "AWS_EXTERNAL_ID"),
|
||||
)
|
||||
)
|
||||
session_tags: Final = _canonical_aws_session_tags(aws_session_tags)
|
||||
|
||||
verbose_logger.debug(
|
||||
"in get credentials\n"
|
||||
|
|
@ -270,7 +325,8 @@ class BaseAWSLLM:
|
|||
"aws_role_name=%s\n"
|
||||
"aws_web_identity_token=[set=%s]\n"
|
||||
"aws_sts_endpoint=%s\n"
|
||||
"aws_external_id=%s",
|
||||
"aws_external_id=%s\n"
|
||||
"aws_session_tags=%s",
|
||||
aws_access_key_id is not None,
|
||||
aws_secret_access_key is not None,
|
||||
aws_session_token is not None,
|
||||
|
|
@ -281,6 +337,7 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token is not None,
|
||||
aws_sts_endpoint,
|
||||
aws_external_id,
|
||||
session_tags,
|
||||
)
|
||||
|
||||
args: Final = {
|
||||
|
|
@ -294,6 +351,7 @@ class BaseAWSLLM:
|
|||
"aws_web_identity_token": aws_web_identity_token,
|
||||
"aws_sts_endpoint": aws_sts_endpoint,
|
||||
"aws_external_id": aws_external_id,
|
||||
"aws_session_tags": session_tags,
|
||||
"ssl_verify": ssl_verify,
|
||||
}
|
||||
|
||||
|
|
@ -336,6 +394,7 @@ class BaseAWSLLM:
|
|||
aws_region_name=aws_region_name,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=session_tags,
|
||||
ssl_verify=ssl_verify,
|
||||
),
|
||||
)
|
||||
|
|
@ -980,6 +1039,7 @@ class BaseAWSLLM:
|
|||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None = None,
|
||||
) -> dict:
|
||||
"""Handle cross-account role assumption for IRSA."""
|
||||
import boto3
|
||||
|
|
@ -1032,16 +1092,9 @@ class BaseAWSLLM:
|
|||
|
||||
# Now assume the target role
|
||||
verbose_logger.debug("Attempting to assume target role: %s with session: %s", aws_role_name, aws_session_name)
|
||||
assume_role_params: Final = {
|
||||
"RoleArn": aws_role_name,
|
||||
"RoleSessionName": aws_session_name,
|
||||
}
|
||||
|
||||
# Add ExternalId parameter if provided
|
||||
if aws_external_id is not None:
|
||||
assume_role_params["ExternalId"] = aws_external_id
|
||||
|
||||
return sts_client_with_creds.assume_role(**assume_role_params)
|
||||
return sts_client_with_creds.assume_role(
|
||||
**_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags)
|
||||
)
|
||||
|
||||
def _handle_irsa_same_account(
|
||||
self,
|
||||
|
|
@ -1051,6 +1104,7 @@ class BaseAWSLLM:
|
|||
aws_sts_endpoint: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_region_name: str | None = None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None = None,
|
||||
) -> dict:
|
||||
"""Handle same-account role assumption for IRSA."""
|
||||
import boto3
|
||||
|
|
@ -1074,16 +1128,9 @@ class BaseAWSLLM:
|
|||
|
||||
# Assume the role
|
||||
verbose_logger.debug("Attempting to assume role: %s with session: %s", aws_role_name, aws_session_name)
|
||||
assume_role_params: Final = {
|
||||
"RoleArn": aws_role_name,
|
||||
"RoleSessionName": aws_session_name,
|
||||
}
|
||||
|
||||
# Add ExternalId parameter if provided
|
||||
if aws_external_id is not None:
|
||||
assume_role_params["ExternalId"] = aws_external_id
|
||||
|
||||
return sts_client.assume_role(**assume_role_params)
|
||||
return sts_client.assume_role(
|
||||
**_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags)
|
||||
)
|
||||
|
||||
def _extract_credentials_and_ttl(self, sts_response: dict) -> tuple[Credentials, int | None]:
|
||||
"""Extract credentials and TTL from STS response.
|
||||
|
|
@ -1118,6 +1165,7 @@ class BaseAWSLLM:
|
|||
aws_region_name: str | None,
|
||||
aws_sts_endpoint: str | None,
|
||||
aws_external_id: str | None,
|
||||
aws_session_tags: tuple[AwsSessionTag, ...] | None,
|
||||
ssl_verify: bool | str | None,
|
||||
) -> tuple[Credentials, int | None]:
|
||||
"""
|
||||
|
|
@ -1144,6 +1192,7 @@ class BaseAWSLLM:
|
|||
aws_region_name=aws_region_name,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
ssl_verify=ssl_verify,
|
||||
)
|
||||
|
||||
|
|
@ -1159,6 +1208,7 @@ class BaseAWSLLM:
|
|||
aws_sts_endpoint: str | None = None,
|
||||
aws_external_id: str | None = None,
|
||||
ssl_verify: bool | str | None = None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None = None,
|
||||
) -> tuple[Credentials, int | None]:
|
||||
"""
|
||||
Authenticate with AWS Role
|
||||
|
|
@ -1189,6 +1239,7 @@ class BaseAWSLLM:
|
|||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
else:
|
||||
sts_response = self._handle_irsa_same_account(
|
||||
|
|
@ -1198,6 +1249,7 @@ class BaseAWSLLM:
|
|||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
ssl_verify=ssl_verify,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
|
||||
return self._extract_credentials_and_ttl(sts_response)
|
||||
|
|
@ -1234,14 +1286,9 @@ class BaseAWSLLM:
|
|||
**sts_client_kwargs,
|
||||
)
|
||||
|
||||
assume_role_params: Final = {
|
||||
"RoleArn": aws_role_name,
|
||||
"RoleSessionName": aws_session_name,
|
||||
}
|
||||
|
||||
# Add ExternalId parameter if provided
|
||||
if aws_external_id is not None:
|
||||
assume_role_params["ExternalId"] = aws_external_id
|
||||
assume_role_params: Final = _assume_role_params(
|
||||
aws_role_name, aws_session_name, aws_external_id, aws_session_tags
|
||||
)
|
||||
|
||||
try:
|
||||
sts_response = sts_client.assume_role(**assume_role_params)
|
||||
|
|
@ -1460,6 +1507,7 @@ class BaseAWSLLM:
|
|||
"aws_bedrock_runtime_endpoint", None
|
||||
) # https://bedrock-runtime.{region_name}.amazonaws.com
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.pop("aws_session_tags", None)
|
||||
|
||||
if bearer_token is not None:
|
||||
return BearerRequestTarget(
|
||||
|
|
@ -1478,6 +1526,7 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
return Boto3CredentialsInfo(
|
||||
credentials=credentials,
|
||||
|
|
@ -1623,6 +1672,7 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token: Final = optional_params.get("aws_web_identity_token", None)
|
||||
aws_sts_endpoint: Final = optional_params.get("aws_sts_endpoint", None)
|
||||
aws_external_id: Final = optional_params.get("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.get("aws_session_tags", None)
|
||||
aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model)
|
||||
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
|
|
@ -1636,6 +1686,7 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
|
||||
sigv4: Final = SigV4Auth(credentials, service_name, aws_region_name)
|
||||
|
|
@ -1668,3 +1719,52 @@ class BaseAWSLLM:
|
|||
request_headers_dict["Authorization"] = incoming_authorization
|
||||
|
||||
return request_headers_dict, request.body
|
||||
|
||||
|
||||
def sign_aws_json_post(
|
||||
get_credentials: Callable[[], Credentials],
|
||||
service_name: str,
|
||||
aws_region_name: str | None,
|
||||
url: str,
|
||||
body: str,
|
||||
headers: Mapping[str, str],
|
||||
) -> AWSPreparedRequest:
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.")
|
||||
|
||||
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
|
||||
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
|
||||
return aws_request.prepare()
|
||||
|
||||
|
||||
_SignParams = ParamSpec("_SignParams")
|
||||
_SignedRequest = TypeVar("_SignedRequest")
|
||||
|
||||
AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing")
|
||||
|
||||
|
||||
async def run_aws_signing(
|
||||
sign: Callable[_SignParams, _SignedRequest],
|
||||
/,
|
||||
*args: _SignParams.args,
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature
|
||||
) -> _SignedRequest:
|
||||
context: Final = contextvars.copy_context()
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs)
|
||||
)
|
||||
|
||||
|
||||
async def sign_request_off_loop_if_aws(
|
||||
provider_config: object,
|
||||
sign_request: Callable[_SignParams, _SignedRequest],
|
||||
/,
|
||||
*args: _SignParams.args,
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
|
||||
) -> _SignedRequest:
|
||||
if isinstance(provider_config, SignsRequestsWithAWS):
|
||||
return await run_aws_signing(sign_request, *args, **kwargs)
|
||||
return sign_request(*args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
|
|
@ -6,6 +6,7 @@ from openai.types.batch import BatchRequestCounts
|
|||
from openai.types.batch import Metadata as OpenAIBatchMetadata
|
||||
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.types.llms.bedrock import AwsSessionTag
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -116,6 +117,7 @@ class BedrockBatchesHandler:
|
|||
aws_web_identity_token: str | None = None,
|
||||
aws_sts_endpoint: str | None = None,
|
||||
aws_external_id: str | None = None,
|
||||
aws_session_tags: Sequence[AwsSessionTag] | None = None,
|
||||
**kwargs: object, # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim
|
||||
) -> "LiteLLMBatch":
|
||||
try:
|
||||
|
|
@ -139,6 +141,7 @@ class BedrockBatchesHandler:
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
|
||||
client: Final = boto3.client(
|
||||
|
|
@ -163,6 +166,7 @@ class BedrockBatchesHandler:
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -283,7 +287,7 @@ class BedrockBatchesHandler:
|
|||
``aws_session_token``, ``aws_profile_name``,
|
||||
``aws_role_name``, ``aws_session_name``,
|
||||
``aws_web_identity_token``, ``aws_sts_endpoint``,
|
||||
``aws_external_id``). Unknown keys are ignored.
|
||||
``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored.
|
||||
|
||||
Returns:
|
||||
``LiteLLMBatch`` shaped like an OpenAI Batch resource.
|
||||
|
|
@ -317,6 +321,7 @@ class BedrockBatchesHandler:
|
|||
aws_web_identity_token=kwargs.get("aws_web_identity_token"),
|
||||
aws_sts_endpoint=kwargs.get("aws_sts_endpoint"),
|
||||
aws_external_id=kwargs.get("aws_external_id"),
|
||||
aws_session_tags=kwargs.get("aws_session_tags"),
|
||||
)
|
||||
|
||||
client: Final = boto3.client(
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
|
|||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
|
||||
from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
|
||||
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
|
||||
|
||||
|
|
@ -136,7 +136,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
data: Final = json.dumps(request_data)
|
||||
|
||||
prepped: Final = self.get_request_headers(
|
||||
prepped: Final = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
|
||||
extra_headers=headers,
|
||||
|
|
@ -206,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
data: Final = json.dumps(request_data)
|
||||
|
||||
prepped: Final = self.get_request_headers(
|
||||
prepped: Final = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
|
||||
extra_headers=headers,
|
||||
|
|
@ -355,6 +357,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None)
|
||||
aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None)
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.pop("aws_session_tags", None)
|
||||
optional_params.pop("aws_region_name", None)
|
||||
|
||||
litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls
|
||||
|
|
@ -373,6 +376,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse`
|
|||
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import types
|
||||
from collections.abc import Mapping
|
||||
|
|
@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig):
|
|||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_openai_gpt_reasoning_model(model: str) -> bool:
|
||||
return re.search(r"openai\.gpt-\d", model) is not None
|
||||
|
||||
def _is_nova_2_model(self, model: str) -> bool:
|
||||
"""
|
||||
Check if the model is a Nova 2 model that supports reasoningConfig.
|
||||
|
|
@ -422,15 +427,15 @@ class AmazonConverseConfig(BaseConfig):
|
|||
"""
|
||||
Handle the reasoning_effort parameter based on the model type.
|
||||
|
||||
- GPT-OSS models: passed through unchanged via additionalModelRequestFields.
|
||||
- OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
|
||||
- GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields.
|
||||
- OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
|
||||
- Nova 2 models: transformed to reasoningConfig.
|
||||
- Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on
|
||||
adaptive Claude 4.6 / 4.7).
|
||||
"""
|
||||
if "gpt-oss" in model:
|
||||
if "gpt-oss" in model or "deepseek" in model:
|
||||
optional_params["reasoning_effort"] = reasoning_effort
|
||||
elif "openai.gpt-5" in model:
|
||||
elif self._is_openai_gpt_reasoning_model(model):
|
||||
reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort}
|
||||
optional_params["reasoning"] = reasoning
|
||||
elif self._is_nova_2_model(model):
|
||||
|
|
@ -509,6 +514,36 @@ class AmazonConverseConfig(BaseConfig):
|
|||
)
|
||||
thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS
|
||||
|
||||
def _is_deepseek_model(self, model: str, base_model: str) -> bool:
|
||||
return "deepseek" in model or "deepseek" in base_model
|
||||
|
||||
def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool:
|
||||
return "deepseek.r1" in model or "deepseek.r1" in base_model
|
||||
|
||||
def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool:
|
||||
"""Whether the model accepts the Anthropic-shaped ``thinking`` request field.
|
||||
|
||||
Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons
|
||||
natively: R1 returns a 400 when the field is sent and V3 silently ignores it.
|
||||
"""
|
||||
if self._is_deepseek_model(model=model, base_model=base_model):
|
||||
return False
|
||||
return (
|
||||
"claude-3-7" in model
|
||||
or "claude-sonnet-4" in model
|
||||
or "claude-opus-4" in model
|
||||
or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
|
||||
or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
|
||||
)
|
||||
|
||||
def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool:
|
||||
"""Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse.
|
||||
|
||||
DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw
|
||||
``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts.
|
||||
"""
|
||||
return self._is_deepseek_r1_model(model=model, base_model=base_model)
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
from litellm.utils import supports_function_calling
|
||||
|
||||
|
|
@ -564,23 +599,20 @@ class AmazonConverseConfig(BaseConfig):
|
|||
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
supported_params.append("tool_choice")
|
||||
|
||||
if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model:
|
||||
if (
|
||||
"gpt-oss" in model
|
||||
or self._is_openai_gpt_reasoning_model(model)
|
||||
or self._is_openai_gpt_reasoning_model(base_model)
|
||||
):
|
||||
supported_params.append("reasoning_effort")
|
||||
elif self._is_deepseek_model(model=model, base_model=base_model):
|
||||
if not self._is_deepseek_r1_model(model=model, base_model=base_model):
|
||||
supported_params.append("reasoning_effort")
|
||||
elif self._is_nova_2_model(model):
|
||||
# Nova 2 models support reasoning_effort (transformed to reasoningConfig)
|
||||
# These models use a different reasoning structure than Anthropic's thinking parameter
|
||||
supported_params.append("reasoning_effort")
|
||||
elif (
|
||||
"claude-3-7" in model
|
||||
or "claude-sonnet-4" in model
|
||||
or "claude-opus-4" in model
|
||||
or "deepseek.r1" in model
|
||||
or supports_reasoning(
|
||||
model=model,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
)
|
||||
or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
|
||||
):
|
||||
elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model):
|
||||
supported_params.append("thinking")
|
||||
supported_params.append("reasoning_effort")
|
||||
supported_params.append("output_config")
|
||||
|
|
@ -872,6 +904,11 @@ class AmazonConverseConfig(BaseConfig):
|
|||
drop_params: bool,
|
||||
) -> dict:
|
||||
is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
|
||||
base_model: Final = BedrockModelInfo.get_base_model(model)
|
||||
drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model)
|
||||
drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param(
|
||||
model=model, base_model=base_model
|
||||
)
|
||||
|
||||
for param, value in non_default_params.items():
|
||||
if param == "response_format" and isinstance(value, dict):
|
||||
|
|
@ -920,7 +957,12 @@ class AmazonConverseConfig(BaseConfig):
|
|||
optional_params["_parallel_tool_use_config"] = {
|
||||
"tool_choice": {"type": "auto", "disable_parallel_tool_use": not value}
|
||||
}
|
||||
if param == "thinking" and "openai.gpt-5" not in model:
|
||||
if param == "thinking" and drop_thinking_param:
|
||||
verbose_logger.debug(
|
||||
"Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.",
|
||||
model,
|
||||
)
|
||||
elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model):
|
||||
if (
|
||||
isinstance(value, dict)
|
||||
and value.get("type") == "adaptive"
|
||||
|
|
@ -946,6 +988,11 @@ class AmazonConverseConfig(BaseConfig):
|
|||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model, optional_params=optional_params, custom_llm_provider="bedrock"
|
||||
)
|
||||
elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param:
|
||||
verbose_logger.debug(
|
||||
"Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.",
|
||||
model,
|
||||
)
|
||||
elif param == "reasoning_effort" and isinstance(value, str):
|
||||
self._handle_reasoning_effort_parameter(
|
||||
model=model, reasoning_effort=value, optional_params=optional_params
|
||||
|
|
@ -1805,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
data=request_data,
|
||||
messages=messages,
|
||||
encoding=encoding,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
|
||||
def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str:
|
||||
|
|
@ -2237,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
data: dict | str,
|
||||
messages: list,
|
||||
encoding,
|
||||
json_mode: bool | None = None,
|
||||
) -> ModelResponse:
|
||||
## LOGGING
|
||||
if logging_obj is not None:
|
||||
|
|
@ -2247,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
|
||||
json_mode: Final[bool | None] = optional_params.get("json_mode", None)
|
||||
resolved_json_mode: Final[bool | None] = (
|
||||
json_mode if json_mode is not None else optional_params.get("json_mode", None)
|
||||
)
|
||||
## RESPONSE OBJECT
|
||||
try:
|
||||
completion_response: Final = ConverseResponseBlock(**response.json())
|
||||
|
|
@ -2339,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
chat_completion_message["content"] = content_str
|
||||
filtered_tools: Final = self._filter_json_mode_tools(
|
||||
json_mode=json_mode,
|
||||
json_mode=resolved_json_mode,
|
||||
tools=tools,
|
||||
chat_completion_message=chat_completion_message,
|
||||
)
|
||||
|
|
@ -2363,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
# When json_mode filtered out all synthetic tool calls the response
|
||||
# is plain content, not a pending tool invocation. Fix finish_reason
|
||||
# so callers (e.g. OpenAI SDK) don't misinterpret it.
|
||||
if json_mode and not filtered_tools and tools:
|
||||
if resolved_json_mode and not filtered_tools and tools:
|
||||
initial_finish_reason = "stop"
|
||||
|
||||
(
|
||||
|
|
|
|||
|
|
@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=encoding,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
elif provider == "twelvelabs":
|
||||
return litellm.AmazonTwelveLabsPegasusConfig().transform_response(
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (
|
|||
"aws_web_identity_token",
|
||||
"aws_sts_endpoint",
|
||||
"aws_external_id",
|
||||
"aws_session_tags",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1663,6 +1664,7 @@ class CommonBatchFilesUtils:
|
|||
aws_web_identity_token=optional_params.get("aws_web_identity_token"),
|
||||
aws_sts_endpoint=optional_params.get("aws_sts_endpoint"),
|
||||
aws_external_id=optional_params.get("aws_external_id"),
|
||||
aws_session_tags=optional_params.get("aws_session_tags"),
|
||||
)
|
||||
|
||||
# Prepare the request data
|
||||
|
|
|
|||
|
|
@ -10,9 +10,10 @@ import httpx
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
|
||||
|
||||
class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
||||
|
|
@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
request_data: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
resolved_model: str,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Handle a CountTokens request using existing LiteLLM patterns.
|
||||
|
|
@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
# Extract api_key for bearer token auth if provided
|
||||
api_key: Final = litellm_params.get("api_key", None)
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
signed_headers, signed_body = self._sign_request(
|
||||
signed_headers, signed_body = await run_aws_signing(
|
||||
self._sign_request,
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
optional_params=litellm_params,
|
||||
|
|
@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
|
||||
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
|
||||
|
||||
response: Final = await async_client.post(
|
||||
endpoint_url,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Handles embedding calls to Bedrock's `/invoke` endpoint
|
|||
import copy
|
||||
import json
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Final, get_args, overload
|
||||
|
||||
import httpx
|
||||
|
|
@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import (
|
|||
)
|
||||
from litellm.types.utils import EmbeddingResponse, LlmProviders
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
|
||||
from .amazon_titan_g1_transformation import AmazonTitanG1Config
|
||||
|
|
@ -41,6 +41,20 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
def _sign_get_request(
|
||||
credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str
|
||||
) -> AWSPreparedRequest:
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers)
|
||||
SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request)
|
||||
return request.prepare()
|
||||
|
||||
|
||||
class BedrockEmbedding(BaseAWSLLM):
|
||||
@overload
|
||||
def _load_credentials(
|
||||
|
|
@ -73,6 +87,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None)
|
||||
aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None)
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.pop("aws_session_tags", None)
|
||||
|
||||
### SET REGION NAME ###
|
||||
if aws_region_name is None:
|
||||
|
|
@ -103,6 +118,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
)
|
||||
return credentials, aws_region_name
|
||||
|
|
@ -342,7 +358,8 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
if extra_headers is not None:
|
||||
headers = {"Content-Type": "application/json", **extra_headers}
|
||||
|
||||
prepped = self.get_request_headers(
|
||||
prepped = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=aws_region_name,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -600,9 +617,6 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
dict: Status response from AWS Bedrock
|
||||
"""
|
||||
|
||||
# Get AWS credentials using the same method as other Bedrock methods
|
||||
credentials, _ = self._load_credentials(kwargs)
|
||||
|
||||
# Get the runtime endpoint
|
||||
endpoint_url, _ = self.get_runtime_endpoint(
|
||||
api_base=None,
|
||||
|
|
@ -619,27 +633,13 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
# Prepare headers for GET request
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
|
||||
# Use AWSRequest directly for GET requests (get_request_headers hardcodes POST)
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
def sign_status_request() -> AWSPreparedRequest:
|
||||
credentials, _ = self._load_credentials(kwargs)
|
||||
return _sign_get_request(
|
||||
credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name
|
||||
)
|
||||
|
||||
# Create AWSRequest with GET method and encoded URL
|
||||
request: Final = AWSRequest(
|
||||
method="GET",
|
||||
url=status_url,
|
||||
data=None, # GET request, no body
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
# Sign the request - SigV4Auth will create canonical string from request URL
|
||||
sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
|
||||
sigv4.add_auth(request)
|
||||
|
||||
# Prepare the request
|
||||
prepped: Final = request.prepare()
|
||||
prepped: Final = await run_aws_signing(sign_status_request)
|
||||
|
||||
# LOGGING
|
||||
if logging_obj is not None:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE
|
|||
from litellm.types.llms.openai import OpenAIRealtimeEvents
|
||||
from litellm.types.realtime import RealtimeResponseTransformInput
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .transformation import BedrockRealtimeConfig
|
||||
|
||||
|
|
@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model)
|
||||
|
||||
credentials: Final = self.get_credentials(
|
||||
credentials: Final = await run_aws_signing(
|
||||
self.get_credentials,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
|
|
@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
"or configure credentials in the environment"
|
||||
),
|
||||
)
|
||||
frozen_credentials: Final = credentials.get_frozen_credentials()
|
||||
frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials)
|
||||
|
||||
# Initialize Bedrock client with aws_sdk_bedrock_runtime
|
||||
config: Final = Config(
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from botocore.exceptions import (
|
|||
ProfileNotFound,
|
||||
)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
|
||||
|
|
@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str:
|
|||
)
|
||||
|
||||
|
||||
class BedrockMantleAuthMixin:
|
||||
class BedrockMantleAuthMixin(SignsRequestsWithAWS):
|
||||
_aws_signer: BaseAWSLLM
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -751,7 +751,9 @@ class AsyncHTTPHandler:
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
stream: bool = False,
|
||||
content: _RequestContent | None = None,
|
||||
follow_redirects: bool | None = None,
|
||||
):
|
||||
_follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT
|
||||
try:
|
||||
if timeout is None:
|
||||
timeout = self.timeout
|
||||
|
|
@ -769,22 +771,30 @@ class AsyncHTTPHandler:
|
|||
timeout=timeout,
|
||||
content=request_content,
|
||||
)
|
||||
response: Final = await self.client.send(req)
|
||||
response: Final = await self.client.send(req, follow_redirects=_follow_redirects)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except (httpx.RemoteProtocolError, httpx.ConnectError):
|
||||
# Retry the request with a new session if there is a connection error
|
||||
new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks)
|
||||
try:
|
||||
return await self.single_connection_post_request(
|
||||
url=url,
|
||||
client=new_client,
|
||||
data=data,
|
||||
retry_data, retry_content = _prepare_request_data_and_content(data, content)
|
||||
retry: Final = new_client.build_request(
|
||||
"PUT",
|
||||
url,
|
||||
data=retry_data,
|
||||
json=json,
|
||||
params=params,
|
||||
headers=headers,
|
||||
stream=stream,
|
||||
timeout=timeout,
|
||||
content=retry_content,
|
||||
)
|
||||
retried: Final = await new_client.send(retry, stream=stream, follow_redirects=_follow_redirects)
|
||||
try:
|
||||
retried.raise_for_status()
|
||||
except httpx.HTTPStatusError as retried_error:
|
||||
await _raise_masked_async_error(retried_error, stream)
|
||||
return retried
|
||||
finally:
|
||||
await new_client.aclose()
|
||||
except httpx.TimeoutException as e:
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import (
|
|||
BaseVectorStoreFilesConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws
|
||||
from litellm.llms.custom_httpx.container_handler import raise_for_error_status
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -637,7 +638,12 @@ class BaseLLMHTTPHandler:
|
|||
headers=request_headers,
|
||||
),
|
||||
)
|
||||
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
|
||||
signed_request: Final = await (
|
||||
run_aws_signing(sign_and_log, transformed)
|
||||
if isinstance(provider_config, SignsRequestsWithAWS)
|
||||
else asyncio.to_thread(sign_and_log, transformed)
|
||||
)
|
||||
return await dispatch_async(*signed_request)
|
||||
|
||||
return transform_then_dispatch()
|
||||
|
||||
|
|
@ -1973,7 +1979,9 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
signed_headers, signed_json_body = provider_config.sign_request(
|
||||
signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
provider_config,
|
||||
provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
|
|
@ -2074,7 +2082,9 @@ class BaseLLMHTTPHandler:
|
|||
max_attempts,
|
||||
)
|
||||
provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body)
|
||||
headers, signed_json_body = provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
provider_config,
|
||||
provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=optional_params_dict,
|
||||
request_data=request_body,
|
||||
|
|
@ -2234,7 +2244,9 @@ class BaseLLMHTTPHandler:
|
|||
stream=stream,
|
||||
)
|
||||
|
||||
headers, signed_json_body = anthropic_messages_provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
anthropic_messages_provider_config,
|
||||
anthropic_messages_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params
|
||||
request_data=request_body,
|
||||
|
|
@ -2910,7 +2922,9 @@ class BaseLLMHTTPHandler:
|
|||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers, signed_body = await sign_request_off_loop_if_aws(
|
||||
responses_api_provider_config,
|
||||
responses_api_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
|
|
@ -4618,7 +4632,9 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers, signed_body = await sign_request_off_loop_if_aws(
|
||||
responses_api_provider_config,
|
||||
responses_api_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
|
|
@ -9845,7 +9861,9 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
all_optional_params: Final[dict[str, object]] = dict(litellm_params)
|
||||
all_optional_params.update(vector_store_search_optional_params or {})
|
||||
headers, signed_json_body = vector_store_provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
vector_store_provider_config,
|
||||
vector_store_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=all_optional_params,
|
||||
request_data=request_body,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,8 @@
|
|||
Common utilities for the DashScope LLM provider.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -16,6 +17,27 @@ if TYPE_CHECKING:
|
|||
)
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
|
||||
DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1"
|
||||
DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks"
|
||||
|
||||
|
||||
def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
|
||||
if api_base is None:
|
||||
return None
|
||||
parsed: Final = urlparse(api_base)
|
||||
host: Final = parsed.hostname or ""
|
||||
on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com")
|
||||
if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH:
|
||||
return None
|
||||
return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}"
|
||||
|
||||
|
||||
def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str:
|
||||
remapped: Final = _rerank_base_for_chat_shaped_base(api_base)
|
||||
if api_base is not None and remapped is None:
|
||||
return api_base
|
||||
return get_secret_str(env_var) or remapped or default_rerank_base
|
||||
|
||||
|
||||
def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import Final
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
|
@ -51,7 +52,9 @@ class QwenAIPlatformRerankConfig(DashScopeRerankConfig):
|
|||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_RERANK") or QWEN_AI_PLATFORM_RERANK_API_BASE
|
||||
return resolve_dashscope_family_rerank_api_base(
|
||||
api_base, "QWEN_AI_PLATFORM_API_BASE_RERANK", QWEN_AI_PLATFORM_RERANK_API_BASE
|
||||
)
|
||||
|
||||
|
||||
class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig):
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import Final
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
|
@ -51,7 +52,9 @@ class QwenCloudRerankConfig(DashScopeRerankConfig):
|
|||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE_RERANK") or QWENCLOUD_RERANK_API_BASE
|
||||
return resolve_dashscope_family_rerank_api_base(
|
||||
api_base, "QWENCLOUD_API_BASE_RERANK", QWENCLOUD_RERANK_API_BASE
|
||||
)
|
||||
|
||||
|
||||
class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig):
|
||||
|
|
|
|||
|
|
@ -12,8 +12,12 @@ Endpoint
|
|||
- https://dashscope.aliyuncs.com/compatible-api/v1/reranks
|
||||
|
||||
Note: chat/embed live under `/compatible-mode/v1/`, but DashScope's rerank
|
||||
route is exposed under `/compatible-api/v1/reranks` per the docs. Override
|
||||
with `DASHSCOPE_API_BASE_RERANK` to point at a different host or path.
|
||||
route is exposed under `/compatible-api/v1/reranks` per the docs. A chat-shaped
|
||||
`.aliyuncs.com/compatible-mode/v1` base reaching this config (the chat default
|
||||
from `get_llm_provider`, or a `DASHSCOPE_API_BASE` env var) is redirected to
|
||||
the same host's rerank route, since `/compatible-mode/v1/reranks` is a dead
|
||||
route on every DashScope host. Override with `DASHSCOPE_API_BASE_RERANK` to
|
||||
point at a different host or path.
|
||||
|
||||
Empirically, qwen3-rerank accepts `return_documents=true` and echoes
|
||||
`results[].document.text` back, even though the public docs list the flag
|
||||
|
|
@ -40,7 +44,7 @@ from litellm.types.rerank import (
|
|||
RerankTokens,
|
||||
)
|
||||
|
||||
from ..common_utils import DashScopeError
|
||||
from ..common_utils import DashScopeError, resolve_dashscope_family_rerank_api_base
|
||||
|
||||
DEFAULT_RERANK_URL: Final = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
|
|
@ -67,9 +71,7 @@ class DashScopeRerankConfig(BaseRerankConfig):
|
|||
return resolved_api_key
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
if api_base is not None:
|
||||
return api_base
|
||||
return get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL
|
||||
return resolve_dashscope_family_rerank_api_base(api_base, "DASHSCOPE_API_BASE_RERANK", DEFAULT_RERANK_URL)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
_should_convert_tool_call_to_json_mode,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
_extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation
|
||||
strip_litellm_internal_message_fields,
|
||||
strip_name_from_message,
|
||||
)
|
||||
|
|
@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues
|
|||
from litellm.types.llms.databricks import (
|
||||
AllDatabricksContentValues,
|
||||
DatabricksChoice,
|
||||
DatabricksDelta,
|
||||
DatabricksFunction,
|
||||
DatabricksMessage,
|
||||
DatabricksResponse,
|
||||
DatabricksTool,
|
||||
)
|
||||
|
|
@ -247,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
api_base = self._get_api_base(api_base)
|
||||
complete_url: Final = f"{api_base}/chat/completions"
|
||||
use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2
|
||||
api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway)
|
||||
url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base
|
||||
complete_url: Final = f"{url_base}/chat/completions"
|
||||
return complete_url
|
||||
|
||||
def get_supported_openai_params(self, model: str | None = None) -> list:
|
||||
|
|
@ -534,6 +539,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
thinking_blocks.append(thinking_block)
|
||||
return reasoning_content, thinking_blocks
|
||||
|
||||
@staticmethod
|
||||
def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None:
|
||||
return delta.get("reasoning_content")
|
||||
|
||||
@staticmethod
|
||||
def resolve_reasoning_and_content(
|
||||
message: DatabricksMessage, block_reasoning_content: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
content_str: Final = DatabricksConfig.extract_content_str(message["content"])
|
||||
if block_reasoning_content is not None:
|
||||
return block_reasoning_content, content_str
|
||||
return _extract_reasoning_content({**message, "content": content_str})
|
||||
|
||||
@staticmethod
|
||||
def extract_citations(
|
||||
content: AllDatabricksContentValues | None,
|
||||
|
|
@ -577,14 +595,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
finish_reason = "stop"
|
||||
|
||||
if translated_message is None:
|
||||
## get the content str
|
||||
content_str = DatabricksConfig.extract_content_str(choice["message"]["content"])
|
||||
|
||||
## get the reasoning content
|
||||
(
|
||||
reasoning_content,
|
||||
block_reasoning_content,
|
||||
thinking_blocks,
|
||||
) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content"))
|
||||
reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content(
|
||||
choice["message"], block_reasoning_content
|
||||
)
|
||||
|
||||
citations = DatabricksConfig.extract_citations(choice["message"].get("content"))
|
||||
|
||||
|
|
@ -738,12 +755,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator):
|
|||
|
||||
# extract the reasoning content
|
||||
(
|
||||
reasoning_content,
|
||||
block_reasoning_content,
|
||||
thinking_blocks,
|
||||
) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content"))
|
||||
|
||||
choice["delta"]["content"] = content_str
|
||||
choice["delta"]["reasoning_content"] = reasoning_content
|
||||
choice["delta"]["reasoning_content"] = (
|
||||
block_reasoning_content
|
||||
if block_reasoning_content is not None
|
||||
else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"])
|
||||
)
|
||||
choice["delta"]["thinking_blocks"] = thinking_blocks
|
||||
translated_choices.append(choice)
|
||||
return ModelResponseStream(
|
||||
|
|
|
|||
|
|
@ -177,19 +177,13 @@ class DatabricksBase:
|
|||
# Default: just litellm
|
||||
return f"litellm/{version}"
|
||||
|
||||
def _get_api_base(self, api_base: str | None) -> str:
|
||||
"""
|
||||
Get the Databricks API base URL.
|
||||
|
||||
If not provided, attempts to get it from the Databricks SDK.
|
||||
"""
|
||||
def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str:
|
||||
if api_base is None:
|
||||
try:
|
||||
from databricks.sdk import WorkspaceClient
|
||||
|
||||
databricks_client: Final = WorkspaceClient()
|
||||
api_base = f"{databricks_client.config.host}/serving-endpoints"
|
||||
return api_base
|
||||
except ImportError:
|
||||
raise DatabricksException(
|
||||
status_code=400,
|
||||
|
|
@ -198,6 +192,18 @@ class DatabricksBase:
|
|||
"or install the databricks-sdk Python library."
|
||||
),
|
||||
)
|
||||
|
||||
if not use_ai_gateway:
|
||||
return api_base
|
||||
|
||||
normalized_api_base: Final = api_base.rstrip("/")
|
||||
if normalized_api_base.endswith("/ai-gateway/mlflow/v1"):
|
||||
return normalized_api_base
|
||||
if normalized_api_base.endswith("/serving-endpoints"):
|
||||
return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1"
|
||||
api_base_parts: Final = urlsplit(normalized_api_base)
|
||||
if api_base_parts.path in ("", "/"):
|
||||
return f"{normalized_api_base}/ai-gateway/mlflow/v1"
|
||||
return api_base
|
||||
|
||||
def _get_oauth_m2m_token(
|
||||
|
|
|
|||
9
litellm/llms/hosted_vllm/image_edit/__init__.py
Normal file
9
litellm/llms/hosted_vllm/image_edit/__init__.py
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
|
||||
from .transformation import HostedVLLMImageEditConfig
|
||||
|
||||
__all__ = ("HostedVLLMImageEditConfig",)
|
||||
|
||||
|
||||
def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig:
|
||||
return HostedVLLMImageEditConfig()
|
||||
43
litellm/llms/hosted_vllm/image_edit/transformation.py
Normal file
43
litellm/llms/hosted_vllm/image_edit/transformation.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"})
|
||||
|
||||
|
||||
class HostedVLLMImageEditConfig(OpenAIImageEditConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract
|
||||
return [ # mutable-ok: BaseImageEditConfig returns list
|
||||
param
|
||||
for param in super().get_supported_openai_params(model)
|
||||
if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT
|
||||
]
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: BaseImageEditConfig contract
|
||||
model: str,
|
||||
api_key: str | None = None,
|
||||
litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract
|
||||
api_base: str | None = None,
|
||||
) -> dict: # mutable-ok: BaseImageEditConfig contract
|
||||
resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key"
|
||||
return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
model: str,
|
||||
api_base: str | None,
|
||||
litellm_params: dict, # mutable-ok: BaseImageEditConfig contract
|
||||
) -> str:
|
||||
resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE")
|
||||
if resolved_api_base is None:
|
||||
raise ValueError(
|
||||
"api_base not set for Hosted VLLM images edits API. "
|
||||
"Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable"
|
||||
)
|
||||
trimmed: Final = resolved_api_base.rstrip("/")
|
||||
if trimmed.endswith("/v1"):
|
||||
return f"{trimmed}/images/edits"
|
||||
return f"{trimmed}/v1/images/edits"
|
||||
|
|
@ -157,9 +157,6 @@ class JinaAIRerankConfig(BaseRerankConfig):
|
|||
billed_units: RerankBilledUnits | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Jina AI reranker is priced at $0.000000018 per token.
|
||||
"""
|
||||
if (
|
||||
model_info is None
|
||||
or "input_cost_per_token" not in model_info
|
||||
|
|
|
|||
|
|
@ -19,10 +19,12 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
_should_convert_tool_call_to_json_mode,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_non_python_regex_patterns,
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
flatten_combinators_and_drop_non_python_regex_patterns,
|
||||
get_tool_call_names,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
tool_with_sanitized_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
async_convert_url_to_base64,
|
||||
|
|
@ -432,7 +434,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
custom_llm_provider, api_base
|
||||
)
|
||||
|
||||
def _flattened_tools_update_for_openai(
|
||||
def _sanitized_tools_update_for_openai(
|
||||
self,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
|
|
@ -440,22 +442,26 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
"""
|
||||
OpenAI's chat completions validator rejects tool `parameters` carrying
|
||||
'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every
|
||||
model family, unlike the Responses API, where GPT-5+ accepts them.
|
||||
model family, unlike the Responses API, where GPT-5+ accepts them, and
|
||||
a `pattern` Python's `re` cannot compile for every model family on both.
|
||||
A custom api_base on the `openai` provider is usually a proxy in front of
|
||||
the same validator, so regexes are dropped there too, while the lossier
|
||||
combinator flattening stays limited to api.openai.com hosts.
|
||||
"""
|
||||
tools: Final = optional_params.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return _NO_TOOLS_UPDATE
|
||||
provider: Final = litellm_params.get("custom_llm_provider")
|
||||
raw_api_base: Final = litellm_params.get("api_base")
|
||||
if not self._targets_openai_hosted_endpoint(
|
||||
provider if isinstance(provider, str) else None,
|
||||
raw_api_base if isinstance(raw_api_base, str) else None,
|
||||
):
|
||||
if not isinstance(tools, list) or provider != "openai":
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
raw_api_base: Final = litellm_params.get("api_base")
|
||||
sanitize: Final = (
|
||||
flatten_combinators_and_drop_non_python_regex_patterns
|
||||
if self._targets_openai_hosted_endpoint(provider, raw_api_base if isinstance(raw_api_base, str) else None)
|
||||
else drop_non_python_regex_patterns
|
||||
)
|
||||
sanitized: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_sanitized_parameters(tool, sanitize) if isinstance(tool, dict) else tool for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
return MappingProxyType({"tools": sanitized})
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
|
|
@ -489,7 +495,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
"model": model,
|
||||
"messages": messages,
|
||||
**optional_params,
|
||||
**self._flattened_tools_update_for_openai(optional_params, litellm_params),
|
||||
**self._sanitized_tools_update_for_openai(optional_params, litellm_params),
|
||||
}
|
||||
|
||||
async def async_transform_request(
|
||||
|
|
@ -521,7 +527,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
"model": model,
|
||||
"messages": transformed_messages,
|
||||
**optional_params,
|
||||
**self._flattened_tools_update_for_openai(optional_params, litellm_params),
|
||||
**self._sanitized_tools_update_for_openai(optional_params, litellm_params),
|
||||
}
|
||||
else:
|
||||
## allow for any object specific behaviour to be handled
|
||||
|
|
|
|||
|
|
@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
|
|||
coerce_stream_holdback_value,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionDeltaToolCall,
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
GenericGuardrailAPIInputs,
|
||||
ModelResponse,
|
||||
|
|
@ -78,7 +80,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
Methods can be overridden to customize behavior for different message formats.
|
||||
"""
|
||||
|
||||
delivers_ended_stream_text_rewrites = True
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
|
||||
"""
|
||||
|
|
@ -610,13 +613,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
deliver_ended_stream_rewrites: bool,
|
||||
) -> None:
|
||||
"""Ended-stream path: rebuild the full response, run the non-streaming
|
||||
output guardrail against it, and (when opted in) write any text rewrite
|
||||
back across the buffered chunks."""
|
||||
output guardrail against it, and (when opted in) write any text or
|
||||
tool-call rewrite back across the buffered chunks."""
|
||||
model_response: Final = cast(
|
||||
ModelResponse,
|
||||
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
|
||||
)
|
||||
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
|
||||
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
|
||||
await self.process_output_response(
|
||||
response=model_response,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
|
|
@ -624,13 +628,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
if deliver_ended_stream_rewrites:
|
||||
await self._write_ended_stream_text_rewrites(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrailed_response=model_response,
|
||||
pre_guardrail_texts=pre_guardrail_texts,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
if not deliver_ended_stream_rewrites:
|
||||
return
|
||||
guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
|
||||
await self._write_ended_stream_text_rewrites(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrailed_response=model_response,
|
||||
pre_guardrail_texts=pre_guardrail_texts,
|
||||
guardrail_name=guardrail_name,
|
||||
)
|
||||
self._write_ended_stream_tool_call_rewrites(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrailed_response=model_response,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_name,
|
||||
)
|
||||
|
||||
def build_stream_error_items(
|
||||
self,
|
||||
|
|
@ -1043,6 +1055,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]:
|
||||
return tuple(
|
||||
(tool_call.function.name, tool_call.function.arguments)
|
||||
for choice in response.choices
|
||||
for tool_call in choice.message.tool_calls or ()
|
||||
if isinstance(tool_call, ChatCompletionMessageToolCall)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _function_tool_call_fragments(
|
||||
responses_so_far: Sequence["ModelResponseStream"],
|
||||
) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
|
||||
"""Group the stream's function tool-call fragments by their tool-call index, in
|
||||
the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
|
||||
only the indices the builder keeps (an id and a name somewhere in the stream)."""
|
||||
fragments: Final = tuple(
|
||||
tool_call
|
||||
for response in responses_so_far
|
||||
for choice in response.choices
|
||||
for tool_call in choice.delta.tool_calls or ()
|
||||
if isinstance(tool_call, ChatCompletionDeltaToolCall)
|
||||
)
|
||||
identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
|
||||
named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
|
||||
return tuple(
|
||||
tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
|
||||
)
|
||||
|
||||
def _write_ended_stream_tool_call_rewrites(
|
||||
self,
|
||||
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
guardrailed_response: "ModelResponse",
|
||||
pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
|
||||
guardrail_name: str,
|
||||
) -> None:
|
||||
"""Write ended-stream guardrail tool-call rewrites back across the buffered
|
||||
chunks: the rewritten name and full arguments land in the tool call's first
|
||||
fragment and the arguments of its later fragments are blanked, mirroring the
|
||||
text write-back. A rewrite on a stream carrying more than one distinct choice
|
||||
index, or whose fragments do not line up with the rebuilt tool calls, is
|
||||
reported as undeliverable, so the pipeline executor discards it and releases
|
||||
the original chunks."""
|
||||
post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
stream_choice_indices: Final = frozenset(
|
||||
choice.index for response in responses_so_far for choice in response.choices
|
||||
)
|
||||
fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
|
||||
if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(guardrail_name)
|
||||
for before, (name, arguments), fragments in zip(
|
||||
pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
|
||||
):
|
||||
if (name, arguments) == before:
|
||||
continue
|
||||
head, *tail = fragments
|
||||
head.function.name = name
|
||||
head.function.arguments = arguments
|
||||
for fragment in tail:
|
||||
fragment.function.arguments = ""
|
||||
|
||||
async def _apply_guardrail_responses_to_output_streaming(
|
||||
self,
|
||||
responses: list["ModelResponseStream"],
|
||||
|
|
|
|||
|
|
@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
|
||||
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
tool_call_dict_from_output_item,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
|
|
@ -84,7 +84,6 @@ from litellm.types.llms.openai import (
|
|||
)
|
||||
from litellm.types.responses.main import (
|
||||
GenericResponseOutputItem,
|
||||
OutputFunctionToolCall,
|
||||
OutputText,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
|
@ -101,6 +100,72 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.openai import ResponseInputParam
|
||||
|
||||
|
||||
class _ToolCallShape(NamedTuple):
|
||||
name: str | None
|
||||
arguments: str
|
||||
|
||||
|
||||
class _ToolCallFunctionFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str | None = None
|
||||
arguments: str = ""
|
||||
|
||||
|
||||
class _ToolCallFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
function: _ToolCallFunctionFields
|
||||
|
||||
|
||||
def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
|
||||
return tuple(
|
||||
_ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
|
||||
for tool_call in tool_calls
|
||||
)
|
||||
|
||||
|
||||
def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
|
||||
payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
|
||||
try:
|
||||
fields: Final = _ToolCallFields.model_validate(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
|
||||
|
||||
|
||||
def _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls: Sequence[object] | None,
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
guardrail_name: str | None,
|
||||
) -> tuple[_ToolCallShape, ...]:
|
||||
if not pre_guardrail_tool_calls:
|
||||
return pre_guardrail_tool_calls
|
||||
if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
|
||||
"leaving the tool call output items unchanged",
|
||||
guardrail_name,
|
||||
"no" if returned_tool_calls is None else len(returned_tool_calls),
|
||||
len(pre_guardrail_tool_calls),
|
||||
)
|
||||
return pre_guardrail_tool_calls
|
||||
returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
|
||||
validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
|
||||
if len(validated_shapes) != len(returned_shapes):
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
|
||||
"leaving the tool call output items unchanged",
|
||||
guardrail_name,
|
||||
)
|
||||
return pre_guardrail_tool_calls
|
||||
return validated_shapes
|
||||
|
||||
|
||||
def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
|
||||
return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
|
||||
|
||||
|
||||
class ResponseOutputEnvelope(TypedDict, total=False):
|
||||
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
|
||||
|
||||
|
|
@ -128,6 +193,20 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
|
||||
_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"function_call": "arguments", "custom_tool_call": "input"}
|
||||
)
|
||||
_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset(
|
||||
{"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"}
|
||||
)
|
||||
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"}
|
||||
)
|
||||
_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset(
|
||||
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
|
||||
)
|
||||
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
|
||||
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"function_call_output": "output", "message": "content"}
|
||||
)
|
||||
|
|
@ -164,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
|
|||
return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts
|
||||
|
||||
|
||||
def _is_function_call_item(item: object) -> bool:
|
||||
return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
|
||||
def _is_tool_call_item(item: object) -> bool:
|
||||
return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES
|
||||
|
||||
|
||||
def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
|
||||
if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
|
||||
return None
|
||||
if isinstance(item, Mapping):
|
||||
return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects
|
||||
return item.model_dump() if isinstance(item, BaseModel) else None
|
||||
|
||||
|
||||
def _is_tool_call_output_item(item: object) -> bool:
|
||||
return _tool_call_output_item_mapping(item) is not None
|
||||
|
||||
|
||||
def _last_message_role(messages: Sequence[object]) -> str | None:
|
||||
|
|
@ -189,7 +280,7 @@ def _provenance_unit_bounds(
|
|||
start_indexes: Final = tuple(
|
||||
index
|
||||
for index in range(len(raw_input))
|
||||
if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
|
||||
if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
|
||||
)
|
||||
return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
|
||||
|
||||
|
|
@ -340,7 +431,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
Methods can be overridden to customize behavior for different message formats.
|
||||
"""
|
||||
|
||||
delivers_ended_stream_text_rewrites = True
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
|
||||
"""
|
||||
|
|
@ -587,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
- response.output is a list of output items
|
||||
- Each output item can be:
|
||||
* GenericResponseOutputItem with a content list of OutputText objects
|
||||
* ResponseFunctionToolCall with tool call data
|
||||
* ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data
|
||||
- Each OutputText object has a text field
|
||||
"""
|
||||
|
||||
|
|
@ -652,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if response_model:
|
||||
inputs["model"] = response_model
|
||||
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -660,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
# Step 3: Map guardrail responses back to original response structure
|
||||
await self._apply_guardrail_responses_to_output(
|
||||
|
|
@ -667,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
responses=guardrailed_texts,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
self._write_tool_call_rewrites_to_output(
|
||||
tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=post_guardrail_tool_calls,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
|
||||
|
||||
|
|
@ -754,6 +857,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if response_model:
|
||||
inputs["model"] = response_model
|
||||
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
|
|
@ -762,6 +866,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
# Write guardrailed texts back into the output items in-place.
|
||||
# final_chunk is a reference into responses_so_far so this
|
||||
|
|
@ -784,6 +893,13 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
stream_events=responses_so_far[:-1],
|
||||
rewrites_by_position=rewrites_by_position,
|
||||
)
|
||||
self._deliver_ended_stream_tool_call_rewrites(
|
||||
responses_so_far=responses_so_far,
|
||||
outputs=outputs,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=post_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
return responses_so_far
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
|
|
@ -894,6 +1010,148 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
continue
|
||||
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
|
||||
|
||||
def _deliver_ended_stream_tool_call_rewrites(
|
||||
self,
|
||||
responses_so_far: Sequence[object],
|
||||
outputs: Sequence[object],
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
guardrail_name: str,
|
||||
) -> None:
|
||||
"""Write ended-stream guardrail tool-call rewrites into the completed
|
||||
envelope's ``function_call`` and ``custom_tool_call`` items and sync the
|
||||
earlier stream events, keyed by ``call_id``. The guardrail sees the
|
||||
envelope's tool calls in output order, which is how a rewritten call
|
||||
finds its ``call_id``; the stream events find their call through the
|
||||
``call_id`` on ``output_item`` events and the ``item_id`` on argument
|
||||
and custom-input events, since an
|
||||
event's ``output_index`` need not match the envelope's (the chat bridge
|
||||
numbers tool calls from 1 while the envelope lists them after the
|
||||
message). A rewrite whose calls do not line up with the envelope, or
|
||||
whose events cannot be found, is reported as undeliverable, so the
|
||||
pipeline executor discards it and releases the original events."""
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
|
||||
call_ids: Final = tuple(
|
||||
call_id
|
||||
for output_item in tool_call_items
|
||||
if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
|
||||
)
|
||||
stream_events: Final = responses_so_far[:-1]
|
||||
call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events)
|
||||
event_call_ids: Final = tuple(
|
||||
self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events
|
||||
)
|
||||
rewrites_by_call_id: Final = MappingProxyType(
|
||||
{
|
||||
call_id: _tool_call_rewrite(before, after)
|
||||
for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if after != before
|
||||
}
|
||||
)
|
||||
unresolved_argument_event: Final = any(
|
||||
call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
|
||||
for event, call_id in zip(stream_events, event_call_ids)
|
||||
)
|
||||
if (
|
||||
len(call_ids) != len(tool_call_items)
|
||||
or len(frozenset(call_ids)) != len(call_ids)
|
||||
or len(call_ids) != len(post_guardrail_tool_calls)
|
||||
or unresolved_argument_event
|
||||
or not rewrites_by_call_id.keys() <= frozenset(event_call_ids)
|
||||
):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(guardrail_name)
|
||||
for output_item, rewrite in (
|
||||
(output_item, rewrites_by_call_id[call_id])
|
||||
for output_item, call_id in zip(tool_call_items, call_ids)
|
||||
if call_id in rewrites_by_call_id
|
||||
):
|
||||
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
delta_replacements: Final = MappingProxyType(
|
||||
{call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
|
||||
)
|
||||
for event, call_id in zip(stream_events, event_call_ids):
|
||||
if call_id not in rewrites_by_call_id:
|
||||
continue
|
||||
match stream_item_field(event, "type"):
|
||||
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES:
|
||||
self._write_event_field(event, "delta", next(delta_replacements[call_id]))
|
||||
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS:
|
||||
self._write_event_field(
|
||||
event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments
|
||||
)
|
||||
case "response.output_item.added":
|
||||
self._write_tool_call_item(
|
||||
stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
|
||||
)
|
||||
case "response.output_item.done":
|
||||
self._write_tool_call_item(
|
||||
stream_item_field(event, "item"),
|
||||
rewrites_by_call_id[call_id].name,
|
||||
rewrites_by_call_id[call_id].arguments,
|
||||
)
|
||||
case _:
|
||||
pass
|
||||
|
||||
def _write_tool_call_rewrites_to_output(
|
||||
self,
|
||||
tool_call_items: Sequence[object],
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
) -> None:
|
||||
if len(tool_call_items) != len(post_guardrail_tool_calls):
|
||||
return
|
||||
for output_item, rewrite in (
|
||||
(output_item, _tool_call_rewrite(before, after))
|
||||
for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if after != before
|
||||
):
|
||||
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
|
||||
items: Final = tuple(
|
||||
stream_item_field(event, "item")
|
||||
for event in stream_events
|
||||
if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
item_id: call_id
|
||||
for item in items
|
||||
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
|
||||
and isinstance(item_id := stream_item_field(item, "id"), str)
|
||||
and isinstance(call_id := stream_item_field(item, "call_id"), str)
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
|
||||
event_type: Final = stream_item_field(event, "type")
|
||||
if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES:
|
||||
item_id: Final = stream_item_field(event, "item_id")
|
||||
return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
|
||||
if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
|
||||
return None
|
||||
item: Final = stream_item_field(event, "item")
|
||||
call_id: Final = stream_item_field(item, "call_id")
|
||||
return (
|
||||
call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
|
||||
if item is None:
|
||||
return
|
||||
if name is not None:
|
||||
OpenAIResponsesHandler._write_event_field(item, "name", name)
|
||||
item_type: Final = stream_item_field(item, "type")
|
||||
if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
|
||||
OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
|
||||
|
||||
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
|
||||
"""
|
||||
Check if the streaming has ended.
|
||||
|
|
@ -920,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
def _completed_response_scan_key(response: object) -> StreamingScanKey:
|
||||
output_items: Final = stream_item_items(response, "output")
|
||||
message_items: Final = tuple(
|
||||
item for item in output_items if stream_item_field(item, "type") != "function_call"
|
||||
item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
|
||||
)
|
||||
return StreamingScanKey(
|
||||
texts=tuple(
|
||||
|
|
@ -932,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
tool_calls=tuple(
|
||||
stream_item_fingerprint(item)
|
||||
for item in output_items
|
||||
if stream_item_field(item, "type") == "function_call"
|
||||
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
|
||||
),
|
||||
stream_ended=True,
|
||||
)
|
||||
|
|
@ -1043,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
Override this method to customize text/image/tool extraction logic.
|
||||
"""
|
||||
|
||||
# Check if this is a tool call (OutputFunctionToolCall)
|
||||
if isinstance(output_item, OutputFunctionToolCall) or (
|
||||
isinstance(output_item, BaseModel)
|
||||
and hasattr(output_item, "type")
|
||||
and getattr(output_item, "type") == "function_call"
|
||||
):
|
||||
tool_call_item: Final = _tool_call_output_item_mapping(output_item)
|
||||
if tool_call_item is not None:
|
||||
if tool_calls_to_check is not None:
|
||||
tool_call_dict = (
|
||||
LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
|
||||
tool_call_item=output_item,
|
||||
index=output_idx,
|
||||
)
|
||||
)
|
||||
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
|
||||
return
|
||||
elif isinstance(output_item, dict) and output_item.get("type") == "function_call":
|
||||
# Handle dict representation of tool call
|
||||
if tool_calls_to_check is not None:
|
||||
# Convert dict to ResponseFunctionToolCall for processing
|
||||
try:
|
||||
tool_call_obj: Final = ResponseFunctionToolCall(**output_item)
|
||||
tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
|
||||
tool_call_item=tool_call_obj,
|
||||
index=output_idx,
|
||||
)
|
||||
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
|
||||
except Exception:
|
||||
pass
|
||||
tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
|
||||
return
|
||||
|
||||
# Handle both GenericResponseOutputItem and dict
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints
|
||||
|
|
@ -15,6 +15,10 @@ from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
|||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
_safe_convert_created_field,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_non_python_regex_patterns,
|
||||
flatten_combinators_and_drop_non_python_regex_patterns,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
|
|
@ -40,7 +44,7 @@ else:
|
|||
|
||||
_NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4")
|
||||
_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
|
||||
_PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
|
||||
_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
|
||||
|
||||
|
||||
|
|
@ -293,7 +297,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
model=model, input=validated_input, tools=tools
|
||||
)
|
||||
object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools)
|
||||
sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai(
|
||||
sanitized_tools: Final = self._sanitized_tool_schemas_for_openai(
|
||||
model=model, tools=object_schema_tools, litellm_params=litellm_params
|
||||
)
|
||||
return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools
|
||||
|
|
@ -378,35 +382,35 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
return item
|
||||
return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item
|
||||
|
||||
def _flatten_tool_schema_combinators_for_openai(
|
||||
def _sanitized_tool_schemas_for_openai(
|
||||
self,
|
||||
model: str,
|
||||
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None:
|
||||
"""Flatten top-level schema combinators only where OpenAI's validator rejects them.
|
||||
"""Rewrite tool schemas only where OpenAI's validator rejects them.
|
||||
|
||||
OpenAI-compatible backends reusing this config (and the ChatGPT backend
|
||||
Codex talks to natively) accept them, and so do GPT-5 and later models,
|
||||
which also call tools better with the union intact. Codex wraps MCP tools
|
||||
inside namespace entries, so nested ``tools`` arrays are walked too.
|
||||
Azure OpenAI shares the validator but names deployments arbitrarily, so
|
||||
the router's declared ``model_info.base_model`` wins over the deployment
|
||||
name and an unrecognized name without one is left untouched.
|
||||
Every model family refuses a ``pattern`` Python's ``re`` cannot compile,
|
||||
while top-level schema combinators are flattened only for the families
|
||||
whose validator rejects them: OpenAI-compatible backends reusing this
|
||||
config (and the ChatGPT backend Codex talks to natively) accept them,
|
||||
and so do GPT-5 and later models, which also call tools better with the
|
||||
union intact. Codex wraps MCP tools inside namespace entries, so nested
|
||||
``tools`` arrays are walked too. Azure OpenAI shares the validator but
|
||||
names deployments arbitrarily, so the router's declared
|
||||
``model_info.base_model`` wins over the deployment name and an
|
||||
unrecognized name without one keeps its combinators.
|
||||
"""
|
||||
if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR:
|
||||
if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR:
|
||||
return tools
|
||||
gate_model: Final = self._combinator_gate_model(model=model, litellm_params=litellm_params)
|
||||
if not self._rejects_top_level_schema_combinators(gate_model):
|
||||
return tools
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
self._flattened_tool_or_passthrough(tool) for tool in tools
|
||||
]
|
||||
return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: spread keeps each tool's shape
|
||||
|
||||
@staticmethod
|
||||
def _flattened_tool_or_passthrough(tool: object) -> object:
|
||||
return OpenAIResponsesAPIConfig._flattened_tool_entry(tool) if isinstance(tool, dict) else tool
|
||||
sanitize: Final = (
|
||||
flatten_combinators_and_drop_non_python_regex_patterns
|
||||
if self._rejects_top_level_schema_combinators(gate_model)
|
||||
else drop_non_python_regex_patterns
|
||||
)
|
||||
sanitized: Final = self._sanitized_tools(tools, sanitize)
|
||||
return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", sanitized) # cast-ok: spread keeps each tool's shape
|
||||
|
||||
@staticmethod
|
||||
def _rejects_top_level_schema_combinators(model: str) -> bool:
|
||||
|
|
@ -421,35 +425,42 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
return base_model if isinstance(base_model, str) and base_model else model
|
||||
|
||||
@staticmethod
|
||||
def _flattened_tool_entry(
|
||||
def _sanitized_tool_entry(
|
||||
entry: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: request tools are JSON dicts
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
flatten_top_level_schema_combinators,
|
||||
)
|
||||
|
||||
sanitize: Callable[[Mapping[str, object]], Mapping[str, object]],
|
||||
) -> Mapping[str, object]:
|
||||
parameters: Final = entry.get("parameters")
|
||||
nested_tools: Final = entry.get("tools")
|
||||
sanitized_parameters: Final = sanitize(parameters) if isinstance(parameters, dict) else parameters
|
||||
sanitized_nested_tools: Final = (
|
||||
OpenAIResponsesAPIConfig._sanitized_tools(nested_tools, sanitize)
|
||||
if isinstance(nested_tools, list)
|
||||
else nested_tools
|
||||
)
|
||||
parameters_update: Final = (
|
||||
MappingProxyType({"parameters": flatten_top_level_schema_combinators(parameters)})
|
||||
if isinstance(parameters, dict)
|
||||
MappingProxyType({"parameters": sanitized_parameters})
|
||||
if sanitized_parameters is not parameters
|
||||
else _NO_TOOL_UPDATE
|
||||
)
|
||||
tools_update: Final = (
|
||||
MappingProxyType({"tools": OpenAIResponsesAPIConfig._flattened_nested_tools(nested_tools)})
|
||||
if isinstance(nested_tools, list)
|
||||
MappingProxyType({"tools": sanitized_nested_tools})
|
||||
if sanitized_nested_tools is not nested_tools
|
||||
else _NO_TOOL_UPDATE
|
||||
)
|
||||
if not parameters_update and not tools_update:
|
||||
return entry
|
||||
return {**entry, **parameters_update, **tools_update} # mutable-ok: request tools are JSON dicts
|
||||
|
||||
@staticmethod
|
||||
def _flattened_nested_tools(
|
||||
nested_tools: Sequence[object],
|
||||
) -> list[object]: # mutable-ok: namespace tools are a JSON list
|
||||
return [ # mutable-ok: namespace tools are a JSON list
|
||||
OpenAIResponsesAPIConfig._flattened_tool_entry(item) if isinstance(item, dict) else item
|
||||
for item in nested_tools
|
||||
def _sanitized_tools(
|
||||
tools: Sequence[object],
|
||||
sanitize: Callable[[Mapping[str, object]], Mapping[str, object]],
|
||||
) -> Sequence[object]:
|
||||
sanitized: Final = [ # mutable-ok: request tools are a JSON list
|
||||
OpenAIResponsesAPIConfig._sanitized_tool_entry(item, sanitize) if isinstance(item, dict) else item
|
||||
for item in tools
|
||||
]
|
||||
return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized
|
||||
|
||||
def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
|
||||
"""
|
||||
|
|
@ -620,15 +631,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
return event_pydantic_model.model_construct(**parsed_chunk)
|
||||
|
||||
@staticmethod
|
||||
def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None:
|
||||
def parse_terminal_event_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesTerminalEvent | None:
|
||||
for chunk_str in reversed(all_chunks):
|
||||
for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent):
|
||||
try:
|
||||
return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response
|
||||
return event_model.model_validate_json(chunk_str.removeprefix("data: "))
|
||||
except ValueError:
|
||||
continue
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None:
|
||||
terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks)
|
||||
return None if terminal_event is None else terminal_event.response
|
||||
|
||||
@staticmethod
|
||||
def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ class SagemakerChatHandler(BaseAWSLLM):
|
|||
aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None)
|
||||
aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None)
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.pop("aws_session_tags", None)
|
||||
|
||||
### SET REGION NAME ###
|
||||
if aws_region_name is None:
|
||||
|
|
@ -63,6 +64,7 @@ class SagemakerChatHandler(BaseAWSLLM):
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
return credentials, aws_region_name
|
||||
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None)
|
||||
aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None)
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
aws_session_tags: Final = optional_params.pop("aws_session_tags", None)
|
||||
|
||||
### SET REGION NAME ###
|
||||
if aws_region_name is None:
|
||||
|
|
@ -86,6 +87,7 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
aws_session_tags=aws_session_tags,
|
||||
)
|
||||
return credentials, aws_region_name
|
||||
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ import random
|
|||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence
|
||||
from concurrent import futures
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
from copy import deepcopy
|
||||
|
|
@ -5398,6 +5398,14 @@ def completion(
|
|||
if dynamic_api_key is not None:
|
||||
api_key = dynamic_api_key
|
||||
# check if user passed in any of the OpenAI optional params
|
||||
bridges_to_responses_api: Final = (
|
||||
responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge
|
||||
)
|
||||
allowed_openai_params: Final[list[str] | None] = (
|
||||
[*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"]
|
||||
if bridges_to_responses_api
|
||||
else kwargs.get("allowed_openai_params")
|
||||
)
|
||||
optional_param_args: Final = {
|
||||
"functions": functions,
|
||||
"function_call": function_call,
|
||||
|
|
@ -5442,7 +5450,7 @@ def completion(
|
|||
"service_tier": service_tier,
|
||||
"store": store,
|
||||
"prompt_cache_key": prompt_cache_key,
|
||||
"allowed_openai_params": kwargs.get("allowed_openai_params"),
|
||||
"allowed_openai_params": allowed_openai_params,
|
||||
"base_model": base_model,
|
||||
}
|
||||
optional_params = get_optional_params(**optional_param_args, **non_default_params)
|
||||
|
|
@ -8587,7 +8595,7 @@ def config_completion(**kwargs):
|
|||
)
|
||||
|
||||
|
||||
def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse:
|
||||
def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse:
|
||||
id: Final = chunks[0]["id"]
|
||||
object: Final = chunks[0]["object"]
|
||||
created: Final = chunks[0]["created"]
|
||||
|
|
@ -8704,10 +8712,11 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o
|
|||
|
||||
def stream_chunk_builder(
|
||||
chunks: list,
|
||||
messages: list | None = None,
|
||||
messages: Sequence | None = None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
logging_obj: Optional["Logging"] = None,
|
||||
count_prompt_tokens: Callable[[], int] | None = None,
|
||||
) -> ModelResponse | TextCompletionResponse | None:
|
||||
try:
|
||||
if chunks is None:
|
||||
|
|
@ -8781,6 +8790,7 @@ def stream_chunk_builder(
|
|||
completion_output=completion_output,
|
||||
messages=messages,
|
||||
reasoning_tokens=0,
|
||||
count_prompt_tokens=count_prompt_tokens,
|
||||
)
|
||||
setattr(response, "usage", usage)
|
||||
|
||||
|
|
@ -8958,6 +8968,7 @@ def stream_chunk_builder(
|
|||
completion_output=completion_output,
|
||||
messages=messages,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
count_prompt_tokens=count_prompt_tokens,
|
||||
)
|
||||
|
||||
setattr(response, "usage", usage)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -179,16 +179,7 @@ def _gateway_dcr_challenge_target(
|
|||
mcp_servers: list[str] | None,
|
||||
client_ip: str | None,
|
||||
) -> str | None:
|
||||
"""The single path-named server this request targets, iff it resolves to a
|
||||
gateway-managed oauth2 server — the one per-server shape the gateway's own keyless
|
||||
DCR flow serves end to end, so the 401 challenge may advertise the per-server
|
||||
protected-resource metadata (whose ``authorization_servers`` names the gateway).
|
||||
|
||||
Multi-server CSV paths, header/path mismatches, unknown names, and every
|
||||
client-forwarded or delegated mode return ``None``: those cells keep their existing
|
||||
challenge (or absence of one), and a challenge is never emitted for a name the
|
||||
public discovery routes would 404, so this reveals exactly the server set the
|
||||
per-server protected-resource metadata already reveals."""
|
||||
"""Resolve a single path target whose sign-in metadata advertises the gateway."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
|
@ -217,7 +208,7 @@ def _is_gateway_dcr_challenge_scope(
|
|||
the caller is not a cold-start DCR client), on the scopes the gateway's keyless
|
||||
flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request
|
||||
(the resource the client configured is still ``/mcp``), or a per-server path whose
|
||||
single target is a gateway-managed oauth2 server. Every other named target keeps
|
||||
single target advertises gateway-owned sign-in. Every other named target keeps
|
||||
its existing behavior, failing closed to the original admission error."""
|
||||
if not _is_litellm_auth_admission_error(exc):
|
||||
return False
|
||||
|
|
@ -236,7 +227,7 @@ def _gateway_dcr_challenge(
|
|||
) -> HTTPException:
|
||||
"""The RFC 9728 challenge pointing the client at the protected-resource metadata
|
||||
matching the scope it requested: the per-server document (same URL spelling the
|
||||
request arrived on) when the single target is a gateway-managed oauth2 server,
|
||||
request arrived on) when the single target advertises gateway-owned sign-in,
|
||||
else the gateway's aggregate document. Either way the client discovers the gateway
|
||||
as its authorization server and starts the same sign-in flow.
|
||||
|
||||
|
|
|
|||
|
|
@ -2310,8 +2310,7 @@ async def _build_oauth_protected_resource_response(
|
|||
it. Only the legacy ``is_oauth_passthrough`` opt-in rewrites ``resource`` to
|
||||
the gateway's own URL so clients present the bearer token back to the gateway.
|
||||
|
||||
An explicitly named gateway-managed oauth2 server (interactive with
|
||||
gateway-vaulted per-user tokens, or M2M) advertises the gateway's own
|
||||
An explicitly named server with gateway-owned sign-in advertises the gateway's own
|
||||
authorization server (``{base}/mcp``): a keyless DCR client that configured the
|
||||
per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint
|
||||
supports and is admitted with a gateway session bearer. The per-server relay
|
||||
|
|
@ -2401,11 +2400,6 @@ async def _build_oauth_protected_resource_response(
|
|||
if obo_response is not None:
|
||||
return obo_response
|
||||
|
||||
# An OBO server with no configured issuer falls through to the gateway default so discovery still
|
||||
# returns metadata; every other non-oauth2 named server 404s to avoid enumeration.
|
||||
if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange:
|
||||
_raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource")
|
||||
|
||||
if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
|
||||
return {
|
||||
"authorization_servers": [f"{request_base_url}/mcp"],
|
||||
|
|
@ -2413,6 +2407,9 @@ async def _build_oauth_protected_resource_response(
|
|||
"scopes_supported": (mcp_server.scopes if mcp_server.scopes else []),
|
||||
}
|
||||
|
||||
if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange:
|
||||
_raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource")
|
||||
|
||||
return {
|
||||
"authorization_servers": [
|
||||
(f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}")
|
||||
|
|
|
|||
|
|
@ -411,7 +411,7 @@ def relative_request_url(request: Request) -> str:
|
|||
|
||||
|
||||
def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None:
|
||||
"""Resolve an RFC 8707 ``resource`` value to the single gateway-managed oauth2 server it
|
||||
"""Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it
|
||||
names, or ``None`` for every other shape: absent, the aggregate resource, a foreign
|
||||
host, an unparseable value, a multi-server path, an unknown name, or any server mode the
|
||||
keyless gateway flow does not serve (whose protected-resource metadata never directs a
|
||||
|
|
@ -443,7 +443,7 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC
|
|||
if len(names) != 1:
|
||||
return None
|
||||
server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0])
|
||||
if server is None or not server.is_gateway_managed_oauth2:
|
||||
if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server):
|
||||
return None
|
||||
return server
|
||||
|
||||
|
|
@ -729,11 +729,15 @@ async def _flow_target(
|
|||
server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
|
||||
if (
|
||||
server is None
|
||||
or not server.is_gateway_managed_oauth2
|
||||
or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server)
|
||||
or not await lookup_server_reachability(flow.user_id, server.server_id)
|
||||
):
|
||||
return "stale", None
|
||||
state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive"
|
||||
state: Final = (
|
||||
"interactive"
|
||||
if server.is_gateway_managed_oauth2 and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
|
||||
else "m2m"
|
||||
)
|
||||
return state, server
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -22,7 +22,16 @@ Response headers returned (all values are masked for safety):
|
|||
x-mcp-debug-auth-resolution
|
||||
Which auth priority was used for the outbound MCP call:
|
||||
``per-request-header``, ``m2m-client-credentials``, ``static-token``,
|
||||
``oauth2-passthrough``, or ``no-auth``.
|
||||
``oauth2-passthrough``, ``stored-user-token``, ``token-exchange``,
|
||||
``id-jag``, ``aws-sigv4``, ``extra-headers``, or ``no-auth``.
|
||||
``unresolved`` means no outcome was available before the first response
|
||||
frame; ``multiple`` means several servers resolved credentials;
|
||||
``not-applicable`` covers stdio; ``resolution-failed`` is a resolver error.
|
||||
|
||||
x-mcp-debug-auth-resolutions
|
||||
For multiple servers, a JSON map of server IDs to resolution labels.
|
||||
At most 32 entries are included; x-mcp-debug-auth-resolutions-truncated
|
||||
is true when additional servers were omitted. No credentials are included.
|
||||
|
||||
x-mcp-debug-outbound-url
|
||||
The upstream MCP server URL that will receive the request.
|
||||
|
|
@ -58,10 +67,16 @@ header is free for OAuth2 discovery::
|
|||
Symptom: ``x-mcp-debug-oauth2-token`` shows ``(none)`` and
|
||||
``x-mcp-debug-auth-resolution`` shows ``no-auth``.
|
||||
|
||||
This means the client didn't go through the OAuth2 flow. Check that:
|
||||
1. The ``Authorization`` header is NOT set as a static header in the client config.
|
||||
2. The ``.well-known/oauth-protected-resource`` endpoint returns valid metadata.
|
||||
3. The MCP server in LiteLLM config has ``auth_type: oauth2``.
|
||||
``no-auth`` means the resolved upstream client carries no authentication.
|
||||
An absent inbound OAuth2 token does not imply the user skipped OAuth: the gateway
|
||||
can retrieve a stored per-user token, reported as ``stored-user-token``.
|
||||
``unresolved`` is used when a stream starts before credential resolution, or a
|
||||
request (such as initialization or a cached tool listing) resolves no credential.
|
||||
Debug reporting does not fetch credentials or delay a streaming frame to resolve them.
|
||||
``extra-headers`` identifies supplied headers that won over the resolver or were
|
||||
the only headers supplied; their values are never inspected to guess a scheme.
|
||||
``per-request-header`` denotes a legacy credential override, including a BYOK
|
||||
credential supplied by the gateway; it does not imply a caller-supplied token.
|
||||
|
||||
**Common issue: M2M token used instead of user token**
|
||||
|
||||
|
|
@ -69,8 +84,8 @@ Symptom: ``x-mcp-debug-auth-resolution`` shows ``m2m-client-credentials``.
|
|||
|
||||
This means the server has ``client_id``/``client_secret``/``token_url``
|
||||
configured and LiteLLM is fetching a machine-to-machine token instead of
|
||||
using the per-user OAuth2 token. If you want per-user tokens, remove the
|
||||
client credentials from the server config.
|
||||
using the per-user OAuth2 token. For gateway-stored per-user tokens,
|
||||
configure ``oauth2_flow: authorization_code``.
|
||||
|
||||
Usage from Claude Code::
|
||||
|
||||
|
|
@ -85,14 +100,16 @@ Usage with curl::
|
|||
http://localhost:4000/mcp/atlassian_mcp
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from starlette.requests import HTTPConnection
|
||||
from starlette.types import Message, Send
|
||||
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
|
||||
|
||||
# Header the client sends to opt into debug mode
|
||||
MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
|
||||
|
|
@ -101,6 +118,83 @@ MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
|
|||
_RESPONSE_HEADER_PREFIX: Final = "x-mcp-debug"
|
||||
|
||||
|
||||
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics"
|
||||
|
||||
|
||||
def record_auth_resolution(server_id: str, source: AuthResolution) -> None:
|
||||
from mcp.server.lowlevel.server import request_ctx
|
||||
|
||||
context: Final[object] = request_ctx.get(None)
|
||||
request: Final[object] = getattr(context, "request", None)
|
||||
if isinstance(request, HTTPConnection):
|
||||
diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY)
|
||||
if isinstance(diagnostics, MCPAuthDiagnostics):
|
||||
diagnostics.record(server_id, source)
|
||||
|
||||
|
||||
class MCPAuthDiagnostics:
|
||||
def __init__(self) -> None:
|
||||
self._outcomes: tuple[tuple[str, AuthResolution], ...] = ()
|
||||
|
||||
def record(self, server_id: str, resolution: AuthResolution) -> None:
|
||||
self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),)
|
||||
|
||||
def resolution(self) -> str:
|
||||
match self._outcomes:
|
||||
case ():
|
||||
return AuthResolution.unresolved.value
|
||||
case ((_, source),):
|
||||
return source.value
|
||||
case _:
|
||||
return AuthResolution.multiple.value
|
||||
|
||||
def headers(self) -> Mapping[str, str]:
|
||||
if len(self._outcomes) <= 1:
|
||||
return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()})
|
||||
return MappingProxyType(
|
||||
{
|
||||
"x-mcp-debug-auth-resolution": AuthResolution.multiple.value,
|
||||
"x-mcp-debug-auth-resolutions": json.dumps(
|
||||
{
|
||||
server_id: source.value for server_id, source in self._outcomes[:32]
|
||||
}, # mutable-ok: JSON encoder requires a concrete dict
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=True,
|
||||
),
|
||||
**(
|
||||
MappingProxyType({"x-mcp-debug-auth-resolutions-truncated": "true"})
|
||||
if len(self._outcomes) > 32
|
||||
else MappingProxyType({})
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _DiagnosticSend:
|
||||
def __init__(self, send: Send, headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]]) -> None:
|
||||
self._send = send
|
||||
self._headers = headers
|
||||
self._resolution = resolution
|
||||
self._start: Message | None = None
|
||||
|
||||
async def __call__(self, message: Message) -> None:
|
||||
if message["type"] == "http.response.start":
|
||||
self._start = message
|
||||
return
|
||||
if self._start is not None:
|
||||
start: Final = self._start
|
||||
self._start = None
|
||||
headers: Final = MappingProxyType({**self._headers, **self._resolution()})
|
||||
await self._send(
|
||||
{ # mutable-ok: ASGI send consumes a mutable message mapping
|
||||
**start,
|
||||
"headers": tuple(start.get("headers", ()))
|
||||
+ tuple((key.encode(), value.encode()) for key, value in headers.items()),
|
||||
}
|
||||
)
|
||||
await self._send(message)
|
||||
|
||||
|
||||
class MCPDebug:
|
||||
"""
|
||||
Static helper class for MCP OAuth2 debug headers.
|
||||
|
|
@ -144,37 +238,6 @@ class MCPDebug:
|
|||
return val.strip().lower() in ("true", "1", "yes")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def resolve_auth_resolution(
|
||||
server: "MCPServer",
|
||||
mcp_auth_header: str | None,
|
||||
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
|
||||
oauth2_headers: dict[str, str] | None,
|
||||
) -> str:
|
||||
"""
|
||||
Determine which auth priority will be used for the outbound MCP call.
|
||||
|
||||
Returns one of: ``per-request-header``, ``m2m-client-credentials``,
|
||||
``static-token``, ``oauth2-passthrough``, or ``no-auth``.
|
||||
"""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
has_server_specific: Final = bool(
|
||||
mcp_server_auth_headers
|
||||
and (
|
||||
mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "")
|
||||
)
|
||||
)
|
||||
if has_server_specific or mcp_auth_header:
|
||||
return "per-request-header"
|
||||
if server.has_client_credentials:
|
||||
return "m2m-client-credentials"
|
||||
if server.authentication_token:
|
||||
return "static-token"
|
||||
if oauth2_headers and server.auth_type == MCPAuth.oauth2:
|
||||
return "oauth2-passthrough"
|
||||
return "no-auth"
|
||||
|
||||
@staticmethod
|
||||
def build_debug_headers(
|
||||
*,
|
||||
|
|
@ -244,12 +307,21 @@ class MCPDebug:
|
|||
return debug
|
||||
|
||||
@staticmethod
|
||||
def wrap_send_with_debug_headers(send: Send, debug_headers: dict[str, str]) -> Send:
|
||||
def wrap_send_with_debug_headers(
|
||||
send: Send,
|
||||
debug_headers: Mapping[str, str],
|
||||
resolution: Callable[[], Mapping[str, str]] | None = None,
|
||||
*,
|
||||
request_method: str | None = None,
|
||||
) -> Send:
|
||||
"""
|
||||
Return a new ASGI ``send`` callable that injects *debug_headers*
|
||||
into the ``http.response.start`` message.
|
||||
"""
|
||||
|
||||
if resolution is not None and request_method == "POST":
|
||||
return _DiagnosticSend(send, debug_headers, resolution)
|
||||
|
||||
async def _send_with_debug(message: Message) -> None:
|
||||
if message["type"] == "http.response.start":
|
||||
headers: Final = list(message.get("headers", []))
|
||||
|
|
@ -266,8 +338,6 @@ class MCPDebug:
|
|||
raw_headers: dict[str, str] | None,
|
||||
scope: dict,
|
||||
mcp_servers: list[str] | None,
|
||||
mcp_auth_header: str | None,
|
||||
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
|
||||
oauth2_headers: dict[str, str] | None,
|
||||
client_ip: str | None,
|
||||
) -> dict[str, str]:
|
||||
|
|
@ -288,16 +358,13 @@ class MCPDebug:
|
|||
|
||||
server_url: str | None = None
|
||||
server_auth_type: str | None = None
|
||||
auth_resolution = "no-auth"
|
||||
auth_resolution: Final = AuthResolution.unresolved.value
|
||||
|
||||
for server_name in mcp_servers or []:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip)
|
||||
if server:
|
||||
server_url = server.url
|
||||
server_auth_type = server.auth_type
|
||||
auth_resolution = MCPDebug.resolve_auth_resolution(
|
||||
server, mcp_auth_header, mcp_server_auth_headers, oauth2_headers
|
||||
)
|
||||
break
|
||||
|
||||
scope_headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope)
|
||||
|
|
|
|||
|
|
@ -80,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
|
|||
raise_classified_list_failure,
|
||||
upstream_auth_challenge,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
MCPPerUserTokenCache,
|
||||
mcp_per_user_token_cache,
|
||||
|
|
@ -108,12 +109,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto
|
|||
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import (
|
||||
LazyPerUserOAuthTokenStore,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import (
|
||||
build_token_exchanger,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
DEFAULT_CREDENTIAL_HEADER,
|
||||
AuthorizationCodeConfig,
|
||||
AuthResolution,
|
||||
ClientCredentialsConfig,
|
||||
CredError,
|
||||
IdJagConfig,
|
||||
|
|
@ -3832,13 +3835,21 @@ class MCPServerManager:
|
|||
(authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any
|
||||
other ``CredError`` onto its public HTTP status; it never returns an error as a value.
|
||||
"""
|
||||
match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
|
||||
case Ok(auth):
|
||||
match await resolve_credentials_with_source(provider, to_subject(user_api_key_auth, subject_token), spec):
|
||||
case Ok(credential):
|
||||
auth: Final = credential.auth
|
||||
# NoOpAuth has no header_name and so never conflicts.
|
||||
header_name: Final[str | None] = getattr(auth, "header_name", None)
|
||||
if header_name is None or not extra_headers:
|
||||
source: Final = (
|
||||
AuthResolution.extra_headers
|
||||
if credential.source == AuthResolution.no_auth and extra_headers
|
||||
else credential.source
|
||||
)
|
||||
record_auth_resolution(server.server_id, source)
|
||||
return auth, extra_headers
|
||||
if not has_header(extra_headers, header_name):
|
||||
record_auth_resolution(server.server_id, credential.source)
|
||||
return auth, extra_headers
|
||||
if isinstance(
|
||||
spec.config,
|
||||
|
|
@ -3853,11 +3864,14 @@ class MCPServerManager:
|
|||
# one-shot 401 refetch is lost with it). Drop only the header the resolved
|
||||
# credential is about to occupy, so a static credential the operator aimed at a
|
||||
# DIFFERENT header still reaches upstream.
|
||||
record_auth_resolution(server.server_id, credential.source)
|
||||
return auth, without_header(extra_headers, header_name)
|
||||
# Other modes: an Authorization already supplied via extra_headers (a forwarded caller
|
||||
# header or static_headers) is intentional and wins; v1 applies those last.
|
||||
record_auth_resolution(server.server_id, AuthResolution.extra_headers)
|
||||
return None, extra_headers
|
||||
case Error(err):
|
||||
record_auth_resolution(server.server_id, AuthResolution.failed)
|
||||
if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig):
|
||||
# authorization_code's missing per-user token -> the per-server browser-OAuth
|
||||
# challenge, built here where the full MCPServer is in hand.
|
||||
|
|
@ -3960,6 +3974,7 @@ class MCPServerManager:
|
|||
Returns:
|
||||
Configured MCP client instance.
|
||||
"""
|
||||
record_auth_resolution(server.server_id, AuthResolution.unresolved)
|
||||
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
|
||||
transport: Final = resolved_server.transport or MCPTransport.sse
|
||||
spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server)
|
||||
|
|
@ -4032,6 +4047,7 @@ class MCPServerManager:
|
|||
env=resolved_env,
|
||||
)
|
||||
|
||||
record_auth_resolution(server.server_id, AuthResolution.not_applicable)
|
||||
return MCPClient(
|
||||
server_url="", # Not used for stdio
|
||||
transport_type=transport,
|
||||
|
|
@ -4086,6 +4102,20 @@ class MCPServerManager:
|
|||
aws_session_name=resolved_server.aws_session_name,
|
||||
)
|
||||
|
||||
legacy_source: Final = (
|
||||
AuthResolution.aws_sigv4
|
||||
if aws_auth is not None
|
||||
else AuthResolution.extra_headers
|
||||
if extra_headers and has_header(extra_headers, auth_header_name or "Authorization")
|
||||
else AuthResolution.per_request_header
|
||||
if mcp_auth_header
|
||||
else AuthResolution.static_token
|
||||
if auth_value
|
||||
else AuthResolution.extra_headers
|
||||
if extra_headers
|
||||
else AuthResolution.no_auth
|
||||
)
|
||||
record_auth_resolution(server.server_id, legacy_source)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger
|
|||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
ApiKeyConfig,
|
||||
AuthorizationCodeConfig,
|
||||
AuthResolution,
|
||||
AuthSpecKind,
|
||||
AwsSigV4Config,
|
||||
Byok,
|
||||
|
|
@ -76,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
|||
NoneConfig,
|
||||
PassthroughConfig,
|
||||
PrivateKeyJwtAuth,
|
||||
ResolvedCredential,
|
||||
ServerSpec,
|
||||
SharedKey,
|
||||
Subject,
|
||||
|
|
@ -448,3 +450,32 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str:
|
|||
|
||||
def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
|
||||
return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet"))
|
||||
|
||||
|
||||
async def resolve_credentials_with_source(
|
||||
provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec
|
||||
) -> Result[ResolvedCredential, CredError]:
|
||||
match await provider.resolve_credentials(subject, server):
|
||||
case Error(err):
|
||||
return Error(err)
|
||||
case Ok(auth):
|
||||
if isinstance(auth, NoOpAuth):
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
|
||||
match server.config:
|
||||
case NoneConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.no_auth))
|
||||
case ApiKeyConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.static_token))
|
||||
case PassthroughConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough))
|
||||
case ClientCredentialsConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.client_credentials))
|
||||
case TokenExchangeConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.token_exchange))
|
||||
case IdJagConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.id_jag))
|
||||
case AuthorizationCodeConfig():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token))
|
||||
case AwsSigV4Config():
|
||||
return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4))
|
||||
assert_never(server.config)
|
||||
|
|
|
|||
|
|
@ -26,10 +26,11 @@ union (see `result.py`), not `expression.Result`.
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Annotated, Final, Literal
|
||||
|
||||
import httpx
|
||||
from expression import case, tag, tagged_union
|
||||
from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
|
||||
from typing_extensions import assert_never
|
||||
|
|
@ -46,6 +47,29 @@ from litellm.types.mcp import (
|
|||
)
|
||||
|
||||
|
||||
class AuthResolution(str, Enum):
|
||||
no_auth = "no-auth"
|
||||
stored_user_token = "stored-user-token"
|
||||
static_token = "static-token"
|
||||
per_request_header = "per-request-header"
|
||||
oauth2_passthrough = "oauth2-passthrough"
|
||||
client_credentials = "m2m-client-credentials"
|
||||
token_exchange = "token-exchange"
|
||||
id_jag = "id-jag"
|
||||
aws_sigv4 = "aws-sigv4"
|
||||
extra_headers = "extra-headers"
|
||||
not_applicable = "not-applicable"
|
||||
unresolved = "unresolved"
|
||||
failed = "resolution-failed"
|
||||
multiple = "multiple"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResolvedCredential:
|
||||
auth: httpx.Auth = field(repr=False)
|
||||
source: AuthResolution
|
||||
|
||||
|
||||
class AuthSpecKind(str, Enum):
|
||||
"""The server's statically-declared upstream-auth mode — derived from its `config`.
|
||||
|
||||
|
|
|
|||
|
|
@ -1234,7 +1234,7 @@ if MCP_AVAILABLE:
|
|||
return client_id, client_secret, scopes
|
||||
|
||||
_STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset(
|
||||
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization)
|
||||
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, MCPAuth.token)
|
||||
)
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -1243,6 +1243,17 @@ if MCP_AVAILABLE:
|
|||
mcp_auth_header: str | None
|
||||
oauth2_headers: dict[str, str] | None
|
||||
|
||||
def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None:
|
||||
if not url:
|
||||
return None
|
||||
try:
|
||||
parsed: Final = httpx.URL(url)
|
||||
except httpx.InvalidURL:
|
||||
return None
|
||||
if parsed.scheme not in ("http", "https") or not parsed.host:
|
||||
return None
|
||||
return parsed.scheme, parsed.host, parsed.port
|
||||
|
||||
def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest:
|
||||
"""
|
||||
Resolve the credentials a not-yet-saved server config carries for a preview call.
|
||||
|
|
@ -1255,7 +1266,19 @@ if MCP_AVAILABLE:
|
|||
MCPRequestHandler,
|
||||
)
|
||||
|
||||
request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request)
|
||||
saved_server: Final = (
|
||||
global_mcp_server_manager.get_mcp_server_by_id(new_mcp_server_request.server_id)
|
||||
if new_mcp_server_request.server_id
|
||||
else None
|
||||
)
|
||||
saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None
|
||||
preview_origin: Final = _preview_origin(new_mcp_server_request.url)
|
||||
may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or (
|
||||
saved_origin is not None and saved_origin == preview_origin
|
||||
)
|
||||
request: Final = (
|
||||
_inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request
|
||||
)
|
||||
mcp_auth_header: Final = (
|
||||
request.credentials.get("auth_value")
|
||||
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
|
||||
|
|
@ -1318,8 +1341,15 @@ if MCP_AVAILABLE:
|
|||
if _oauth2_flow == "client_credentials" and not request.token_url:
|
||||
_oauth2_flow = None
|
||||
|
||||
# Static previews inherit credentials before this step, but must not resolve back to
|
||||
# the saved record during client creation and discard the edited connection settings.
|
||||
preview_server_id: Final = (
|
||||
""
|
||||
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none)
|
||||
else request.server_id or ""
|
||||
)
|
||||
server_model: Final = MCPServer(
|
||||
server_id=request.server_id or "",
|
||||
server_id=preview_server_id,
|
||||
name=request.alias or request.server_name or "",
|
||||
url=request.url,
|
||||
transport=request.transport,
|
||||
|
|
|
|||
|
|
@ -49,7 +49,11 @@ from litellm.proxy._experimental.mcp_server.mcp_context import (
|
|||
_mcp_gateway_server_name,
|
||||
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import (
|
||||
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY,
|
||||
MCPAuthDiagnostics,
|
||||
MCPDebug,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
||||
_redact_mcp_resource_url,
|
||||
get_passthrough_www_authenticate,
|
||||
|
|
@ -4472,13 +4476,15 @@ if MCP_AVAILABLE:
|
|||
raw_headers=raw_headers,
|
||||
scope=dict(scope),
|
||||
mcp_servers=mcp_servers,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
mcp_server_auth_headers=mcp_server_auth_headers,
|
||||
oauth2_headers=oauth2_headers,
|
||||
client_ip=_client_ip,
|
||||
)
|
||||
if _debug_headers:
|
||||
send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers)
|
||||
diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None
|
||||
if diagnostics is not None:
|
||||
scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics
|
||||
send = MCPDebug.wrap_send_with_debug_headers(
|
||||
send, _debug_headers, diagnostics.headers, request_method=scope.get("method")
|
||||
)
|
||||
|
||||
# Ensure session managers are initialized
|
||||
if not _SESSION_MANAGERS_INITIALIZED:
|
||||
|
|
|
|||
|
|
@ -2836,6 +2836,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"UI username/password login. Default is False."
|
||||
),
|
||||
)
|
||||
disable_env_credential_login: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"If True, disables signing in to the Admin UI with the environment credentials: "
|
||||
"UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback "
|
||||
"means env-credential login is always live by default). Database users with passwords "
|
||||
"are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password "
|
||||
"before enabling, or nobody can sign in to the UI. A locked-out admin can still "
|
||||
"administer the proxy over the API with the master key, and can unset this setting "
|
||||
"and restart the proxy to restore env-credential login. Default is False."
|
||||
),
|
||||
)
|
||||
disable_budget_reservation: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
|
|
@ -3714,6 +3726,15 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
|
|||
],
|
||||
)
|
||||
|
||||
pointfive: CallbackOnUI = CallbackOnUI(
|
||||
litellm_callback_name="pointfive",
|
||||
ui_callback_name="PointFive",
|
||||
litellm_callback_params=[ # mutable-ok: the registry field is typed list
|
||||
"POINTFIVE_API_KEY",
|
||||
"POINTFIVE_API_URL",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class SpendLogsRouterMetadata(TypedDict):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -94,6 +94,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
team_membership_auth_cache_key,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.db_lookup_gate import db_lookup_gate
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.guardrails.tool_name_extraction import (
|
||||
TOOL_CAPABLE_CALL_TYPES,
|
||||
|
|
@ -3450,36 +3451,37 @@ async def _fetch_key_object_from_db_with_reconnect(
|
|||
"""
|
||||
Fetch key object from DB and retry once if a DB connection error can be healed.
|
||||
"""
|
||||
try:
|
||||
return await prisma_client.get_data(
|
||||
token=hashed_token,
|
||||
table_name="combined_view",
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
if PrismaDBExceptionHandler.is_database_transport_error(e):
|
||||
did_reconnect = False
|
||||
if hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0)
|
||||
if not isinstance(auth_reconnect_timeout, (int, float)):
|
||||
auth_reconnect_timeout = 2.0
|
||||
auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1)
|
||||
if not isinstance(auth_reconnect_lock_timeout, (int, float)):
|
||||
auth_reconnect_lock_timeout = 0.1
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason="auth_get_key_object_lookup_failure",
|
||||
timeout_seconds=auth_reconnect_timeout,
|
||||
lock_timeout_seconds=auth_reconnect_lock_timeout,
|
||||
)
|
||||
if did_reconnect:
|
||||
return await prisma_client.get_data(
|
||||
token=hashed_token,
|
||||
table_name="combined_view",
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
raise
|
||||
async with db_lookup_gate.current():
|
||||
try:
|
||||
return await prisma_client.get_data(
|
||||
token=hashed_token,
|
||||
table_name="combined_view",
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
if PrismaDBExceptionHandler.is_database_transport_error(e):
|
||||
did_reconnect = False
|
||||
if hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0)
|
||||
if not isinstance(auth_reconnect_timeout, (int, float)):
|
||||
auth_reconnect_timeout = 2.0
|
||||
auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1)
|
||||
if not isinstance(auth_reconnect_lock_timeout, (int, float)):
|
||||
auth_reconnect_lock_timeout = 0.1
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason="auth_get_key_object_lookup_failure",
|
||||
timeout_seconds=auth_reconnect_timeout,
|
||||
lock_timeout_seconds=auth_reconnect_lock_timeout,
|
||||
)
|
||||
if did_reconnect:
|
||||
return await prisma_client.get_data(
|
||||
token=hashed_token,
|
||||
table_name="combined_view",
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from litellm.litellm_core_utils.url_utils import (
|
|||
provider_url_destination_candidates,
|
||||
validate_url,
|
||||
)
|
||||
from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
|
|
@ -316,6 +317,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
|
|||
"aws_profile_name",
|
||||
"aws_session_name",
|
||||
"aws_external_id",
|
||||
"aws_session_tags",
|
||||
"vertex_credentials",
|
||||
# Azure managed-identity / federated-auth token. The Azure provider
|
||||
# transformer reads ``azure_ad_token`` (top-level or via
|
||||
|
|
@ -2003,9 +2005,20 @@ def get_model_from_request(
|
|||
bedrock_model: Final = _model_from_bedrock_route(route)
|
||||
return model if bedrock_model is None else bedrock_model
|
||||
|
||||
if route.lower().startswith(("/azure/", "/azure_ai/")):
|
||||
azure_model: Final = _router_model_from_azure_route(route, llm_router)
|
||||
return model if azure_model is None else azure_model
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None:
|
||||
if llm_router is None:
|
||||
return None
|
||||
endpoint: Final = re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE)
|
||||
return azure_router_model_in_endpoint(endpoint, frozenset(llm_router.get_model_names()))
|
||||
|
||||
|
||||
def _model_from_bedrock_route(route: str) -> str | None:
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_extract_model_from_bedrock_endpoint,
|
||||
|
|
|
|||
|
|
@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]:
|
|||
return ui_username, ui_password
|
||||
|
||||
|
||||
def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool:
|
||||
ui_username, ui_password = get_ui_credentials(master_key)
|
||||
return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
|
||||
password.encode("utf-8"), ui_password.encode("utf-8")
|
||||
)
|
||||
|
||||
|
||||
def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool:
|
||||
"""Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed.
|
||||
|
||||
Two settings can turn it off: `disable_env_credential_login` unconditionally, and
|
||||
`disable_password_login_when_sso_enabled` as a side effect, since its gate rejects
|
||||
every username/password login before the env comparison runs. Feeds both the
|
||||
`authenticate_user` gate and the Admin UI warning banner, so the banner never nags
|
||||
about a login path that is already unreachable.
|
||||
"""
|
||||
if general_settings.get("disable_env_credential_login") is True:
|
||||
return False
|
||||
if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class LoginResult:
|
||||
"""Result object containing authentication data from login."""
|
||||
|
||||
|
|
@ -129,7 +152,8 @@ async def authenticate_user(
|
|||
master_key: Master key for the proxy (required)
|
||||
prisma_client: Prisma database client (optional)
|
||||
general_settings: Proxy general_settings, checked for
|
||||
`disable_password_login_when_sso_enabled`
|
||||
`disable_password_login_when_sso_enabled` and
|
||||
`disable_env_credential_login`
|
||||
|
||||
Returns:
|
||||
LoginResult: Object containing authentication data
|
||||
|
|
@ -170,8 +194,6 @@ async def authenticate_user(
|
|||
code=500,
|
||||
)
|
||||
|
||||
ui_username, ui_password = get_ui_credentials(master_key)
|
||||
|
||||
# Check if we can find the `username` in the db. On the UI, users can enter username=their email
|
||||
_user_row: LiteLLM_UserTable | None = None
|
||||
user_role: (
|
||||
|
|
@ -197,8 +219,8 @@ async def authenticate_user(
|
|||
- Login with UI_USERNAME and UI_PASSWORD
|
||||
- Login with Invite Link `user_email` and `password` combination
|
||||
"""
|
||||
if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
|
||||
password.encode("utf-8"), ui_password.encode("utf-8")
|
||||
if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
|
||||
username, password, master_key
|
||||
):
|
||||
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
|
||||
user_role = LitellmUserRoles.PROXY_ADMIN
|
||||
|
|
@ -340,8 +362,13 @@ async def authenticate_user(
|
|||
code=401,
|
||||
)
|
||||
else:
|
||||
env_credentials_hint: Final = (
|
||||
"\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file"
|
||||
if is_env_credential_login_enabled(general_settings)
|
||||
else ""
|
||||
)
|
||||
raise ProxyException(
|
||||
message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file",
|
||||
message=f"Invalid credentials used to access UI.{env_credentials_hint}",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="invalid_credentials",
|
||||
code=401,
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
_safe_set_request_parsed_body,
|
||||
populate_request_with_path_params,
|
||||
)
|
||||
from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group
|
||||
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
|
|
@ -183,6 +184,44 @@ def _get_model_from_request_context(
|
|||
)
|
||||
|
||||
|
||||
_CLAUDE_MODEL_ROUTES: Final = frozenset(
|
||||
f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses")
|
||||
)
|
||||
_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized"
|
||||
|
||||
|
||||
async def _normalize_claude_model(
|
||||
request_data: dict, valid_token: UserAPIKeyAuth, request: Request | None, route: str
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj
|
||||
|
||||
if route not in _CLAUDE_MODEL_ROUTES or llm_router is None:
|
||||
return
|
||||
if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True:
|
||||
return
|
||||
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router)
|
||||
if not isinstance(requested, str) or requested != request_data.get("model"):
|
||||
return
|
||||
if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"):
|
||||
return
|
||||
settings: Final = await proxy_config.get_hierarchical_router_settings(
|
||||
user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
aliases: Final = settings.get("model_group_alias") if isinstance(settings, Mapping) else None
|
||||
source: Final = claude_code_requested_group(
|
||||
requested, llm_router, valid_token.team_id, (valid_token.aliases, valid_token.team_model_aliases, aliases)
|
||||
)
|
||||
if request is not None:
|
||||
request.scope[_CLAUDE_MODEL_NORMALIZED] = True
|
||||
if source is None:
|
||||
return
|
||||
request_data["model"] = source
|
||||
_safe_set_request_parsed_body(request=request, parsed_body=request_data)
|
||||
if request is not None:
|
||||
request._json = request_data
|
||||
request._body = orjson.dumps(request_data)
|
||||
|
||||
|
||||
def _get_model_names_for_budget_checks(
|
||||
model: str | list[str] | None,
|
||||
) -> list[str]:
|
||||
|
|
@ -2768,6 +2807,7 @@ async def _authorize_authenticated_request(
|
|||
"""
|
||||
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
|
||||
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
|
||||
await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route)
|
||||
|
||||
# Single authorization point. Builder paths MUST NOT call common_checks.
|
||||
# Route through the same exception handler the builder uses so
|
||||
|
|
@ -3131,6 +3171,7 @@ async def _enforce_key_and_fallback_model_access(
|
|||
Key-level model allowlist and client fallbacks (same as standard auth).
|
||||
Not included in common_checks — common_checks enforces team/user/project model access only.
|
||||
"""
|
||||
await _normalize_claude_model(request_data, valid_token, request, route)
|
||||
config: Final = valid_token.config
|
||||
|
||||
if config != {}:
|
||||
|
|
|
|||
|
|
@ -490,7 +490,7 @@ lite codex exec "summarize the repo"
|
|||
|
||||
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
|
||||
|
||||
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
|
||||
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
|
||||
|
||||
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
|
||||
|
||||
|
|
@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
|
|||
|
||||
### Route Every Claude Code Session Through the Proxy
|
||||
|
||||
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
|
||||
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
|
||||
|
||||
Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
|
||||
|
||||
|
|
@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
|
|||
lite --base-url https://your-proxy.example.com login --config-claude
|
||||
```
|
||||
|
||||
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
|
||||
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
|
||||
|
||||
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
|
||||
|
||||
Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
|
||||
|
||||
#### Configuring Claude Code Once, With a Virtual Key or Your Login
|
||||
|
||||
`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
|
||||
lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto
|
||||
claude
|
||||
```
|
||||
|
||||
With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-<UTF-8 hex of the group name>` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control
|
||||
|
||||
Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt
|
||||
|
||||
What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request
|
||||
|
||||
### QA Complexity-Based Auto-Routing Against Your Real Proxy
|
||||
|
||||
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
|
||||
|
|
@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t
|
|||
|
||||
The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
|
||||
|
||||
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
|
||||
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
|
||||
|
||||
You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import subprocess
|
|||
import sys
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
|
|
@ -12,10 +13,12 @@ import requests
|
|||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
|
||||
from .claude_settings import claude_settings_path, lite_api_key_helper_configured
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
from .pi import (
|
||||
LITELLM_PROXY_API_KEY_ENV,
|
||||
PI_PROVIDER_NAME,
|
||||
ListingFailure,
|
||||
PiSyncError,
|
||||
fetch_model_ids,
|
||||
fetch_model_limits,
|
||||
|
|
@ -83,6 +86,8 @@ def build_agent_env(
|
|||
base_url: str,
|
||||
api_key: str,
|
||||
profiles: frozenset[str],
|
||||
*,
|
||||
export_anthropic_token: bool = True,
|
||||
) -> dict[str, str]:
|
||||
"""Return a copy of base_env wired to route the agent through the proxy.
|
||||
|
||||
|
|
@ -97,12 +102,19 @@ def build_agent_env(
|
|||
proxy's /v1/models; likewise left alone when already set.
|
||||
pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY
|
||||
from its synced models.json provider entry.
|
||||
|
||||
With export_anthropic_token=False the bearer is left out (and any inherited
|
||||
one dropped) so Claude Code asks its configured apiKeyHelper instead; Claude
|
||||
Code prefers ANTHROPIC_AUTH_TOKEN over the helper and warns when both are set.
|
||||
"""
|
||||
env: Final = dict(base_env)
|
||||
root: Final = base_url.rstrip("/")
|
||||
if PROFILE_ANTHROPIC in profiles:
|
||||
env[ANTHROPIC_BASE_URL_ENV] = root
|
||||
env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key
|
||||
if export_anthropic_token:
|
||||
env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key
|
||||
else:
|
||||
env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None)
|
||||
env.pop(ANTHROPIC_API_KEY_ENV, None)
|
||||
if ENABLE_TOOL_SEARCH_ENV not in env:
|
||||
env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE
|
||||
|
|
@ -165,7 +177,9 @@ def prepare_pi(
|
|||
"""
|
||||
ids: Final = fetch_model_ids(base_url, api_key, get=get)
|
||||
if isinstance(ids, PiSyncError):
|
||||
raise AgentRunError(ids.message)
|
||||
raise AgentRunError(
|
||||
f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message
|
||||
)
|
||||
limits: Final = fetch_model_limits(base_url, api_key, get=get)
|
||||
path: Final = models_json_path(base_env)
|
||||
error: Final = sync_models_json(path, base_url, ids, limits)
|
||||
|
|
@ -460,6 +474,7 @@ def run_agent(
|
|||
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
|
||||
reattach_terminal: Callable[[], None] | None = None,
|
||||
preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS),
|
||||
export_anthropic_token: bool = True,
|
||||
) -> None:
|
||||
"""Validate, wire the environment, and hand off to the agent.
|
||||
|
||||
|
|
@ -491,7 +506,9 @@ def run_agent(
|
|||
|
||||
env: Final = MappingProxyType(
|
||||
{
|
||||
**build_agent_env(env_before_sync, base_url, api_key, profiles),
|
||||
**build_agent_env(
|
||||
env_before_sync, base_url, api_key, profiles, export_anthropic_token=export_anthropic_token
|
||||
),
|
||||
**(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
|
||||
}
|
||||
)
|
||||
|
|
@ -529,14 +546,26 @@ def resolve_api_key(ctx: click.Context) -> str:
|
|||
_SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy."
|
||||
|
||||
|
||||
def _helper_supplies_token(
|
||||
ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str], settings_path: Path
|
||||
) -> bool:
|
||||
if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"):
|
||||
return False
|
||||
return lite_api_key_helper_configured(base_url, settings_path)
|
||||
|
||||
|
||||
def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None:
|
||||
ctx_obj: Final[CliContextObj] = ctx.obj
|
||||
base_url: Final = ctx_obj["base_url"]
|
||||
started_interactive: Final = _is_interactive()
|
||||
api_key: Final = resolve_api_key(ctx)
|
||||
|
||||
display_name, _ = agent_profile(binary)
|
||||
display_name, profiles = agent_profile(binary)
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles, settings_path)
|
||||
click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}")
|
||||
if helper_supplies_token:
|
||||
click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {settings_path}")
|
||||
|
||||
try:
|
||||
run_agent(
|
||||
|
|
@ -545,6 +574,7 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify
|
|||
[binary, *args],
|
||||
skip_verify=skip_verify,
|
||||
reattach_terminal=(_restore_controlling_terminal if started_interactive else None),
|
||||
export_anthropic_token=not helper_supplies_token,
|
||||
)
|
||||
except AgentRunError as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
|
|
@ -40,10 +41,16 @@ from litellm.litellm_core_utils.cli_token_utils import (
|
|||
)
|
||||
|
||||
from .claude_settings import (
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
SETTINGS_FILE_OWNERS,
|
||||
STARTING_MODEL_ROLE,
|
||||
ApiKeyHelper,
|
||||
ClaudeSettingsError,
|
||||
write_claude_settings,
|
||||
KeepModel,
|
||||
claude_settings_path,
|
||||
configure_claude_settings,
|
||||
configure_state_path,
|
||||
refuse_while_owned,
|
||||
resolve_api_key_helper,
|
||||
settings_file_owners,
|
||||
)
|
||||
from .pkce_login import (
|
||||
Http,
|
||||
|
|
@ -778,13 +785,24 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
|
|||
|
||||
|
||||
def _configure_claude_code(base_url: str) -> None:
|
||||
"""Point Claude Code at base_url by patching ~/.claude/settings.json."""
|
||||
"""Point Claude Code at base_url by patching the settings.json it reads, undoable with `lite unconfigure claude`."""
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
try:
|
||||
write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
|
||||
configure_claude_settings(
|
||||
base_url,
|
||||
ApiKeyHelper(resolve_api_key_helper(base_url)),
|
||||
KeepModel(),
|
||||
settings_path,
|
||||
configure_state_path(settings_path),
|
||||
settings_file_owners(settings_path),
|
||||
)
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
|
||||
click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
|
||||
click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
|
||||
click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.")
|
||||
click.echo(
|
||||
"Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. "
|
||||
f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}."
|
||||
)
|
||||
|
||||
|
||||
def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None:
|
||||
|
|
@ -853,6 +871,12 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
|
|||
|
||||
ctx_obj: Final[CliContextObj] = ctx.obj
|
||||
base_url: Final = ctx_obj["base_url"]
|
||||
if config_claude:
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
try:
|
||||
refuse_while_owned(settings_path, settings_file_owners(settings_path))
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}")
|
||||
|
||||
try:
|
||||
if pkce:
|
||||
|
|
|
|||
|
|
@ -14,11 +14,13 @@ from ..claude_settings import (
|
|||
AUTOROUTE_BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ClaudeSettingsError,
|
||||
StaticToken,
|
||||
load_json_or_empty,
|
||||
merge_claude_settings,
|
||||
)
|
||||
from ..up import BackupRecord as ClaudeBackupRecord
|
||||
from ..up import restore_claude_settings, write_backup
|
||||
from .config import master_key_from_config
|
||||
from .config import AUTOROUTER_MODEL_NAME, master_key_from_config
|
||||
from .process import (
|
||||
CONFIG_PATH,
|
||||
DEFAULT_AUTOROUTE_PORT,
|
||||
|
|
@ -37,7 +39,6 @@ from .process import (
|
|||
terminate,
|
||||
write_pid_record,
|
||||
)
|
||||
from .settings import merge_claude_settings_static_token
|
||||
from .wizard import run_configure_wizard
|
||||
|
||||
_GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
|
@ -156,7 +157,9 @@ def up(port: int) -> None:
|
|||
ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
|
||||
AUTOROUTE_BACKUP_PATH,
|
||||
)
|
||||
merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key)
|
||||
merged: Final = merge_claude_settings(
|
||||
original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME
|
||||
)
|
||||
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
with secure_create(CLAUDE_SETTINGS_PATH) as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
|
|
|
|||
|
|
@ -1,51 +0,0 @@
|
|||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from .config import AUTOROUTER_MODEL_NAME
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
|
||||
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
|
||||
# Force every one of Claude Code's own model tiers to request the auto-router by name.
|
||||
# Router's auto-router registry is keyed by the literal requested model string
|
||||
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
|
||||
# model_name can never work as a catch-all -- these overrides are what actually makes
|
||||
# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
|
||||
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
|
||||
"ANTHROPIC_DEFAULT_SONNET_MODEL",
|
||||
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
|
||||
"ANTHROPIC_DEFAULT_OPUS_MODEL",
|
||||
)
|
||||
|
||||
|
||||
def merge_claude_settings_static_token(
|
||||
settings: dict[str, JsonValue], base_url: str, auth_token: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to a local ephemeral proxy with a static token.
|
||||
|
||||
Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
|
||||
remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the
|
||||
locally persisted autoroute master key, so a plain env var is simpler and correct. Any
|
||||
existing apiKeyHelper is cleared so it can't fight with the static token.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final[dict[str, JsonValue]] = {
|
||||
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
|
||||
**base_env,
|
||||
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
|
||||
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
|
||||
**{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
|
||||
}
|
||||
env.pop(ANTHROPIC_API_KEY_KEY, None)
|
||||
merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env}
|
||||
merged.pop(API_KEY_HELPER_KEY, None)
|
||||
return merged
|
||||
|
||||
|
||||
__all__ = ["merge_claude_settings_static_token"]
|
||||
|
|
@ -1,37 +1,71 @@
|
|||
"""Shared handling of Claude Code's ~/.claude/settings.json.
|
||||
|
||||
`lite up` patches this file temporarily and restores it on exit; `lite login
|
||||
--config-claude` patches it persistently. Both need the same merge and the same
|
||||
apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
|
||||
live here rather than in either command module.
|
||||
`lite up` and `lite autoroute up` patch this file temporarily and restore it on
|
||||
exit; `lite login --config-claude` and `lite configure claude` patch it
|
||||
persistently and record how to undo it. All of them need the same merge and the
|
||||
same apiKeyHelper command, and `up` already imports from `auth`, so the shared
|
||||
parts live here rather than in any one command module.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import shlex
|
||||
import shutil
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.private_json import write_private_json
|
||||
from litellm.litellm_core_utils.private_json import (
|
||||
commit_staged_json,
|
||||
discard_staged_json,
|
||||
ensure_private_dir,
|
||||
stage_private_json,
|
||||
)
|
||||
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
MODEL_KEY: Final = "model"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
|
||||
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
|
||||
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
|
||||
ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
|
||||
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
|
||||
"ANTHROPIC_DEFAULT_SONNET_MODEL",
|
||||
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
|
||||
"ANTHROPIC_DEFAULT_OPUS_MODEL",
|
||||
"ANTHROPIC_DEFAULT_FABLE_MODEL",
|
||||
)
|
||||
OWNED_ENV_KEYS: Final = (
|
||||
ENABLE_TOOL_SEARCH_KEY,
|
||||
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY,
|
||||
ANTHROPIC_BASE_URL_KEY,
|
||||
ANTHROPIC_AUTH_TOKEN_KEY,
|
||||
ANTHROPIC_API_KEY_KEY,
|
||||
)
|
||||
OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY)
|
||||
OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS)
|
||||
_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY))
|
||||
_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY)
|
||||
_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}"
|
||||
STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on"
|
||||
|
||||
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
|
||||
CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR"
|
||||
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
|
||||
CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -55,6 +89,129 @@ class ClaudeSettingsError(Exception):
|
|||
"""Raised for any user-actionable failure while reading or writing Claude Code settings."""
|
||||
|
||||
|
||||
def claude_settings_path(environ: Mapping[str, str]) -> Path:
|
||||
"""The settings.json Claude Code reads: under CLAUDE_CONFIG_DIR when set, else ~/.claude/settings.json."""
|
||||
config_dir: Final = environ.get(CLAUDE_CONFIG_DIR_ENV, "")
|
||||
if not config_dir:
|
||||
return CLAUDE_SETTINGS_PATH
|
||||
return Path(config_dir).expanduser() / "settings.json"
|
||||
|
||||
|
||||
def _is_default_settings_file(settings_path: Path) -> bool:
|
||||
return settings_path.resolve() == CLAUDE_SETTINGS_PATH.resolve()
|
||||
|
||||
|
||||
def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]:
|
||||
"""The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file."""
|
||||
return SETTINGS_FILE_OWNERS if _is_default_settings_file(settings_path) else ()
|
||||
|
||||
|
||||
def configure_state_path(settings_path: Path) -> Path:
|
||||
"""The receipt describing settings_path: the default file keeps CONFIGURE_STATE_PATH, and any other file
|
||||
(a CLAUDE_CONFIG_DIR) gets its own beside it, keyed by its resolved path, so two settings files never
|
||||
share one undo record."""
|
||||
if _is_default_settings_file(settings_path):
|
||||
return CONFIGURE_STATE_PATH
|
||||
digest: Final = hashlib.sha256(str(settings_path.resolve()).encode()).hexdigest()
|
||||
return CONFIGURE_STATE_PATH.parent / CONFIGURE_STATE_PATH.stem / f"{digest}.json"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StaticToken:
|
||||
"""A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN."""
|
||||
|
||||
token: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ApiKeyHelper:
|
||||
"""A `lite auth print-token` command Claude Code runs per request, so a login renews in place."""
|
||||
|
||||
command: str
|
||||
|
||||
|
||||
ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KeepModel:
|
||||
"""Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login)."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnpinModel:
|
||||
"""Let go of a `model` an earlier configure pinned; one the user set themselves stays."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StartOn:
|
||||
"""Pin the top-level `model`, the row Claude Code starts on."""
|
||||
|
||||
model: str
|
||||
|
||||
|
||||
ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn
|
||||
|
||||
|
||||
class OwnedValue(BaseModel):
|
||||
"""What one key held at a moment in time; `present=False` is an absent key, not a null one."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
present: bool
|
||||
value: JsonValue = None
|
||||
|
||||
|
||||
class ConfigureReceipt(BaseModel):
|
||||
"""What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key).
|
||||
|
||||
Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the
|
||||
value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and
|
||||
carries the earlier one otherwise, so a key the user edited in between stops matching and is left
|
||||
alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier
|
||||
snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the
|
||||
repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot
|
||||
was captured beside, so a credential is only ever put back next to the server it was issued for.
|
||||
No fingerprint is a second copy of a token.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
file_existed: bool
|
||||
env_present: bool
|
||||
env_was_object: bool
|
||||
previous: Mapping[str, OwnedValue]
|
||||
written: Mapping[str, str]
|
||||
endpoints: Mapping[str, OwnedValue]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WithheldCredential:
|
||||
"""A credential left removed: captured beside `endpoint`, while the restored file points elsewhere."""
|
||||
|
||||
key: str
|
||||
endpoint: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnconfigureOutcome:
|
||||
"""Keys whose value unconfigure changed back, keys the user changed since and so were left as they
|
||||
are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the
|
||||
URL points back), and whether no settings file remains."""
|
||||
|
||||
restored: tuple[str, ...]
|
||||
kept: tuple[str, ...]
|
||||
withheld: tuple[WithheldCredential, ...] = ()
|
||||
file_removed: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Claim:
|
||||
previous: OwnedValue
|
||||
written: str | None
|
||||
endpoint: OwnedValue | None
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
try:
|
||||
content: Final = path.read_bytes() if path.exists() else b""
|
||||
|
|
@ -70,29 +227,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
|||
)
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]:
|
||||
raw_env: Final = settings.get(ENV_KEY)
|
||||
if raw_env is None:
|
||||
return MappingProxyType({})
|
||||
if not isinstance(raw_env, dict):
|
||||
raise ClaudeSettingsError(
|
||||
f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.'
|
||||
)
|
||||
return raw_env
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
|
||||
defaults to true because Claude Code turns tool search off when
|
||||
ANTHROPIC_BASE_URL is not a first-party Anthropic host, and
|
||||
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker
|
||||
is filled from the proxy's /v1/models; existing values of both are left
|
||||
alone. Every other key is preserved untouched.
|
||||
|
||||
def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
"""Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a
|
||||
purely local check, so commands run it before any login prompt or request."""
|
||||
for owner in owners:
|
||||
if owner.backup_path.exists():
|
||||
raise ClaudeSettingsError(
|
||||
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
|
||||
f"{owner.backup_path}) and will restore it when it stops. "
|
||||
f"Run `{owner.stop_command}` first, then retry."
|
||||
)
|
||||
|
||||
|
||||
def _write_target(settings_path: Path) -> Path:
|
||||
"""Write through a symlinked settings.json rather than replacing the link, which would silently
|
||||
detach a file symlinked into a dotfiles repo."""
|
||||
try:
|
||||
return settings_path.resolve() if settings_path.is_symlink() else settings_path
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e
|
||||
|
||||
|
||||
def _stage(path: Path, document: Mapping[str, object]) -> str:
|
||||
try:
|
||||
return stage_private_json(str(path), document)
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not write {path}: {e}") from e
|
||||
|
||||
|
||||
def _land(
|
||||
path: Path,
|
||||
staged: str | None,
|
||||
also_discard: Sequence[str | None] = (),
|
||||
commit: Callable[[str, str], None] = commit_staged_json,
|
||||
) -> None:
|
||||
"""Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a
|
||||
filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are
|
||||
discarded, so no temp file holding a token is left behind."""
|
||||
try:
|
||||
if staged is None:
|
||||
path.unlink(missing_ok=True)
|
||||
else:
|
||||
commit(staged, str(path))
|
||||
except OSError as e:
|
||||
for other in also_discard:
|
||||
if other is not None:
|
||||
discard_staged_json(other)
|
||||
raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue],
|
||||
base_url: str,
|
||||
credential: ClaudeCredential,
|
||||
default_model: str | None = None,
|
||||
tier_model: str | None = None,
|
||||
) -> Mapping[str, JsonValue]:
|
||||
"""Return a new settings mapping wired to route Claude Code through the proxy.
|
||||
|
||||
A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper;
|
||||
the other credential slots are removed either way, since Claude Code given two credentials may
|
||||
send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their
|
||||
defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts
|
||||
on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one
|
||||
group. Apart from those tier keys, exactly OWNED_PATHS are touched.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = {
|
||||
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
|
||||
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE,
|
||||
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
|
||||
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
|
||||
}
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
current_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
|
||||
chain(
|
||||
(
|
||||
(ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE),
|
||||
(ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE),
|
||||
),
|
||||
((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS),
|
||||
((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),),
|
||||
((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (),
|
||||
((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None),
|
||||
)
|
||||
)
|
||||
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
|
||||
chain(
|
||||
((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)),
|
||||
((ENV_KEY, env),),
|
||||
((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (),
|
||||
((MODEL_KEY, default_model),) if default_model is not None else (),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
|
||||
|
|
@ -121,56 +353,273 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
|
|||
return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
|
||||
|
||||
|
||||
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
"""Persistently point Claude Code at base_url, preserving every unrelated setting.
|
||||
def lite_api_key_helper_configured(base_url: str, settings_path: Path) -> bool:
|
||||
"""Whether settings_path already carries the apiKeyHelper `lite login --config-claude` writes for base_url.
|
||||
|
||||
Refuses while any owner holds a backup: each restores its backup when it
|
||||
stops, which would silently undo this write.
|
||||
Only an exact match counts: a helper for another proxy, a hand-written one, or
|
||||
settings that cannot be read leave the caller on the env-token path.
|
||||
"""
|
||||
for owner in owners:
|
||||
if owner.backup_path.exists():
|
||||
raise ClaudeSettingsError(
|
||||
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
|
||||
f"{owner.backup_path}) and will restore it when it stops. "
|
||||
f"Run `{owner.stop_command}` first, then retry."
|
||||
)
|
||||
normalized_base_url: Final = base_url.rstrip("/")
|
||||
api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
|
||||
existing: Final = load_json_or_empty(settings_path)
|
||||
raw_env: Final = existing.get(ENV_KEY)
|
||||
if raw_env is not None and not isinstance(raw_env, dict):
|
||||
raise ClaudeSettingsError(
|
||||
f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
|
||||
"Fix or remove it, then retry."
|
||||
)
|
||||
merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
|
||||
# os.replace() swaps the symlink itself for a regular file, silently detaching a
|
||||
# settings.json that is symlinked into a dotfiles repo. There is no backup to undo
|
||||
# that here, unlike `lite up`, so write through to the link's target instead.
|
||||
target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
|
||||
try:
|
||||
write_private_json(str(target), merged)
|
||||
configured_helper: Final = load_json_or_empty(settings_path).get(API_KEY_HELPER_KEY)
|
||||
return configured_helper == resolve_api_key_helper(base_url.rstrip("/"))
|
||||
except ClaudeSettingsError:
|
||||
return False
|
||||
|
||||
|
||||
def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue:
|
||||
return OwnedValue(present=key in container, value=container.get(key))
|
||||
|
||||
|
||||
def _fingerprint(owned: OwnedValue) -> str:
|
||||
return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest()
|
||||
|
||||
|
||||
def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
||||
raw_env: Final = settings.get(ENV_KEY)
|
||||
return raw_env if isinstance(raw_env, dict) else MappingProxyType({})
|
||||
|
||||
|
||||
def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue:
|
||||
section, _, key = path.rpartition(".")
|
||||
return _owned(_env(settings) if section else settings, key)
|
||||
|
||||
|
||||
def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
|
||||
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
|
||||
chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ())
|
||||
)
|
||||
|
||||
|
||||
def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
|
||||
"""`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes."""
|
||||
section, _, key = path.rpartition(".")
|
||||
if not section:
|
||||
return _with_key(settings, key, owned)
|
||||
return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned)))
|
||||
|
||||
|
||||
def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]:
|
||||
return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings)
|
||||
|
||||
|
||||
def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool:
|
||||
"""Whether the key still holds what a configure wrote (a key no configure ever changed is never ours)."""
|
||||
return receipt.written.get(path) == _fingerprint(_lookup(settings, path))
|
||||
|
||||
|
||||
def _claim(
|
||||
path: str,
|
||||
current: Mapping[str, JsonValue],
|
||||
merged: Mapping[str, JsonValue],
|
||||
earlier: ConfigureReceipt | None,
|
||||
url_now: OwnedValue,
|
||||
) -> _Claim:
|
||||
"""What this configure records for one key; see ConfigureReceipt for the rules."""
|
||||
before, after = _lookup(current, path), _lookup(merged, path)
|
||||
carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None
|
||||
return _Claim(
|
||||
previous=before if carried is None else carried.previous.get(path, before),
|
||||
written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)),
|
||||
endpoint=None
|
||||
if path not in _CREDENTIAL_PATHS
|
||||
else (url_now if carried is None else carried.endpoints.get(path, url_now)),
|
||||
)
|
||||
|
||||
|
||||
def _receipt(
|
||||
current: Mapping[str, JsonValue],
|
||||
merged: Mapping[str, JsonValue],
|
||||
earlier: ConfigureReceipt | None,
|
||||
file_exists: bool,
|
||||
) -> ConfigureReceipt:
|
||||
url_now: Final = _lookup(current, _BASE_URL_PATH)
|
||||
claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS})
|
||||
return ConfigureReceipt(
|
||||
file_existed=file_exists if earlier is None else earlier.file_existed,
|
||||
env_present=ENV_KEY in current if earlier is None else earlier.env_present,
|
||||
env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object,
|
||||
previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}),
|
||||
written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}),
|
||||
endpoints=MappingProxyType(
|
||||
{path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None:
|
||||
if not state_path.exists():
|
||||
return None
|
||||
try:
|
||||
return ConfigureReceipt.model_validate_json(state_path.read_bytes())
|
||||
except (OSError, ValidationError) as e:
|
||||
raise ClaudeSettingsError(
|
||||
f"{state_path} is not a readable `lite configure claude` receipt ({e}). "
|
||||
"Remove it and edit Claude Code's settings by hand if they still point at the proxy."
|
||||
) from e
|
||||
|
||||
|
||||
def configure_claude_settings(
|
||||
base_url: str,
|
||||
credential: ClaudeCredential,
|
||||
model: ModelChoice,
|
||||
settings_path: Path,
|
||||
state_path: Path,
|
||||
owners: Sequence[SettingsFileOwner],
|
||||
commit: Callable[[str, str], None] = commit_staged_json,
|
||||
) -> None:
|
||||
"""Persistently route Claude Code through base_url, recording how to undo it.
|
||||
|
||||
Both files are staged before either is committed, so a full disk or a read-only directory fails
|
||||
before anything changes. The two commits are still two renames: a receipt rename that fails
|
||||
discards the staged settings, and a settings rename that fails after the receipt landed puts the
|
||||
earlier receipt back (or removes the new one), so the receipt on disk never describes settings
|
||||
that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an
|
||||
earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login).
|
||||
"""
|
||||
refuse_while_owned(settings_path, owners)
|
||||
current: Final = load_json_or_empty(settings_path)
|
||||
_env_object(current, settings_path)
|
||||
earlier: Final = read_configure_receipt(state_path)
|
||||
existing: Final = (
|
||||
_with(current, MODEL_KEY, earlier.previous[MODEL_KEY])
|
||||
if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier)
|
||||
else current
|
||||
)
|
||||
merged: Final = merge_claude_settings(
|
||||
existing, base_url, credential, model.model if isinstance(model, StartOn) else None
|
||||
)
|
||||
receipt: Final = _receipt(current, merged, earlier, settings_path.exists())
|
||||
target: Final = _write_target(settings_path)
|
||||
try:
|
||||
ensure_private_dir(state_path.parent)
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
|
||||
raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e
|
||||
staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json"))
|
||||
try:
|
||||
staged_settings: Final = _stage(target, merged)
|
||||
except ClaudeSettingsError:
|
||||
discard_staged_json(staged_receipt)
|
||||
raise
|
||||
_land(state_path, staged_receipt, (staged_settings,), commit)
|
||||
try:
|
||||
_land(target, staged_settings, commit=commit)
|
||||
except ClaudeSettingsError as settings_error:
|
||||
try:
|
||||
_land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json")))
|
||||
except ClaudeSettingsError as receipt_error:
|
||||
raise ClaudeSettingsError(
|
||||
f"{settings_error} The receipt at {state_path} now describes settings that were not written and "
|
||||
f"could not be put back either ({receipt_error}); remove it before retrying."
|
||||
) from settings_error
|
||||
raise
|
||||
|
||||
|
||||
def _endpoint_text(endpoint: OwnedValue) -> str:
|
||||
if not endpoint.present:
|
||||
return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)"
|
||||
return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value)
|
||||
|
||||
|
||||
def unconfigure_claude_settings(
|
||||
settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner]
|
||||
) -> UnconfigureOutcome:
|
||||
"""Undo `lite configure claude`: put back every key still holding what configure wrote, leave the
|
||||
rest alone, and withhold a credential the restored file would send to a different server than it
|
||||
was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish)."""
|
||||
refuse_while_owned(settings_path, owners)
|
||||
receipt: Final = read_configure_receipt(state_path)
|
||||
if receipt is None:
|
||||
raise ClaudeSettingsError(
|
||||
f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo."
|
||||
)
|
||||
current: Final = load_json_or_empty(settings_path)
|
||||
_env_object(current, settings_path)
|
||||
ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt))
|
||||
kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present)
|
||||
put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours}))
|
||||
url_after: Final = _lookup(put_back, _BASE_URL_PATH)
|
||||
withheld: Final = tuple(
|
||||
WithheldCredential(path, _endpoint_text(receipt.endpoints[path]))
|
||||
for path in _CREDENTIAL_PATHS
|
||||
if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after
|
||||
)
|
||||
absent: Final = OwnedValue(present=False)
|
||||
trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld}))
|
||||
settings: Final = (
|
||||
trimmed
|
||||
if _env(trimmed) or receipt.env_was_object
|
||||
else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None))
|
||||
)
|
||||
target: Final = _write_target(settings_path)
|
||||
file_removed: Final = not settings and not (receipt.file_existed and target.exists())
|
||||
kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy
|
||||
receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}})
|
||||
if withheld
|
||||
else None
|
||||
)
|
||||
staged_settings: Final = None if file_removed else _stage(target, settings)
|
||||
try:
|
||||
staged_receipt: Final = (
|
||||
None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json"))
|
||||
)
|
||||
except ClaudeSettingsError:
|
||||
if staged_settings is not None:
|
||||
discard_staged_json(staged_settings)
|
||||
raise
|
||||
_land(target, staged_settings, (staged_receipt,))
|
||||
_land(state_path, staged_receipt)
|
||||
return UnconfigureOutcome(
|
||||
restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)),
|
||||
kept=kept,
|
||||
withheld=withheld,
|
||||
file_removed=file_removed,
|
||||
)
|
||||
|
||||
|
||||
__all__ = (
|
||||
"ANTHROPIC_API_KEY_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN_KEY",
|
||||
"ANTHROPIC_BASE_URL_KEY",
|
||||
"ANTHROPIC_DEFAULT_MODEL_ENV_KEYS",
|
||||
"API_KEY_HELPER_KEY",
|
||||
"AUTOROUTE_BACKUP_PATH",
|
||||
"BACKUP_PATH",
|
||||
"CLAUDE_CONFIG_DIR_ENV",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"CONFIGURE_STATE_PATH",
|
||||
"ENABLE_GATEWAY_MODEL_DISCOVERY_KEY",
|
||||
"ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE",
|
||||
"ENABLE_TOOL_SEARCH_KEY",
|
||||
"ENABLE_TOOL_SEARCH_VALUE",
|
||||
"ENV_KEY",
|
||||
"MODEL_KEY",
|
||||
"OWNED_ENV_KEYS",
|
||||
"OWNED_PATHS",
|
||||
"OWNED_TOP_LEVEL_KEYS",
|
||||
"SETTINGS_FILE_OWNERS",
|
||||
"STARTING_MODEL_ROLE",
|
||||
"ApiKeyHelper",
|
||||
"ClaudeCredential",
|
||||
"ClaudeSettingsError",
|
||||
"ConfigureReceipt",
|
||||
"KeepModel",
|
||||
"ModelChoice",
|
||||
"OwnedValue",
|
||||
"SettingsFileOwner",
|
||||
"StartOn",
|
||||
"StaticToken",
|
||||
"UnconfigureOutcome",
|
||||
"UnpinModel",
|
||||
"WithheldCredential",
|
||||
"claude_settings_path",
|
||||
"configure_claude_settings",
|
||||
"configure_state_path",
|
||||
"lite_api_key_helper_configured",
|
||||
"load_json_or_empty",
|
||||
"merge_claude_settings",
|
||||
"read_configure_receipt",
|
||||
"refuse_while_owned",
|
||||
"resolve_api_key_helper",
|
||||
"write_claude_settings",
|
||||
"settings_file_owners",
|
||||
"unconfigure_claude_settings",
|
||||
)
|
||||
|
|
|
|||
292
litellm/proxy/client/cli/commands/configure.py
Normal file
292
litellm/proxy/client/cli/commands/configure.py
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import click
|
||||
from InquirerPy import inquirer
|
||||
from InquirerPy.base.control import Choice
|
||||
|
||||
from litellm.proxy.common_utils.model_listing_utils import (
|
||||
CLAUDE_CODE_CLIENT,
|
||||
CLAUDE_CODE_PICKER_PATTERN,
|
||||
GATEWAY_CLIENT_HEADER,
|
||||
)
|
||||
|
||||
from .auth import CliContextObj, context_secret_vault, get_stored_api_key
|
||||
from .claude_settings import (
|
||||
STARTING_MODEL_ROLE,
|
||||
ApiKeyHelper,
|
||||
ClaudeCredential,
|
||||
ClaudeSettingsError,
|
||||
ModelChoice,
|
||||
StartOn,
|
||||
StaticToken,
|
||||
UnconfigureOutcome,
|
||||
UnpinModel,
|
||||
claude_settings_path,
|
||||
configure_claude_settings,
|
||||
configure_state_path,
|
||||
refuse_while_owned,
|
||||
resolve_api_key_helper,
|
||||
settings_file_owners,
|
||||
unconfigure_claude_settings,
|
||||
)
|
||||
from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing
|
||||
from .up import ensure_fresh_login
|
||||
|
||||
_LISTED_MODELS_SHOWN: Final = 20
|
||||
_CLAUDE_TARGET: Final = "claude"
|
||||
_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),)
|
||||
_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default"
|
||||
_CLAUDE_CODE_VIEW: Final = MappingProxyType(
|
||||
{"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT}
|
||||
)
|
||||
_MODEL_OPTION_HELP: Final = (
|
||||
f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, "
|
||||
"Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude "
|
||||
"Code's sub-agent or background tiers; `lite autoroute up` is the mode that does."
|
||||
)
|
||||
|
||||
|
||||
def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]:
|
||||
"""The credential to write and the key to check the proxy with.
|
||||
|
||||
An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes
|
||||
into settings.json as a static token. Without one, the stored `lite login` credential is used
|
||||
the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a
|
||||
day and renews in place there; a missing or stale login is refreshed first, as `lite up` does.
|
||||
"""
|
||||
ctx_obj: Final[CliContextObj] = ctx.obj
|
||||
explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key"))
|
||||
if explicit:
|
||||
return StaticToken(explicit), explicit
|
||||
base_url: Final = ctx_obj["base_url"]
|
||||
ensure_fresh_login(ctx)
|
||||
stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
|
||||
if not stored:
|
||||
raise ClaudeSettingsError("Login did not produce a usable token.")
|
||||
return ApiKeyHelper(resolve_api_key_helper(base_url)), stored
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Listing:
|
||||
models: tuple[ListedModel, ...]
|
||||
|
||||
@property
|
||||
def ids(self) -> tuple[str, ...]:
|
||||
return tuple(model.id for model in self.models)
|
||||
|
||||
|
||||
def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, _Listing]:
|
||||
"""Every configure path begins the same way: the local ownership check first, so a `lite up`
|
||||
session is refused before any login prompt or request, then the credential, then the listing."""
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
try:
|
||||
refuse_while_owned(settings_path, settings_file_owners(settings_path))
|
||||
credential, key = resolve_credential(ctx, api_key)
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
return credential, _listed_models(ctx.obj["base_url"], key)
|
||||
|
||||
|
||||
def _listing_error(base_url: str, error: PiSyncError) -> str:
|
||||
"""The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question."""
|
||||
if error.kind is ListingFailure.REJECTED:
|
||||
return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key."
|
||||
if error.kind is ListingFailure.UNREACHABLE:
|
||||
return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?"
|
||||
if error.kind is ListingFailure.EMPTY:
|
||||
return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model."
|
||||
return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy."
|
||||
|
||||
|
||||
def _listed_models(base_url: str, key: str) -> _Listing:
|
||||
listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW)
|
||||
if isinstance(listed, PiSyncError):
|
||||
raise click.ClickException(_listing_error(base_url, listed))
|
||||
return _Listing(listed)
|
||||
|
||||
|
||||
def _starting_model(model: str, listing: _Listing) -> str | None:
|
||||
source: Final = next((listed.id for listed in listing.models if listed.source_model == model), None)
|
||||
return source or next((listed.id for listed in listing.models if listed.id == model), None)
|
||||
|
||||
|
||||
def _model_choice(model: str | None) -> ModelChoice:
|
||||
return StartOn(model) if model is not None else UnpinModel()
|
||||
|
||||
|
||||
def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Listing, model: str | None) -> None:
|
||||
ctx_obj: Final[CliContextObj] = ctx.obj
|
||||
base_url: Final = ctx_obj["base_url"]
|
||||
listed: Final = listing.ids
|
||||
starting: Final = _starting_model(model, listing) if model is not None else None
|
||||
if model is not None and starting is None:
|
||||
shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN])
|
||||
more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else ""
|
||||
raise click.ClickException(
|
||||
f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}."
|
||||
)
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
try:
|
||||
configure_claude_settings(
|
||||
base_url,
|
||||
credential,
|
||||
_model_choice(starting),
|
||||
settings_path,
|
||||
configure_state_path(settings_path),
|
||||
settings_file_owners(settings_path),
|
||||
)
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model))
|
||||
click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.")
|
||||
|
||||
click.echo(
|
||||
"Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN."
|
||||
if isinstance(credential, StaticToken)
|
||||
else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it."
|
||||
)
|
||||
click.echo(
|
||||
f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model."
|
||||
if starting is not None
|
||||
else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or "
|
||||
"pass --model to start on a proxy model."
|
||||
)
|
||||
click.echo(
|
||||
f"/model will list all {len(listed)} of the proxy's models."
|
||||
if in_picker == len(listed)
|
||||
else f"/model will list {in_picker} of the proxy's {len(listed)} models: Claude Code shows only ids containing "
|
||||
"'claude' or 'anthropic', and this proxy does not list the rest under such names."
|
||||
)
|
||||
click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.")
|
||||
if isinstance(credential, StaticToken) and settings_path.is_symlink():
|
||||
click.echo(
|
||||
f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in "
|
||||
"that file; keep it out of version control.",
|
||||
err=True,
|
||||
)
|
||||
|
||||
|
||||
def _pick_targets() -> tuple[str, ...]:
|
||||
picked: Final = inquirer.checkbox(
|
||||
message="Which agents should route through LiteLLM?",
|
||||
choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS],
|
||||
validate=lambda chosen: len(chosen) > 0,
|
||||
invalid_message="Pick at least one.",
|
||||
).execute()
|
||||
return tuple(str(value) for value in picked)
|
||||
|
||||
|
||||
def _pick_model(listed: Sequence[str]) -> str | None:
|
||||
picked: Final = inquirer.fuzzy(
|
||||
message="Model Claude Code starts on (type to filter; /model switches any time):",
|
||||
choices=[_KEEP_DEFAULT_MODEL, *listed],
|
||||
).execute()
|
||||
return None if picked == _KEEP_DEFAULT_MODEL else str(picked)
|
||||
|
||||
|
||||
def interactive_configure(
|
||||
ctx: click.Context,
|
||||
pick_targets: Callable[[], tuple[str, ...]] = _pick_targets,
|
||||
pick_model: Callable[[Sequence[str]], str | None] = _pick_model,
|
||||
) -> None:
|
||||
"""`lite configure` with no agent named: ask which agents to wire and which model to pin."""
|
||||
targets: Final = pick_targets()
|
||||
if _CLAUDE_TARGET not in targets:
|
||||
return
|
||||
credential, listing = _start(ctx, None)
|
||||
_apply_claude(
|
||||
ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models))
|
||||
)
|
||||
|
||||
|
||||
@click.group(name="configure", invoke_without_command=True)
|
||||
@click.pass_context
|
||||
def configure_group(ctx: click.Context) -> None:
|
||||
"""Persistently route a coding agent through your LiteLLM proxy.
|
||||
|
||||
With no agent named, asks which agents to wire and which proxy model to pin.
|
||||
"""
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
if not sys.stdin.isatty():
|
||||
raise click.ClickException(
|
||||
"`lite configure` asks questions, so it needs a terminal. Non-interactively, run "
|
||||
"`lite configure claude --api-key <key> --model <model>`."
|
||||
)
|
||||
interactive_configure(ctx)
|
||||
|
||||
|
||||
@click.group(name="unconfigure")
|
||||
def unconfigure_group() -> None:
|
||||
"""Undo `lite configure` for a coding agent."""
|
||||
|
||||
|
||||
@configure_group.command(name="claude")
|
||||
@click.option(
|
||||
"--api-key",
|
||||
"api_key",
|
||||
default=None,
|
||||
help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / "
|
||||
"LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.",
|
||||
)
|
||||
@click.option("--model", default=None, help=_MODEL_OPTION_HELP)
|
||||
@click.pass_context
|
||||
def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None:
|
||||
"""Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`.
|
||||
|
||||
Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a
|
||||
static token, or your `lite login` through apiKeyHelper), and gateway model discovery so
|
||||
/model lists the proxy's models; --model picks the one Claude Code starts on. Every other
|
||||
setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back.
|
||||
Assumes the proxy is already running.
|
||||
"""
|
||||
credential, listing = _start(ctx, api_key)
|
||||
_apply_claude(ctx, credential, listing, model)
|
||||
|
||||
|
||||
@unconfigure_group.command(name="claude")
|
||||
def unconfigure_claude() -> None:
|
||||
"""Return Claude Code's settings to what they were before `lite configure claude`.
|
||||
|
||||
Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are
|
||||
put back; anything you changed since is left as it is and named in the output.
|
||||
"""
|
||||
settings_path: Final = claude_settings_path(os.environ)
|
||||
state_path: Final = configure_state_path(settings_path)
|
||||
try:
|
||||
outcome: Final = unconfigure_claude_settings(settings_path, state_path, settings_file_owners(settings_path))
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
_report_unconfigure(settings_path, state_path, outcome)
|
||||
|
||||
|
||||
def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None:
|
||||
"""Say what unconfigure did, naming only keys whose value it changed."""
|
||||
if outcome.file_removed:
|
||||
click.echo(
|
||||
f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys."
|
||||
)
|
||||
elif outcome.restored:
|
||||
click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.")
|
||||
else:
|
||||
click.echo(f"Nothing in {settings_path} was still ours to restore.")
|
||||
if outcome.kept:
|
||||
click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.")
|
||||
if outcome.withheld:
|
||||
click.echo(
|
||||
"Left removed, since the file now points at a different server than they were issued for: "
|
||||
+ "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld)
|
||||
+ f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` "
|
||||
"again to put them back, or delete that file to drop them."
|
||||
)
|
||||
|
||||
|
||||
__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group")
|
||||
|
|
@ -10,21 +10,40 @@ import os
|
|||
import tempfile
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import Annotated, Final
|
||||
|
||||
import requests
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator
|
||||
from pydantic.types import StringConstraints
|
||||
|
||||
PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
|
||||
PI_PROVIDER_NAME: Final = "litellm"
|
||||
LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
|
||||
_REJECTED_STATUSES: Final = frozenset((401, 403))
|
||||
|
||||
|
||||
class ListingFailure(StrEnum):
|
||||
"""Why a proxy could not be listed, decided once where the HTTP outcome is classified.
|
||||
|
||||
`unreachable` means no response at all; the other kinds prove the proxy answered, so callers
|
||||
must not suggest checking whether it is running.
|
||||
"""
|
||||
|
||||
UNREACHABLE = "unreachable"
|
||||
REJECTED = "rejected"
|
||||
BAD_BODY = "bad_body"
|
||||
EMPTY = "empty"
|
||||
OTHER = "other"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PiSyncError:
|
||||
message: str
|
||||
status: int | None = None
|
||||
kind: ListingFailure | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -33,12 +52,25 @@ class ModelLimits:
|
|||
max_tokens: int | None
|
||||
|
||||
|
||||
class _Model(BaseModel):
|
||||
id: str
|
||||
_NonEmptyString = Annotated[str, StringConstraints(min_length=1)]
|
||||
|
||||
|
||||
class ListedModel(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: _NonEmptyString
|
||||
source_model: _NonEmptyString | None = None
|
||||
|
||||
|
||||
class _ModelList(BaseModel):
|
||||
data: tuple[_Model, ...]
|
||||
data: tuple[ListedModel, ...]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def unique_id_mappings(self) -> "_ModelList":
|
||||
mappings: Final = frozenset((model.id, model.source_model or model.id) for model in self.data)
|
||||
if len(frozenset(model.id for model in self.data)) != len(mappings):
|
||||
raise ValueError("model ids must not map to multiple source models")
|
||||
return self
|
||||
|
||||
|
||||
class _ModelGroup(BaseModel):
|
||||
|
|
@ -51,31 +83,47 @@ class _ModelGroupList(BaseModel):
|
|||
data: tuple[_ModelGroup, ...]
|
||||
|
||||
|
||||
def fetch_model_listing(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
*,
|
||||
get: Callable[..., requests.Response] = requests.get,
|
||||
headers: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> tuple[ListedModel, ...] | PiSyncError:
|
||||
url: Final = base_url.rstrip("/") + "/v1/models"
|
||||
try:
|
||||
resp: Final = get(
|
||||
url,
|
||||
headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict
|
||||
timeout=10,
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE)
|
||||
if resp.status_code != 200:
|
||||
return PiSyncError(
|
||||
f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.",
|
||||
resp.status_code,
|
||||
ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER,
|
||||
)
|
||||
try:
|
||||
listing: Final = _ModelList.model_validate(resp.json())
|
||||
except (ValueError, ValidationError) as e:
|
||||
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY)
|
||||
models: Final = tuple(dict.fromkeys(listing.data))
|
||||
if not models:
|
||||
return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY)
|
||||
return models
|
||||
|
||||
|
||||
def fetch_model_ids(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
*,
|
||||
get: Callable[..., requests.Response] = requests.get,
|
||||
headers: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> tuple[str, ...] | PiSyncError:
|
||||
url: Final = base_url.rstrip("/") + "/v1/models"
|
||||
try:
|
||||
resp: Final = get(
|
||||
url,
|
||||
headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict
|
||||
timeout=10,
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
return PiSyncError(f"Could not list models from the proxy: {e}")
|
||||
if resp.status_code != 200:
|
||||
return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
|
||||
try:
|
||||
listing: Final = _ModelList.model_validate(resp.json())
|
||||
except (ValueError, ValidationError) as e:
|
||||
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
|
||||
ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
|
||||
if not ids:
|
||||
return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
|
||||
return ids
|
||||
listed: Final = fetch_model_listing(base_url, api_key, get=get, headers=headers)
|
||||
return listed if isinstance(listed, PiSyncError) else tuple(dict.fromkeys(model.id for model in listed))
|
||||
|
||||
|
||||
_NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({})
|
||||
|
|
@ -200,10 +248,13 @@ __all__ = (
|
|||
"LITELLM_PROXY_API_KEY_ENV",
|
||||
"PI_CONFIG_DIR_ENV",
|
||||
"PI_PROVIDER_NAME",
|
||||
"ListedModel",
|
||||
"ListingFailure",
|
||||
"ModelLimits",
|
||||
"PiSyncError",
|
||||
"fetch_model_ids",
|
||||
"fetch_model_limits",
|
||||
"fetch_model_listing",
|
||||
"models_json_path",
|
||||
"provider_block",
|
||||
"sync_models_json",
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_
|
|||
from .claude_settings import (
|
||||
BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ApiKeyHelper,
|
||||
ClaudeSettingsError,
|
||||
load_json_or_empty,
|
||||
merge_claude_settings,
|
||||
|
|
@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool:
|
|||
return token_data is not None and token_data.get("refresh_token") is not None
|
||||
|
||||
|
||||
def _ensure_fresh_login(ctx: click.Context) -> None:
|
||||
def ensure_fresh_login(ctx: click.Context) -> None:
|
||||
ctx_obj: Final[CliContextObj] = ctx.obj
|
||||
base_url: Final = ctx_obj["base_url"].rstrip("/")
|
||||
vault: Final = context_secret_vault(ctx)
|
||||
|
|
@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None:
|
|||
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
|
||||
ctx.invoke(login, pkce=pkce)
|
||||
if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault):
|
||||
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
|
||||
raise UpError("Login did not produce a usable token.")
|
||||
|
||||
|
||||
def _restore_and_report() -> None:
|
||||
|
|
@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None:
|
|||
base_url: Final = ctx.obj["base_url"]
|
||||
|
||||
try:
|
||||
_ensure_fresh_login(ctx)
|
||||
ensure_fresh_login(ctx)
|
||||
api_key: Final = resolve_api_key(ctx)
|
||||
verify_proxy_key(base_url, api_key)
|
||||
|
||||
|
|
@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None:
|
|||
)
|
||||
|
||||
CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
|
||||
merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
|
||||
merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper))
|
||||
with open(CLAUDE_SETTINGS_PATH, "w") as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except (AgentRunError, ClaudeSettingsError) as e:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue