mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_bedrock_sol_anthropic_1hr
This commit is contained in:
commit
4b3e82b8a3
315 changed files with 15746 additions and 2549 deletions
|
|
@ -84,7 +84,7 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1810
|
||||
"limit": 1808
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44530
|
||||
"limit": 44528
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38808
|
||||
"limit": 38804
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19829
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30356
|
||||
"limit": 30355
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
|
|
@ -135,12 +135,12 @@
|
|||
"limit": 21
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"limit": 139
|
||||
"limit": 138
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 545
|
||||
"limit": 544
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 146
|
||||
"limit": 145
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -157,6 +157,9 @@ COST_DESCRIPTIONS: dict[str, str] = {
|
|||
"input_cost_per_token": "USD per prompt token.",
|
||||
"output_cost_per_token": "USD per generated token.",
|
||||
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
|
||||
"google_maps_grounding_cost_per_query": (
|
||||
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
|
||||
),
|
||||
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
|
||||
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
|
||||
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""
|
||||
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
|
||||
Cost tracking is handled automatically by the get-responses call.
|
||||
Cost tracking is handled by the get-responses call, which prices normally only because the
|
||||
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
|
||||
same route are non-inference and free.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -9,12 +11,14 @@ from typing import TYPE_CHECKING, Dict, Optional, cast
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
STALE_OBJECT_CLEANUP_BATCH_SIZE,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -113,7 +117,8 @@ class CheckResponsesCost:
|
|||
Check if background responses are complete and track their cost.
|
||||
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
|
||||
- Query the provider to check if response is complete
|
||||
- Cost is automatically tracked by the get-responses call
|
||||
- Cost is tracked by the get-responses call, billed because the poll is stamped
|
||||
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
- Mark responses in a terminal state as complete in the database
|
||||
"""
|
||||
try:
|
||||
|
|
@ -153,6 +158,7 @@ class CheckResponsesCost:
|
|||
# Prepare metadata with model information for cost tracking
|
||||
litellm_metadata = {
|
||||
"user_api_key_user_id": job.created_by or "default-user-id",
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
}
|
||||
|
||||
# Add model information if available
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import asyncio
|
|||
import datetime
|
||||
import inspect
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.caching import InMemoryCache
|
||||
from litellm.caching.caching import S3Cache
|
||||
from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
update_response_metadata,
|
||||
)
|
||||
|
|
@ -124,6 +125,29 @@ def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") ->
|
|||
return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
|
||||
|
||||
|
||||
_PENDING_CACHE_WRITES: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs to pending write tasks
|
||||
|
||||
|
||||
async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], Awaitable[None]]) -> None:
|
||||
try:
|
||||
await write_factory()
|
||||
except asyncio.CancelledError:
|
||||
try:
|
||||
await asyncio.wait_for(write_factory(), timeout=CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
|
||||
except Exception as flush_error: # noqa: BLE001 # shutdown flush failures are logged, never raised
|
||||
verbose_logger.warning(
|
||||
"LiteLLM Cache: pending cache write failed during event loop shutdown: %s", flush_error
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]":
|
||||
task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory))
|
||||
_PENDING_CACHE_WRITES.add(task)
|
||||
task.add_done_callback(_PENDING_CACHE_WRITES.discard)
|
||||
return task
|
||||
|
||||
|
||||
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
|
||||
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
|
||||
return request_kwargs.get("cache_key", None)
|
||||
|
|
@ -983,6 +1007,7 @@ class LLMCachingHandler:
|
|||
|
||||
if litellm.cache is None:
|
||||
return
|
||||
cache: Final = litellm.cache
|
||||
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
|
|
@ -1004,24 +1029,24 @@ class LLMCachingHandler:
|
|||
):
|
||||
if (
|
||||
isinstance(result, EmbeddingResponse)
|
||||
and litellm.cache is not None
|
||||
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
|
||||
and not isinstance(cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
|
||||
):
|
||||
asyncio.create_task(
|
||||
litellm.cache.async_add_cache_pipeline(
|
||||
create_cache_write_task(
|
||||
lambda: cache.async_add_cache_pipeline(
|
||||
result, dynamic_cache_object=self.dual_cache, **new_kwargs
|
||||
)
|
||||
)
|
||||
else:
|
||||
asyncio.create_task(
|
||||
litellm.cache.async_add_cache(
|
||||
result.model_dump_json(),
|
||||
result_json: Final = result.model_dump_json()
|
||||
create_cache_write_task(
|
||||
lambda: cache.async_add_cache(
|
||||
result_json,
|
||||
dynamic_cache_object=self.dual_cache,
|
||||
**new_kwargs,
|
||||
)
|
||||
)
|
||||
else:
|
||||
asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs))
|
||||
create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
|
||||
|
||||
def sync_set_cache(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -435,7 +435,7 @@ class RedisCache(BaseCache):
|
|||
"""
|
||||
if key is None:
|
||||
return key
|
||||
if self.namespace is not None and not key.startswith(self.namespace):
|
||||
if self.namespace and not key.startswith(self.namespace + ":"):
|
||||
key = self.namespace + ":" + key
|
||||
|
||||
return key
|
||||
|
|
|
|||
|
|
@ -381,6 +381,7 @@ AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_
|
|||
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
|
||||
AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96))
|
||||
REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
|
||||
CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
|
||||
REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
|
||||
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
|
||||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))
|
||||
|
|
@ -1363,8 +1364,6 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
|
|||
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
|
||||
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
|
||||
AUTO_ROUTED_REQUEST_METADATA_KEY: Final = "_auto_routed_request"
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD: Final = "router_model_name"
|
||||
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
||||
|
|
@ -1813,6 +1812,43 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
|
|||
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
|
||||
|
||||
# A retrieved response replays the usage of the call that created it, so pricing these
|
||||
# read/management routes like inference bills the same tokens twice.
|
||||
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"get_responses",
|
||||
"aget_responses",
|
||||
"delete_responses",
|
||||
"adelete_responses",
|
||||
"cancel_responses",
|
||||
"acancel_responses",
|
||||
"list_input_items",
|
||||
"alist_input_items",
|
||||
"vector_store_create",
|
||||
"avector_store_create",
|
||||
"vector_store_retrieve",
|
||||
"avector_store_retrieve",
|
||||
"vector_store_list",
|
||||
"avector_store_list",
|
||||
"vector_store_update",
|
||||
"avector_store_update",
|
||||
"vector_store_delete",
|
||||
"avector_store_delete",
|
||||
"vector_store_file_create",
|
||||
"avector_store_file_create",
|
||||
"vector_store_file_list",
|
||||
"avector_store_file_list",
|
||||
"vector_store_file_retrieve",
|
||||
"avector_store_file_retrieve",
|
||||
"vector_store_file_content",
|
||||
"avector_store_file_content",
|
||||
"vector_store_file_update",
|
||||
"avector_store_file_update",
|
||||
"vector_store_file_delete",
|
||||
"avector_store_file_delete",
|
||||
}
|
||||
)
|
||||
|
||||
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
|
||||
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
|
||||
# spend under the table's composite unique constraint.
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
## File for 'response_cost' calculation in Logging
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
||||
|
|
@ -591,6 +592,7 @@ def cost_per_token(
|
|||
prompt_characters=prompt_characters,
|
||||
completion_characters=completion_characters,
|
||||
usage=usage_block,
|
||||
service_tier=service_tier,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
elif cost_router == "cost_per_token":
|
||||
|
|
@ -794,14 +796,27 @@ def _select_model_name_for_cost_calc(
|
|||
and custom_llm_provider is not None
|
||||
and not _model_contains_known_llm_provider(return_model)
|
||||
): # add provider prefix if not already present, to match model_cost
|
||||
if region_name is not None:
|
||||
return_model = f"{custom_llm_provider}/{region_name}/{return_model}"
|
||||
else:
|
||||
return_model = f"{custom_llm_provider}/{return_model}"
|
||||
provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}"
|
||||
return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name)
|
||||
|
||||
return return_model
|
||||
|
||||
|
||||
def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str:
|
||||
"""Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the
|
||||
registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already
|
||||
resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069)."""
|
||||
segments: Final = model.split("/")
|
||||
if "/".join(segments[1:]) in litellm.model_cost:
|
||||
return model
|
||||
head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1
|
||||
head: Final = "/".join(segments[:head_len])
|
||||
tail: Final = segments[head_len:]
|
||||
strippable: Final = next((index for index, segment in enumerate(tail) if segment in LlmProvidersSet), len(tail))
|
||||
candidates: Final = (f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1))
|
||||
return next((candidate for candidate in candidates if candidate in litellm.model_cost), model)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _model_contains_known_llm_provider(model: str) -> bool:
|
||||
"""
|
||||
|
|
@ -832,9 +847,11 @@ def _get_response_model(completion_response: object) -> str | None:
|
|||
_GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = {
|
||||
# ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc.
|
||||
"ON_DEMAND_PRIORITY": "priority",
|
||||
# FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc.
|
||||
# FLEX / BATCH / ON_DEMAND_FLEX maps to "flex" — selects input_cost_per_token_flex, etc.
|
||||
# Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX, not FLEX.
|
||||
"FLEX": "flex",
|
||||
"BATCH": "flex",
|
||||
"ON_DEMAND_FLEX": "flex",
|
||||
# ON_DEMAND is standard pricing — no service_tier suffix applied
|
||||
"ON_DEMAND": None,
|
||||
}
|
||||
|
|
@ -849,9 +866,9 @@ def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
|
|||
|
||||
trafficType values seen in practice
|
||||
------------------------------------
|
||||
ON_DEMAND -> standard pricing (service_tier = None)
|
||||
ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority")
|
||||
FLEX / BATCH -> batch/flex pricing (service_tier = "flex")
|
||||
ON_DEMAND -> standard pricing (service_tier = None)
|
||||
ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority")
|
||||
FLEX / BATCH / ON_DEMAND_FLEX -> batch/flex pricing (service_tier = "flex")
|
||||
"""
|
||||
if traffic_type is None:
|
||||
return None
|
||||
|
|
@ -2357,6 +2374,64 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
|
|||
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
|
||||
|
||||
|
||||
def _candidate_realtime_token_costs(
|
||||
model_name: str,
|
||||
combined_usage_object: Usage,
|
||||
custom_llm_provider: str,
|
||||
data_residency: str | None,
|
||||
) -> tuple[float, float] | None:
|
||||
try:
|
||||
return generic_cost_per_token(
|
||||
model=model_name,
|
||||
usage=combined_usage_object,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) -> bool:
|
||||
entries: Final = (
|
||||
litellm.model_cost.get(model_name),
|
||||
litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"),
|
||||
)
|
||||
return any(
|
||||
entry is not None and any("cost_per" in field and value is not None for field, value in entry.items())
|
||||
for entry in entries
|
||||
)
|
||||
|
||||
|
||||
def _first_priced_realtime_token_costs(
|
||||
potential_model_names: Sequence[str | None],
|
||||
combined_usage_object: Usage,
|
||||
custom_llm_provider: str,
|
||||
data_residency: str | None,
|
||||
) -> tuple[float, float]:
|
||||
candidate_costs: Final = (
|
||||
(model_name, costs)
|
||||
for model_name in potential_model_names
|
||||
if model_name is not None
|
||||
and (
|
||||
costs := _candidate_realtime_token_costs(
|
||||
model_name=model_name,
|
||||
combined_usage_object=combined_usage_object,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
return next(
|
||||
(
|
||||
costs
|
||||
for model_name, costs in candidate_costs
|
||||
if sum(costs) > 0 or _cost_map_entry_declares_pricing(model_name, custom_llm_provider)
|
||||
),
|
||||
(0.0, 0.0),
|
||||
)
|
||||
|
||||
|
||||
def handle_realtime_stream_cost_calculation(
|
||||
results: OpenAIRealtimeStreamList,
|
||||
combined_usage_object: Usage,
|
||||
|
|
@ -2381,24 +2456,12 @@ def handle_realtime_stream_cost_calculation(
|
|||
potential_model_names.append(received_model)
|
||||
|
||||
potential_model_names.append(litellm_model_name)
|
||||
input_cost_per_token = 0.0
|
||||
output_cost_per_token = 0.0
|
||||
|
||||
for model_name in potential_model_names:
|
||||
try:
|
||||
if model_name is None:
|
||||
continue
|
||||
_input_cost_per_token, _output_cost_per_token = generic_cost_per_token(
|
||||
model=model_name,
|
||||
usage=combined_usage_object,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
input_cost_per_token += _input_cost_per_token
|
||||
output_cost_per_token += _output_cost_per_token
|
||||
break # exit if we find a valid model
|
||||
input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs(
|
||||
potential_model_names=potential_model_names,
|
||||
combined_usage_object=combined_usage_object,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
transcription_cost: Final = (
|
||||
handle_realtime_transcription_cost_calculation(
|
||||
results=results,
|
||||
|
|
|
|||
|
|
@ -8,6 +8,14 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def _completion_response_cost(model_response: "ModelResponse") -> float | None:
|
||||
hidden_params: Final = getattr(model_response, "_hidden_params", None)
|
||||
if not isinstance(hidden_params, dict):
|
||||
return None
|
||||
response_cost: Final = hidden_params.get("response_cost")
|
||||
return response_cost if isinstance(response_cost, float) else None
|
||||
|
||||
|
||||
class SpeechToCompletionBridgeTransformationHandler:
|
||||
def transform_request(
|
||||
self,
|
||||
|
|
@ -123,4 +131,6 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
|
||||
# Create an httpx.Response object
|
||||
response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers)
|
||||
return HttpxBinaryResponseContent(response)
|
||||
binary_response: Final = HttpxBinaryResponseContent(response)
|
||||
binary_response.set_response_cost(_completion_response_cost(model_response))
|
||||
return binary_response
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import base64
|
|||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from datetime import timedelta
|
||||
from importlib import metadata
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
import httpx
|
||||
|
|
@ -21,6 +22,18 @@ try:
|
|||
streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
MCP_STREAMABLE_HTTP_REQUIREMENT: Final = "mcp>=1.28.1"
|
||||
|
||||
|
||||
def missing_streamable_http_client_error() -> ImportError:
|
||||
return ImportError(
|
||||
f"MCP streamable HTTP transport requires {MCP_STREAMABLE_HTTP_REQUIREMENT}, but the installed "
|
||||
f"mcp {metadata.version('mcp')} does not provide streamable_http_client. "
|
||||
"Fix with: pip install 'litellm[mcp]' (or upgrade mcp directly: pip install -U mcp)"
|
||||
)
|
||||
|
||||
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import (
|
||||
|
|
@ -323,7 +336,7 @@ class MCPClient:
|
|||
)
|
||||
# HTTP transport (default)
|
||||
if streamable_http_client is None:
|
||||
raise ImportError("streamable_http_client is not available. Please install mcp with HTTP support.")
|
||||
raise missing_streamable_http_client_error()
|
||||
headers = self._get_auth_headers()
|
||||
httpx_client_factory = self._create_httpx_client_factory()
|
||||
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
PassThroughEndpointLogging,
|
||||
|
|
@ -65,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
|
|||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
request_body: dict,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
hidden_params: dict[str, Any] | None = None,
|
||||
):
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
|
|
@ -72,6 +74,10 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
|
|||
self.start_time = datetime.now()
|
||||
self.collected_chunks: list[bytes] = []
|
||||
self.model = model
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
self.endpoint_type: Final = (
|
||||
EndpointType.GEMINI if custom_llm_provider == litellm.LlmProviders.GEMINI.value else EndpointType.VERTEX_AI
|
||||
)
|
||||
self._hidden_params: dict[str, Any] = hidden_params or {}
|
||||
|
||||
async def _handle_async_streaming_logging(
|
||||
|
|
@ -89,7 +95,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
|
|||
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
|
||||
url_route="/v1/generateContent",
|
||||
request_body=self.request_body or {},
|
||||
endpoint_type=EndpointType.VERTEX_AI,
|
||||
endpoint_type=self.endpoint_type,
|
||||
start_time=self.start_time,
|
||||
raw_bytes=self.collected_chunks,
|
||||
end_time=end_time,
|
||||
|
|
@ -118,13 +124,13 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
|
|||
litellm_logging_obj=logging_obj,
|
||||
request_body=request_body or {},
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
self.response = response
|
||||
self.model = model
|
||||
self.generate_content_provider_config = generate_content_provider_config
|
||||
self.litellm_metadata = litellm_metadata
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
# Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
|
||||
# large inlineData payloads (e.g. image/jpeg) intact within one event.
|
||||
self.stream_iterator = response.iter_lines()
|
||||
|
|
@ -169,13 +175,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
|
|||
litellm_logging_obj=logging_obj,
|
||||
request_body=request_body or {},
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
self.response = response
|
||||
self.model = model
|
||||
self.generate_content_provider_config = generate_content_provider_config
|
||||
self.litellm_metadata = litellm_metadata
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
# Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
|
||||
# large inlineData payloads (e.g. image/jpeg) intact within one event.
|
||||
self.stream_iterator = response.aiter_lines()
|
||||
|
|
|
|||
|
|
@ -104,6 +104,13 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool:
|
|||
return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES
|
||||
|
||||
|
||||
# Set by a caller whose message list is not the one that goes upstream -- today the
|
||||
# Responses API layer, whose `instructions` only becomes a system message further down.
|
||||
# Tells this hook to hand role-targeted points to the pass holding the final messages
|
||||
# rather than spending them on a list that is still missing some of their targets.
|
||||
CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_points"
|
||||
|
||||
|
||||
class AnthropicCacheControlHook(CustomPromptManagement):
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
|
|
@ -128,6 +135,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
- non_default_params: dict - params with any global cache controls
|
||||
"""
|
||||
# Extract cache control injection points
|
||||
carry_unmatched: Final = bool(non_default_params.pop(CARRY_UNMATCHED_MESSAGE_POINTS, False))
|
||||
injection_points: Final[list[CacheControlInjectionPoint]] = non_default_params.pop(
|
||||
"cache_control_injection_points", []
|
||||
)
|
||||
|
|
@ -161,12 +169,25 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
non_default_params.get("prompt_cache_options"),
|
||||
)
|
||||
)
|
||||
# A provisional message list defers every role-targeted point to the pass holding
|
||||
# the final one: a role with no message here may have one there, and settling all
|
||||
# of them in one pass is what lets config order decide the shared breakpoint
|
||||
# budget. An ordinal names a different message once a later layer builds its own
|
||||
# list, so it is placed here or not at all.
|
||||
carried_message_points: Final[Sequence[CacheControlMessageInjectionPoint]] = (
|
||||
tuple(point for point in message_points if point.get("index") is None) if carry_unmatched else ()
|
||||
)
|
||||
applied_message_points: Final[Sequence[CacheControlMessageInjectionPoint]] = (
|
||||
tuple(point for point in message_points if point.get("index") is not None)
|
||||
if carry_unmatched
|
||||
else tuple(message_points)
|
||||
)
|
||||
reserved_blocks: Final = (
|
||||
1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0
|
||||
)
|
||||
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
|
||||
processed_messages = self._apply_message_injections(
|
||||
points=message_points,
|
||||
points=applied_message_points,
|
||||
messages=processed_messages,
|
||||
max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
|
||||
openai_dialect=openai_dialect,
|
||||
|
|
@ -177,10 +198,15 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
):
|
||||
non_default_params.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
|
||||
|
||||
# Pass through non-message injection points for provider-specific handling
|
||||
if remaining_points:
|
||||
# Points this pass did not place: non-message ones for the provider transform, and
|
||||
# the deferred role-targeted ones. Deferring is what reaches the Responses API's
|
||||
# `instructions`, which is only a system message once the bridge builds one. The
|
||||
# judged stamp is what makes it safe: the next pass must not re-judge points
|
||||
# against messages this pass already marked (see `_should_stand_down`).
|
||||
carried_points: Final[Sequence[CacheControlInjectionPoint]] = (*remaining_points, *carried_message_points)
|
||||
if carried_points:
|
||||
non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(
|
||||
remaining_points
|
||||
carried_points
|
||||
)
|
||||
|
||||
return model, processed_messages, non_default_params
|
||||
|
|
@ -218,7 +244,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
|
||||
@staticmethod
|
||||
def _apply_message_injections(
|
||||
points: list[CacheControlMessageInjectionPoint],
|
||||
points: Sequence[CacheControlMessageInjectionPoint],
|
||||
messages: list[AllMessageValues],
|
||||
max_blocks: int,
|
||||
openai_dialect: bool = False,
|
||||
|
|
|
|||
|
|
@ -220,6 +220,12 @@
|
|||
"ui_name": "Host URL",
|
||||
"description": "Langfuse host URL (default: https://cloud.langfuse.com)",
|
||||
"required": false
|
||||
},
|
||||
"langfuse_environment": {
|
||||
"type": "text",
|
||||
"ui_name": "Tracing Environment",
|
||||
"description": "Langfuse tracing environment (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "Langfuse v2 Logging Integration"
|
||||
|
|
@ -247,6 +253,12 @@
|
|||
"ui_name": "Host URL",
|
||||
"description": "Langfuse host URL (default: https://cloud.langfuse.com)",
|
||||
"required": false
|
||||
},
|
||||
"langfuse_environment": {
|
||||
"type": "text",
|
||||
"ui_name": "Tracing Environment",
|
||||
"description": "Langfuse tracing environment (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "Langfuse v3 OTEL Logging Integration"
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ class CustomBatchLogger(CustomLogger):
|
|||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def periodic_flush(self):
|
||||
async def periodic_flush(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval)
|
||||
|
|
|
|||
|
|
@ -62,12 +62,16 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom
|
|||
if dotprompt_content and not prompt_data and not prompt_file:
|
||||
prompt_data = _get_prompt_data_from_dotprompt_content(dotprompt_content)
|
||||
|
||||
from .prompt_manager import strip_version_suffix
|
||||
|
||||
registration_prompt_id: Final = prompt_id or strip_version_suffix(prompt_spec.prompt_id) or prompt_spec.prompt_id
|
||||
|
||||
try:
|
||||
dot_prompt_manager: Final = DotpromptManager(
|
||||
prompt_directory=prompt_directory,
|
||||
prompt_data=prompt_data,
|
||||
prompt_file=prompt_file,
|
||||
prompt_id=prompt_id,
|
||||
prompt_id=registration_prompt_id,
|
||||
)
|
||||
|
||||
return dot_prompt_manager
|
||||
|
|
|
|||
|
|
@ -96,7 +96,7 @@ class DotpromptManager(CustomPromptManagement):
|
|||
if prompt_id is None:
|
||||
return False
|
||||
try:
|
||||
return prompt_id in self.prompt_manager.list_prompts()
|
||||
return self.prompt_manager.get_prompt(prompt_id) is not None
|
||||
except Exception:
|
||||
# If there's any error accessing prompts, don't run prompt management
|
||||
return False
|
||||
|
|
@ -209,6 +209,8 @@ class DotpromptManager(CustomPromptManagement):
|
|||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
|
||||
async def async_get_chat_completion_prompt(
|
||||
|
|
|
|||
|
|
@ -11,6 +11,13 @@ from jinja2 import DictLoader, select_autoescape
|
|||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
|
||||
|
||||
def strip_version_suffix(prompt_id: str) -> str | None:
|
||||
base, separator, version = prompt_id.rpartition(".v")
|
||||
if separator and base and version.isdigit():
|
||||
return base
|
||||
return None
|
||||
|
||||
|
||||
class PromptTemplate:
|
||||
"""Represents a single prompt template with metadata and content."""
|
||||
|
||||
|
|
@ -124,11 +131,13 @@ class PromptManager:
|
|||
"content": "template content",
|
||||
"metadata": {"model": "gpt-4", "temperature": 0.7, ...}
|
||||
} + prompt_id
|
||||
"""
|
||||
if prompt_id:
|
||||
prompt_data = {prompt_id: prompt_data}
|
||||
|
||||
for prompt_id, prompt_info in prompt_data.items():
|
||||
A dict carrying a "content" key is a single flat template registered under
|
||||
prompt_id; anything else is treated as already keyed by template ID.
|
||||
"""
|
||||
keyed_prompts: Final = {prompt_id: prompt_data} if prompt_id and "content" in prompt_data else prompt_data
|
||||
|
||||
for template_id, prompt_info in keyed_prompts.items():
|
||||
try:
|
||||
content = prompt_info.get("content", "")
|
||||
metadata = prompt_info.get("metadata", {})
|
||||
|
|
@ -136,11 +145,11 @@ class PromptManager:
|
|||
template = PromptTemplate(
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
template_id=prompt_id,
|
||||
template_id=template_id,
|
||||
)
|
||||
self.prompts[prompt_id] = template
|
||||
self.prompts[template_id] = template
|
||||
except Exception:
|
||||
# Optional: print(f"Error loading prompt from JSON: {prompt_id}")
|
||||
# Optional: print(f"Error loading prompt from JSON: {template_id}")
|
||||
pass
|
||||
|
||||
def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate:
|
||||
|
|
@ -272,8 +281,12 @@ class PromptManager:
|
|||
if versioned_id in self.prompts:
|
||||
return self.prompts[versioned_id]
|
||||
|
||||
# Fall back to base prompt_id
|
||||
return self.prompts.get(prompt_id)
|
||||
direct_match: Final = self.prompts.get(prompt_id)
|
||||
if direct_match is not None:
|
||||
return direct_match
|
||||
|
||||
base_prompt_id: Final = strip_version_suffix(prompt_id)
|
||||
return self.prompts.get(base_prompt_id) if base_prompt_id else None
|
||||
|
||||
def list_prompts(self) -> list[str]:
|
||||
"""Get a list of all available prompt IDs."""
|
||||
|
|
|
|||
|
|
@ -416,17 +416,8 @@ class GenericPromptManager(CustomPromptManagement):
|
|||
tools=tools,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
ignore_prompt_manager_model=(
|
||||
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_optional_params=(
|
||||
ignore_prompt_manager_optional_params
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
|
|
@ -457,17 +448,8 @@ class GenericPromptManager(CustomPromptManagement):
|
|||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
ignore_prompt_manager_model=(
|
||||
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_optional_params=(
|
||||
ignore_prompt_manager_optional_params
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
#### What this does ####
|
||||
# On success, logs events to Langfuse
|
||||
import inspect
|
||||
import os
|
||||
import traceback
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
|
|
@ -21,6 +22,9 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
reconstruct_model_name,
|
||||
safe_deep_copy,
|
||||
)
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
validate_langfuse_environment_value,
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
|
@ -140,6 +144,7 @@ class LangFuseLogger:
|
|||
langfuse_public_key=None,
|
||||
langfuse_secret=None,
|
||||
langfuse_host=None,
|
||||
langfuse_environment: str | None = None,
|
||||
flush_interval=1,
|
||||
allow_env_credentials: bool = True,
|
||||
):
|
||||
|
|
@ -159,6 +164,10 @@ class LangFuseLogger:
|
|||
if not (self.langfuse_host.startswith("http://") or self.langfuse_host.startswith("https://")):
|
||||
# add http:// if unset, assume communicating over private network - e.g. render
|
||||
self.langfuse_host = "http://" + self.langfuse_host
|
||||
_env_override: Final = str(langfuse_environment).strip() if langfuse_environment is not None else None
|
||||
self.langfuse_environment = _env_override or os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
|
||||
if self.langfuse_environment:
|
||||
validate_langfuse_environment_value(self.langfuse_environment)
|
||||
self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
|
||||
self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")
|
||||
self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(flush_interval)
|
||||
|
|
@ -182,6 +191,8 @@ class LangFuseLogger:
|
|||
}
|
||||
self.langfuse_sdk_version: str = langfuse.version.__version__
|
||||
|
||||
if "environment" in inspect.signature(Langfuse.__init__).parameters:
|
||||
parameters["environment"] = self.langfuse_environment
|
||||
if Version(self.langfuse_sdk_version) >= Version("2.6.0"):
|
||||
parameters["sdk_integration"] = "litellm"
|
||||
self.Langfuse: Langfuse = self.safe_init_langfuse_client(parameters)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import os
|
||||
|
||||
"""
|
||||
This file contains the LangFuseHandler class
|
||||
|
||||
|
|
@ -108,6 +110,7 @@ class LangFuseHandler:
|
|||
langfuse_public_key=credentials.get("langfuse_public_key"),
|
||||
langfuse_secret=credentials.get("langfuse_secret") or credentials.get("langfuse_secret_key"),
|
||||
langfuse_host=credentials.get("langfuse_host"),
|
||||
langfuse_environment=credentials.get("langfuse_environment"),
|
||||
allow_env_credentials=credentials.get("langfuse_host") is None,
|
||||
)
|
||||
in_memory_dynamic_logger_cache.set_cache(
|
||||
|
|
@ -135,8 +138,29 @@ class LangFuseHandler:
|
|||
or standard_callback_dynamic_params.get("langfuse_secret_key"),
|
||||
langfuse_public_key=standard_callback_dynamic_params.get("langfuse_public_key"),
|
||||
langfuse_host=standard_callback_dynamic_params.get("langfuse_host"),
|
||||
langfuse_environment=LangFuseHandler._meaningful_dynamic_environment(standard_callback_dynamic_params),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _meaningful_dynamic_environment(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> str | None:
|
||||
"""Return the per-request environment only when it changes behavior.
|
||||
|
||||
Empty/whitespace values and values equal to the deployment-wide
|
||||
LANGFUSE_TRACING_ENVIRONMENT fallback are treated as absent so an
|
||||
environment-only override that matches the default does not mint a
|
||||
duplicate SDK client (each client costs threads and counts against
|
||||
MAX_LANGFUSE_INITIALIZED_CLIENTS).
|
||||
"""
|
||||
raw = standard_callback_dynamic_params.get("langfuse_environment")
|
||||
if raw is None:
|
||||
return None
|
||||
value = str(raw).strip()
|
||||
if not value or value == os.getenv("LANGFUSE_TRACING_ENVIRONMENT"):
|
||||
return None
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _dynamic_langfuse_credentials_are_passed(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
|
|
@ -153,6 +177,7 @@ class LangFuseHandler:
|
|||
or standard_callback_dynamic_params.get("langfuse_public_key") is not None
|
||||
or standard_callback_dynamic_params.get("langfuse_secret") is not None
|
||||
or standard_callback_dynamic_params.get("langfuse_secret_key") is not None
|
||||
or LangFuseHandler._meaningful_dynamic_environment(standard_callback_dynamic_params) is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -231,7 +231,10 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
from litellm.integrations.arize._utils import safe_set_attribute
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
langfuse_environment: Final = os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
|
||||
dynamic_params: Final = kwargs.get("standard_callback_dynamic_params")
|
||||
langfuse_environment: Final = (
|
||||
dynamic_params.get("langfuse_environment") if dynamic_params else None
|
||||
) or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
|
||||
if langfuse_environment:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
|
|
|
|||
395
litellm/integrations/newrelic/newrelic_metrics.py
Normal file
395
litellm/integrations/newrelic/newrelic_metrics.py
Normal file
|
|
@ -0,0 +1,395 @@
|
|||
"""
|
||||
New Relic Metric API Integration - sends per-team cost/usage metrics to /metric/v1
|
||||
|
||||
NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-api/introduction-metric-api/
|
||||
|
||||
`async_log_success_event` / `async_log_failure_event` queue one record per request;
|
||||
at flush the queue is aggregated by (team, model group, model, provider, status)
|
||||
into count/summary metrics. `interval.ms` is the real window between flushes,
|
||||
computed at flush time.
|
||||
|
||||
Team-scoped by construction: the ingest key is injected explicitly and there is
|
||||
deliberately no environment-variable fallback, so a team's metrics are never sent
|
||||
with the proxy operator's credentials (mirrors ``allow_env_credentials=False`` on
|
||||
the Datadog team logger).
|
||||
|
||||
Error policy on flush: 4xx drops the batch (a retry would fail identically; 403
|
||||
is a permanent credential failure), 5xx/network re-queues capped at
|
||||
``max_queue_size`` records with the oldest dropped.
|
||||
|
||||
For batching specific details see CustomBatchLogger class
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gzip
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from math import ceil
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from httpx import HTTPStatusError, Response
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.integrations.newrelic import (
|
||||
NEWRELIC_DEFAULT_REGION,
|
||||
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN,
|
||||
NEWRELIC_METRIC_COMPLETION_TOKENS,
|
||||
NEWRELIC_METRIC_COST_USD,
|
||||
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
|
||||
NEWRELIC_METRIC_PROMPT_TOKENS,
|
||||
NEWRELIC_METRIC_REQUEST_DURATION_MS,
|
||||
NEWRELIC_METRIC_REQUESTS,
|
||||
NEWRELIC_METRIC_TOTAL_TOKENS,
|
||||
NEWRELIC_METRICS_MAX_BATCH_SIZE,
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
|
||||
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
|
||||
NewRelicCountMetric,
|
||||
NewRelicMetric,
|
||||
NewRelicMetricCommon,
|
||||
NewRelicMetricEnvelope,
|
||||
NewRelicMetricRecord,
|
||||
NewRelicSummaryMetric,
|
||||
NewRelicSummaryValue,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
# 408 (request timeout) and 429 (rate limit) are transient client errors the
|
||||
# Metric API expects a retry on, unlike 400/403 which a retry would only repeat.
|
||||
_RETRYABLE_CLIENT_STATUSES: Final = frozenset({408, 429})
|
||||
|
||||
|
||||
def resolve_newrelic_metric_endpoint(newrelic_region: str | None) -> str:
|
||||
if not newrelic_region:
|
||||
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
|
||||
endpoint: Final = NEWRELIC_METRIC_ENDPOINT_BY_REGION.get(newrelic_region.lower())
|
||||
if endpoint is None:
|
||||
verbose_logger.warning(
|
||||
"New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.",
|
||||
newrelic_region,
|
||||
", ".join(sorted(NEWRELIC_METRIC_ENDPOINT_BY_REGION)),
|
||||
)
|
||||
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
|
||||
return endpoint
|
||||
|
||||
|
||||
def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) -> NewRelicMetricRecord:
|
||||
metadata: Final = standard_logging_object.get("metadata")
|
||||
team_id: Final = ((metadata.get("user_api_key_team_id") or metadata.get("team_id")) if metadata else None) or ""
|
||||
team_alias: Final = (
|
||||
(metadata.get("user_api_key_team_alias") or metadata.get("team_alias")) if metadata else None
|
||||
) or ""
|
||||
return NewRelicMetricRecord(
|
||||
team_id=team_id,
|
||||
team_alias=team_alias,
|
||||
model_group=standard_logging_object.get("model_group") or "",
|
||||
model=standard_logging_object.get("model") or "",
|
||||
custom_llm_provider=standard_logging_object.get("custom_llm_provider") or "",
|
||||
status=str(standard_logging_object.get("status") or "success"),
|
||||
response_cost=float(standard_logging_object.get("response_cost") or 0.0),
|
||||
prompt_tokens=int(standard_logging_object.get("prompt_tokens") or 0),
|
||||
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
|
||||
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
|
||||
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
|
||||
)
|
||||
|
||||
|
||||
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
|
||||
first: Final = bucket_records[0]
|
||||
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
|
||||
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
|
||||
for key, value in (
|
||||
("team_id", first.team_id),
|
||||
("team_alias", first.team_alias),
|
||||
("model_group", first.model_group),
|
||||
("model", first.model),
|
||||
("custom_llm_provider", first.custom_llm_provider),
|
||||
("status", first.status),
|
||||
)
|
||||
if value
|
||||
}
|
||||
durations: Final = tuple(record.duration_ms for record in bucket_records)
|
||||
counts: Final[tuple[tuple[str, float], ...]] = (
|
||||
(NEWRELIC_METRIC_REQUESTS, float(len(bucket_records))),
|
||||
(NEWRELIC_METRIC_COST_USD, sum(record.response_cost for record in bucket_records)),
|
||||
(NEWRELIC_METRIC_PROMPT_TOKENS, float(sum(record.prompt_tokens for record in bucket_records))),
|
||||
(NEWRELIC_METRIC_COMPLETION_TOKENS, float(sum(record.completion_tokens for record in bucket_records))),
|
||||
(NEWRELIC_METRIC_TOTAL_TOKENS, float(sum(record.total_tokens for record in bucket_records))),
|
||||
)
|
||||
count_metrics: Final[tuple[NewRelicMetric, ...]] = tuple(
|
||||
NewRelicCountMetric(name=name, type="count", value=value, attributes=attributes) for name, value in counts
|
||||
)
|
||||
summary_metric: Final = NewRelicSummaryMetric(
|
||||
name=NEWRELIC_METRIC_REQUEST_DURATION_MS,
|
||||
type="summary",
|
||||
value=NewRelicSummaryValue(
|
||||
count=len(durations),
|
||||
sum=sum(durations),
|
||||
min=min(durations),
|
||||
max=max(durations),
|
||||
),
|
||||
attributes=attributes,
|
||||
)
|
||||
return (*count_metrics, summary_metric)
|
||||
|
||||
|
||||
def build_metric_payload(
|
||||
records: tuple[NewRelicMetricRecord, ...],
|
||||
*,
|
||||
window_start: float,
|
||||
now: float,
|
||||
) -> tuple[NewRelicMetricEnvelope, ...]:
|
||||
"""Aggregates records into one Metric API envelope for the flush window."""
|
||||
interval_ms: Final = max(1, int((now - window_start) * 1000))
|
||||
bucket_keys: Final = tuple(dict.fromkeys(record.bucket_key for record in records))
|
||||
metrics: Final = tuple(
|
||||
metric
|
||||
for key in bucket_keys
|
||||
for metric in _bucket_metrics(tuple(record for record in records if record.bucket_key == key))
|
||||
)
|
||||
common: Final[NewRelicMetricCommon] = {
|
||||
"timestamp": int(window_start * 1000),
|
||||
"interval.ms": interval_ms,
|
||||
}
|
||||
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
|
||||
|
||||
|
||||
class NewRelicMetricsLogger(CustomBatchLogger):
|
||||
def __init__(
|
||||
self,
|
||||
newrelic_api_key: str,
|
||||
newrelic_region: str | None = None,
|
||||
) -> None:
|
||||
if not newrelic_api_key:
|
||||
raise ValueError(
|
||||
"newrelic_api_key is required for NewRelicMetricsLogger; "
|
||||
"team-scoped metrics never fall back to environment credentials"
|
||||
)
|
||||
self.newrelic_api_key: Final = newrelic_api_key
|
||||
self.metric_api_url: Final = resolve_newrelic_metric_endpoint(newrelic_region)
|
||||
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
self._stopped: bool = False
|
||||
self._drain_lock = asyncio.Lock()
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
self.flush_lock = asyncio.Lock()
|
||||
super().__init__(
|
||||
flush_lock=self.flush_lock,
|
||||
batch_size=NEWRELIC_METRICS_MAX_BATCH_SIZE,
|
||||
max_queue_size=NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Ends the periodic flush loop; called on DynamicLoggingCache eviction.
|
||||
|
||||
Schedules one final drain of anything still queued, so eviction never
|
||||
silently discards records. Guarded so it can never raise into the
|
||||
cache's eviction path.
|
||||
"""
|
||||
self._stopped = True
|
||||
try:
|
||||
asyncio.get_running_loop().create_task(self._final_drain())
|
||||
except Exception: # noqa: BLE001 # no running loop / shutdown; the periodic loop's final drain still runs
|
||||
verbose_logger.debug("New Relic Metrics: could not schedule final drain on stop()", exc_info=True)
|
||||
|
||||
async def _drain_with_retry(self) -> None:
|
||||
"""Deliver everything queued on a stopped logger, or drop it with a log.
|
||||
|
||||
A stopped logger has no periodic loop left, so every post-stop path
|
||||
funnels through here. ``_drain_lock`` serializes drains: a callback that
|
||||
appends and starts its own drain queues behind the running one instead
|
||||
of racing it. Each pass attempts the whole current queue in
|
||||
``batch_size`` chunks, unlike the periodic path it does not stop at the
|
||||
first failing chunk, so a persistently failing head never starves the
|
||||
tail. Only after ``_MAX_DRAIN_PASSES`` against a permanently failing
|
||||
destination is the remainder dropped, and then only the records that were
|
||||
queued when this drain began, so every dropped record got the full retry
|
||||
budget: a record a callback appended mid-drain is not in that snapshot,
|
||||
so it is left for its own serialized drain rather than dropped after
|
||||
fewer attempts, and is never stranded.
|
||||
"""
|
||||
async with self._drain_lock:
|
||||
attempted: Final = tuple(self.log_queue)
|
||||
for _pass in range(NEWRELIC_METRICS_MAX_DRAIN_PASSES):
|
||||
await self._drain_flush_once()
|
||||
if not self.log_queue:
|
||||
return
|
||||
if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1:
|
||||
await asyncio.sleep(2**_pass)
|
||||
async with self.flush_lock:
|
||||
tried_ids: Final = frozenset(id(record) for record in attempted)
|
||||
survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids)
|
||||
dropped: Final = len(self.log_queue) - len(survivors)
|
||||
if dropped:
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: dropping %s records after %s drain passes",
|
||||
dropped,
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
|
||||
)
|
||||
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
|
||||
|
||||
async def _drain_flush_once(self) -> None:
|
||||
"""Attempt every queued record once, in ``batch_size`` chunks, without
|
||||
stopping at the first failing chunk so a persistently failing head does
|
||||
not starve the tail (the periodic ``flush_queue`` deliberately stops
|
||||
instead). Takes the queue under ``flush_lock`` and re-queues only the
|
||||
chunks a 5xx/network error left undelivered, so records a concurrent
|
||||
request appends during the sends survive for the next pass."""
|
||||
async with self.flush_lock:
|
||||
pending: Final = tuple(self.log_queue)
|
||||
window_start: Final = self.last_flush_time
|
||||
self.last_flush_time = time.time()
|
||||
del self.log_queue[:]
|
||||
if not pending:
|
||||
return
|
||||
chunks: Final = tuple(
|
||||
pending[start : start + self.batch_size] for start in range(0, len(pending), self.batch_size)
|
||||
)
|
||||
delivered: Final = tuple([await self._classify_and_send(chunk, window_start) for chunk in chunks])
|
||||
failed: Final = tuple(record for chunk, ok in zip(chunks, delivered) for record in (() if ok else chunk))
|
||||
if failed:
|
||||
self._requeue(failed)
|
||||
|
||||
async def _final_drain(self) -> None:
|
||||
await self._drain_with_retry()
|
||||
|
||||
async def periodic_flush(self) -> None:
|
||||
while not self._stopped:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
if self._stopped:
|
||||
break
|
||||
await self.flush_queue()
|
||||
await self._final_drain()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
try:
|
||||
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
|
||||
except Exception as e: # noqa: BLE001 # logging must never break the request path
|
||||
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
try:
|
||||
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
|
||||
except Exception as e: # noqa: BLE001 # logging must never break the request path
|
||||
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def _log_async_event(self, standard_logging_object: StandardLoggingPayload | None) -> None:
|
||||
if standard_logging_object is None:
|
||||
raise ValueError("standard_logging_object not found in kwargs")
|
||||
self.log_queue.append(_metric_record_from_payload(standard_logging_object))
|
||||
if self._stopped:
|
||||
# A stopped logger has no periodic loop left; an in-flight callback
|
||||
# that appends after the eviction drain delivers its own record.
|
||||
await self._drain_with_retry()
|
||||
return
|
||||
if len(self.log_queue) >= self.batch_size:
|
||||
await self.flush_queue()
|
||||
|
||||
async def flush_queue(self) -> None:
|
||||
async with self.flush_lock:
|
||||
window_start: Final = self.last_flush_time
|
||||
self.last_flush_time = time.time()
|
||||
queued: Final = len(self.log_queue)
|
||||
if not queued:
|
||||
return
|
||||
verbose_logger.debug("New Relic Metrics: Flushing %s queued records", queued)
|
||||
# Bounded by what is queued now: records appended mid-flush belong to
|
||||
# the next window, and looping until empty would never end under load.
|
||||
for _chunk in range(ceil(queued / self.batch_size)):
|
||||
if not await self.async_send_batch(window_start=window_start):
|
||||
return
|
||||
|
||||
async def async_send_batch(self, window_start: float | None = None) -> bool:
|
||||
"""Sends the oldest ``batch_size`` records only, so a queue grown past that
|
||||
by re-queues cannot breach the Metric API data point cap in one request.
|
||||
Returns False once a chunk fails and is re-queued, so the caller stops."""
|
||||
if not self.log_queue:
|
||||
return False
|
||||
|
||||
batch_to_send: Final[tuple[NewRelicMetricRecord, ...]] = tuple(self.log_queue[: self.batch_size])
|
||||
del self.log_queue[: len(batch_to_send)]
|
||||
|
||||
delivered: Final = await self._classify_and_send(
|
||||
batch_to_send, window_start if window_start is not None else self.last_flush_time
|
||||
)
|
||||
if not delivered:
|
||||
self._requeue(batch_to_send)
|
||||
return delivered
|
||||
|
||||
async def _classify_and_send(self, batch: tuple[NewRelicMetricRecord, ...], window_start: float) -> bool:
|
||||
"""Send one chunk and classify the outcome, never touching the queue.
|
||||
Returns True when the batch is done with (delivered on any 2xx, or a 4xx
|
||||
a retry would only repeat, 403 being a permanent bad-key rejection), and
|
||||
False when a 5xx or network error means the caller should re-queue it.
|
||||
|
||||
``AsyncHTTPHandler.post`` raises ``HTTPStatusError`` on any non-2xx, so a
|
||||
4xx never returns a response here; the status is read off the raised
|
||||
error to keep the client-error path (drop) distinct from 5xx (retry)."""
|
||||
payload: Final = build_metric_payload(records=batch, window_start=window_start, now=time.time())
|
||||
try:
|
||||
status = (
|
||||
await self.async_send_compressed_data(payload)
|
||||
).status_code # rebind-ok: reassigned from the raised HTTPStatusError below
|
||||
except HTTPStatusError as e:
|
||||
status = e.response.status_code
|
||||
except Exception as e: # noqa: BLE001 # transport/network failure re-queues the batch
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: network error sending %s records, will retry - %s",
|
||||
len(batch),
|
||||
e,
|
||||
)
|
||||
return False
|
||||
|
||||
if 200 <= status < 300:
|
||||
return True
|
||||
|
||||
if 400 <= status < 500 and status not in _RETRYABLE_CLIENT_STATUSES:
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: %s from Metric API%s, dropping %s records.",
|
||||
status,
|
||||
" (permanent credential failure: invalid or revoked team ingest key)" if status == 403 else "",
|
||||
len(batch),
|
||||
)
|
||||
return True
|
||||
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: %s from Metric API, will retry %s records",
|
||||
status,
|
||||
len(batch),
|
||||
)
|
||||
return False
|
||||
|
||||
def _requeue(self, batch: tuple[NewRelicMetricRecord, ...]) -> None:
|
||||
"""Prepends ``batch`` in place (never by assignment: records appended by
|
||||
concurrent requests during the flush await must survive), keeping
|
||||
chronological order so the cap drops the oldest records first."""
|
||||
self.log_queue[:0] = batch
|
||||
overflow: Final = len(self.log_queue) - self.max_queue_size
|
||||
if overflow > 0:
|
||||
del self.log_queue[:overflow]
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: retry queue exceeded max_queue_size=%s; dropped %s oldest records.",
|
||||
self.max_queue_size,
|
||||
overflow,
|
||||
)
|
||||
|
||||
async def async_send_compressed_data(self, payload: tuple[NewRelicMetricEnvelope, ...]) -> Response:
|
||||
compressed_data: Final = gzip.compress(safe_dumps(payload).encode("utf-8"))
|
||||
headers: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/json",
|
||||
"Content-Encoding": "gzip",
|
||||
"Api-Key": self.newrelic_api_key,
|
||||
}
|
||||
)
|
||||
return await self.async_client.post(
|
||||
url=self.metric_api_url,
|
||||
data=compressed_data,
|
||||
headers=headers,
|
||||
)
|
||||
90
litellm/integrations/newrelic/newrelic_team_handler.py
Normal file
90
litellm/integrations/newrelic/newrelic_team_handler.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
"""
|
||||
New Relic Team Handler
|
||||
|
||||
Used to get the NewRelicMetricsLogger for a given request.
|
||||
Handles Key/Team Based New Relic metrics, following the same pattern as DataDogHandler.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
|
||||
|
||||
from .newrelic_metrics import NewRelicMetricsLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
|
||||
|
||||
|
||||
class NewRelicLoggingConfig(TypedDict):
|
||||
newrelic_api_key: ReadOnly[str | None]
|
||||
newrelic_region: ReadOnly[str | None]
|
||||
|
||||
|
||||
class NewRelicHandler:
|
||||
@staticmethod
|
||||
def get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
|
||||
) -> NewRelicMetricsLogger:
|
||||
"""
|
||||
Get a team-scoped NewRelicMetricsLogger for a given request.
|
||||
|
||||
Resolves and caches per-team NewRelicMetricsLogger instances using
|
||||
DynamicLoggingCache, keyed by the team's New Relic credentials. Each unique
|
||||
set of credentials gets its own logger instance with its own batch/flush loop.
|
||||
|
||||
Note: This handler is only called when a team-scoped newrelic_api_key is
|
||||
present. The trace logger for the ``newrelic`` callback (OTel v2 / legacy
|
||||
agent) is managed separately by _init_custom_logger_compatible_class via
|
||||
_in_memory_loggers.
|
||||
"""
|
||||
_credentials: Final = NewRelicHandler.get_dynamic_newrelic_logging_config(
|
||||
standard_callback_dynamic_params=standard_callback_dynamic_params,
|
||||
)
|
||||
|
||||
temp_newrelic_logger = in_memory_dynamic_logger_cache.get_cache(
|
||||
credentials=_credentials, service_name="newrelic"
|
||||
)
|
||||
|
||||
if temp_newrelic_logger is None:
|
||||
temp_newrelic_logger = NewRelicHandler._create_newrelic_logger_from_credentials(
|
||||
credentials=_credentials,
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
|
||||
return temp_newrelic_logger
|
||||
|
||||
@staticmethod
|
||||
def _create_newrelic_logger_from_credentials(
|
||||
credentials: NewRelicLoggingConfig,
|
||||
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
|
||||
) -> NewRelicMetricsLogger:
|
||||
newrelic_logger: Final = NewRelicMetricsLogger(
|
||||
newrelic_api_key=credentials.get("newrelic_api_key") or "",
|
||||
newrelic_region=credentials.get("newrelic_region"),
|
||||
)
|
||||
in_memory_dynamic_logger_cache.set_cache(
|
||||
credentials=credentials,
|
||||
service_name="newrelic",
|
||||
logging_obj=newrelic_logger,
|
||||
)
|
||||
verbose_logger.debug("New Relic: Created and cached new NewRelicMetricsLogger for team-scoped credentials")
|
||||
return newrelic_logger
|
||||
|
||||
@staticmethod
|
||||
def get_dynamic_newrelic_logging_config(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> NewRelicLoggingConfig:
|
||||
return NewRelicLoggingConfig(
|
||||
newrelic_api_key=standard_callback_dynamic_params.get("newrelic_api_key"),
|
||||
newrelic_region=standard_callback_dynamic_params.get("newrelic_region"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _dynamic_newrelic_credentials_are_passed(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
return standard_callback_dynamic_params.get("newrelic_api_key") is not None
|
||||
|
|
@ -22,6 +22,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
|
|||
)
|
||||
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
|
||||
from litellm.integrations.otel.model.semconv import Metric
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.litellm_core_utils.service_tier_utils import (
|
||||
|
|
@ -1643,7 +1644,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
if self._operation_duration_histogram:
|
||||
self._operation_duration_histogram.record(duration_s, attributes=common_attrs)
|
||||
if response_obj and (usage := response_obj.get("usage")) and self._token_usage_histogram:
|
||||
if (
|
||||
self._token_usage_histogram
|
||||
and response_obj
|
||||
and not is_unbilled_non_inference_call_from_params(kwargs.get("call_type"), params, response_obj)
|
||||
and (usage := response_obj.get("usage"))
|
||||
):
|
||||
in_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"}
|
||||
out_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"}
|
||||
self._token_usage_histogram.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
|
||||
|
|
@ -1719,6 +1725,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
if not self._time_per_output_token_histogram:
|
||||
return
|
||||
|
||||
if is_unbilled_non_inference_call_from_params(
|
||||
kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj
|
||||
):
|
||||
return
|
||||
|
||||
# Get completion tokens from response_obj
|
||||
completion_tokens = None
|
||||
if response_obj and (usage := response_obj.get("usage")):
|
||||
|
|
@ -2049,6 +2060,26 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# serialise to JSON once so set_attribute never coerces.
|
||||
guardrail_span.set_attribute("guardrail_violation_categories", safe_dumps(violation_categories))
|
||||
|
||||
# Billable usage counters and USD cost stamped by the provider hook
|
||||
# (e.g. Azure Prompt Shield text records, Bedrock policy units).
|
||||
guardrail_usage = guardrail_information.get("guardrail_usage")
|
||||
if guardrail_usage is not None:
|
||||
guardrail_span.set_attribute("guardrail_usage", safe_dumps(guardrail_usage))
|
||||
guardrail_cost = guardrail_information.get("guardrail_cost")
|
||||
if guardrail_cost is not None:
|
||||
self.safe_set_attribute(
|
||||
span=guardrail_span,
|
||||
key="guardrail_cost",
|
||||
value=guardrail_cost,
|
||||
)
|
||||
guardrail_cost_in_spend = guardrail_information.get("guardrail_cost_in_spend")
|
||||
if isinstance(guardrail_cost_in_spend, bool):
|
||||
self.safe_set_attribute(
|
||||
span=guardrail_span,
|
||||
key="guardrail_cost_in_spend",
|
||||
value=guardrail_cost_in_spend,
|
||||
)
|
||||
|
||||
self._set_team_attributes_from_kwargs(guardrail_span, kwargs)
|
||||
|
||||
guardrail_span.end(end_time=self._to_ns(end_time_datetime))
|
||||
|
|
@ -2468,7 +2499,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload)
|
||||
|
||||
usage: Final = response_obj and response_obj.get("usage")
|
||||
usage: Final = (
|
||||
response_obj.get("usage")
|
||||
if response_obj
|
||||
and not is_unbilled_non_inference_call_from_params(
|
||||
kwargs.get("call_type"), litellm_params, response_obj
|
||||
)
|
||||
else None
|
||||
)
|
||||
if usage:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
|
|
|
|||
|
|
@ -136,6 +136,9 @@ class GenAIMapper:
|
|||
LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id,
|
||||
LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template,
|
||||
LiteLLM.GUARDRAIL_DETECTION_METHOD: lambda d: d.detection_method,
|
||||
LiteLLM.GUARDRAIL_USAGE: lambda d: d.usage_json,
|
||||
LiteLLM.GUARDRAIL_COST: lambda d: d.cost,
|
||||
LiteLLM.GUARDRAIL_COST_IN_SPEND: lambda d: d.cost_in_spend,
|
||||
}
|
||||
|
||||
_SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = {
|
||||
|
|
|
|||
|
|
@ -190,6 +190,15 @@ class GuardrailSpanData:
|
|||
guardrail_id: str | None = None
|
||||
policy_template: str | None = None
|
||||
detection_method: str | None = None
|
||||
# Provider-reported billable usage counters (JSON-serialized) and the USD cost
|
||||
# priced from them by the provider hook (``guardrail_usage`` /
|
||||
# ``guardrail_cost`` on ``StandardLoggingGuardrailInformation``).
|
||||
usage_json: str | None = None
|
||||
cost: float | None = None
|
||||
# Whether ``cost`` participates in the request's billed spend (absent means
|
||||
# billed, the default; False means report-only). Mirrors
|
||||
# ``guardrail_cost_in_spend`` so trace consumers can avoid double-counting.
|
||||
cost_in_spend: bool | None = None
|
||||
# Set when the guardrail intervened/blocked or failed, so the emitter marks
|
||||
# the span ERROR — a blocking guardrail is an error outcome for that span.
|
||||
error: SpanError | None = None
|
||||
|
|
@ -209,6 +218,8 @@ class GuardrailSpanData:
|
|||
get: Final = cast(Mapping[str, object], entry).get
|
||||
status: Final = as_str(get("guardrail_status"))
|
||||
response: Final = get("guardrail_response")
|
||||
usage: Final = get("guardrail_usage")
|
||||
in_spend: Final = get("guardrail_cost_in_spend")
|
||||
error: Final = (
|
||||
SpanError(error_type=status, message=as_str(get("guardrail_action")))
|
||||
if status in cls._ERROR_STATUSES
|
||||
|
|
@ -231,6 +242,9 @@ class GuardrailSpanData:
|
|||
guardrail_id=as_str(get("guardrail_id")),
|
||||
policy_template=as_str(get("policy_template")),
|
||||
detection_method=as_str(get("detection_method")),
|
||||
usage_json=_json_or_none(usage) if usage is not None else None,
|
||||
cost=as_float(get("guardrail_cost")),
|
||||
cost_in_spend=in_spend if isinstance(in_spend, bool) else None,
|
||||
error=error,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ class GenAIOperation(str, Enum):
|
|||
EXECUTE_TOOL = "execute_tool" # MCP tool-call spans
|
||||
LITELLM_VECTOR_STORE_MANAGEMENT = "litellm.vector_store_management"
|
||||
LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management"
|
||||
LITELLM_RESPONSES_MANAGEMENT = "litellm.responses_management"
|
||||
LITELLM_MODERATION = "litellm.moderation"
|
||||
|
||||
|
||||
|
|
@ -307,6 +308,15 @@ class LiteLLM:
|
|||
GUARDRAIL_ID: Final = "litellm.guardrail.id"
|
||||
GUARDRAIL_POLICY_TEMPLATE: Final = "litellm.guardrail.policy_template"
|
||||
GUARDRAIL_DETECTION_METHOD: Final = "litellm.guardrail.detection_method"
|
||||
# Provider-reported billable usage counters, JSON-serialized into one value.
|
||||
GUARDRAIL_USAGE: Final = "litellm.guardrail.usage"
|
||||
# Numeric USD cost of the guardrail invocation; lives under the litellm.cost.*
|
||||
# namespace (COST_PREFIX) beside the LLM call's litellm.cost.total.
|
||||
GUARDRAIL_COST: Final = "litellm.cost.guardrail"
|
||||
# Whether litellm.cost.guardrail is already inside litellm.cost.total (True,
|
||||
# the billed default) or reported alongside it (False) — without this a trace
|
||||
# consumer cannot tell whether adding the two double-counts.
|
||||
GUARDRAIL_COST_IN_SPEND: Final = "litellm.guardrail.cost_in_spend"
|
||||
SERVICE_NAME: Final = "litellm.service.name"
|
||||
SERVICE_CALL_TYPE: Final = "litellm.service.call_type"
|
||||
PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms"
|
||||
|
|
@ -374,6 +384,14 @@ _OPERATION_BY_CALL_TYPE: Final[dict[str, GenAIOperation]] = {
|
|||
"aembedding": GenAIOperation.EMBEDDINGS,
|
||||
"responses": GenAIOperation.CHAT,
|
||||
"aresponses": GenAIOperation.CHAT,
|
||||
"get_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"aget_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"delete_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"adelete_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"cancel_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"acancel_responses": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"list_input_items": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"alist_input_items": GenAIOperation.LITELLM_RESPONSES_MANAGEMENT,
|
||||
"image_generation": GenAIOperation.GENERATE_CONTENT,
|
||||
"aimage_generation": GenAIOperation.GENERATE_CONTENT,
|
||||
"moderation": GenAIOperation.LITELLM_MODERATION,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.integrations.otel.model.semconv import (
|
|||
resolve_provider,
|
||||
)
|
||||
from litellm.integrations.otel.model.utils import to_seconds
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
|
||||
|
|
@ -198,16 +199,21 @@ class GenAIMetricRecorder:
|
|||
) -> None:
|
||||
common_attrs: Final = self._filter_attributes(self._bounded_attributes(kwargs))
|
||||
duration_s: Final = (end_time - start_time).total_seconds()
|
||||
usage_is_replayed: Final = is_unbilled_non_inference_call_from_params(
|
||||
kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj
|
||||
)
|
||||
|
||||
self._metrics.operation_duration.record(duration_s, attributes=common_attrs)
|
||||
self._record_token_usage(response_obj, common_attrs)
|
||||
if not usage_is_replayed:
|
||||
self._record_token_usage(response_obj, common_attrs)
|
||||
|
||||
cost: Final = kwargs.get("response_cost")
|
||||
if cost:
|
||||
self._metrics.token_cost.record(cost, attributes=common_attrs)
|
||||
|
||||
self._record_time_to_first_token(kwargs, common_attrs)
|
||||
self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs)
|
||||
if not usage_is_replayed:
|
||||
self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs)
|
||||
self._record_response_duration(kwargs, end_time, common_attrs)
|
||||
|
||||
def record_failure(
|
||||
|
|
|
|||
|
|
@ -19,6 +19,19 @@ class PromptManagementClient(TypedDict):
|
|||
completed_messages: list[AllMessageValues] | None
|
||||
|
||||
|
||||
def resolve_prompt_manager_ignore_flags(
|
||||
prompt_spec: PromptSpec | None,
|
||||
ignore_prompt_manager_model: bool | None,
|
||||
ignore_prompt_manager_optional_params: bool | None,
|
||||
) -> tuple[bool, bool]:
|
||||
spec_params: Final = prompt_spec.litellm_params if prompt_spec is not None else None
|
||||
return (
|
||||
bool(ignore_prompt_manager_model) or bool(spec_params is not None and spec_params.ignore_prompt_manager_model),
|
||||
bool(ignore_prompt_manager_optional_params)
|
||||
or bool(spec_params is not None and spec_params.ignore_prompt_manager_optional_params),
|
||||
)
|
||||
|
||||
|
||||
class PromptManagementBase(ABC):
|
||||
@property
|
||||
@abstractmethod
|
||||
|
|
@ -182,13 +195,18 @@ class PromptManagementBase(ABC):
|
|||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
resolved_ignore_model, resolved_ignore_optional_params = resolve_prompt_manager_ignore_flags(
|
||||
prompt_spec=prompt_spec,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
return self.post_compile_prompt_processing(
|
||||
prompt_template=prompt_template,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
model=model,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
ignore_prompt_manager_model=resolved_ignore_model,
|
||||
ignore_prompt_manager_optional_params=resolved_ignore_optional_params,
|
||||
)
|
||||
|
||||
async def async_get_chat_completion_prompt(
|
||||
|
|
@ -224,11 +242,16 @@ class PromptManagementBase(ABC):
|
|||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
resolved_ignore_model, resolved_ignore_optional_params = resolve_prompt_manager_ignore_flags(
|
||||
prompt_spec=prompt_spec,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
return self.post_compile_prompt_processing(
|
||||
prompt_template=prompt_template,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
model=model,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
ignore_prompt_manager_model=resolved_ignore_model,
|
||||
ignore_prompt_manager_optional_params=resolved_ignore_optional_params,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,19 +3,25 @@ Helper functions for health check calls.
|
|||
"""
|
||||
|
||||
import base64
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
# Minimal PDF for health checks - base64 encoded 1-page PDF with just "test"
|
||||
TEST_PDF_URL = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="
|
||||
|
||||
# Minimal image for health checks - base64 encoded 512x512 solid-gray PNG
|
||||
TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFlklEQVR42u3VMQEAAAzCMKQjHQ97l0jo0xSAlyIBgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAUgAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAADcDrctaAb6XeXAAAAAASUVORK5CYII="
|
||||
# Minimal image for health checks - base64 encoded 512x512 blue circle on a white background PNG
|
||||
TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAJk0lEQVR42u3VQREAIRADwVWCOmTjBVzwSLorCri6nbkAVBpPACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAZdY+HgEBgIRr/meeGgGA8EMvDAgAuPh6gACAi68HCAA4+mKAAICjLwYIALj7SoAAgLuvBAgA7r4pAQKAu29KgADg7psSIAA4/SYDCADuvikBAoDTbzKAAOD0mwwgADj9JgMIAE6/yQACgNNvMoAA4PSbDCAAOP0mAwgATr/JAAKA668BIAA4/TKAAOD0mwwgALj+pgEIAE6/yQACgOtvGoAA4PSbDCAAuP6mAQgATr/JAAKA628agADg9JsMIAC4/qYBCACuv2kAAoDrbxqAAOD0mwwgALj+pgEIAK6/aQACgOtvGoAA4PqbBiAArr+ZBiAATr+ZDCAArr+ZBiAArr+ZBiAArr+ZBiAArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggAAmAmAAKA62+mAQKA62+mAQKA62/m1xYAXH/TAAQA1980AAHA9TcNQAAEwEwAEADX30wDEADX30wDEADX30wDEADX30wDEAABMBMABMD1N9MABMD1N9MABEAAzAQAAXD9zTQAAXD9zTRAABAAMwEQAFx/Mw0QAFx/Mw0QAATATAAEANffTAMEANffTAMEAAEwEwABwPU30wABQADMBEAAXH8z0wABcP3NTAMEQADMTAAEwPU3Mw0QAAEwMwEQANffTAMQAAEwEwAEwPU30wAEQADMBAABcP3NNAABEAAzAUAAXH8zDUAABMBMABAA199MAwQAATATAAHA9TfTAAFAAMwEQABw/c00QAAQADMBEAAEwEwABMD1NzMNEAABMDMBEADX38w0QAAEwMwEQAAEwMwEQABcfzPTAAEQADMTAAEQADMTAAFw/c1MAwRAAMxMAARAAMxMAATA9TczDRAAATATAARAAMwEAAFw/c00AAEQADMBQAAEwEwAEADX30wDBAABMBMAAUAAzARAABAAMwEQAFx/Mw0QAAEwMwEQAAEwMwEQAAEwMwEQANffzDRAAATAzARAAATAzARAAATAzARAAATAzARAAFx/M9MAARAAMxMAARAAMxMAARAAMxMAARAAMxMAAXD9zUwDBEAAzEwABEAAzAQAARAAMwFAAATATAAQAAEwEwABQADMBEAAEAAzARAAXH8zDRAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAA19/MNEAANMDM9UcABMBMABAAATATAAHwBAJgJgACgACYCYAAIABmAiAA+IvMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAANMDMXH8BEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzAUAABMBMABAADTBz/REAATATAAFAAMwEQAAQADMBEAAEwEwABAANMHP9BUAAzEwABEAAzEwABEAAzEwABEAAzEwABEADzMz1FwABMDMBEAABMDMBEAABMDMBEAANMDPXXwAEwMwEQAAEwMwEQAAEwMwEQAA0wMxcfwEQADMTAAEQADMBQAA0wMz1RwAEwEwAEAABMBMABEADzFx/AUAAzARAABAAMwEQADTAzPUXAATATAAEAAEwEwABQAPMXH8BEAAzEwABEAAzEwAB0AAzc/0FQADMTAAEQAPMzPUXAAEwMwEQAAEwMwEQAA0wM9dfAATAzARAADTAzFx/ARAAMwFAADTAzPVHAATATAAQAA0wc/0RAAEwEwAEQAPMXH8EQADMBAAB0AAz118AEAAzARAANMDM9RcABMBMAAQADTBz/QUAATATAAFAA8xcfwFAA8xcfwFAAMwEQADQADPXXwAEwMwEQAA0wMxcfwHQADNz/QVAAMxMAARAA8xcfwRAA8xcfwRAAMwEAAHQADPXHwHQADPXHwEQADMBQAA0wMz1RwA0wMz1RwAEwEwAEAANMHP9BQANMHP9BQANMHP9BQANMHP9BQABMBMAAUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzPVHABAAEwAEAA0w1x8BQAPM9UcA0ABz/REANMBcfwQADTDXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwGQATOnHwHQADPXHwHQADPXXwDQADPXXwDQADPXXwDQADPXXwCQAXP6EQA0wFx/BAANMNcfAUADzPVHAJABc/oRADTAXH8EABkwpx8BQAPM9UcAkAFz+hEANMBcfwQAGTCnHwFAA8z1RwCQAXP6EQBkwJx+BAANMNcfAUAGzOlHAJABc/oRAGTA6QcBQAacfhAAZMDpRwBABpx+BABkwOlHAEAGnH4EAJTA3UcAQAacfgQAlMDdRwBACdx9BACUwN1HAEAJ3H0EAJTA3UcAQAwcfQQAmmLgsyIA0NIDHw4BgJYe+DQIAISHwVMjAJDQDI+AAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACACAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAIAABNHpialFcmLajuAAAAAElFTkSuQmCC"
|
||||
|
||||
|
||||
IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
|
||||
"Add a small yellow star in the top right corner of this simple drawing of a blue circle on a white background"
|
||||
)
|
||||
|
||||
|
||||
def get_image_file_for_health_check() -> bytes:
|
||||
|
|
@ -121,6 +127,17 @@ class HealthCheckHelpers:
|
|||
else:
|
||||
return await litellm.acompletion(**model_params)
|
||||
|
||||
@staticmethod
|
||||
async def _image_edit_health_check(edit_request: Callable[[], Awaitable["ImageResponse"]]) -> "ImageResponse":
|
||||
import litellm
|
||||
|
||||
try:
|
||||
return await edit_request()
|
||||
except litellm.BadRequestError as e:
|
||||
if isinstance(e, litellm.ContentPolicyViolationError) or "moderation_blocked" in str(e):
|
||||
return litellm.ImageResponse()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def get_mode_handlers(
|
||||
model: str,
|
||||
|
|
@ -195,10 +212,12 @@ class HealthCheckHelpers:
|
|||
**_filter_model_params(model_params=model_params),
|
||||
prompt=prompt,
|
||||
),
|
||||
"image_edit": lambda: litellm.aimage_edit(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
image=get_image_file_for_health_check(),
|
||||
prompt=prompt or "test",
|
||||
"image_edit": lambda: HealthCheckHelpers._image_edit_health_check(
|
||||
edit_request=lambda: litellm.aimage_edit(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
image=get_image_file_for_health_check(),
|
||||
prompt=IMAGE_EDIT_HEALTH_CHECK_PROMPT,
|
||||
),
|
||||
),
|
||||
"video_generation": lambda: litellm.avideo_generation(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import re
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -45,12 +46,29 @@ def validate_no_callback_env_reference(param: str, value: object, *, source: str
|
|||
_raise_env_reference_error(param, source=source)
|
||||
|
||||
|
||||
# Langfuse rejects events whose environment does not match this pattern
|
||||
# (lowercase alphanumerics, hyphens, underscores; no "langfuse" prefix).
|
||||
# Validating here fails fast at config/init time instead of silently
|
||||
# dropping every trace server-side.
|
||||
LANGFUSE_ENVIRONMENT_PATTERN: Final = r"^(?!langfuse)[a-z0-9-_]+$"
|
||||
|
||||
|
||||
def validate_langfuse_environment_value(value: str) -> None:
|
||||
if not re.match(LANGFUSE_ENVIRONMENT_PATTERN, value):
|
||||
raise ValueError(
|
||||
f"Invalid langfuse_environment {value!r}: must be lowercase "
|
||||
"alphanumerics/hyphens/underscores and must not start with "
|
||||
f"'langfuse' (pattern {LANGFUSE_ENVIRONMENT_PATTERN})"
|
||||
)
|
||||
|
||||
|
||||
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
|
||||
_supported_callback_params: Final[tuple[str, ...]] = (
|
||||
"langfuse_public_key",
|
||||
"langfuse_secret",
|
||||
"langfuse_secret_key",
|
||||
"langfuse_host",
|
||||
"langfuse_environment",
|
||||
"langfuse_prompt_version",
|
||||
"langsmith_api_key",
|
||||
"langsmith_project",
|
||||
|
|
|
|||
|
|
@ -20,8 +20,8 @@ from __future__ import annotations
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.types.utils import InternalCallOrigin
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
|
||||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
|
|
@ -45,6 +45,60 @@ budget-checked like the request that spawned it. Everything else on the parent's
|
|||
be a lie on a sub-call that runs after it returned."""
|
||||
|
||||
|
||||
def is_background_response(response: object) -> bool:
|
||||
"""Whether a retrieved object is a response created with ``background=true``.
|
||||
|
||||
Such a create returns ``status="queued"`` and no usage at all, so nothing has billed the
|
||||
job by the time anyone reads it back. Accepts the response as a mapping or a model,
|
||||
because the callers hold it in both shapes.
|
||||
"""
|
||||
if isinstance(response, Mapping):
|
||||
return response.get("background") is True
|
||||
return getattr(response, "background", None) is True
|
||||
|
||||
|
||||
def is_unbilled_non_inference_call(
|
||||
call_type: str | None,
|
||||
metadata: Mapping[str, object] | None,
|
||||
response: object,
|
||||
) -> bool:
|
||||
"""A read/management route priced at zero, because the usage it reports belongs to the
|
||||
call that created the object it just read.
|
||||
|
||||
Retrieving a background response is the exception, and the enterprise cost poller's read
|
||||
is the same exception seen from the other side: that job's create billed nothing, so its
|
||||
retrieval is the only place the spend is ever visible. Pricing those at zero would lose
|
||||
the spend rather than deduplicate it.
|
||||
"""
|
||||
if call_type not in NON_INFERENCE_CALL_TYPES:
|
||||
return False
|
||||
if is_background_response(response):
|
||||
return False
|
||||
if metadata is None:
|
||||
return True
|
||||
return metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
|
||||
def is_unbilled_non_inference_call_from_params(
|
||||
call_type: str | None,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
response: object,
|
||||
) -> bool:
|
||||
""":func:`is_unbilled_non_inference_call` for callers holding raw ``litellm_params``.
|
||||
|
||||
The call-type membership test runs first so that inference traffic, which is every
|
||||
request in a normal workload, never pays for the metadata merge behind it.
|
||||
"""
|
||||
if call_type not in NON_INFERENCE_CALL_TYPES:
|
||||
return False
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
metadata: Final = (
|
||||
StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) if litellm_params is not None else None
|
||||
)
|
||||
return is_unbilled_non_inference_call(call_type, metadata, response)
|
||||
|
||||
|
||||
def sanitize_user_api_key_auth(auth: object) -> object:
|
||||
"""Copy of the auth object with its budget reservation removed; the cost callback
|
||||
falls back to reading the reservation from inside the auth object."""
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ from litellm.integrations.mlflow import MlflowLogger
|
|||
from litellm.integrations.sqs import SQSLogger
|
||||
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
|
||||
cost_breakdown_with_guardrail,
|
||||
guardrail_information_cost,
|
||||
|
|
@ -612,37 +613,60 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
processed_list: Final[list[str | Callable | CustomLogger]] = []
|
||||
for callback in callback_list:
|
||||
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
|
||||
# For callbacks that support team-scoped credentials (e.g. datadog),
|
||||
# pass only the relevant dynamic params as custom_logger_init_args.
|
||||
_custom_logger_init_args: dict | None = None
|
||||
if callback == "datadog":
|
||||
# dd_* params are blocked from standard_callback_dynamic_params
|
||||
# (request-level security); only the proxy-stamped team/key
|
||||
# callback vars are admin-configured and trusted.
|
||||
_custom_logger_init_args = {k: v for k, v in self._trusted_callback_vars if k.startswith("dd_")}
|
||||
|
||||
callback_class = _init_custom_logger_compatible_class(
|
||||
callback,
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
custom_logger_init_args=_custom_logger_init_args,
|
||||
)
|
||||
if callback_class is not None:
|
||||
processed_list.append(callback_class)
|
||||
for callback_instance in self._resolve_dynamic_callback_string(callback):
|
||||
processed_list.append(callback_instance)
|
||||
|
||||
# If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks
|
||||
if dynamic_callbacks_type == "success":
|
||||
if self.dynamic_async_success_callbacks is None:
|
||||
self.dynamic_async_success_callbacks = []
|
||||
self.dynamic_async_success_callbacks.append(callback_class)
|
||||
self.dynamic_async_success_callbacks.append(callback_instance)
|
||||
elif dynamic_callbacks_type == "failure":
|
||||
if self.dynamic_async_failure_callbacks is None:
|
||||
self.dynamic_async_failure_callbacks = []
|
||||
self.dynamic_async_failure_callbacks.append(callback_class)
|
||||
self.dynamic_async_failure_callbacks.append(callback_instance)
|
||||
else:
|
||||
processed_list.append(callback)
|
||||
return processed_list
|
||||
|
||||
def _resolve_dynamic_callback_string(self, callback: str) -> "tuple[CustomLogger, ...]":
|
||||
"""
|
||||
Resolve a known callback name to the logger instance(s) it dispatches to.
|
||||
|
||||
For callbacks that support team-scoped credentials (datadog, newrelic),
|
||||
only the proxy-stamped team/key callback vars are passed as
|
||||
custom_logger_init_args: dd_*/newrelic_* params are blocked from
|
||||
standard_callback_dynamic_params (request-level security), so the
|
||||
trusted-vars channel is the only way credentials reach a per-team logger.
|
||||
"""
|
||||
_trusted_var_prefix: Final = "dd_" if callback == "datadog" else "newrelic_" if callback == "newrelic" else None
|
||||
_custom_logger_init_args: Final[dict | None] = (
|
||||
{k: v for k, v in self._trusted_callback_vars if k.startswith(_trusted_var_prefix)}
|
||||
if _trusted_var_prefix is not None
|
||||
else None
|
||||
)
|
||||
|
||||
callback_class: Final = _init_custom_logger_compatible_class(
|
||||
callback,
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
custom_logger_init_args=_custom_logger_init_args,
|
||||
)
|
||||
if callback_class is None:
|
||||
return ()
|
||||
|
||||
# With team creds, "newrelic" resolves to the per-team METRICS logger;
|
||||
# resolve the name again without creds so the trace logger (OTel v2 /
|
||||
# legacy agent) keeps receiving this request.
|
||||
_newrelic_trace_class: Final = (
|
||||
_init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
|
||||
if callback == "newrelic" and _custom_logger_init_args and _custom_logger_init_args.get("newrelic_api_key")
|
||||
else None
|
||||
)
|
||||
if _newrelic_trace_class is not None and _newrelic_trace_class is not callback_class:
|
||||
return (callback_class, _newrelic_trace_class)
|
||||
return (callback_class,)
|
||||
|
||||
def initialize_standard_callback_dynamic_params(self, kwargs: dict | None = None) -> StandardCallbackDynamicParams:
|
||||
"""
|
||||
Initialize the standard callback dynamic params from the kwargs
|
||||
|
|
@ -1586,11 +1610,16 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if cache_hit is True:
|
||||
return 0.0
|
||||
|
||||
if is_unbilled_non_inference_call(
|
||||
self.call_type, StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params), result
|
||||
):
|
||||
return 0.0
|
||||
|
||||
transformed_result: Final = self._generate_content_result_as_model_response(result)
|
||||
if transformed_result is not None:
|
||||
result = transformed_result
|
||||
|
||||
if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"):
|
||||
if isinstance(result, (BaseModel, HttpxBinaryResponseContent)) and hasattr(result, "_hidden_params"):
|
||||
hidden_params: Final = getattr(result, "_hidden_params", {})
|
||||
if (
|
||||
"response_cost" in hidden_params and hidden_params["response_cost"] is not None
|
||||
|
|
@ -4636,6 +4665,19 @@ def _init_custom_logger_compatible_class(
|
|||
_in_memory_loggers.append(gitlab_logger)
|
||||
return gitlab_logger
|
||||
elif logging_integration == "newrelic":
|
||||
if custom_logger_init_args.get("newrelic_api_key"):
|
||||
# Team-scoped credentials: per-team METRICS logger, isolated per
|
||||
# credential set via DynamicLoggingCache. The trace logger for
|
||||
# this name stays on the global path below.
|
||||
from litellm.integrations.newrelic.newrelic_team_handler import (
|
||||
NewRelicHandler,
|
||||
)
|
||||
|
||||
return NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=custom_logger_init_args,
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
|
||||
_v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
|
||||
if _v2 is not None:
|
||||
return _v2
|
||||
|
|
@ -5057,7 +5099,7 @@ class StandardLoggingPayloadSetup:
|
|||
return messages
|
||||
|
||||
@staticmethod
|
||||
def merge_litellm_metadata(litellm_params: dict) -> dict:
|
||||
def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict:
|
||||
"""
|
||||
Merge both litellm_metadata and metadata from litellm_params.
|
||||
|
||||
|
|
@ -5819,7 +5861,7 @@ def get_standard_logging_object_payload(
|
|||
cache_hit: Final = kwargs.get("cache_hit", False)
|
||||
# Extract usage as a plain dict, avoiding Pydantic round-trip
|
||||
raw_usage_dict: Final = StandardLoggingPayloadSetup.get_usage_as_dict(
|
||||
response_obj=response_obj,
|
||||
response_obj=None if is_unbilled_non_inference_call(call_type, metadata, response_obj) else response_obj,
|
||||
combined_usage_object=cast(Usage | None, kwargs.get("combined_usage_object")),
|
||||
)
|
||||
usage_dict: Final = (
|
||||
|
|
|
|||
|
|
@ -21,11 +21,13 @@ class GuardrailCostEntry(BaseModel):
|
|||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
guardrail_cost: float | None = None
|
||||
# ``bool | None`` because the TypedDict sanctions None; None means "not set"
|
||||
# and keeps the default billed behavior, so a None-carrying entry must not
|
||||
# fail union validation and silently zero a sibling entry's real cost.
|
||||
guardrail_cost_in_spend: bool | None = True
|
||||
|
||||
|
||||
GuardrailInformationShape = tuple[GuardrailCostEntry, ...] | GuardrailCostEntry | None
|
||||
|
||||
_GUARDRAIL_INFORMATION_ADAPTER: Final[TypeAdapter[GuardrailInformationShape]] = TypeAdapter(GuardrailInformationShape)
|
||||
_GUARDRAIL_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailCostEntry]] = TypeAdapter(GuardrailCostEntry)
|
||||
|
||||
|
||||
def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None:
|
||||
|
|
@ -47,23 +49,55 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str
|
|||
return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items())
|
||||
|
||||
|
||||
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"
|
||||
|
||||
|
||||
def azure_prompt_shield_guardrail_cost(
|
||||
usage_units: Mapping[str, int],
|
||||
cost_tier: str | None,
|
||||
price_per_1000_text_records: float | None,
|
||||
) -> float | None:
|
||||
"""USD cost of an Azure Prompt Shield invocation from its text-record count.
|
||||
|
||||
Returns 0.0 on the free tier, ``text_records * price / 1000`` when a price is
|
||||
configured, and None when pricing is not configured (usage-only tracking).
|
||||
"""
|
||||
if cost_tier == "free":
|
||||
return 0.0
|
||||
if price_per_1000_text_records is None:
|
||||
return None
|
||||
return usage_units.get(AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT, 0) * price_per_1000_text_records / 1000.0
|
||||
|
||||
|
||||
def _billable_entry_cost(entry: GuardrailCostEntry) -> float:
|
||||
if entry.guardrail_cost_in_spend is False:
|
||||
return 0.0
|
||||
cost: Final = entry.guardrail_cost
|
||||
if cost is None or not math.isfinite(cost) or cost <= 0.0:
|
||||
return 0.0
|
||||
return cost
|
||||
|
||||
|
||||
def guardrail_information_cost(guardrail_information: object) -> float:
|
||||
def _validated_entry_cost(raw: object) -> float:
|
||||
"""Billable cost of one raw ``guardrail_information`` entry.
|
||||
|
||||
Validated per entry so one malformed entry (e.g. a custom hook stamping a
|
||||
non-boolean ``guardrail_cost_in_spend``) prices to 0.0 by itself instead of
|
||||
failing a whole-payload validation and silently zeroing a sibling entry's
|
||||
real billable cost."""
|
||||
try:
|
||||
parsed: Final = _GUARDRAIL_INFORMATION_ADAPTER.validate_python(guardrail_information)
|
||||
except ValidationError:
|
||||
return _billable_entry_cost(_GUARDRAIL_COST_ENTRY_ADAPTER.validate_python(raw))
|
||||
except ValidationError as e:
|
||||
verbose_logger.warning("Ignoring malformed guardrail_information entry for guardrail cost: %s", e)
|
||||
return 0.0
|
||||
if parsed is None:
|
||||
|
||||
|
||||
def guardrail_information_cost(guardrail_information: object) -> float:
|
||||
if guardrail_information is None:
|
||||
return 0.0
|
||||
if isinstance(parsed, GuardrailCostEntry):
|
||||
return _billable_entry_cost(parsed)
|
||||
return sum(_billable_entry_cost(entry) for entry in parsed)
|
||||
if isinstance(guardrail_information, (list, tuple)):
|
||||
return sum(_validated_entry_cost(entry) for entry in guardrail_information)
|
||||
return _validated_entry_cost(guardrail_information)
|
||||
|
||||
|
||||
def cost_breakdown_with_guardrail(cost_breakdown: CostBreakdown | None, guardrail_cost: float) -> CostBreakdown | None:
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from typing import Any, Final, Literal
|
|||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
|
||||
from litellm.types.llms.openai import (
|
||||
FileSearchTool,
|
||||
ResponsesAPIResponse,
|
||||
|
|
@ -64,11 +64,17 @@ class StandardBuiltInToolCostTracking:
|
|||
"""
|
||||
standard_built_in_tools_params = standard_built_in_tools_params or {}
|
||||
|
||||
google_maps_grounding_cost: Final = StandardBuiltInToolCostTracking._handle_google_maps_grounding_cost(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
# Handle web search
|
||||
if StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
|
||||
response_object=response_object, usage=usage
|
||||
):
|
||||
return StandardBuiltInToolCostTracking._handle_web_search_cost(
|
||||
return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_web_search_cost(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
|
|
@ -78,19 +84,56 @@ class StandardBuiltInToolCostTracking:
|
|||
|
||||
# Handle file search
|
||||
if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response_object):
|
||||
return StandardBuiltInToolCostTracking._handle_file_search_cost(
|
||||
return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_file_search_cost(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
standard_built_in_tools_params=standard_built_in_tools_params,
|
||||
)
|
||||
|
||||
# Handle Azure assistant features
|
||||
return StandardBuiltInToolCostTracking._handle_azure_assistant_costs(
|
||||
return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_azure_assistant_costs(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
standard_built_in_tools_params=standard_built_in_tools_params,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model_info(model: str, custom_llm_provider: str | None) -> tuple[ModelInfo | None, str | None]:
|
||||
direct: Final = StandardBuiltInToolCostTracking._safe_get_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if direct is not None:
|
||||
return direct, custom_llm_provider or direct["litellm_provider"]
|
||||
if "/" not in model:
|
||||
return None, custom_llm_provider
|
||||
by_prefix: Final = StandardBuiltInToolCostTracking._safe_get_model_info(model=model)
|
||||
if by_prefix is None:
|
||||
return None, custom_llm_provider
|
||||
return by_prefix, by_prefix["litellm_provider"]
|
||||
|
||||
@staticmethod
|
||||
def _handle_google_maps_grounding_cost(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
usage: Usage | None,
|
||||
) -> float:
|
||||
from litellm.llms import get_cost_for_google_maps_grounding_request
|
||||
from litellm.llms.gemini.cost_calculator import google_maps_grounding_requests
|
||||
|
||||
if usage is None or google_maps_grounding_requests(usage) is None:
|
||||
return 0.0
|
||||
model_info, resolved_provider = StandardBuiltInToolCostTracking._resolve_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if model_info is None or resolved_provider is None:
|
||||
return 0.0
|
||||
return (
|
||||
get_cost_for_google_maps_grounding_request(
|
||||
custom_llm_provider=resolved_provider, usage=usage, model_info=model_info
|
||||
)
|
||||
or 0.0
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _handle_web_search_cost(
|
||||
model: str,
|
||||
|
|
@ -102,29 +145,21 @@ class StandardBuiltInToolCostTracking:
|
|||
"""Handle web search cost calculation."""
|
||||
from litellm.llms import get_cost_for_web_search_request
|
||||
|
||||
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
|
||||
# A provider-prefixed model (e.g. gemini/gemini-3.1-flash-lite) may not map under the
|
||||
# request's custom_llm_provider. _resolve_model_info re-resolves from the prefix and adopts
|
||||
# that provider so the cost is routed and priced with the model_info that was actually
|
||||
# resolved, instead of feeding a re-resolved model into the original provider's calculator.
|
||||
model_info, resolved_provider = StandardBuiltInToolCostTracking._resolve_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
|
||||
# A provider-prefixed model (e.g. gemini/gemini-3.1-flash-lite) may not map under the
|
||||
# request's custom_llm_provider. Re-resolve from the prefix and adopt that provider so the
|
||||
# cost is routed and priced with the model_info that was actually resolved, instead of
|
||||
# feeding a re-resolved model into the original provider's calculator.
|
||||
if model_info is None and "/" in model:
|
||||
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(model=model)
|
||||
if model_info is not None:
|
||||
custom_llm_provider = model_info["litellm_provider"]
|
||||
|
||||
if custom_llm_provider is None and model_info is not None:
|
||||
custom_llm_provider = model_info["litellm_provider"]
|
||||
|
||||
resolved_usage: Final = StandardBuiltInToolCostTracking._usage_with_anthropic_web_search(
|
||||
usage=usage, response_object=response_object
|
||||
)
|
||||
|
||||
if model_info is not None and resolved_usage is not None and custom_llm_provider is not None:
|
||||
if model_info is not None and resolved_usage is not None and resolved_provider is not None:
|
||||
result: Final = get_cost_for_web_search_request(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
custom_llm_provider=resolved_provider,
|
||||
usage=resolved_usage,
|
||||
model_info=model_info,
|
||||
)
|
||||
|
|
@ -333,7 +368,7 @@ class StandardBuiltInToolCostTracking:
|
|||
get_anthropic_web_search_requests_from_response,
|
||||
)
|
||||
|
||||
if usage is not None and (_get_web_search_requests(getattr(usage, "server_tool_use", None)) is not None):
|
||||
if usage is not None and (get_web_search_requests(getattr(usage, "server_tool_use", None)) is not None):
|
||||
return usage
|
||||
web_search_requests: Final = get_anthropic_web_search_requests_from_response(response_object)
|
||||
if web_search_requests is None:
|
||||
|
|
@ -381,7 +416,7 @@ class StandardBuiltInToolCostTracking:
|
|||
# Anthropic Claude (direct API and Vertex AI) uses server_tool_use.web_search_requests.
|
||||
# Without this check, Claude ModelResponse always falls through to return False
|
||||
# and _handle_web_search_cost() is never called.
|
||||
if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
|
||||
if hasattr(usage, "server_tool_use") and get_web_search_requests(usage.server_tool_use) is not None:
|
||||
return True
|
||||
# xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
|
||||
# answer with no url_citation annotations has no other chat-path signal
|
||||
|
|
@ -396,7 +431,7 @@ class StandardBuiltInToolCostTracking:
|
|||
elif usage is not None:
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
and get_web_search_requests(usage.server_tool_use) is not None
|
||||
or (
|
||||
hasattr(usage, "prompt_tokens_details")
|
||||
and usage.prompt_tokens_details is not None
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ def _get_token_detail_value(details: object, key: str) -> int | None:
|
|||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _get_web_search_requests(server_tool_use: Any) -> int | None:
|
||||
def get_web_search_requests(server_tool_use: Any) -> int | None:
|
||||
"""
|
||||
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
|
||||
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
|
||||
|
|
@ -889,11 +889,22 @@ def generic_cost_per_token(
|
|||
total_details: Final = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
|
||||
has_double_counting: Final = (cache_hit > 0 or cache_creation > 0) and total_details > usage.prompt_tokens
|
||||
|
||||
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
|
||||
if has_double_counting:
|
||||
# cached and per-modality counts are both subsets of prompt_tokens and may overlap, so a
|
||||
# modality can only bill what the cache did not already cover or the overlap is billed twice
|
||||
uncached_budget: Final = max(usage.prompt_tokens - cache_hit - cache_creation, 0)
|
||||
billable_audio: Final = min(audio_tokens, uncached_budget)
|
||||
billable_image: Final = min(image_tokens, uncached_budget - billable_audio)
|
||||
billable_video: Final = min(video_tokens, uncached_budget - billable_audio - billable_image)
|
||||
prompt_tokens_details["audio_tokens"] = billable_audio
|
||||
prompt_tokens_details["image_tokens"] = billable_image
|
||||
prompt_tokens_details["video_tokens"] = billable_video
|
||||
prompt_tokens_details["text_tokens"] = uncached_budget - billable_audio - billable_image - billable_video
|
||||
elif text_tokens == 0 and prompt_tokens_details["image_count"] == 0:
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
text_tokens = max(text_tokens, 0)
|
||||
prompt_tokens_details["text_tokens"] = text_tokens
|
||||
prompt_tokens_details["text_tokens"] = max(
|
||||
usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
|
||||
)
|
||||
|
||||
(
|
||||
prompt_base_cost,
|
||||
|
|
@ -1063,15 +1074,17 @@ def get_token_type_cost_breakdown(
|
|||
reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
|
||||
|
||||
# Reasoning is billed at the selected tier's reasoning rate for tiered models,
|
||||
# else at the explicit per-reasoning-token rate when the model defines one,
|
||||
# otherwise at the standard output-token rate - this mirrors how the total
|
||||
# completion cost is computed, so the breakdown can never diverge from it.
|
||||
# else at the service-tier-aware per-reasoning-token rate - this mirrors how the
|
||||
# total completion cost is computed, so the breakdown can never diverge from it.
|
||||
tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage)
|
||||
flat_reasoning_rate: Final = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None)
|
||||
reasoning_rate: Final = (
|
||||
tiered_reasoning_rate
|
||||
if tiered_reasoning_rate is not None
|
||||
else (flat_reasoning_rate if flat_reasoning_rate is not None else completion_base_cost)
|
||||
else _resolve_reasoning_token_cost(
|
||||
model_info=model_info,
|
||||
service_tier=service_tier,
|
||||
completion_base_cost=completion_base_cost,
|
||||
)
|
||||
)
|
||||
reasoning_cost = float(reasoning_tokens) * reasoning_rate
|
||||
|
||||
|
|
|
|||
|
|
@ -178,7 +178,7 @@ def update_response_metadata(
|
|||
- response._hidden_params["litellm_overhead_time_ms"]
|
||||
- response.response_time_ms
|
||||
"""
|
||||
if result is None:
|
||||
if result is None or not hasattr(result, "_hidden_params"):
|
||||
return
|
||||
|
||||
metadata: Final = ResponseMetadata(result)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
import asyncio
|
||||
import atexit
|
||||
import contextvars
|
||||
import inspect
|
||||
import logging
|
||||
from collections.abc import Coroutine, Iterator
|
||||
from typing import Final
|
||||
|
|
@ -53,6 +54,7 @@ class LoggingWorker:
|
|||
self._queue: asyncio.Queue[LoggingTask] | None = None
|
||||
self._worker_task: asyncio.Task | None = None
|
||||
self._running_tasks: set[asyncio.Task] = set()
|
||||
self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks
|
||||
self._sem: asyncio.Semaphore | None = None
|
||||
self._bound_loop: asyncio.AbstractEventLoop | None = None
|
||||
self._last_aggressive_clear_time: float = 0.0
|
||||
|
|
@ -61,6 +63,38 @@ class LoggingWorker:
|
|||
# Register cleanup handler to flush remaining events on exit
|
||||
atexit.register(self._flush_on_exit)
|
||||
|
||||
def _track_dequeued(self, task: LoggingTask) -> None:
|
||||
self._dequeued_tasks[id(task)] = task
|
||||
|
||||
def _untrack_dequeued(self, task: LoggingTask) -> None:
|
||||
self._dequeued_tasks.pop(id(task), None)
|
||||
|
||||
def _unstarted_dequeued_tasks(self) -> tuple[LoggingTask, ...]:
|
||||
return tuple(
|
||||
task
|
||||
for task in self._dequeued_tasks.values()
|
||||
if inspect.getcoroutinestate(task["coroutine"]) == inspect.CORO_CREATED
|
||||
)
|
||||
|
||||
def _requeue_unstarted_dequeued(self, new_queue: "asyncio.Queue[LoggingTask]") -> int:
|
||||
revived: Final = self._unstarted_dequeued_tasks()
|
||||
self._dequeued_tasks.clear()
|
||||
for index, revived_task in enumerate(revived):
|
||||
try:
|
||||
new_queue.put_nowait(revived_task)
|
||||
except asyncio.QueueFull:
|
||||
for leftover in revived[index:]:
|
||||
self._track_dequeued(leftover)
|
||||
return index
|
||||
return len(revived)
|
||||
|
||||
def _run_coroutine_silently(self, loop: asyncio.AbstractEventLoop, coroutine: Coroutine) -> bool:
|
||||
try:
|
||||
loop.run_until_complete(asyncio.wait_for(coroutine, timeout=self.timeout))
|
||||
except (Exception, asyncio.CancelledError): # noqa: BLE001 # atexit flush must never break the user's program
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _drain_pending(queue: "asyncio.Queue[LoggingTask]") -> tuple[LoggingTask, ...]:
|
||||
"""Pop every task still queued, without awaiting them, so they can be moved to another queue."""
|
||||
|
|
@ -90,10 +124,12 @@ class LoggingWorker:
|
|||
new_queue: Final[asyncio.Queue[LoggingTask]] = asyncio.Queue(maxsize=self.max_queue_size)
|
||||
for carried_task in carried_over:
|
||||
new_queue.put_nowait(carried_task)
|
||||
if carried_over:
|
||||
revived_count: Final = self._requeue_unstarted_dequeued(new_queue)
|
||||
if carried_over or revived_count:
|
||||
verbose_logger.warning(
|
||||
"LoggingWorker: event loop changed; carried %d pending logging task(s) onto the new loop",
|
||||
"LoggingWorker: event loop changed; carried %d pending and revived %d dequeued logging task(s) onto the new loop",
|
||||
len(carried_over),
|
||||
revived_count,
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
|
||||
|
|
@ -129,6 +165,7 @@ class LoggingWorker:
|
|||
except Exception as e:
|
||||
verbose_logger.exception("LoggingWorker error: %s", e)
|
||||
finally:
|
||||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
finally:
|
||||
# Always release semaphore, even if queue is None
|
||||
|
|
@ -146,6 +183,7 @@ class LoggingWorker:
|
|||
await self._sem.acquire()
|
||||
try:
|
||||
task = await self._queue.get()
|
||||
self._track_dequeued(task)
|
||||
# Track each spawned coroutine so we can cancel on shutdown.
|
||||
processing_task = asyncio.create_task(self._process_log_task(task, self._sem))
|
||||
self._running_tasks.add(processing_task)
|
||||
|
|
@ -298,9 +336,10 @@ class LoggingWorker:
|
|||
extracted_tasks: Final = []
|
||||
for _ in range(items_to_extract):
|
||||
try:
|
||||
extracted_tasks.append(self._queue.get_nowait())
|
||||
extracted_tasks.append(extracted := self._queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
self._track_dequeued(extracted)
|
||||
|
||||
return extracted_tasks
|
||||
|
||||
|
|
@ -318,6 +357,7 @@ class LoggingWorker:
|
|||
|
||||
# Add new task to extracted tasks to process directly
|
||||
if new_task is not None:
|
||||
self._track_dequeued(new_task)
|
||||
extracted_tasks.append(new_task)
|
||||
|
||||
# Process extracted tasks directly
|
||||
|
|
@ -343,6 +383,7 @@ class LoggingWorker:
|
|||
# Suppress errors during processing to ensure we keep going
|
||||
pass
|
||||
finally:
|
||||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
|
||||
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
|
||||
|
|
@ -486,11 +527,12 @@ class LoggingWorker:
|
|||
self._safe_log("debug", "[LoggingWorker] atexit: No queue initialized")
|
||||
return
|
||||
|
||||
if self._queue.empty():
|
||||
unstarted_dequeued: Final = self._unstarted_dequeued_tasks()
|
||||
if self._queue.empty() and not unstarted_dequeued:
|
||||
self._safe_log("debug", "[LoggingWorker] atexit: Queue is empty")
|
||||
return
|
||||
|
||||
queue_size: Final = self._queue.qsize()
|
||||
queue_size: Final = self._queue.qsize() + len(unstarted_dequeued)
|
||||
self._safe_log("info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...")
|
||||
|
||||
# Create a new event loop since the original is closed
|
||||
|
|
@ -509,6 +551,16 @@ class LoggingWorker:
|
|||
previous_raise_exceptions: Final = logging.raiseExceptions
|
||||
logging.raiseExceptions = False
|
||||
try:
|
||||
for pending in unstarted_dequeued:
|
||||
if (
|
||||
processed >= MAX_ITERATIONS_TO_CLEAR_QUEUE
|
||||
or loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE
|
||||
):
|
||||
break
|
||||
if self._run_coroutine_silently(loop, pending["coroutine"]):
|
||||
processed += 1
|
||||
self._untrack_dequeued(pending)
|
||||
|
||||
while not self._queue.empty() and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE:
|
||||
if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
|
||||
self._safe_log(
|
||||
|
|
@ -526,11 +578,8 @@ class LoggingWorker:
|
|||
# Note: We run the coroutine directly, not via create_task,
|
||||
# since we're in a new event loop context
|
||||
try:
|
||||
loop.run_until_complete(task["coroutine"])
|
||||
processed += 1
|
||||
except Exception:
|
||||
# Silent failure to not break user's program
|
||||
pass
|
||||
if self._run_coroutine_silently(loop, task["coroutine"]):
|
||||
processed += 1
|
||||
finally:
|
||||
# Clear reference to prevent memory leaks
|
||||
task = None
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in MirroredPricingParams.model_
|
|||
"cache_creation_input_token_cost_above_1hr",
|
||||
"cache_creation_input_token_cost_above_200k_tokens",
|
||||
"cache_read_input_token_cost_above_200k_tokens",
|
||||
"google_maps_grounding_cost_per_query",
|
||||
)
|
||||
# tiered_pricing is emptied rather than zeroed: its tiers outrank the zeros written beside
|
||||
# them, so a zero here would leave the cost map's tiers billing the traffic the reserved
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import json
|
|||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
|
||||
|
||||
from ...caching import InMemoryCache
|
||||
|
|
@ -46,6 +47,15 @@ class LangfuseInMemoryCache(InMemoryCache):
|
|||
_created_langfuse_logger.Langfuse.flush()
|
||||
_created_langfuse_logger.Langfuse.shutdown()
|
||||
|
||||
# Loggers with a periodic flush task (e.g. NewRelicMetricsLogger) expose
|
||||
# stop() so eviction actually ends the task instead of leaking it.
|
||||
_evicted_stop: Final = getattr(self.cache_dict[key], "stop", None)
|
||||
if callable(_evicted_stop):
|
||||
try:
|
||||
_evicted_stop()
|
||||
except Exception: # noqa: BLE001 # a failing stop() must not block eviction
|
||||
verbose_logger.debug("DynamicLoggingCache: stop() raised during eviction", exc_info=True)
|
||||
|
||||
#########################################################
|
||||
# Call parent class to remove key from cache
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -173,6 +173,27 @@ def attach_cache_creation_token_details(
|
|||
return prompt_tokens_details.model_copy(update={"cache_creation_token_details": cache_creation_token_details})
|
||||
|
||||
|
||||
def apply_grounding_request_counts(
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None,
|
||||
web_search_requests: int | None,
|
||||
google_maps_grounding_requests: int | None,
|
||||
) -> PromptTokensDetailsWrapper | None:
|
||||
updates: Final = MappingProxyType(
|
||||
{
|
||||
field: value
|
||||
for field, value in (
|
||||
("web_search_requests", web_search_requests),
|
||||
("google_maps_grounding_requests", google_maps_grounding_requests),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
if not updates:
|
||||
return prompt_tokens_details
|
||||
counted: Final = prompt_tokens_details if prompt_tokens_details is not None else PromptTokensDetailsWrapper()
|
||||
return counted.model_copy(update=updates)
|
||||
|
||||
|
||||
class ChunkProcessor:
|
||||
def __init__(self, chunks: list, messages: list | None = None):
|
||||
self.chunks = self._sort_chunks(chunks)
|
||||
|
|
@ -778,6 +799,7 @@ class ChunkProcessor:
|
|||
|
||||
server_tool_use: ServerToolUse | None = None
|
||||
web_search_requests: int | None = None
|
||||
google_maps_grounding_requests: int | None = None
|
||||
completion_tokens_details: CompletionTokensDetails | None = None
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
|
||||
# Anthropic emits the cache-creation TTL breakdown (5m/1h split) only on
|
||||
|
|
@ -827,6 +849,13 @@ class ChunkProcessor:
|
|||
)
|
||||
if chunk_web_search_requests is not None:
|
||||
web_search_requests = chunk_web_search_requests
|
||||
chunk_google_maps_grounding_requests: int | None = getattr(
|
||||
usage_chunk_dict["prompt_tokens_details"],
|
||||
"google_maps_grounding_requests",
|
||||
None,
|
||||
)
|
||||
if chunk_google_maps_grounding_requests is not None:
|
||||
google_maps_grounding_requests = chunk_google_maps_grounding_requests
|
||||
|
||||
prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details
|
||||
|
||||
|
|
@ -852,6 +881,7 @@ class ChunkProcessor:
|
|||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
server_tool_use=server_tool_use,
|
||||
web_search_requests=web_search_requests,
|
||||
google_maps_grounding_requests=google_maps_grounding_requests,
|
||||
completion_tokens_details=completion_tokens_details,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
cost=cost,
|
||||
|
|
@ -939,6 +969,7 @@ class ChunkProcessor:
|
|||
|
||||
server_tool_use: Final[ServerToolUse | None] = calculated_usage_per_chunk["server_tool_use"]
|
||||
web_search_requests: Final[int | None] = calculated_usage_per_chunk["web_search_requests"]
|
||||
google_maps_grounding_requests: Final[int | None] = calculated_usage_per_chunk["google_maps_grounding_requests"]
|
||||
completion_tokens_details: Final[CompletionTokensDetails | None] = calculated_usage_per_chunk[
|
||||
"completion_tokens_details"
|
||||
]
|
||||
|
|
@ -998,13 +1029,11 @@ class ChunkProcessor:
|
|||
|
||||
if server_tool_use is not None:
|
||||
returned_usage.server_tool_use = server_tool_use
|
||||
if web_search_requests is not None:
|
||||
if returned_usage.prompt_tokens_details is None:
|
||||
returned_usage.prompt_tokens_details = PromptTokensDetailsWrapper(
|
||||
web_search_requests=web_search_requests
|
||||
)
|
||||
else:
|
||||
returned_usage.prompt_tokens_details.web_search_requests = web_search_requests
|
||||
returned_usage.prompt_tokens_details = apply_grounding_request_counts(
|
||||
returned_usage.prompt_tokens_details,
|
||||
web_search_requests,
|
||||
google_maps_grounding_requests,
|
||||
)
|
||||
|
||||
if cost is not None:
|
||||
setattr(returned_usage, "cost", cost)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,21 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import ModelInfo, Usage
|
||||
|
||||
|
||||
def get_cost_for_google_maps_grounding_request(
|
||||
custom_llm_provider: str, usage: "Usage", model_info: "ModelInfo"
|
||||
) -> float | None:
|
||||
"""
|
||||
Get the cost of Grounding with Google Maps for a given model. Only Gemini models on the
|
||||
Gemini API and Vertex AI can populate the Maps grounding counter, so every other provider
|
||||
returns None.
|
||||
"""
|
||||
if custom_llm_provider != "gemini" and not custom_llm_provider.startswith("vertex_ai"):
|
||||
return None
|
||||
from .gemini.cost_calculator import cost_per_google_maps_grounding_request
|
||||
|
||||
return cost_per_google_maps_grounding_request(usage=usage, model_info=model_info)
|
||||
|
||||
|
||||
def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", model_info: "ModelInfo") -> float | None:
|
||||
"""
|
||||
Get the cost for a web search request for a given model.
|
||||
|
|
|
|||
|
|
@ -8,9 +8,9 @@ from typing import TYPE_CHECKING, Final, Optional
|
|||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
_get_web_search_requests,
|
||||
generic_cost_per_token,
|
||||
get_provider_specific_geo_multiplier,
|
||||
get_web_search_requests,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -104,7 +104,7 @@ def get_cost_for_anthropic_web_search(
|
|||
|
||||
if usage is None:
|
||||
return 0.0
|
||||
web_search_requests: Final = _get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
web_search_requests: Final = get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
if web_search_requests is None:
|
||||
return 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -99,6 +99,7 @@ from litellm.types.llms.anthropic import (
|
|||
ContextManagementResponse,
|
||||
MessageBlockDelta,
|
||||
MessageDelta,
|
||||
ServerToolUsage,
|
||||
StreamingContentBlockDeltaType,
|
||||
UsageDelta,
|
||||
UsageIteration,
|
||||
|
|
@ -1354,10 +1355,24 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
return explicit_value
|
||||
return cls._first_positive_prompt_tokens_detail_value(usage, ("cache_creation_tokens", "cache_write_tokens"))
|
||||
|
||||
@classmethod
|
||||
def _get_web_search_request_count(cls, usage: Usage) -> int:
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
get_web_search_requests,
|
||||
)
|
||||
|
||||
from_server_tool_use: Final = cls._positive_int(
|
||||
get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
)
|
||||
if from_server_tool_use > 0:
|
||||
return from_server_tool_use
|
||||
return cls._first_positive_prompt_tokens_detail_value(usage, ("web_search_requests",))
|
||||
|
||||
@classmethod
|
||||
def _translate_openai_usage_to_anthropic_usage_delta(cls, usage: Usage) -> UsageDelta:
|
||||
cache_read_input_tokens: Final = cls._get_cache_read_input_tokens(usage)
|
||||
cache_creation_input_tokens: Final = cls._get_cache_creation_input_tokens(usage)
|
||||
web_search_requests: Final = cls._get_web_search_request_count(usage)
|
||||
input_tokens: Final = max(
|
||||
(usage.prompt_tokens or 0) - cache_read_input_tokens - cache_creation_input_tokens,
|
||||
0,
|
||||
|
|
@ -1371,6 +1386,11 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens
|
||||
if cache_read_input_tokens > 0:
|
||||
usage_delta["cache_read_input_tokens"] = cache_read_input_tokens
|
||||
if web_search_requests > 0:
|
||||
return UsageDelta(
|
||||
**usage_delta,
|
||||
server_tool_use=ServerToolUsage(web_search_requests=web_search_requests),
|
||||
)
|
||||
return usage_delta
|
||||
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ from litellm.types.llms.openai import (
|
|||
OpenAIMessageContentListBlock,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CacheCreationTokenDetails,
|
||||
ChatCompletionMessageToolCall,
|
||||
CompletionTokensDetailsWrapper,
|
||||
Function,
|
||||
|
|
@ -1807,6 +1808,26 @@ class AmazonConverseConfig(BaseConfig):
|
|||
thinking_blocks_list.append(_redacted_block)
|
||||
return thinking_blocks_list
|
||||
|
||||
@staticmethod
|
||||
def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenDetails | None":
|
||||
"""Split ``cacheDetails`` into 5m/1h buckets, or ``None`` unless the split fully
|
||||
accounts for ``cacheWriteInputTokens``, since a partial or unrecognized-ttl
|
||||
breakdown would understate the cache-write cost.
|
||||
|
||||
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html
|
||||
"""
|
||||
cache_details: Final = usage.get("cacheDetails")
|
||||
if not cache_details:
|
||||
return None
|
||||
tokens_5m: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
|
||||
tokens_1h: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
|
||||
if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0):
|
||||
return None
|
||||
return CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=tokens_5m,
|
||||
ephemeral_1h_input_tokens=tokens_1h,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None:
|
||||
"""Converse omits thinking tokens from its usage block; they only arrive under
|
||||
|
|
@ -1878,6 +1899,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
prompt_tokens_details: Final = PromptTokensDetailsWrapper(
|
||||
cached_tokens=cache_read_input_tokens,
|
||||
cache_creation_tokens=cache_creation_input_tokens,
|
||||
cache_creation_token_details=self._parse_cache_details(usage),
|
||||
text_tokens=raw_input_tokens,
|
||||
)
|
||||
estimated_reasoning_tokens: Final = (
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import ssl
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
|
|
@ -18,6 +19,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
_get_httpx_client,
|
||||
get_ssl_configuration,
|
||||
)
|
||||
from litellm.types.llms.openai import FileTypes
|
||||
from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders
|
||||
|
|
@ -56,7 +58,11 @@ class BaseLLMAIOHTTPHandler:
|
|||
|
||||
# Create a transport using AsyncHTTPHandler's logic
|
||||
try:
|
||||
self.transport = AsyncHTTPHandler._create_aiohttp_transport()
|
||||
ssl_config: Final = get_ssl_configuration()
|
||||
self.transport = AsyncHTTPHandler._create_aiohttp_transport(
|
||||
ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
|
||||
ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None,
|
||||
)
|
||||
self._owns_transport = True
|
||||
return self.transport
|
||||
except Exception:
|
||||
|
|
@ -79,20 +85,19 @@ class BaseLLMAIOHTTPHandler:
|
|||
|
||||
def _create_client_session_with_transport(self) -> ClientSession:
|
||||
"""Create a new client session using transport or connector configuration."""
|
||||
connector: Final = self._get_connector()
|
||||
if self.transport is None:
|
||||
connector: Final = self._get_connector()
|
||||
if connector:
|
||||
return aiohttp.ClientSession(connector=connector)
|
||||
|
||||
if self.transport and hasattr(self.transport, "_get_valid_client_session"):
|
||||
# Use transport's session creation if available
|
||||
session = self.transport._get_valid_client_session()
|
||||
return session
|
||||
elif connector:
|
||||
# Use provided connector
|
||||
session = aiohttp.ClientSession(connector=connector)
|
||||
return session
|
||||
else:
|
||||
# Default session creation
|
||||
session = aiohttp.ClientSession()
|
||||
return session
|
||||
transport: Final = self.transport or self._get_or_create_transport()
|
||||
if transport is not None and hasattr(transport, "_get_valid_client_session"):
|
||||
try:
|
||||
return transport._get_valid_client_session()
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
return aiohttp.ClientSession()
|
||||
|
||||
def _get_async_client_session(self, dynamic_client_session: ClientSession | None = None) -> ClientSession:
|
||||
if dynamic_client_session:
|
||||
|
|
|
|||
|
|
@ -2,16 +2,17 @@
|
|||
Translates from OpenAI's `/v1/chat/completions` to DeepSeek's `/v1/chat/completions`
|
||||
"""
|
||||
|
||||
from collections.abc import Coroutine
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from typing import Any, Final, Literal, cast, overload
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
handle_messages_with_content_list_to_str_conversion,
|
||||
convert_content_list_to_str,
|
||||
extract_search_results_text,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.utils import supports_reasoning
|
||||
from litellm.utils import supports_reasoning, supports_vision
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
|
|
@ -117,13 +118,98 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
|
|||
self, messages: list[AllMessageValues], model: str, is_async: bool = False
|
||||
) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]:
|
||||
"""
|
||||
DeepSeek does not support content in list format.
|
||||
DeepSeek vision models accept image_url content blocks in user
|
||||
messages (https://api-docs.deepseek.com/guides/vision), so those
|
||||
content lists are forwarded as-is, with any search_results text
|
||||
appended as a trailing text block. Every other message keeps the
|
||||
historical string collapse (which also folds search_results text
|
||||
into string content); a list with no extractable text stays
|
||||
unchanged, matching what DeepSeek historically received.
|
||||
"""
|
||||
messages = handle_messages_with_content_list_to_str_conversion(messages)
|
||||
forward_images: Final = any(
|
||||
isinstance(message.get("content"), list) for message in messages
|
||||
) and supports_vision(model=model, custom_llm_provider="deepseek")
|
||||
transformed: Final = [ # mutable-ok: provider messages must stay JSON-array lists the base transform mutates
|
||||
self._forward_or_collapse_content(message=message, forward_images=forward_images) for message in messages
|
||||
]
|
||||
|
||||
if is_async:
|
||||
return super()._transform_messages(messages=messages, model=model, is_async=True)
|
||||
return super()._transform_messages(messages=transformed, model=model, is_async=True)
|
||||
else:
|
||||
return super()._transform_messages(messages=messages, model=model, is_async=False)
|
||||
return super()._transform_messages(messages=transformed, model=model, is_async=False)
|
||||
|
||||
def _forward_or_collapse_content(self, message: AllMessageValues, forward_images: bool) -> AllMessageValues:
|
||||
"""
|
||||
Returns the vision-forwardable message with any search_results text
|
||||
appended as a text block; every other message keeps the historical
|
||||
string collapse, which extracts the text from a content list and
|
||||
folds search_results text into string content.
|
||||
"""
|
||||
content: Final = message.get("content")
|
||||
if (
|
||||
forward_images
|
||||
and isinstance(content, list)
|
||||
and self._is_vision_forwardable_content(message=message, content=content)
|
||||
):
|
||||
return self._with_search_results_text_block(message=message, content=content)
|
||||
collapsed: Final = convert_content_list_to_str(message=message)
|
||||
if not collapsed or collapsed == content:
|
||||
return message
|
||||
collapsed_message: Final = {**message, "content": collapsed} # mutable-ok: wire messages are plain JSON dicts
|
||||
return cast(AllMessageValues, collapsed_message) # cast-ok: TypedDict spread narrows to dict
|
||||
|
||||
def _is_vision_forwardable_content(self, message: AllMessageValues, content: Sequence[object]) -> bool:
|
||||
"""
|
||||
True only for a user message whose content list holds well-formed
|
||||
text and image_url blocks with at least one image; a block missing
|
||||
its payload falls back to the string collapse instead of crashing
|
||||
or reaching the wire malformed. The model capability gate lives in
|
||||
the caller.
|
||||
"""
|
||||
if message.get("role") != "user":
|
||||
return False
|
||||
if not all(self._is_forwardable_block(block) for block in content):
|
||||
return False
|
||||
return any(isinstance(block, dict) and block.get("type") == "image_url" for block in content)
|
||||
|
||||
@staticmethod
|
||||
def _is_forwardable_block(block: object) -> bool:
|
||||
"""A dict block typed text or image_url that carries its payload."""
|
||||
if not isinstance(block, dict):
|
||||
return False
|
||||
block_type: Final = block.get("type")
|
||||
if block_type == "image_url":
|
||||
return DeepSeekChatConfig._is_image_url_payload(block.get("image_url"))
|
||||
if block_type == "text":
|
||||
return isinstance(block.get("text"), str)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_image_url_payload(payload: object) -> bool:
|
||||
"""A url string or an object carrying one, per the OpenAI image_url shape."""
|
||||
if isinstance(payload, str):
|
||||
return bool(payload)
|
||||
if not isinstance(payload, Mapping):
|
||||
return False
|
||||
url: Final = payload.get("url")
|
||||
return isinstance(url, str) and bool(url)
|
||||
|
||||
def _with_search_results_text_block(self, message: AllMessageValues, content: Sequence[object]) -> AllMessageValues:
|
||||
"""
|
||||
Appends the message's search_results text as a trailing text block,
|
||||
keeping the context that the string collapse used to fold in, and
|
||||
drops the non-OpenAI search_results key from the wire message.
|
||||
"""
|
||||
message_fields: Final = cast(Mapping[str, object], message) # cast-ok: search_results is not on the TypedDicts
|
||||
search_text: Final = extract_search_results_text(message_fields.get("search_results"))
|
||||
if not search_text:
|
||||
return message
|
||||
forwarded_content: Final = [*content, {"type": "text", "text": search_text}] # mutable-ok: JSON-array content
|
||||
forwarded: Final = { # mutable-ok: wire messages are plain JSON dicts
|
||||
**{key: value for key, value in message_fields.items() if key != "search_results"},
|
||||
"content": forwarded_content,
|
||||
}
|
||||
return cast(AllMessageValues, forwarded) # cast-ok: TypedDict spread narrows to dict
|
||||
|
||||
def _thinking_mode_active(self, model: str, optional_params: dict) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -13,6 +13,13 @@ class FireworksAIException(BaseLLMException):
|
|||
|
||||
|
||||
def get_fireworks_session_id(litellm_params: dict) -> str | None:
|
||||
"""
|
||||
Session id to send as `x-session-affinity`, or None when the caller gave none.
|
||||
|
||||
Deliberately does not fall back to `litellm_trace_id`: that is generated per
|
||||
request (`str(uuid.uuid4())` when absent), so using it pins every request to a
|
||||
different Fireworks node and prompt caching never hits.
|
||||
"""
|
||||
params: Final = litellm_params
|
||||
for key in ("litellm_session_id", "session_id"):
|
||||
value = params.get(key)
|
||||
|
|
@ -23,9 +30,6 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
|
|||
value = metadata.get("session_id")
|
||||
if value:
|
||||
return str(value)
|
||||
value = params.get("litellm_trace_id")
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -39,25 +39,69 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa
|
|||
``model_info`` when available, falling back to $0.035 for models not
|
||||
yet updated in the pricing JSON.
|
||||
"""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper
|
||||
|
||||
_DEFAULT_COST: Final = 35e-3
|
||||
search_costs: Final = model_info.get("search_context_cost_per_query") or {}
|
||||
_cost: Final = search_costs.get("search_context_size_medium", _DEFAULT_COST)
|
||||
|
||||
number_of_web_search_requests = 0
|
||||
if (
|
||||
usage is not None
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
):
|
||||
number_of_web_search_requests = usage.prompt_tokens_details.web_search_requests
|
||||
requests_from_prompt_details: Final = (
|
||||
usage.prompt_tokens_details.web_search_requests
|
||||
if (
|
||||
usage is not None
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
)
|
||||
else None
|
||||
)
|
||||
requests_from_server_tool_use: Final = get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
number_of_web_search_requests: Final = requests_from_prompt_details or requests_from_server_tool_use or 0
|
||||
|
||||
# per_prompt billing: clamp to 1 (flat fee per grounded API call)
|
||||
billing_mode: Final = model_info.get("web_search_billing_unit") or "per_prompt"
|
||||
if number_of_web_search_requests > 0 and billing_mode == "per_prompt":
|
||||
number_of_web_search_requests = 1
|
||||
billable_requests: Final = (
|
||||
1 if (number_of_web_search_requests > 0 and billing_mode == "per_prompt") else number_of_web_search_requests
|
||||
)
|
||||
|
||||
return _cost * number_of_web_search_requests
|
||||
return _cost * billable_requests
|
||||
|
||||
|
||||
GOOGLE_MAPS_GROUNDING_DEFAULT_COST_PER_QUERY: Final = 14e-3
|
||||
GOOGLE_MAPS_GROUNDING_DEFAULT_COST_PER_PROMPT: Final = 25e-3
|
||||
|
||||
|
||||
def google_maps_grounding_requests(usage: "Usage | None") -> int | None:
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper
|
||||
|
||||
details: Final = usage.prompt_tokens_details if usage is not None else None
|
||||
if not isinstance(details, PromptTokensDetailsWrapper) or not hasattr(details, "google_maps_grounding_requests"):
|
||||
return None
|
||||
return details.google_maps_grounding_requests
|
||||
|
||||
|
||||
def cost_per_google_maps_grounding_request(usage: "Usage", model_info: "ModelInfo") -> float:
|
||||
"""
|
||||
Calculates the cost of Grounding with Google Maps.
|
||||
|
||||
Billing follows ``web_search_billing_unit`` in model_info the same way Google Search grounding
|
||||
does: ``"per_query"`` (Gemini 3.x) multiplies the executed Maps queries, ``"per_prompt"``
|
||||
(default, Gemini 2.x) charges one flat fee per grounded prompt.
|
||||
|
||||
The rate comes from ``google_maps_grounding_cost_per_query`` in ``model_info``, falling back
|
||||
to Google's list price for that billing unit when the pricing JSON has no entry yet.
|
||||
"""
|
||||
requests: Final = google_maps_grounding_requests(usage)
|
||||
if not requests or requests <= 0:
|
||||
return 0.0
|
||||
billing_mode: Final = model_info.get("web_search_billing_unit") or "per_prompt"
|
||||
default_cost: Final = (
|
||||
GOOGLE_MAPS_GROUNDING_DEFAULT_COST_PER_QUERY
|
||||
if billing_mode == "per_query"
|
||||
else GOOGLE_MAPS_GROUNDING_DEFAULT_COST_PER_PROMPT
|
||||
)
|
||||
configured_cost: Final = model_info.get("google_maps_grounding_cost_per_query")
|
||||
cost: Final = default_cost if configured_cost is None else configured_cost
|
||||
billed_requests: Final = requests if billing_mode == "per_query" else 1
|
||||
return cost * billed_requests
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ This file contains the transformation logic for the Gemini realtime API.
|
|||
|
||||
import json
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final, cast
|
||||
|
||||
import litellm
|
||||
|
|
@ -72,6 +73,28 @@ MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Final[dict[str, OpenAIRealtimeEventTypes | Res
|
|||
_KNOWN_GEMINI_TOP_LEVEL_KEYS: Final[set] = {map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT}
|
||||
|
||||
|
||||
OPENAI_STOCK_REALTIME_VOICES: Final[frozenset[str]] = frozenset(
|
||||
{"alloy", "ash", "ballad", "cedar", "coral", "echo", "marin", "sage", "shimmer", "verse"}
|
||||
)
|
||||
|
||||
|
||||
def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None:
|
||||
"""Build the Gemini Live speechConfig for a client-requested voice.
|
||||
|
||||
OpenAI stock voice names have no Gemini equivalent and Gemini Live closes
|
||||
the session on an unknown voice, so they are dropped with a warning and
|
||||
the model keeps its default voice. Every other name is forwarded verbatim.
|
||||
"""
|
||||
if isinstance(voice, str) and voice.lower() in OPENAI_STOCK_REALTIME_VOICES:
|
||||
verbose_logger.warning(
|
||||
"Gemini Realtime: voice %s is an OpenAI voice with no Gemini equivalent; "
|
||||
"dropping it so the session keeps the model's default voice.",
|
||||
voice,
|
||||
)
|
||||
return None
|
||||
return VertexGeminiConfig()._map_audio_params({"voice": voice})
|
||||
|
||||
|
||||
class GeminiRealtimeConfig(BaseRealtimeConfig):
|
||||
_TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping
|
||||
|
||||
|
|
@ -282,12 +305,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
automaticActivityDetection=transformed_audio_activity_config
|
||||
)
|
||||
elif key == "voice":
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
vertex_gemini_config = VertexGeminiConfig()
|
||||
speech_config = vertex_gemini_config._map_audio_params({"voice": value})
|
||||
speech_config = _gemini_live_speech_config(value)
|
||||
if speech_config:
|
||||
optional_params["generationConfig"]["speechConfig"] = speech_config
|
||||
if len(optional_params["generationConfig"]) == 0:
|
||||
|
|
@ -365,10 +383,6 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
entry: Final = GeminiRealtimeConfig._model_cost_entry(model)
|
||||
return bool(entry.get("gemini_native_audio") or entry.get("gemini_audio_only_live"))
|
||||
|
||||
@staticmethod
|
||||
def _is_native_audio_model(model: str) -> bool:
|
||||
return bool(GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio"))
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_modalities(model: str, modalities: list[Any]) -> list[str]:
|
||||
"""Map unsupported TEXT responseModalities to AUDIO for audio-only Live models."""
|
||||
|
|
@ -384,7 +398,6 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
|
||||
@staticmethod
|
||||
def _finalize_gemini_live_setup(model: str, setup: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Drop fields Gemini Live native-audio rejects on ``setup``."""
|
||||
generation_config: Final = setup.get("generationConfig")
|
||||
if isinstance(generation_config, dict):
|
||||
modalities: Final = generation_config.get("responseModalities")
|
||||
|
|
@ -392,8 +405,6 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
generation_config["responseModalities"] = GeminiRealtimeConfig._coerce_response_modalities(
|
||||
model, modalities
|
||||
)
|
||||
if GeminiRealtimeConfig._is_native_audio_model(model):
|
||||
generation_config.pop("speechConfig", None)
|
||||
return setup
|
||||
|
||||
def _handle_session_update(
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
MiniMax Anthropic transformation config - extends AnthropicConfig for MiniMax's Anthropic-compatible API
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
from typing import Any, Final # noqa: TID251 # override below must mirror the legacy base signature
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
|
||||
|
|
@ -49,6 +49,26 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig):
|
|||
"""
|
||||
return api_base or get_secret_str("MINIMAX_API_BASE") or "https://api.minimax.io/anthropic/v1/messages"
|
||||
|
||||
def validate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: mirrors the legacy base override signature
|
||||
model: str,
|
||||
messages: list[Any], # mutable-ok: mirrors the legacy base override signature
|
||||
optional_params: dict, # mutable-ok: mirrors the legacy base override signature
|
||||
litellm_params: dict, # mutable-ok: mirrors the legacy base override signature
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]: # mutable-ok: mirrors the legacy base override signature
|
||||
return super().validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=self.get_api_key(api_key=api_key),
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ def cost_per_character(
|
|||
usage: Usage,
|
||||
prompt_characters: float | None = None,
|
||||
completion_characters: float | None = None,
|
||||
service_tier: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
|
|
@ -74,6 +75,8 @@ def cost_per_character(
|
|||
- custom_llm_provider: str, "vertex_ai-*"
|
||||
- prompt_characters: float, the number of input characters
|
||||
- completion_characters: float, the number of output characters
|
||||
- service_tier: optional tier derived from Gemini trafficType
|
||||
("priority" for ON_DEMAND_PRIORITY, "flex" for FLEX/batch).
|
||||
- vertex_location: the Vertex AI location serving the request; non-global
|
||||
locations apply the model's regional-endpoint uplift multiplier
|
||||
|
||||
|
|
@ -92,6 +95,7 @@ def cost_per_character(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
|
|
@ -123,6 +127,7 @@ def cost_per_character(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
)
|
||||
|
||||
## CALCULATE OUTPUT COST
|
||||
|
|
@ -131,6 +136,7 @@ def cost_per_character(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
)
|
||||
else:
|
||||
completion_tokens: Final = usage.completion_tokens
|
||||
|
|
@ -162,6 +168,7 @@ def cost_per_character(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
)
|
||||
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
|
|
|
|||
56
litellm/llms/vertex_ai/gemini/grounding_requests.py
Normal file
56
litellm/llms/vertex_ai/gemini/grounding_requests.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GroundingRequests:
|
||||
web_search_requests: int | None
|
||||
google_maps_grounding_requests: int | None
|
||||
|
||||
def has_billable_grounding(self) -> bool:
|
||||
return bool(self.web_search_requests or self.google_maps_grounding_requests)
|
||||
|
||||
|
||||
def _chunk_kinds(item: Mapping[str, object]) -> frozenset[str]:
|
||||
chunks: Final = item.get("groundingChunks")
|
||||
if not isinstance(chunks, list):
|
||||
return frozenset()
|
||||
return frozenset(kind for chunk in chunks if isinstance(chunk, Mapping) for kind in chunk)
|
||||
|
||||
|
||||
def _queries(item: Mapping[str, object]) -> frozenset[str]:
|
||||
queries: Final = item.get("webSearchQueries")
|
||||
if not isinstance(queries, list):
|
||||
return frozenset()
|
||||
return frozenset(query for query in queries if isinstance(query, str) and query)
|
||||
|
||||
|
||||
def _is_maps_item(item: Mapping[str, object]) -> bool:
|
||||
return "maps" in _chunk_kinds(item) or bool(item.get("googleMapsWidgetContextToken"))
|
||||
|
||||
|
||||
def _attributes_queries_to_maps(item: Mapping[str, object]) -> bool:
|
||||
return _is_maps_item(item) and "web" not in _chunk_kinds(item)
|
||||
|
||||
|
||||
def calculate_grounding_requests(grounding_metadata: Sequence[Mapping[str, object]]) -> GroundingRequests:
|
||||
"""Billable grounding requests across candidates, counting each distinct query once.
|
||||
|
||||
Duplicate queries within and across grounding metadata items collapse to the
|
||||
distinct-query count (#36377), and empty strings are ignored. Maps grounding is
|
||||
floored at one request whenever a candidate carries maps chunks or a widget token,
|
||||
since per-prompt billing charges the prompt even when no query is reported.
|
||||
"""
|
||||
items: Final = tuple(item for item in grounding_metadata if isinstance(item, Mapping))
|
||||
web_queries: Final = frozenset(
|
||||
query for item in items if not _attributes_queries_to_maps(item) for query in _queries(item)
|
||||
)
|
||||
maps_queries: Final = frozenset(
|
||||
query for item in items if _attributes_queries_to_maps(item) for query in _queries(item)
|
||||
)
|
||||
has_maps: Final = any(_is_maps_item(item) for item in items)
|
||||
return GroundingRequests(
|
||||
web_search_requests=len(web_queries) or None,
|
||||
google_maps_grounding_requests=max(len(maps_queries), 1) if has_maps else None,
|
||||
)
|
||||
|
|
@ -89,6 +89,7 @@ from ..common_utils import (
|
|||
supports_response_json_schema,
|
||||
)
|
||||
from ..vertex_llm_base import VertexBase
|
||||
from .grounding_requests import calculate_grounding_requests
|
||||
from .transformation import (
|
||||
_gemini_convert_messages_with_history,
|
||||
async_transform_request_body,
|
||||
|
|
@ -1717,14 +1718,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
completion_response: GenerateContentResponseBody | BidiGenerateContentServerMessage,
|
||||
) -> bool:
|
||||
"""
|
||||
Whether the response used Grounding with Google Search, detected via
|
||||
groundingMetadata.webSearchQueries (an actual web search was performed).
|
||||
Whether the response used Grounding with Google Search or Grounding with Google Maps,
|
||||
detected via groundingMetadata.webSearchQueries (an actual web search was performed) or
|
||||
groundingMetadata.groundingChunks[].maps (a Maps lookup was performed).
|
||||
|
||||
Google bills grounding-with-Google-Search retrieved tokens separately (a per-request /
|
||||
per-query search fee) and excludes them from input token billing, unlike URL context /
|
||||
File Search / code execution whose tool-use tokens are charged at the input token rate.
|
||||
URL context also emits groundingMetadata (with groundingChunks but no webSearchQueries),
|
||||
so presence of groundingMetadata alone is not a sufficient signal.
|
||||
Google bills both groundings separately (a per-request / per-query fee) and excludes their
|
||||
retrieved tokens from input token billing, unlike URL context / File Search / code execution
|
||||
whose tool-use tokens are charged at the input token rate. URL context also emits
|
||||
groundingMetadata (with web groundingChunks but no webSearchQueries), so presence of
|
||||
groundingMetadata alone is not a sufficient signal.
|
||||
See https://ai.google.dev/gemini-api/docs/pricing and
|
||||
https://github.com/BerriAI/litellm/discussions/33198
|
||||
"""
|
||||
|
|
@ -1732,7 +1734,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return False
|
||||
for candidate in completion_response["candidates"] or []:
|
||||
grounding_metadata, _, _, _ = VertexGeminiConfig._extract_candidate_metadata(candidate)
|
||||
if VertexGeminiConfig._calculate_web_search_requests(grounding_metadata):
|
||||
if calculate_grounding_requests(grounding_metadata).has_billable_grounding():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
@ -1979,16 +1981,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
|
||||
@staticmethod
|
||||
def _calculate_web_search_requests(grounding_metadata: list[dict]) -> int | None:
|
||||
web_search_requests: int | None = None
|
||||
return calculate_grounding_requests(grounding_metadata).web_search_requests
|
||||
|
||||
if grounding_metadata and isinstance(grounding_metadata, list) and len(grounding_metadata) > 0:
|
||||
for grounding_metadata_item in grounding_metadata:
|
||||
web_search_queries = grounding_metadata_item.get("webSearchQueries")
|
||||
if web_search_queries and web_search_requests:
|
||||
web_search_requests += len([q for q in web_search_queries if q])
|
||||
elif web_search_queries:
|
||||
web_search_requests = len([q for q in web_search_queries if q])
|
||||
return web_search_requests
|
||||
@staticmethod
|
||||
def _set_grounding_usage_counters(usage: Usage, grounding_metadata: Sequence[Mapping[str, object]]) -> None:
|
||||
grounding_requests: Final = calculate_grounding_requests(grounding_metadata)
|
||||
details: Final = cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details)
|
||||
if grounding_requests.web_search_requests is not None:
|
||||
details.web_search_requests = grounding_requests.web_search_requests
|
||||
if grounding_requests.google_maps_grounding_requests is not None:
|
||||
details.google_maps_grounding_requests = grounding_requests.google_maps_grounding_requests
|
||||
|
||||
@staticmethod
|
||||
def _create_streaming_choice(
|
||||
|
|
@ -2454,9 +2456,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
|
||||
usage: Final = VertexGeminiConfig._calculate_usage(completion_response=completion_response)
|
||||
|
||||
web_search_requests: Final = VertexGeminiConfig._calculate_web_search_requests(grounding_metadata)
|
||||
if web_search_requests is not None:
|
||||
cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details).web_search_requests = web_search_requests
|
||||
VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)
|
||||
|
||||
setattr(model_response, "usage", usage)
|
||||
|
||||
|
|
@ -3221,9 +3221,7 @@ class ModelResponseIterator:
|
|||
completion_response=processed_chunk,
|
||||
)
|
||||
|
||||
web_search_requests: Final = VertexGeminiConfig._calculate_web_search_requests(grounding_metadata)
|
||||
if web_search_requests is not None:
|
||||
cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details).web_search_requests = web_search_requests
|
||||
VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)
|
||||
|
||||
traffic_type: Final = processed_chunk.get("usageMetadata", {}).get("trafficType")
|
||||
if traffic_type:
|
||||
|
|
|
|||
|
|
@ -8013,7 +8013,7 @@ def speech(
|
|||
|
||||
if max_retries is None:
|
||||
max_retries = litellm.num_retries or openai.DEFAULT_MAX_RETRIES
|
||||
litellm_params_dict: Final = get_litellm_params(**kwargs)
|
||||
litellm_params_dict: Final = get_litellm_params(metadata=metadata, api_key=api_key or dynamic_api_key, **kwargs)
|
||||
|
||||
# Get provider-specific text-to-speech config and map parameters
|
||||
text_to_speech_provider_config = ProviderConfigManager.get_provider_text_to_speech_config(
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -46,6 +46,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
|
||||
_run_centralized_common_checks,
|
||||
user_api_key_auth,
|
||||
)
|
||||
|
|
@ -429,7 +430,10 @@ class MCPRequestHandler:
|
|||
# An explicit x-litellm-api-key is always a LiteLLM credential, even
|
||||
# for a delegated server, so validate it: identity / spend / rate
|
||||
# limits resolve and any stored upstream token can be forwarded.
|
||||
validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request)
|
||||
validated_user_api_key_auth = await user_api_key_auth(
|
||||
api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}",
|
||||
request=request,
|
||||
)
|
||||
elif MCPRequestHandler._target_servers_delegate_auth_to_upstream(
|
||||
path=request_route,
|
||||
mcp_servers=mcp_servers,
|
||||
|
|
|
|||
|
|
@ -1600,7 +1600,7 @@ async def refresh_user_oauth_token(
|
|||
) -> OAuthCredentialPayload | None:
|
||||
"""Attempt to refresh a per-user OAuth2 token using its stored refresh_token.
|
||||
|
||||
POSTs to ``server.token_url`` with ``grant_type=refresh_token``.
|
||||
POSTs to ``server.effective_token_url`` with ``grant_type=refresh_token``.
|
||||
|
||||
On success: persists the new credential via ``store_user_oauth_credential``
|
||||
and returns the updated payload dict.
|
||||
|
|
@ -1609,7 +1609,7 @@ async def refresh_user_oauth_token(
|
|||
stale credential and triggering re-authentication.
|
||||
"""
|
||||
refresh_token: Final[str | None] = cred.get("refresh_token")
|
||||
token_url: Final[str | None] = getattr(server, "token_url", None)
|
||||
token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None)
|
||||
server_id: Final[str] = getattr(server, "server_id", "")
|
||||
client_id: Final[str | None] = getattr(server, "client_id", None)
|
||||
client_secret: Final[str | None] = getattr(server, "client_secret", None)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import html as _html
|
|||
import json
|
||||
import secrets
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
|
@ -663,6 +663,26 @@ def _endpoint_not_configured_detail(
|
|||
)
|
||||
|
||||
|
||||
async def _server_with_oauth_endpoints(
|
||||
mcp_server: MCPServer,
|
||||
needed_endpoint: Callable[[MCPServer], str | None],
|
||||
) -> MCPServer:
|
||||
"""Join deferred OAuth discovery only when the endpoint this caller needs is still missing.
|
||||
|
||||
Admin-entered endpoints live on ``configured_*`` after an anchored issuer empties the
|
||||
resolved fields. A caller whose needed endpoint already resolves never awaits discovery
|
||||
and cannot 503 over a leftover pin. A server still missing it joins the deferred task;
|
||||
no slot is a no-op and the caller 400s.
|
||||
"""
|
||||
if needed_endpoint(mcp_server) is not None:
|
||||
return mcp_server
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
return await global_mcp_server_manager.ensure_oauth_metadata_discovered(mcp_server)
|
||||
|
||||
|
||||
def _raise_unless_oauth2_discovery_server(
|
||||
mcp_server: MCPServer | None,
|
||||
mcp_server_name: str | None,
|
||||
|
|
@ -697,7 +717,7 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
|
|||
returns directly to the client's redirect URI without transiting the gateway. Gateway-side
|
||||
redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit
|
||||
arm, where the upstream only knows the gateway's own callback."""
|
||||
return mcp_server.is_dcr_bridge and bool(mcp_server.registration_url) and not mcp_server.client_id
|
||||
return mcp_server.is_dcr_bridge and bool(mcp_server.effective_registration_url) and not mcp_server.client_id
|
||||
|
||||
|
||||
def _require_s256_pkce(
|
||||
|
|
@ -745,7 +765,7 @@ def _redirect_to_upstream_authorize(
|
|||
**({"scope": scope_value} if scope_value else {}),
|
||||
**({"resource": upstream_resource} if upstream_resource else {}),
|
||||
}
|
||||
parsed_auth_url: Final = urlparse(mcp_server.authorization_url or "")
|
||||
parsed_auth_url: Final = urlparse(mcp_server.effective_authorization_url or "")
|
||||
merged_params: Final = {**dict(parse_qsl(parsed_auth_url.query)), **passthrough_params}
|
||||
return RedirectResponse(urlunparse(parsed_auth_url._replace(query=urlencode(merged_params))))
|
||||
|
||||
|
|
@ -812,18 +832,19 @@ async def authorize_with_server(
|
|||
ephemeral_dcr_client: "EphemeralDcrClient | None" = None,
|
||||
):
|
||||
_raise_if_not_oauth2(mcp_server)
|
||||
if mcp_server.authorization_url is None:
|
||||
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint)
|
||||
if resolved_server.effective_authorization_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=_endpoint_not_configured_detail(
|
||||
mcp_server,
|
||||
resolved_server,
|
||||
"authorization url",
|
||||
"set Authorization URL and Token URL manually",
|
||||
"set Issuer to discover them from the identity provider (RFC 8414)",
|
||||
),
|
||||
)
|
||||
|
||||
if mcp_server.is_dcr_bridge:
|
||||
if resolved_server.is_dcr_bridge:
|
||||
# Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated,
|
||||
# now-non-optional pair to the upstream authorize; the short-circuit arm keeps
|
||||
# calling this for its enforcement side effect, then falls through to the gateway
|
||||
|
|
@ -832,9 +853,9 @@ async def authorize_with_server(
|
|||
# A gateway-minted ephemeral client is registered against {base}/callback, so its
|
||||
# flow must run the short-circuit arm; the relay arm is only for clients that
|
||||
# registered themselves through the front door and hold their own redirect binding.
|
||||
if _dcr_bridge_relays_client_registration(mcp_server) and ephemeral_dcr_client is None:
|
||||
if _dcr_bridge_relays_client_registration(resolved_server) and ephemeral_dcr_client is None:
|
||||
return _redirect_to_upstream_authorize(
|
||||
mcp_server=mcp_server,
|
||||
mcp_server=resolved_server,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
|
|
@ -860,7 +881,7 @@ async def authorize_with_server(
|
|||
# litellm key, so the browser session is the only identity source; without one there is nothing to
|
||||
# bind, so send the user through login first. Every other oauth2 server keeps the identity-less state.
|
||||
litellm_user_id: str | None = None
|
||||
if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate:
|
||||
if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate:
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
_user_id_from_session_cookie,
|
||||
)
|
||||
|
|
@ -870,7 +891,7 @@ async def authorize_with_server(
|
|||
return _redirect_to_litellm_login(request)
|
||||
denial: Final = await _bridge_authorize_access_denial(
|
||||
litellm_user_id=litellm_user_id,
|
||||
mcp_server=mcp_server,
|
||||
mcp_server=resolved_server,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
)
|
||||
|
|
@ -884,7 +905,7 @@ async def authorize_with_server(
|
|||
code_challenge_method=code_challenge_method,
|
||||
client_redirect_uri=redirect_uri,
|
||||
litellm_user_id=litellm_user_id,
|
||||
mcp_server_id=mcp_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None,
|
||||
mcp_server_id=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None,
|
||||
dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None,
|
||||
dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None,
|
||||
dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method
|
||||
|
|
@ -894,26 +915,26 @@ async def authorize_with_server(
|
|||
relay_state: Final = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES)
|
||||
|
||||
params: Final = {
|
||||
"client_id": mcp_server.client_id if mcp_server.client_id else client_id,
|
||||
"client_id": resolved_server.client_id if resolved_server.client_id else client_id,
|
||||
"redirect_uri": f"{request_base_url}/callback",
|
||||
"state": relay_state,
|
||||
"response_type": response_type or "code",
|
||||
}
|
||||
if scope:
|
||||
params["scope"] = scope
|
||||
elif mcp_server.scopes:
|
||||
params["scope"] = " ".join(mcp_server.scopes)
|
||||
elif resolved_server.scopes:
|
||||
params["scope"] = " ".join(resolved_server.scopes)
|
||||
|
||||
if code_challenge:
|
||||
params["code_challenge"] = code_challenge
|
||||
if code_challenge_method:
|
||||
params["code_challenge_method"] = code_challenge_method
|
||||
|
||||
upstream_resource: Final = resolve_upstream_resource(mcp_server)
|
||||
upstream_resource: Final = resolve_upstream_resource(resolved_server)
|
||||
if upstream_resource:
|
||||
params["resource"] = upstream_resource
|
||||
|
||||
parsed_auth_url: Final = urlparse(mcp_server.authorization_url)
|
||||
parsed_auth_url: Final = urlparse(resolved_server.effective_authorization_url)
|
||||
existing_params: Final = dict(parse_qsl(parsed_auth_url.query))
|
||||
existing_params.update(params)
|
||||
final_url: Final = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params)))
|
||||
|
|
@ -946,11 +967,13 @@ async def exchange_token_with_server(
|
|||
if grant_type not in ("authorization_code", "refresh_token"):
|
||||
raise HTTPException(status_code=400, detail="Unsupported grant_type")
|
||||
|
||||
if mcp_server.token_url is None:
|
||||
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _token_flow_needed_endpoint)
|
||||
token_url: Final = resolved_server.effective_token_url
|
||||
if token_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=_endpoint_not_configured_detail(
|
||||
mcp_server,
|
||||
resolved_server,
|
||||
"token url",
|
||||
"set Token URL manually",
|
||||
"set Issuer to discover it from the identity provider (RFC 8414)",
|
||||
|
|
@ -965,16 +988,16 @@ async def exchange_token_with_server(
|
|||
# recovered from a sealed code) must authenticate the way its own registration was granted,
|
||||
# not the way the server row is configured; callers that carry no method keep the row's method
|
||||
# as before.
|
||||
resolved_client_id: Final = mcp_server.client_id if mcp_server.client_id else client_id
|
||||
resolved_client_secret: Final = mcp_server.client_secret if mcp_server.client_id else client_secret
|
||||
resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id
|
||||
resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret
|
||||
resolved_auth_method: Final = (
|
||||
mcp_server.token_endpoint_auth_method
|
||||
if mcp_server.client_id
|
||||
else (client_token_endpoint_auth_method or mcp_server.token_endpoint_auth_method)
|
||||
resolved_server.token_endpoint_auth_method
|
||||
if resolved_server.client_id
|
||||
else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method)
|
||||
)
|
||||
try:
|
||||
token_request: Final = build_upstream_oauth2_token_request(
|
||||
mcp_server,
|
||||
resolved_server,
|
||||
auth_method=resolved_auth_method,
|
||||
client_id=resolved_client_id,
|
||||
client_secret=resolved_client_secret,
|
||||
|
|
@ -987,14 +1010,14 @@ async def exchange_token_with_server(
|
|||
bridge_upstream_refresh: SecretStr | None = None
|
||||
bridge_upstream_scope: str | None = None
|
||||
refresh_request_scope: str | None = None
|
||||
is_bridge: Final = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge
|
||||
is_bridge: Final = resolved_server.is_oauth_delegate and resolved_server.is_dcr_bridge
|
||||
|
||||
if grant_type == "refresh_token":
|
||||
# Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed
|
||||
# identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange
|
||||
# sends the upstream token and never the envelope. A failure returns without touching the upstream.
|
||||
if is_bridge:
|
||||
prepared_refresh: Final = await _prepare_bridge_refresh(mcp_server, refresh_token)
|
||||
prepared_refresh: Final = await _prepare_bridge_refresh(resolved_server, refresh_token)
|
||||
if not isinstance(prepared_refresh, _BridgeRefreshReady):
|
||||
return _bridge_mint_error_response(prepared_refresh)
|
||||
bridge_mint_ready = prepared_refresh.ready
|
||||
|
|
@ -1031,13 +1054,13 @@ async def exchange_token_with_server(
|
|||
# A raw upstream code (scripted path) opens to None and the code is used as-is.
|
||||
bridge_identity = open_bridge_authorization_code(code)
|
||||
if bridge_identity is not None:
|
||||
if bridge_identity.mcp_server_id != mcp_server.server_id:
|
||||
if bridge_identity.mcp_server_id != resolved_server.server_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Authorization code was issued for a different MCP server",
|
||||
)
|
||||
code = bridge_identity.upstream_code
|
||||
bridge_token_relay: Final = _dcr_bridge_relays_client_registration(mcp_server)
|
||||
bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
|
||||
if bridge_token_relay and not redirect_uri:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
|
|
@ -1059,7 +1082,7 @@ async def exchange_token_with_server(
|
|||
# Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or
|
||||
# the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code.
|
||||
if is_bridge:
|
||||
prepared: Final = await _prepare_bridge_mint(request, mcp_server, bridge_identity)
|
||||
prepared: Final = await _prepare_bridge_mint(request, resolved_server, bridge_identity)
|
||||
if not isinstance(prepared, _BridgeMintReady):
|
||||
return _bridge_mint_error_response(prepared)
|
||||
bridge_mint_ready = prepared
|
||||
|
|
@ -1067,7 +1090,7 @@ async def exchange_token_with_server(
|
|||
async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
|
||||
try:
|
||||
response: Final = await async_client.post(
|
||||
mcp_server.token_url,
|
||||
token_url,
|
||||
headers={"Accept": "application/json", **token_request.headers},
|
||||
data=token_data,
|
||||
)
|
||||
|
|
@ -1076,8 +1099,8 @@ async def exchange_token_with_server(
|
|||
except httpx.HTTPStatusError as exc:
|
||||
fault: Final = classify_upstream_token_rejection(
|
||||
exc.response,
|
||||
credential_source=_token_credential_source(mcp_server),
|
||||
log_context=mcp_server.server_id,
|
||||
credential_source=_token_credential_source(resolved_server),
|
||||
log_context=resolved_server.server_id,
|
||||
)
|
||||
upstream_rejected_bridge_refresh: Final = (
|
||||
is_bridge
|
||||
|
|
@ -1090,7 +1113,7 @@ async def exchange_token_with_server(
|
|||
"bridge refresh: the upstream rejected the sealed refresh token for server=%s with "
|
||||
"invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client "
|
||||
"re-runs authorization_code rather than an opaque upstream error",
|
||||
mcp_server.server_id,
|
||||
resolved_server.server_id,
|
||||
)
|
||||
return _bridge_mint_error_response("invalid_refresh")
|
||||
return render_token_fault(fault)
|
||||
|
|
@ -1103,22 +1126,22 @@ async def exchange_token_with_server(
|
|||
|
||||
# Validate token response against server-configured rules before any storage.
|
||||
# This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc.
|
||||
if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict):
|
||||
if resolved_server.token_validation and isinstance(resolved_server.token_validation, dict):
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules=mcp_server.token_validation,
|
||||
server_id=mcp_server.server_id,
|
||||
validation_rules=resolved_server.token_validation,
|
||||
server_id=resolved_server.server_id,
|
||||
)
|
||||
|
||||
# Store server-side when the server is configured for per-user OAuth and
|
||||
# the calling client has provided a valid LiteLLM identity.
|
||||
# Errors are non-fatal: the token is still returned to the client.
|
||||
if mcp_server.needs_user_oauth_token:
|
||||
if resolved_server.needs_user_oauth_token:
|
||||
user_id: Final = await _extract_user_id_from_request(request)
|
||||
if user_id:
|
||||
try:
|
||||
await _store_per_user_token_server_side(
|
||||
server=mcp_server,
|
||||
server=resolved_server,
|
||||
user_id=user_id,
|
||||
token_response=token_response,
|
||||
)
|
||||
|
|
@ -1126,7 +1149,7 @@ async def exchange_token_with_server(
|
|||
verbose_logger.warning(
|
||||
"exchange_token_with_server: server-side storage failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
mcp_server.server_id,
|
||||
resolved_server.server_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
|
|
@ -1136,7 +1159,7 @@ async def exchange_token_with_server(
|
|||
"requires the stored token, so the client will be challenged with 401 on reconnect. "
|
||||
"Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), "
|
||||
"or store it via POST /mcp/server/{id}/oauth-user-credential.",
|
||||
mcp_server.server_id,
|
||||
resolved_server.server_id,
|
||||
)
|
||||
|
||||
# A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the
|
||||
|
|
@ -1147,7 +1170,9 @@ async def exchange_token_with_server(
|
|||
token_response = {**token_response, "scope": refresh_request_scope}
|
||||
# Phase 3: seal the upstream grant into the client-held envelope; failures map through the same
|
||||
# OAuth-shaped response as the phase-1 preconditions.
|
||||
minted: Final = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc))
|
||||
minted: Final = _finish_bridge_mint(
|
||||
bridge_mint_ready, resolved_server, token_response, datetime.now(timezone.utc)
|
||||
)
|
||||
return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted)
|
||||
|
||||
raw_access_token: Final = token_response.get("access_token") if isinstance(token_response, dict) else None
|
||||
|
|
@ -1551,7 +1576,8 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) ->
|
|||
bounded by the server count even when the request origin varies) so parallel authorize requests
|
||||
cannot each register an upstream client; the cache stamps nothing onto the server record and
|
||||
correctness never depends on it because the sealed state carries the client through the flow."""
|
||||
if mcp_server.registration_url is None:
|
||||
registration_url: Final = mcp_server.effective_registration_url
|
||||
if registration_url is None:
|
||||
return None
|
||||
request_base_url: Final = get_request_base_url(request)
|
||||
cache_key: Final = f"mcp_ephemeral_dcr_client:{mcp_server.server_id}:{request_base_url}"
|
||||
|
|
@ -1571,7 +1597,7 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) ->
|
|||
"token_endpoint_auth_method": "none",
|
||||
}
|
||||
response: Final = await _post_dcr_registration(
|
||||
registration_url=mcp_server.registration_url,
|
||||
registration_url=registration_url,
|
||||
register_data=register_data,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
|
|
@ -1617,7 +1643,7 @@ async def resolve_ephemeral_dcr_client(
|
|||
usable to generate orphan IdP clients)."""
|
||||
if not (mcp_server.is_true_passthrough or (mcp_server.is_oauth_delegate and not mcp_server.is_dcr_bridge)):
|
||||
return None
|
||||
if mcp_server.authorization_url is None:
|
||||
if mcp_server.effective_authorization_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="MCP server authorization url is not set",
|
||||
|
|
@ -1627,6 +1653,29 @@ async def resolve_ephemeral_dcr_client(
|
|||
return await mint_ephemeral_dcr_client(request, mcp_server)
|
||||
|
||||
|
||||
def _register_flow_needed_endpoint(mcp_server: MCPServer) -> str | None:
|
||||
"""The register flow's deferred-discovery join gate. A DCR bridge with no admin-configured
|
||||
client can only register callers through the upstream's registration endpoint
|
||||
(``_oauth_endpoints_unresolved`` keeps its discovery slot armed for exactly this shape), so
|
||||
the flow must keep joining discovery while registration is still missing instead of silently
|
||||
degrading to the dummy short-circuit. Every other shape only needs the authorization url."""
|
||||
if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None:
|
||||
return None
|
||||
return mcp_server.effective_authorization_url
|
||||
|
||||
|
||||
def _token_flow_needed_endpoint(mcp_server: MCPServer) -> str | None:
|
||||
"""The token exchange's deferred-discovery join gate. The exchange's relay-vs-callback arm
|
||||
(:func:`_dcr_bridge_relays_client_registration`) reads the registration url, so a clientless
|
||||
DCR bridge rebuilt without its discovered registration endpoint must keep joining discovery
|
||||
even when the token url already resolves; skipping it would select the gateway-callback arm
|
||||
and the upstream would reject the code over a redirect_uri mismatch. Every other shape only
|
||||
needs the token url."""
|
||||
if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None:
|
||||
return None
|
||||
return mcp_server.effective_token_url
|
||||
|
||||
|
||||
async def register_client_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -1661,21 +1710,23 @@ async def register_client_with_server(
|
|||
):
|
||||
return dummy_return
|
||||
|
||||
if mcp_server.authorization_url is None:
|
||||
resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint)
|
||||
if resolved_server.effective_authorization_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=_endpoint_not_configured_detail(
|
||||
mcp_server,
|
||||
resolved_server,
|
||||
"authorization url",
|
||||
"set Authorization URL and Token URL manually",
|
||||
"set Issuer to discover them from the identity provider (RFC 8414)",
|
||||
),
|
||||
)
|
||||
|
||||
if mcp_server.registration_url is None:
|
||||
registration_url: Final = resolved_server.effective_registration_url
|
||||
if registration_url is None:
|
||||
return dummy_return
|
||||
|
||||
bridge_relay: Final = _dcr_bridge_relays_client_registration(mcp_server)
|
||||
bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
|
||||
if bridge_relay and not client_redirect_uris:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
|
|
@ -1690,15 +1741,17 @@ async def register_client_with_server(
|
|||
"token_endpoint_auth_method": token_endpoint_auth_method or ("none" if bridge_relay else ""),
|
||||
}
|
||||
response: Final = await _post_dcr_registration(
|
||||
registration_url=mcp_server.registration_url,
|
||||
registration_url=registration_url,
|
||||
register_data=register_data,
|
||||
server_id=mcp_server.server_id,
|
||||
server_id=resolved_server.server_id,
|
||||
)
|
||||
|
||||
token_response = response.json()
|
||||
|
||||
if persist_credentials and not bridge_relay:
|
||||
persistence_result = await _persist_dcr_client_registration(mcp_server, token_response, current_redirect_uri)
|
||||
persistence_result = await _persist_dcr_client_registration(
|
||||
resolved_server, token_response, current_redirect_uri
|
||||
)
|
||||
if persistence_result == "reused":
|
||||
return dummy_return
|
||||
|
||||
|
|
@ -1755,17 +1808,10 @@ async def authorize(
|
|||
lookup_name: Final[str | None] = mcp_server_name or client_id
|
||||
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = (
|
||||
await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
|
||||
if lookup_name
|
||||
else None
|
||||
global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None
|
||||
)
|
||||
if mcp_server is None and mcp_server_name is None:
|
||||
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
mcp_server = (
|
||||
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
|
||||
if unresolved_server is not None
|
||||
else None
|
||||
)
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
_raise_if_not_oauth2(mcp_server)
|
||||
|
|
@ -1846,14 +1892,9 @@ async def token_endpoint(
|
|||
|
||||
lookup_name: Final = mcp_server_name or client_id
|
||||
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
|
||||
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip)
|
||||
if mcp_server is None and mcp_server_name is None:
|
||||
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
mcp_server = (
|
||||
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
|
||||
if unresolved_server is not None
|
||||
else None
|
||||
)
|
||||
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(status_code=404, detail="MCP server not found")
|
||||
return await exchange_token_with_server(
|
||||
|
|
@ -2684,10 +2725,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
|
|||
return await register_aggregate_client(request=request, request_body=data)
|
||||
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if resolved:
|
||||
resolved_server: Final = await global_mcp_server_manager.ensure_oauth_metadata_discovered(resolved)
|
||||
return await register_client_with_server(
|
||||
request=request,
|
||||
mcp_server=resolved_server,
|
||||
mcp_server=resolved,
|
||||
client_name=data.get("client_name", ""),
|
||||
grant_types=data.get("grant_types", []),
|
||||
response_types=data.get("response_types", []),
|
||||
|
|
@ -2697,10 +2737,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
|
|||
)
|
||||
return dummy_return
|
||||
|
||||
mcp_server: Final = await global_mcp_server_manager.get_resolved_mcp_server_by_name(
|
||||
mcp_server_name,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip)
|
||||
if mcp_server is None:
|
||||
return dummy_return
|
||||
return await register_client_with_server(
|
||||
|
|
|
|||
|
|
@ -523,7 +523,7 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool:
|
|||
# can come from resource discovery, so a server that resolved its endpoints but no scopes is
|
||||
# still unresolved for its flow.
|
||||
return True
|
||||
if server.is_dcr_bridge and not server.client_id and server.registration_url is None:
|
||||
if server.is_dcr_bridge and not server.client_id and server.effective_registration_url is None:
|
||||
# A DCR bridge with no admin-configured client can only register callers through the
|
||||
# upstream's registration endpoint, so a build that resolved the authorize and token
|
||||
# endpoints but not registration_endpoint (partial metadata) is still unresolved for its
|
||||
|
|
@ -535,8 +535,8 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool:
|
|||
return _flow_endpoints_missing(
|
||||
server.auth_type,
|
||||
MCPServerManager.effective_oauth2_flow(server),
|
||||
server.authorization_url,
|
||||
server.token_url,
|
||||
server.effective_authorization_url,
|
||||
server.effective_token_url,
|
||||
server.token_exchange_endpoint,
|
||||
)
|
||||
|
||||
|
|
@ -6205,14 +6205,6 @@ class MCPServerManager:
|
|||
return server
|
||||
return None
|
||||
|
||||
async def get_resolved_mcp_server_by_name(
|
||||
self,
|
||||
server_name: str,
|
||||
client_ip: str | None = None,
|
||||
) -> MCPServer | None:
|
||||
server: Final = self.get_mcp_server_by_name(server_name, client_ip=client_ip)
|
||||
return await self.ensure_oauth_metadata_discovered(server) if server is not None else None
|
||||
|
||||
def get_filtered_registry(self, client_ip: str | None = None) -> dict[str, MCPServer]:
|
||||
"""
|
||||
Get registry filtered by client IP access control.
|
||||
|
|
|
|||
|
|
@ -67,7 +67,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
rest of the identity rather than stored in a key."""
|
||||
material: Final = "\x00".join(
|
||||
(
|
||||
server.token_url or "",
|
||||
server.effective_token_url or "",
|
||||
server.client_id or "",
|
||||
server.client_secret or "",
|
||||
" ".join(server.scopes or ()),
|
||||
|
|
@ -82,7 +82,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
|
||||
@staticmethod
|
||||
def _has_client_credentials_config(server: "MCPServer") -> bool:
|
||||
return bool(server.client_id and server.client_secret and server.token_url)
|
||||
return bool(server.client_id and server.client_secret and server.effective_token_url)
|
||||
|
||||
async def async_get_token(self, server: "MCPServer") -> str | None:
|
||||
"""Return a valid access token, fetching or refreshing as needed.
|
||||
|
|
@ -112,19 +112,20 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
return token
|
||||
|
||||
async def _fetch_token(self, server: "MCPServer") -> tuple[str, int]:
|
||||
"""POST to ``token_url`` with ``grant_type=client_credentials``.
|
||||
"""POST to ``effective_token_url`` with ``grant_type=client_credentials``.
|
||||
|
||||
Returns ``(access_token, ttl_seconds)`` where ttl accounts for the
|
||||
expiry buffer so the cache entry expires before the real token does.
|
||||
"""
|
||||
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
|
||||
if not server.client_id or not server.client_secret or not server.token_url:
|
||||
token_url: Final = server.effective_token_url
|
||||
if not server.client_id or not server.client_secret or not token_url:
|
||||
raise ValueError(
|
||||
f"MCP server '{server.server_id}' missing required OAuth2 fields: "
|
||||
f"client_id={bool(server.client_id)}, "
|
||||
f"client_secret={bool(server.client_secret)}, "
|
||||
f"token_url={bool(server.token_url)}"
|
||||
f"token_url={bool(token_url)}"
|
||||
)
|
||||
|
||||
token_request: Final = build_upstream_oauth2_token_request(
|
||||
|
|
@ -146,7 +147,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
)
|
||||
|
||||
try:
|
||||
response: Final = await client.post(server.token_url, data=data, headers=token_request.headers or None)
|
||||
response: Final = await client.post(token_url, data=data, headers=token_request.headers or None)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -142,7 +142,7 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
|
|||
config=ClientCredentialsConfig(
|
||||
client_id=server.client_id,
|
||||
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
|
||||
token_url=server.token_url,
|
||||
token_url=server.effective_token_url,
|
||||
scopes=tuple(server.scopes or ()),
|
||||
audience=server.audience,
|
||||
upstream_resource=resolve_upstream_resource(server),
|
||||
|
|
@ -163,7 +163,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
|||
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
|
||||
forwarded only when the operator set it; a missing one is omitted, not derived.
|
||||
"""
|
||||
endpoint: Final = server.token_exchange_endpoint or server.token_url
|
||||
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
|
||||
if not server.client_id or not server.client_secret:
|
||||
return None
|
||||
profile: Final[Literal["rfc8693", "entra_obo"]] = (
|
||||
|
|
|
|||
|
|
@ -88,7 +88,10 @@ class AuthorizationCodeRefresher:
|
|||
if token.refresh_token is None:
|
||||
return None
|
||||
server: Final = self._server_lookup(server_id)
|
||||
if server is None or not server.token_url:
|
||||
if server is None:
|
||||
return None
|
||||
token_url: Final = server.effective_token_url
|
||||
if not token_url:
|
||||
return None
|
||||
|
||||
try:
|
||||
|
|
@ -106,7 +109,7 @@ class AuthorizationCodeRefresher:
|
|||
"refresh_token": token.refresh_token,
|
||||
**token_request.body,
|
||||
}
|
||||
body: Final = await self._token_endpoint(server.token_url, form, token_request.headers)
|
||||
body: Final = await self._token_endpoint(token_url, form, token_request.headers)
|
||||
if body is None:
|
||||
return None
|
||||
access_token: Final = body.get("access_token")
|
||||
|
|
|
|||
|
|
@ -8,8 +8,8 @@ omits each feature's routes until the feature is warmed.
|
|||
|
||||
import asyncio
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
|
|
@ -397,11 +397,27 @@ def _make_warmup_router(app: "FastAPI") -> "APIRouter":
|
|||
return router
|
||||
|
||||
|
||||
def inject_lazy_stubs(schema: dict) -> dict:
|
||||
"""Inject openapi entries for unloaded features. Uses the snapshot file
|
||||
when available (full route info), otherwise falls back to a single
|
||||
placeholder per feature. Any failure logs and returns the schema unchanged
|
||||
so /openapi.json never 500s on a cosmetic injection bug."""
|
||||
def loaded_lazy_modules(app: "FastAPI") -> frozenset[str]:
|
||||
"""The set of lazy feature modules whose routers are actually registered
|
||||
on this app (tracked by _force_load), empty before the middleware ever ran.
|
||||
sys.modules is the wrong signal: boot code imports several feature modules
|
||||
(mcp_management, cloudzero, vantage, config_overrides) without mounting
|
||||
their routers, and their stubs must still be injected."""
|
||||
loaded: Final = getattr(app.state, "lazy_loaded", None)
|
||||
if not isinstance(loaded, set):
|
||||
return frozenset()
|
||||
return frozenset(m for m in loaded if isinstance(m, str))
|
||||
|
||||
|
||||
def inject_lazy_stubs(
|
||||
schema: dict,
|
||||
loaded_modules: AbstractSet[str],
|
||||
features: tuple[LazyFeature, ...] = LAZY_FEATURES,
|
||||
) -> dict:
|
||||
"""Inject openapi entries for features not in loaded_modules. Uses the
|
||||
snapshot file when available (full route info), otherwise falls back to a
|
||||
single placeholder per feature. Any failure logs and returns the schema
|
||||
unchanged so /openapi.json never 500s on a cosmetic injection bug."""
|
||||
try:
|
||||
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
|
||||
|
||||
|
|
@ -409,8 +425,8 @@ def inject_lazy_stubs(schema: dict) -> dict:
|
|||
paths: Final = schema.setdefault("paths", {})
|
||||
schemas: Final = schema.setdefault("components", {}).setdefault("schemas", {})
|
||||
|
||||
for feat in LAZY_FEATURES:
|
||||
if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
|
||||
for feat in features:
|
||||
if feat.module_path in loaded_modules and not feat.persistent_swagger_stub:
|
||||
continue
|
||||
|
||||
fragment = (snapshot or {}).get(feat.name)
|
||||
|
|
|
|||
|
|
@ -23538,7 +23538,7 @@
|
|||
"paths": {
|
||||
"/prompts": {
|
||||
"post": {
|
||||
"description": "Create a new prompt\n\n\ud83d\udc49 [Prompt docs](https://docs.litellm.ai/docs/proxy/prompt_management)\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/prompts\" \\\n -H \"Authorization: Bearer <your_api_key>\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"prompt_id\": \"my_prompt\",\n \"litellm_params\": {\n \"prompt_id\": \"json_prompt\",\n \"prompt_integration\": \"dotprompt\",\n ### EITHER prompt_directory OR prompt_data MUST BE PROVIDED\n \"prompt_directory\": \"/path/to/dotprompt/folder\",\n \"prompt_data\": {\"json_prompt\": {\"content\": \"This is a prompt\", \"metadata\": {\"model\": \"gpt-4\"}}}\n },\n \"prompt_info\": {\n \"prompt_type\": \"config\"\n }\n }'\n```",
|
||||
"description": "Create a new prompt\n\n\ud83d\udc49 [Prompt docs](https://docs.litellm.ai/docs/proxy/prompt_management)\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/prompts\" \\\n -H \"Authorization: Bearer <your_api_key>\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"prompt_id\": \"my_prompt\",\n \"litellm_params\": {\n \"prompt_id\": \"my_prompt\",\n \"prompt_integration\": \"dotprompt\",\n \"prompt_data\": {\"content\": \"This is a prompt\", \"metadata\": {\"model\": \"gpt-4\"}}\n },\n \"prompt_info\": {\n \"prompt_type\": \"config\"\n }\n }'\n```",
|
||||
"operationId": "create_prompt_prompts_post",
|
||||
"requestBody": {
|
||||
"content": {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
|
|||
from litellm._uuid import uuid
|
||||
from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
validate_langfuse_environment_value,
|
||||
validate_no_callback_env_reference,
|
||||
)
|
||||
from litellm.types.integrations.compression_interception import (
|
||||
|
|
@ -2027,6 +2028,8 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
|
|||
raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}")
|
||||
callback_vars[key] = str(value)
|
||||
validate_no_callback_env_reference(key, callback_vars[key], source="key/team callback metadata")
|
||||
if key == "langfuse_environment":
|
||||
validate_langfuse_environment_value(callback_vars[key])
|
||||
return values
|
||||
|
||||
|
||||
|
|
@ -2507,6 +2510,17 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"are skipped for on-demand GET /health as well as the background health loop."
|
||||
),
|
||||
)
|
||||
model_list_healthy_only: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing "
|
||||
"deployments are all unhealthy, for every caller, without needing `healthy_only=true` "
|
||||
"per request. Requires `background_health_checks: true`, and keeps deployment health "
|
||||
"state cached without turning on `enable_health_check_routing`, so routing is "
|
||||
"unaffected. With no health state nothing is hidden. Hiding is presentation-only, a "
|
||||
"hidden model can still be called."
|
||||
),
|
||||
)
|
||||
alerting: list | None = Field(
|
||||
None,
|
||||
description="List of alerting integrations. Today, just slack - `alerting: ['slack']`",
|
||||
|
|
|
|||
|
|
@ -2588,13 +2588,29 @@ async def _delete_cache_key_object(
|
|||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
):
|
||||
"""
|
||||
Evict one key object, best-effort, matching `delete_cache_team_object` and
|
||||
`delete_cache_key_objects`.
|
||||
|
||||
Every caller runs this after its own write has already committed, and the in-memory entry is
|
||||
dropped before the Redis round trip. Letting a cache-backend error raise here therefore reports
|
||||
failure for work that succeeded without making the cache any less stale; the leftover Redis
|
||||
entry expires at its TTL either way.
|
||||
"""
|
||||
key: Final = hashed_token
|
||||
|
||||
user_api_key_cache.delete_cache(key=key)
|
||||
try:
|
||||
user_api_key_cache.delete_cache(key=key)
|
||||
|
||||
## UPDATE REDIS CACHE ##
|
||||
if proxy_logging_obj is not None:
|
||||
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
|
||||
## UPDATE REDIS CACHE ##
|
||||
if proxy_logging_obj is not None:
|
||||
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: a cache error must not fail a committed write
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to invalidate cached key entry %s; a stale key object may be served until its TTL expires: %s",
|
||||
key,
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
async def delete_cache_key_objects(
|
||||
|
|
|
|||
|
|
@ -21,14 +21,13 @@ import litellm
|
|||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
AUTO_ROUTED_REQUEST_METADATA_KEY,
|
||||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
|
||||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
LITELLM_DETAILED_TIMING,
|
||||
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
|
||||
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD,
|
||||
STREAM_SSE_DATA_PREFIX,
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES,
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS,
|
||||
|
|
@ -39,6 +38,7 @@ from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
|
|||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
get_supported_openai_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
||||
|
|
@ -1302,15 +1302,51 @@ def _uncached_input_cost(
|
|||
return input_cost - (cache_read_cost or 0.0) - (cache_creation_cost or 0.0)
|
||||
|
||||
|
||||
_ZERO_COST_BREAKDOWN: Final = CostBreakdownHeaderValues(
|
||||
original_cost=0.0,
|
||||
discount_amount=0.0,
|
||||
margin_total_amount=0.0,
|
||||
margin_percent=0.0,
|
||||
input_cost=0.0,
|
||||
output_cost=0.0,
|
||||
tool_usage_cost=0.0,
|
||||
)
|
||||
"""The component split a call priced at zero advertises, so a client reading the cost headers off a
|
||||
read or management route still finds the whole family rather than a partially populated one."""
|
||||
|
||||
|
||||
def _totals_to_zero(response_cost: float | str | None) -> bool:
|
||||
"""Whether the total these headers carry is zero, counting a total no route ever priced as one.
|
||||
|
||||
A component split is only reported as zero alongside a total that agrees with it, so a read
|
||||
that did price normally never advertises a real total beside an all-zero split.
|
||||
"""
|
||||
if response_cost is None or response_cost == "":
|
||||
return True
|
||||
try:
|
||||
return float(response_cost) == 0.0
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def _get_cost_breakdown_from_logging_obj(
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None,
|
||||
response_cost: float | str | None = None,
|
||||
) -> CostBreakdownHeaderValues:
|
||||
"""Extract discount, margin, and per-component cost information from logging object's cost breakdown."""
|
||||
"""Extract discount, margin, and per-component cost information from logging object's cost breakdown.
|
||||
|
||||
A non-inference call that priced at zero never records a breakdown, so its components are
|
||||
reported as zero here. Any such call that did price normally (retrieving a background response,
|
||||
and the cost poller's read of one) reports the breakdown it stored, or nothing at all when the
|
||||
breakdown has not landed yet.
|
||||
"""
|
||||
if not litellm_logging_obj or not hasattr(litellm_logging_obj, "cost_breakdown"):
|
||||
return CostBreakdownHeaderValues()
|
||||
|
||||
cost_breakdown: Final = litellm_logging_obj.cost_breakdown
|
||||
if not cost_breakdown:
|
||||
if litellm_logging_obj.call_type in NON_INFERENCE_CALL_TYPES and _totals_to_zero(response_cost):
|
||||
return _ZERO_COST_BREAKDOWN
|
||||
return CostBreakdownHeaderValues()
|
||||
|
||||
return CostBreakdownHeaderValues(
|
||||
|
|
@ -1459,7 +1495,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
exclude_values: Final = {"", None, "None"}
|
||||
hidden_params = hidden_params or {}
|
||||
|
||||
cost_breakdown: Final = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=litellm_logging_obj)
|
||||
cost_breakdown: Final = _get_cost_breakdown_from_logging_obj(
|
||||
litellm_logging_obj=litellm_logging_obj, response_cost=response_cost
|
||||
)
|
||||
|
||||
# Calculate updated spend for header (include current response_cost)
|
||||
current_spend: Final = user_api_key_dict.spend or 0.0
|
||||
|
|
@ -2036,54 +2074,6 @@ class ProxyBaseLLMRequestProcessing:
|
|||
return deployment
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_router_selected_model_name(
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None,
|
||||
) -> str | None:
|
||||
"""Model group an auto-routing strategy selected, or None if none fired.
|
||||
|
||||
The marker and ``deployment_model_name`` are written by different bucket
|
||||
resolvers (``get_or_create_metadata_bucket`` vs
|
||||
``_get_router_metadata_variable_name``), so they can land in different
|
||||
buckets on the same request. Resolve each across both.
|
||||
"""
|
||||
litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if not isinstance(litellm_params, dict):
|
||||
return None
|
||||
buckets: Final = tuple(
|
||||
bucket for key in ("litellm_metadata", "metadata") if isinstance(bucket := litellm_params.get(key), dict)
|
||||
)
|
||||
if not any(bucket.get(AUTO_ROUTED_REQUEST_METADATA_KEY) is True for bucket in buckets):
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
model_group
|
||||
for bucket in buckets
|
||||
if isinstance(model_group := bucket.get("deployment_model_name"), str) and model_group
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def set_router_selected_model_field(
|
||||
*,
|
||||
response_obj: object,
|
||||
router_model_name: str | None,
|
||||
) -> None:
|
||||
if not router_model_name:
|
||||
return
|
||||
if isinstance(response_obj, dict):
|
||||
response_obj[ROUTER_MODEL_NAME_RESPONSE_FIELD] = router_model_name
|
||||
return
|
||||
try:
|
||||
setattr(response_obj, ROUTER_MODEL_NAME_RESPONSE_FIELD, router_model_name)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
verbose_proxy_logger.debug(
|
||||
"Could not set %s on response object of type %s",
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD,
|
||||
type(response_obj),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _response_cost_from_logging_obj(
|
||||
*,
|
||||
|
|
@ -2582,20 +2572,21 @@ class ProxyBaseLLMRequestProcessing:
|
|||
log_context=f"litellm_call_id={logging_obj.litellm_call_id}",
|
||||
return_raw_model_name=_should_return_raw_model_name(self.data),
|
||||
)
|
||||
self.set_router_selected_model_field(
|
||||
response_obj=response,
|
||||
router_model_name=self.get_router_selected_model_name(logging_obj),
|
||||
)
|
||||
|
||||
hidden_params = get_hidden_params_dict(response) # get any updated response headers
|
||||
additional_headers = hidden_params.get("additional_headers", {}) or {}
|
||||
|
||||
recover_response_cost: Final = not response_cost and hidden_params.get("response_cost") is None
|
||||
llm_cost_for_headers: Final = (
|
||||
computed_cost_for_headers: Final = (
|
||||
self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) or ""
|
||||
if recover_response_cost
|
||||
else response_cost
|
||||
)
|
||||
llm_cost_for_headers: Final = (
|
||||
0.0
|
||||
if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response)
|
||||
else computed_cost_for_headers
|
||||
)
|
||||
_, request_metadata_bucket = get_or_create_metadata_bucket(self.data)
|
||||
guardrail_cost_for_headers: Final = guardrail_information_cost(
|
||||
request_metadata_bucket.get("standard_logging_guardrail_information")
|
||||
|
|
|
|||
|
|
@ -14,11 +14,36 @@ _NEWRELIC_VAR_PREFIX: Final = "newrelic_"
|
|||
|
||||
|
||||
def callback_config_error(callback_name: str | None, callback_vars: Mapping[str, str] | None) -> str | None:
|
||||
if callback_name != _NEWRELIC_CALLBACK or not callback_vars:
|
||||
if not callback_vars:
|
||||
return None
|
||||
env_error: Final = _langfuse_environment_error(callback_vars)
|
||||
if env_error is not None:
|
||||
return env_error
|
||||
if callback_name != _NEWRELIC_CALLBACK:
|
||||
return None
|
||||
return _newrelic_config_error(callback_vars)
|
||||
|
||||
|
||||
def _langfuse_environment_error(callback_vars: Mapping[str, str]) -> str | None:
|
||||
"""Reject langfuse_environment values Langfuse ingestion would drop.
|
||||
|
||||
Accepting an invalid value here would 200 the config write and then
|
||||
silently lose every trace for that key/team at request time.
|
||||
"""
|
||||
value: Final = callback_vars.get("langfuse_environment")
|
||||
if value is None:
|
||||
return None
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
validate_langfuse_environment_value,
|
||||
)
|
||||
|
||||
try:
|
||||
validate_langfuse_environment_value(value)
|
||||
except ValueError as e:
|
||||
return str(e)
|
||||
return None
|
||||
|
||||
|
||||
def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None:
|
||||
"""Validate every ``logging`` entry of a team/key metadata payload."""
|
||||
if not metadata:
|
||||
|
|
|
|||
79
litellm/proxy/common_utils/healthy_model_filter.py
Normal file
79
litellm/proxy/common_utils/healthy_model_filter.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""Opt-in health filtering shared by the model listing endpoints.
|
||||
|
||||
`/v1/models`, `GET /v1/models/{id}` and `/v1/model/info` hide models whose
|
||||
backing deployments are all marked unhealthy by background health checks, either
|
||||
per request via `healthy_only=true` or proxy-wide via
|
||||
`general_settings.model_list_healthy_only: true`. Both are opt-in: with neither
|
||||
set the listings are returned unfiltered and no health lookup runs at all.
|
||||
|
||||
The proxy-wide setting is what an operator turns on so every client (UI, SDK,
|
||||
raw API) sees only reachable models without having to pass the query parameter.
|
||||
It also makes the background health check loop keep the deployment health cache
|
||||
populated, so `background_health_checks: true` is the only other setting needed.
|
||||
The per-request parameter reads that same cache, so on its own it needs the
|
||||
cache to be filled by either this setting or `enable_health_check_routing`.
|
||||
|
||||
Filtering is presentation-only and always fails open: it answers "should this
|
||||
model be advertised?", never "should a request for it be attempted?". A hidden
|
||||
model stays callable, and an absent, stale or empty health state hides nothing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
MODEL_LIST_HEALTHY_ONLY_SETTING: Final = "model_list_healthy_only"
|
||||
|
||||
|
||||
def is_healthy_only_listing_default(general_settings: Mapping[str, object]) -> bool:
|
||||
"""Whether `model_list_healthy_only` filters every listing on this proxy.
|
||||
|
||||
Only a real `true` counts, so a quoted YAML value never silently starts
|
||||
hiding models. This also tells the background health check loop to keep the
|
||||
deployment health cache populated, which is the state the filter reads.
|
||||
"""
|
||||
return general_settings.get(MODEL_LIST_HEALTHY_ONLY_SETTING, False) is True
|
||||
|
||||
|
||||
def is_healthy_only_enabled(
|
||||
healthy_only: bool | None,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> bool:
|
||||
"""Whether the health filter applies to this request.
|
||||
|
||||
The per-request `healthy_only=true` and the proxy-wide
|
||||
`model_list_healthy_only` setting are independent opt-ins: either one turns
|
||||
the filter on, and a request cannot turn the proxy-wide setting back off
|
||||
(`healthy_only=false` is the unset default, indistinguishable from absent).
|
||||
"""
|
||||
if healthy_only:
|
||||
return True
|
||||
return is_healthy_only_listing_default(general_settings)
|
||||
|
||||
|
||||
async def get_hidden_unhealthy_model_names(
|
||||
healthy_only: bool | None,
|
||||
general_settings: Mapping[str, object],
|
||||
llm_router: Router | None,
|
||||
) -> set[str]:
|
||||
"""Model names to hide from a listing, empty when the filter is off.
|
||||
|
||||
Empty is also the fail-open answer whenever the router cannot report health
|
||||
(no router, no background health checks, stale state, `allowed_fails_policy`
|
||||
configured), so callers apply it unconditionally and simply hide nothing.
|
||||
"""
|
||||
if llm_router is None or not is_healthy_only_enabled(healthy_only, general_settings):
|
||||
return set()
|
||||
unhealthy_names: Final = await llm_router.async_get_fully_unhealthy_model_names()
|
||||
if not unhealthy_names:
|
||||
verbose_proxy_logger.debug(
|
||||
"healthy-only model listing is enabled but no unhealthy deployment state is "
|
||||
"available (requires background_health_checks); returning unfiltered model list"
|
||||
)
|
||||
return unhealthy_names
|
||||
|
|
@ -16,6 +16,10 @@ if TYPE_CHECKING:
|
|||
# Azure Content Safety APIs have a 10,000 character limit per request.
|
||||
AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000
|
||||
|
||||
# Azure Content Safety bills text in 1,000-character "text records"; a submitted
|
||||
# chunk of N characters consumes ceil(N / 1000) text records.
|
||||
AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000
|
||||
|
||||
|
||||
class AzureGuardrailBase:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,7 +3,10 @@
|
|||
Azure Prompt Shield Native Guardrail Integrationfor LiteLLM
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast
|
||||
import math
|
||||
from collections.abc import Mapping, MutableMapping
|
||||
from contextvars import ContextVar
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NoReturn, cast
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -12,14 +15,24 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
|
||||
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT,
|
||||
azure_prompt_shield_guardrail_cost,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import (
|
||||
CallTypesLiteral,
|
||||
GenericGuardrailAPIInputs,
|
||||
GuardrailTracingDetail,
|
||||
)
|
||||
|
||||
from .base import AzureGuardrailBase
|
||||
from .base import AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH, AzureGuardrailBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_prompt_shield import (
|
||||
AzurePromptShieldGuardrailResponse,
|
||||
|
|
@ -27,6 +40,77 @@ if TYPE_CHECKING:
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
||||
# Per-invocation billing counters. A ContextVar rather than request metadata: the
|
||||
# decorator can swap out ``request_data``, metadata is client-forgeable, and
|
||||
# concurrent guardrails run in separate tasks with their own context copy.
|
||||
_billing_usage_stash: Final[ContextVar[dict[str, int] | None]] = ContextVar( # mutable-ok: task-local stash
|
||||
"azure_prompt_shield_billing_usage", default=None
|
||||
)
|
||||
|
||||
|
||||
def _resolved_secret_value(value: object) -> object:
|
||||
"""Resolve ``os.environ/<VAR>`` references the way guardrail api_key/api_base
|
||||
are resolved; any other value passes through unchanged. A reference that
|
||||
resolves to nothing raises instead of silently disabling pricing, so an
|
||||
intended-paid deployment fails fast rather than starting in usage-only mode."""
|
||||
if isinstance(value, str) and value.startswith("os.environ/"):
|
||||
resolved: Final = get_secret_str(value)
|
||||
if resolved is None or not resolved.strip():
|
||||
raise ValueError(f"Azure Prompt Shield: {value!r} resolves to an unset or blank environment variable")
|
||||
return resolved
|
||||
return value
|
||||
|
||||
|
||||
def _updated_param(litellm_params: "LitellmParams | dict", key: str) -> object: # mutable-ok: DB dict
|
||||
"""Read one param from a Mapping or a pydantic object, including pydantic
|
||||
extras (cost_tier / price_per_1000_text_records live there), which the base
|
||||
class ``vars()`` loop never sees."""
|
||||
if isinstance(litellm_params, Mapping):
|
||||
return litellm_params.get(key)
|
||||
return getattr(litellm_params, key, None)
|
||||
|
||||
|
||||
def _resolved_cost_tier(raw: object) -> str | None:
|
||||
"""Normalize the configured cost_tier to 'free' / 'paid' / None."""
|
||||
value: Final = _resolved_secret_value(raw)
|
||||
if value is None or (isinstance(value, str) and not value.strip()):
|
||||
return None
|
||||
tier: Final = str(value).strip().lower()
|
||||
if tier not in ("free", "paid"):
|
||||
raise ValueError(f"Azure Prompt Shield: cost_tier must be 'free' or 'paid', got {value!r}")
|
||||
return tier
|
||||
|
||||
|
||||
def _resolved_price(raw: object, cost_tier: str | None) -> float | None:
|
||||
"""Normalize price_per_1000_text_records and validate it against the tier.
|
||||
|
||||
A 'paid' tier requires a positive price so a misconfigured deployment fails at
|
||||
startup instead of silently reporting a wrong cost; an omitted price with no
|
||||
tier means usage-only tracking (no cost estimate)."""
|
||||
value: Final = _resolved_secret_value(raw)
|
||||
price: Final = _price_from_value(value)
|
||||
if cost_tier == "paid" and (price is None or price <= 0):
|
||||
raise ValueError("Azure Prompt Shield: cost_tier 'paid' requires a positive price_per_1000_text_records")
|
||||
return price
|
||||
|
||||
|
||||
def _price_from_value(value: object) -> float | None:
|
||||
"""Parse a resolved price value into a float; None for an unset/blank value."""
|
||||
if value is None or (isinstance(value, str) and not value.strip()):
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||||
raise TypeError(f"Azure Prompt Shield: price_per_1000_text_records must be a number, got {value!r}")
|
||||
try:
|
||||
price: Final = float(value)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Azure Prompt Shield: price_per_1000_text_records must be a number, got {value!r}") from e
|
||||
if not math.isfinite(price) or price < 0:
|
||||
raise ValueError(
|
||||
f"Azure Prompt Shield: price_per_1000_text_records must be a finite, non-negative number, got {value!r}"
|
||||
)
|
||||
return price
|
||||
|
||||
|
||||
class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrail):
|
||||
"""
|
||||
LiteLLM Built-in Guardrail for Azure Content Safety Guardrail (Prompt Shield).
|
||||
|
|
@ -61,9 +145,20 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
# Plain (non-Final) attributes: ``update_in_memory_litellm_params``
|
||||
# re-resolves them when the guardrail is updated in place.
|
||||
self.cost_tier: str | None = _resolved_cost_tier(kwargs.get("cost_tier"))
|
||||
self.price_per_1000_text_records: float | None = _resolved_price(
|
||||
kwargs.get("price_per_1000_text_records"), self.cost_tier
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Initialized Azure Prompt Shield Guardrail: %s", guardrail_name)
|
||||
|
||||
async def async_make_request(self, user_prompt: str) -> "AzurePromptShieldGuardrailResponse":
|
||||
async def async_make_request(
|
||||
self,
|
||||
user_prompt: str,
|
||||
usage_accumulator: MutableMapping[str, int], # mutable-ok: callee-filled accumulator
|
||||
) -> "AzurePromptShieldGuardrailResponse":
|
||||
"""
|
||||
Make a request to the Azure Prompt Shield API.
|
||||
|
||||
|
|
@ -71,6 +166,13 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
that respect the Azure Content Safety 10 000-character limit. Each
|
||||
chunk is analysed independently; an attack in *any* chunk raises
|
||||
an HTTPException immediately.
|
||||
|
||||
``usage_accumulator`` collects billable usage per SUBMITTED chunk:
|
||||
``requests`` (Azure API calls), ``input_characters``, and
|
||||
``text_records`` (ceil(chunk_chars / 1000), Azure's billing unit).
|
||||
A chunk that triggers an intervention was still submitted and billed,
|
||||
so it is counted before the block is raised; chunks after it are
|
||||
never submitted and never counted.
|
||||
"""
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_prompt_shield import (
|
||||
AzurePromptShieldGuardrailRequestBody,
|
||||
|
|
@ -89,6 +191,12 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
|
||||
last_response = cast(AzurePromptShieldGuardrailResponse, response_json)
|
||||
|
||||
usage_accumulator["requests"] = usage_accumulator.get("requests", 0) + 1
|
||||
usage_accumulator["input_characters"] = usage_accumulator.get("input_characters", 0) + len(chunk)
|
||||
usage_accumulator[AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT] = usage_accumulator.get(
|
||||
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT, 0
|
||||
) + math.ceil(len(chunk) / AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH)
|
||||
|
||||
if last_response["userPromptAnalysis"].get("attackDetected"):
|
||||
verbose_proxy_logger.warning(
|
||||
"Azure Prompt Shield: Attack detected in chunk of length %d",
|
||||
|
|
@ -114,9 +222,14 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
for text in inputs.get("texts") or ():
|
||||
if text:
|
||||
await self.async_make_request(user_prompt=text)
|
||||
_billing_usage_stash.set(None)
|
||||
usage: Final[dict[str, int]] = {} # mutable-ok: per-invocation billing accumulator
|
||||
try:
|
||||
for text in inputs.get("texts") or ():
|
||||
if text:
|
||||
await self.async_make_request(user_prompt=text, usage_accumulator=usage)
|
||||
finally:
|
||||
self._record_billing_usage(usage)
|
||||
return inputs
|
||||
|
||||
@log_guardrail_information
|
||||
|
|
@ -132,6 +245,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
|
||||
Raises HTTPException if content should be blocked.
|
||||
"""
|
||||
_billing_usage_stash.set(None)
|
||||
verbose_proxy_logger.debug(
|
||||
"Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
|
||||
call_type,
|
||||
|
|
@ -144,13 +258,132 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
|
|||
|
||||
if user_prompt:
|
||||
verbose_proxy_logger.debug("Azure Prompt Shield: User prompt: %s", user_prompt)
|
||||
await self.async_make_request(
|
||||
user_prompt=user_prompt,
|
||||
)
|
||||
usage: Final[dict[str, int]] = {} # mutable-ok: per-invocation billing accumulator
|
||||
try:
|
||||
await self.async_make_request(
|
||||
user_prompt=user_prompt,
|
||||
usage_accumulator=usage,
|
||||
)
|
||||
finally:
|
||||
self._record_billing_usage(usage)
|
||||
else:
|
||||
verbose_proxy_logger.warning("Azure Prompt Shield: No user prompt found")
|
||||
return None
|
||||
|
||||
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | dict") -> None: # mutable-ok: DB dict
|
||||
"""Apply updated params in place, re-resolving billing and credentials.
|
||||
|
||||
Pricing is read via ``_updated_param`` (the values are pydantic extras, and
|
||||
the immediate PUT sync hands this method the raw DB dict). Pricing and any
|
||||
``os.environ/`` credential references are validated and resolved BEFORE any
|
||||
state is mutated, so an invalid update leaves the running guardrail
|
||||
untouched and a raw reference never overwrites a resolved credential.
|
||||
"""
|
||||
cost_tier: Final = _resolved_cost_tier(_updated_param(litellm_params, "cost_tier"))
|
||||
price: Final = _resolved_price(_updated_param(litellm_params, "price_per_1000_text_records"), cost_tier)
|
||||
resolved_credentials: dict[str, object] = {} # mutable-ok: staged before mutation
|
||||
for cred_key in ("api_key", "api_base"):
|
||||
cred_value = _updated_param(litellm_params, cred_key)
|
||||
if isinstance(cred_value, str) and cred_value.startswith("os.environ/"):
|
||||
resolved_credentials[cred_key] = _resolved_secret_value(cred_value)
|
||||
if isinstance(litellm_params, Mapping):
|
||||
for key, value in litellm_params.items():
|
||||
setattr(self, key, resolved_credentials.get(key, value))
|
||||
else:
|
||||
super().update_in_memory_litellm_params(litellm_params)
|
||||
for cred_key, cred_value in resolved_credentials.items():
|
||||
setattr(self, cred_key, cred_value)
|
||||
self.cost_tier = cost_tier
|
||||
self.price_per_1000_text_records = price
|
||||
|
||||
def _record_billing_usage(self, usage: Mapping[str, int]) -> None:
|
||||
"""Stash this invocation's usage counters for the ``_process_*`` call the
|
||||
decorator runs next in the same asyncio task; overwrites any leftover."""
|
||||
_billing_usage_stash.set(dict(usage) if usage else None) # mutable-ok: fresh snapshot, popped by _process_*
|
||||
|
||||
def _pop_billing_tracing_detail(self) -> GuardrailTracingDetail | None:
|
||||
"""Build the billing tracing detail from the stashed usage counters, priced
|
||||
with the configured tier/price. ``guardrail_cost_in_spend=False`` keeps the
|
||||
estimated cost out of ``response_cost`` and budget enforcement: Azure
|
||||
guardrail cost is reported on logs, OTEL spans, and the UI, never billed
|
||||
against team/user/key budgets (LIT-5917)."""
|
||||
usage: Final = _billing_usage_stash.get()
|
||||
_billing_usage_stash.set(None)
|
||||
if not usage:
|
||||
return None
|
||||
cost: Final = azure_prompt_shield_guardrail_cost(
|
||||
usage_units=usage,
|
||||
cost_tier=self.cost_tier,
|
||||
price_per_1000_text_records=self.price_per_1000_text_records,
|
||||
)
|
||||
if cost is None:
|
||||
return GuardrailTracingDetail(guardrail_usage=usage)
|
||||
return GuardrailTracingDetail(
|
||||
guardrail_usage=usage,
|
||||
guardrail_cost=cost,
|
||||
guardrail_cost_in_spend=False,
|
||||
)
|
||||
|
||||
def _process_response(
|
||||
self,
|
||||
response: dict | None, # mutable-ok: matches CustomGuardrail._process_response signature
|
||||
request_data: dict, # mutable-ok: matches CustomGuardrail._process_response signature
|
||||
start_time: float | None = None,
|
||||
end_time: float | None = None,
|
||||
duration: float | None = None,
|
||||
event_type: GuardrailEventHooks | None = None,
|
||||
original_inputs: dict | None = None, # mutable-ok: matches CustomGuardrail._process_response signature
|
||||
) -> dict | None: # mutable-ok: matches CustomGuardrail._process_response return
|
||||
"""Override to attach the Azure billing tracing detail (usage counters and
|
||||
estimated cost) and the ``azure`` provider label to the recorded guardrail
|
||||
information. Follows the OpenAI moderation override pattern
|
||||
(openai/moderations.py)."""
|
||||
guardrail_response: Final[dict | str] = ( # mutable-ok: mirrors CustomGuardrail._process_response
|
||||
("mask" if self._inputs_were_modified(original_inputs, response) else "allow")
|
||||
if original_inputs is not None and isinstance(response, dict)
|
||||
else ({} if response is None else response) # mutable-ok: empty placeholder, never mutated
|
||||
)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=guardrail_response,
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
duration=duration,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
event_type=event_type,
|
||||
guardrail_provider="azure",
|
||||
tracing_detail=self._pop_billing_tracing_detail(),
|
||||
)
|
||||
return response
|
||||
|
||||
def _process_error(
|
||||
self,
|
||||
e: Exception,
|
||||
request_data: dict, # mutable-ok: matches CustomGuardrail._process_error signature
|
||||
start_time: float | None = None,
|
||||
end_time: float | None = None,
|
||||
duration: float | None = None,
|
||||
event_type: GuardrailEventHooks | None = None,
|
||||
) -> NoReturn:
|
||||
"""Override to attach the Azure billing tracing detail to the blocked/error
|
||||
guardrail record; a chunk that triggered an intervention was still submitted
|
||||
to (and billed by) Azure, so its usage is recorded on this path too."""
|
||||
guardrail_status: Final = (
|
||||
"guardrail_intervened" if self._is_guardrail_intervention(e) else "guardrail_failed_to_respond"
|
||||
)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=e,
|
||||
request_data=request_data,
|
||||
guardrail_status=guardrail_status,
|
||||
duration=duration,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
event_type=event_type,
|
||||
guardrail_provider="azure",
|
||||
tracing_detail=self._pop_billing_tracing_detail(),
|
||||
)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type["GuardrailConfigModel"] | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -785,11 +785,30 @@ class InMemoryGuardrailHandler:
|
|||
return None
|
||||
|
||||
# Remove from memory if exists (also removes from callbacks)
|
||||
previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
|
||||
previous_source: Final = self._sources.get(guardrail_id, source)
|
||||
if guardrail_id in self.IN_MEMORY_GUARDRAILS:
|
||||
self.delete_in_memory_guardrail(guardrail_id)
|
||||
|
||||
# Initialize fresh (will add new callback to litellm.callbacks)
|
||||
return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
|
||||
# Initialize fresh (will add new callback to litellm.callbacks). If the new
|
||||
# params are invalid (a raising guardrail __init__), restore the previous
|
||||
# instance instead of leaving the guardrail silently removed: a guardrail
|
||||
# that was enforcing must never fail open because an update was bad.
|
||||
try:
|
||||
return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
|
||||
except Exception:
|
||||
if previous_guardrail is not None:
|
||||
verbose_proxy_logger.exception(
|
||||
"Reinitializing guardrail %s with updated params failed; restoring the previous configuration",
|
||||
guardrail_id,
|
||||
)
|
||||
try:
|
||||
self.initialize_guardrail(
|
||||
guardrail=previous_guardrail, config_file_path=config_file_path, source=previous_source
|
||||
)
|
||||
except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks
|
||||
verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id)
|
||||
raise
|
||||
|
||||
def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,11 +6,15 @@ import random
|
|||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
logger: Final = logging.getLogger(__name__)
|
||||
from litellm.constants import (
|
||||
BACKGROUND_HEALTH_CHECK_MAX_TOKENS,
|
||||
|
|
@ -18,16 +22,29 @@ from litellm.constants import (
|
|||
DEFAULT_HEALTH_CHECK_PROMPT,
|
||||
HEALTH_CHECK_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model
|
||||
from litellm.router_utils.auto_router_model_naming import (
|
||||
StrategyRouterDependency,
|
||||
classify_strategy_router_model,
|
||||
strategy_router_dependencies,
|
||||
)
|
||||
|
||||
ILLEGAL_DISPLAY_PARAMS: Final = [
|
||||
"messages",
|
||||
"api_key",
|
||||
"prompt",
|
||||
"input",
|
||||
"client_secret",
|
||||
"azure_ad_token",
|
||||
"azure_username",
|
||||
"azure_password",
|
||||
"vertex_credentials",
|
||||
"vertex_ai_credentials",
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
"aws_web_identity_token",
|
||||
"extra_headers",
|
||||
"headers",
|
||||
"exception", # internal; not JSON-serializable, never for display
|
||||
"litellm_metadata", # internal tracking metadata with auth objects; not for display
|
||||
]
|
||||
|
|
@ -151,7 +168,7 @@ def health_check_filter_kwargs_from_general_settings(
|
|||
|
||||
|
||||
def filter_deployments_by_id(
|
||||
model_list: list,
|
||||
model_list: Sequence[Mapping[str, object]],
|
||||
) -> list:
|
||||
seen_ids: Final = set()
|
||||
filtered_deployments: Final = []
|
||||
|
|
@ -183,12 +200,240 @@ async def run_with_timeout(task, timeout):
|
|||
return {"error": "Timeout exceeded", "exception": timeout_exception}
|
||||
|
||||
|
||||
def _skips_health_checks(deployment: Mapping[str, object]) -> bool:
|
||||
info: Final = deployment.get("model_info")
|
||||
return bool(info.get("disable_background_health_check", False)) if isinstance(info, Mapping) else False
|
||||
|
||||
|
||||
def _health_check_eligible(
|
||||
model_list: Sequence[Mapping[str, object]], skip_disabled: bool
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""Deployments this run is allowed to contact.
|
||||
|
||||
The one eligibility gate, applied to the requested set and to the pool a router's
|
||||
dependencies are drawn from alike, so an opted-out deployment cannot re-enter through a
|
||||
router that depends on it.
|
||||
"""
|
||||
return tuple(x for x in model_list if not (skip_disabled and _skips_health_checks(x)))
|
||||
|
||||
|
||||
def _deployment_model(deployment: Mapping[str, object]) -> str | None:
|
||||
params: Final = deployment.get("litellm_params")
|
||||
return params.get("model") if isinstance(params, Mapping) else None
|
||||
|
||||
|
||||
def _narrow_to_target(
|
||||
model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""Narrow to the requested deployment. An id matching nothing keeps the whole list."""
|
||||
if model_id is not None:
|
||||
by_id: Final = tuple(x for x in model_list if _deployment_id(x) == model_id)
|
||||
return by_id or tuple(model_list)
|
||||
if model is None:
|
||||
return tuple(model_list)
|
||||
by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model)
|
||||
return by_param or tuple(x for x in model_list if x.get("model_name") == model)
|
||||
|
||||
|
||||
def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool:
|
||||
"""True for strategy-router deployments."""
|
||||
model: Final[object] = litellm_params.get("model", "")
|
||||
return isinstance(model, str) and classify_strategy_router_model(model) is not None
|
||||
|
||||
|
||||
def _is_marker(deployment: Mapping[str, object]) -> bool:
|
||||
params: Final = deployment.get("litellm_params")
|
||||
return isinstance(params, Mapping) and _is_strategy_router_deployment(params)
|
||||
|
||||
|
||||
def _deployment_id(deployment: Mapping[str, object]) -> str | None:
|
||||
info: Final = deployment.get("model_info")
|
||||
ident: Final = info.get("id") if isinstance(info, Mapping) else None
|
||||
return str(ident) if ident else None
|
||||
|
||||
|
||||
def _resolved_deployment_ids(router: "Router", model_name: str) -> frozenset[str] | None:
|
||||
"""Deployment ids backing `model_name`, or None when the name resolves to nothing.
|
||||
|
||||
`get_model_list` composes every channel the request path itself uses (exact name,
|
||||
model_group_alias, routing groups, wildcards); a mirror of any one channel would call a
|
||||
working tier broken. An alias whose target is gone resolves to nothing, which fails a
|
||||
request exactly like an unknown name.
|
||||
"""
|
||||
resolved: Final = router.get_model_list(model_name=model_name)
|
||||
if not resolved:
|
||||
return None
|
||||
return frozenset(ident for entry in resolved if (ident := _deployment_id(entry)))
|
||||
|
||||
|
||||
def _dependency_failure(
|
||||
dependency: StrategyRouterDependency,
|
||||
router: "Router",
|
||||
unhealthy_ids: frozenset[str],
|
||||
) -> str | None:
|
||||
"""Why this dependency makes its router unable to serve, or None when it does not.
|
||||
|
||||
A name reds its router only when *every* deployment behind it is known unhealthy. One
|
||||
replica this run never judged, hidden from the caller or opted out of health checks, can
|
||||
still serve what the dead one drops, so partial evidence leaves the verdict green.
|
||||
"""
|
||||
resolved: Final = _resolved_deployment_ids(router, dependency.model_name)
|
||||
if resolved is None:
|
||||
return f"{dependency.role} model '{dependency.model_name}' matches no deployment on this proxy"
|
||||
if not resolved or not resolved <= unhealthy_ids:
|
||||
return None
|
||||
return f"{dependency.role} model '{dependency.model_name}' has no healthy deployment"
|
||||
|
||||
|
||||
def _strategy_router_dependency_error(
|
||||
deployment: Mapping[str, object],
|
||||
router: "Router",
|
||||
unhealthy_ids: frozenset[str],
|
||||
) -> str | None:
|
||||
"""The first dependency fault that makes this router unable to serve, if any."""
|
||||
params: Final = deployment.get("litellm_params")
|
||||
if not isinstance(params, Mapping):
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
failure
|
||||
for dependency in strategy_router_dependencies(params)
|
||||
if (failure := _dependency_failure(dependency, router, unhealthy_ids))
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _deployments_by_id(
|
||||
universe: Sequence[Mapping[str, object]], ids: frozenset[str]
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""The deployments for `ids`, one row per id.
|
||||
|
||||
Reuses the requested set's own dedupe rule, so an alias that duplicates a row cannot get
|
||||
it probed twice or split a single id's verdict across two disagreeing results.
|
||||
"""
|
||||
matched: Final = tuple(d for d in universe if (uid := _deployment_id(d)) and uid in ids)
|
||||
return tuple(filter_deployments_by_id(model_list=matched))
|
||||
|
||||
|
||||
def _dependency_deployments_to_probe(
|
||||
checked: Sequence[Mapping[str, object]],
|
||||
universe: Sequence[Mapping[str, object]],
|
||||
router: "Router",
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""Deployments backing the checked routers' dependencies that are not already checked.
|
||||
|
||||
Empty on a full-list run, which therefore gains no probe; it is the targeted
|
||||
`/health?model_id=<router>` call the dashboard makes per deployment that needs them,
|
||||
since a router's verdict is a statement about models the request never named. Drawn from
|
||||
`universe`, the caller's access-filtered list, so no deployment is probed that the caller
|
||||
was not already granted. Expansion follows routers through routers, one hop per round,
|
||||
because a child router's own models must be probed for the parent to fail; stopping when
|
||||
a round adds nothing is what makes a router cycle terminate.
|
||||
"""
|
||||
checked_ids: Final = frozenset(cid for d in checked if (cid := _deployment_id(d)))
|
||||
reached = checked_ids # rebind-ok: the sweep's cursor, one hop wider per round
|
||||
frontier = tuple(checked) # rebind-ok: the routers whose dependencies the next round expands
|
||||
for _ in range(len(universe)):
|
||||
names = frozenset(
|
||||
dependency.model_name
|
||||
for deployment in frontier
|
||||
if isinstance(params := deployment.get("litellm_params"), Mapping)
|
||||
for dependency in strategy_router_dependencies(params)
|
||||
)
|
||||
fresh_ids = (
|
||||
frozenset(ident for name in names for ident in (_resolved_deployment_ids(router, name) or ())) - reached
|
||||
)
|
||||
if not fresh_ids:
|
||||
break
|
||||
frontier = _deployments_by_id(universe, fresh_ids)
|
||||
reached = reached | fresh_ids
|
||||
return _deployments_by_id(universe, reached - checked_ids)
|
||||
|
||||
|
||||
def _strategy_router_verdicts(
|
||||
healthy_endpoints: Sequence[Mapping[str, object]],
|
||||
unhealthy_endpoints: Sequence[Mapping[str, object]],
|
||||
checked: Sequence[Mapping[str, object]],
|
||||
router: "Router",
|
||||
) -> Mapping[str, str]:
|
||||
"""The dependency fault, per model id, for every strategy router that cannot serve.
|
||||
|
||||
A marker is filed healthy by `_run_model_health_check` returning `{}`, which says only
|
||||
that nothing was probed. This is where that placeholder becomes a verdict, derived from
|
||||
this run's own results rather than a re-probe or a cache that is empty unless
|
||||
`enable_health_check_routing` is on. A marker never fails a probe of its own, so verdicts
|
||||
settle over rounds, each feeding the last round's reds back in as unhealthy; without that
|
||||
the parent of a red child would stay green. Bounded by the marker count, which is what
|
||||
makes a router cycle terminate green rather than spin.
|
||||
"""
|
||||
by_id: Final = MappingProxyType({i: d for d in checked if (i := _deployment_id(d))})
|
||||
markers: Final = MappingProxyType(
|
||||
{
|
||||
marker_id: by_id[marker_id]
|
||||
for endpoint in healthy_endpoints
|
||||
if isinstance(marker_id := endpoint.get("model_id"), str) and marker_id in by_id
|
||||
if _is_marker(by_id[marker_id])
|
||||
}
|
||||
)
|
||||
probe_failures: Final = frozenset(
|
||||
ident for endpoint in unhealthy_endpoints if isinstance(ident := endpoint.get("model_id"), str)
|
||||
)
|
||||
settled: Mapping[str, str] = MappingProxyType({}) # rebind-ok: the fixed point, a round's verdicts at a time
|
||||
for _ in range(len(markers)):
|
||||
fresh = MappingProxyType(
|
||||
{
|
||||
marker_id: error
|
||||
for marker_id, deployment in markers.items()
|
||||
if marker_id not in settled
|
||||
if (error := _strategy_router_dependency_error(deployment, router, probe_failures | frozenset(settled)))
|
||||
}
|
||||
)
|
||||
if not fresh:
|
||||
break
|
||||
settled = MappingProxyType({**settled, **fresh})
|
||||
return settled
|
||||
|
||||
|
||||
def _finalize_strategy_router_endpoints(
|
||||
healthy_endpoints: Sequence[Mapping[str, object]],
|
||||
unhealthy_endpoints: Sequence[Mapping[str, object]],
|
||||
checked: Sequence[Mapping[str, object]],
|
||||
router: "Router | None",
|
||||
dependency_probes: Sequence[Mapping[str, object]],
|
||||
) -> tuple[Sequence[Mapping[str, object]], Sequence[Mapping[str, object]]]:
|
||||
"""Apply router verdicts, then drop the deployments probed only to reach them.
|
||||
|
||||
The probes exist to judge the routers that depend on them; reporting them would answer a
|
||||
targeted request with deployments the caller never asked about.
|
||||
"""
|
||||
verdicts: Final = (
|
||||
_strategy_router_verdicts(healthy_endpoints, unhealthy_endpoints, checked, router)
|
||||
if router is not None
|
||||
else MappingProxyType({})
|
||||
)
|
||||
dropped: Final = frozenset(i for d in dependency_probes if (i := _deployment_id(d)))
|
||||
|
||||
def keep(endpoint: Mapping[str, object]) -> bool:
|
||||
model_id: Final = endpoint.get("model_id")
|
||||
return not (isinstance(model_id, str) and model_id in dropped)
|
||||
|
||||
def verdict_for(endpoint: Mapping[str, object]) -> str | None:
|
||||
model_id: Final = endpoint.get("model_id")
|
||||
return verdicts.get(model_id) if isinstance(model_id, str) else None
|
||||
|
||||
kept_healthy: Final = tuple(e for e in healthy_endpoints if keep(e))
|
||||
return (
|
||||
tuple(e for e in kept_healthy if verdict_for(e) is None),
|
||||
tuple(e for e in unhealthy_endpoints if keep(e))
|
||||
+ tuple(
|
||||
dict(e, error=error) # mutable-ok: the /health payload must stay a plain JSON-serializable dict
|
||||
for e in kept_healthy
|
||||
if (error := verdict_for(e)) is not None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _run_model_health_check(model: dict):
|
||||
litellm_params = model["litellm_params"]
|
||||
model_info: Final = model.get("model_info", {})
|
||||
|
|
@ -531,6 +776,7 @@ async def perform_health_check(
|
|||
max_concurrency: int | None = None,
|
||||
instrumentation_context: dict | None = None,
|
||||
health_check_skip_disabled_background_models: bool = False,
|
||||
router: "Router | None" = None,
|
||||
):
|
||||
"""
|
||||
Perform a health check on the system.
|
||||
|
|
@ -567,23 +813,9 @@ async def perform_health_check(
|
|||
|
||||
cycle_start_time: Final = time.monotonic()
|
||||
requested_model_count: Final = len(model_list)
|
||||
|
||||
# Filter by model_id first so a single deployment is checked when id is specified
|
||||
if model_id is not None:
|
||||
_by_id: Final = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id]
|
||||
if _by_id:
|
||||
model_list = _by_id
|
||||
elif model is not None:
|
||||
_new_model_list = [x for x in model_list if x["litellm_params"]["model"] == model]
|
||||
if _new_model_list == []:
|
||||
_new_model_list = [x for x in model_list if x["model_name"] == model]
|
||||
model_list = _new_model_list
|
||||
|
||||
if health_check_skip_disabled_background_models:
|
||||
model_list = [
|
||||
x for x in model_list if not (x.get("model_info") or {}).get("disable_background_health_check", False)
|
||||
]
|
||||
if not model_list:
|
||||
skip_disabled: Final = health_check_skip_disabled_background_models
|
||||
narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id), skip_disabled)
|
||||
if not narrowed:
|
||||
if instrumentation_enabled:
|
||||
logger.debug(
|
||||
"health_check_cycle_skipped source=%s cycle_id=%s reason=no_models_after_filter",
|
||||
|
|
@ -592,11 +824,16 @@ async def perform_health_check(
|
|||
)
|
||||
return [], [], {}
|
||||
|
||||
post_filter_model_count: Final = len(model_list)
|
||||
model_list = filter_deployments_by_id(
|
||||
model_list=model_list
|
||||
) # filter duplicate deployments (e.g. when model alias'es are used)
|
||||
deduped_model_count: Final = len(model_list)
|
||||
post_filter_model_count: Final = len(narrowed)
|
||||
requested: Final = filter_deployments_by_id(model_list=narrowed)
|
||||
deduped_model_count: Final = len(requested)
|
||||
|
||||
dependency_probes: Final = (
|
||||
_dependency_deployments_to_probe(requested, _health_check_eligible(model_list, skip_disabled), router)
|
||||
if router is not None
|
||||
else ()
|
||||
)
|
||||
checked: Final = requested + list(dependency_probes) # mutable-ok: _perform_health_check takes a list
|
||||
|
||||
if instrumentation_enabled:
|
||||
logger.debug(
|
||||
|
|
@ -613,15 +850,20 @@ async def perform_health_check(
|
|||
|
||||
try:
|
||||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
probed_healthy,
|
||||
probed_unhealthy,
|
||||
exceptions_by_model_id,
|
||||
) = await _perform_health_check(
|
||||
model_list,
|
||||
checked,
|
||||
details,
|
||||
max_concurrency=max_concurrency,
|
||||
instrumentation_context=instrumentation_context,
|
||||
)
|
||||
graded_healthy, graded_unhealthy = _finalize_strategy_router_endpoints(
|
||||
probed_healthy, probed_unhealthy, checked, router, dependency_probes
|
||||
)
|
||||
healthy_endpoints: Final = list(graded_healthy)
|
||||
unhealthy_endpoints: Final = list(graded_unhealthy)
|
||||
except Exception:
|
||||
if instrumentation_enabled:
|
||||
logger.exception(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
|
@ -12,6 +12,9 @@ from litellm.constants import (
|
|||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy.health_check import perform_health_check
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
|
||||
class SharedHealthCheckManager:
|
||||
"""
|
||||
|
|
@ -185,6 +188,7 @@ class SharedHealthCheckManager:
|
|||
details: bool = True,
|
||||
max_concurrency: int | None = None,
|
||||
health_check_skip_disabled_background_models: bool = False,
|
||||
router: "Router | None" = None,
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
|
||||
"""
|
||||
Perform health check with shared state coordination.
|
||||
|
|
@ -235,6 +239,7 @@ class SharedHealthCheckManager:
|
|||
details=details,
|
||||
max_concurrency=max_concurrency,
|
||||
health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
|
||||
router=router,
|
||||
)
|
||||
|
||||
# Cache the results
|
||||
|
|
@ -254,6 +259,7 @@ class SharedHealthCheckManager:
|
|||
details=details,
|
||||
max_concurrency=max_concurrency,
|
||||
health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
|
||||
router=router,
|
||||
)
|
||||
|
||||
# Lock not acquired — poll for cached results until the lock
|
||||
|
|
@ -309,6 +315,7 @@ class SharedHealthCheckManager:
|
|||
details=details,
|
||||
max_concurrency=max_concurrency,
|
||||
health_check_skip_disabled_background_models=health_check_skip_disabled_background_models,
|
||||
router=router,
|
||||
)
|
||||
|
||||
async def is_health_check_in_progress(self) -> bool:
|
||||
|
|
|
|||
|
|
@ -1113,6 +1113,7 @@ async def health_endpoint(
|
|||
user_id=user_api_key_dict.user_id,
|
||||
model_id=model_id,
|
||||
max_concurrency=health_check_concurrency,
|
||||
router=llm_router,
|
||||
**_hc_filter,
|
||||
)
|
||||
return _post_process(router_result)
|
||||
|
|
|
|||
|
|
@ -195,7 +195,7 @@ def _models_this_test_can_call(config: RequestComplexityRouterConfig) -> tuple[s
|
|||
model
|
||||
for model in (
|
||||
config.classifier_llm_config.model
|
||||
if config.classifier_type == "llm" and config.classifier_llm_config is not None
|
||||
if config.uses_llm_classifier and config.classifier_llm_config is not None
|
||||
else None,
|
||||
config.embedding_model if config.semantic_keyword_matching else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ All /budget management endpoints
|
|||
|
||||
#### BUDGET TABLE MANAGEMENT ####
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
|
@ -178,13 +179,17 @@ async def update_budget(
|
|||
else {}
|
||||
)
|
||||
|
||||
response: Final = await BudgetRepository(prisma_client).table.update(
|
||||
where={"budget_id": budget_obj.budget_id},
|
||||
data={
|
||||
budget_obj_jsonified: Final[Mapping[str, object]] = jsonify_object(
|
||||
{
|
||||
**budget_obj.model_dump(exclude_unset=True),
|
||||
**recomputed_reset_at,
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
response: Final = await BudgetRepository(prisma_client).table.update(
|
||||
where={"budget_id": budget_obj.budget_id},
|
||||
data=budget_obj_jsonified,
|
||||
)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -252,6 +252,38 @@ def _raise_on_strategy_router_write_violation(
|
|||
)
|
||||
|
||||
|
||||
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
|
||||
_REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm")
|
||||
|
||||
|
||||
def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None:
|
||||
"""Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml.
|
||||
|
||||
Off by default, so deployments keep adding models without limits. When
|
||||
``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added
|
||||
without both rpm and tpm set to a positive value is rejected rather than stored
|
||||
unbounded (or effectively excluded from routing by a zero/negative limit).
|
||||
"""
|
||||
if not enforced:
|
||||
return
|
||||
missing: Final = tuple(
|
||||
field
|
||||
for field in _REQUIRED_RATE_LIMIT_FIELDS
|
||||
if (value := getattr(litellm_params, field)) is None or value <= 0
|
||||
)
|
||||
if not missing:
|
||||
return
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"{' and '.join(missing)} must be set to a positive value when "
|
||||
f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings"
|
||||
),
|
||||
type=ProxyErrorTypes.validation_error.value,
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
param=f"litellm_params.{missing[0]}",
|
||||
)
|
||||
|
||||
|
||||
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
|
||||
|
||||
|
||||
|
|
@ -326,9 +358,10 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
|
|||
raise HTTPException(status_code=400, detail=error)
|
||||
|
||||
|
||||
# The mirrored per-token pricing fields plus the three remaining fields
|
||||
# Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is
|
||||
# what that back-fill targets, so a field left out here is one a PTU deployment still bills.
|
||||
# The mirrored per-token pricing fields plus the remaining rates the public cost map or a
|
||||
# provider default would otherwise supply (the cache back-fills, the Maps grounding rate). An
|
||||
# unset field falls back to those sources, so a field left out here is one a PTU deployment
|
||||
# still bills.
|
||||
# tiered_pricing is the one mirrored field that is a table of ranges, not a rate, so it is stored
|
||||
# empty (see _PTU_EMPTIED_PRICING_FIELDS): its tiers outrank the zeros written beside them, so
|
||||
# dropping it would leave the cost map's tiers billing the traffic the reserved capacity covers.
|
||||
|
|
@ -1748,6 +1781,11 @@ async def add_new_model(
|
|||
existing_params=None,
|
||||
)
|
||||
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
litellm_params=model_params.litellm_params,
|
||||
enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)),
|
||||
)
|
||||
|
||||
model_response: prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None = None
|
||||
# update DB
|
||||
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ This is an enterprise feature and requires a premium license.
|
|||
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from itertools import chain
|
||||
|
|
@ -2375,6 +2376,37 @@ async def get_group(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
def _new_team_request_with_defaults(
|
||||
team_id: str,
|
||||
team_alias: str | None,
|
||||
members_with_roles: Sequence[Member],
|
||||
) -> NewTeamRequest:
|
||||
"""Build the SCIM group's team request, applying litellm.default_team_params
|
||||
(including models) the same way SSO auto-created teams do."""
|
||||
default_params: Final = litellm.default_team_params
|
||||
defaults: Final[Mapping[str, object]] = (
|
||||
deepcopy(default_params)
|
||||
if isinstance(default_params, dict)
|
||||
else default_params.model_dump(exclude_none=True)
|
||||
if default_params is not None
|
||||
else {}
|
||||
)
|
||||
default_metadata: Final = defaults.get("metadata")
|
||||
metadata: Final = {
|
||||
**(default_metadata if isinstance(default_metadata, dict) else {}),
|
||||
SCIM_MANAGED_TEAM_METADATA_KEY: True,
|
||||
}
|
||||
return NewTeamRequest.model_validate(
|
||||
{
|
||||
**defaults,
|
||||
"team_id": team_id,
|
||||
"team_alias": team_alias,
|
||||
"members_with_roles": members_with_roles,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@scim_router.post(
|
||||
"/Groups",
|
||||
response_model=SCIMGroup,
|
||||
|
|
@ -2412,11 +2444,10 @@ async def create_group(
|
|||
|
||||
# Create team in database
|
||||
created_team: Final = await new_team(
|
||||
data=NewTeamRequest(
|
||||
data=_new_team_request_with_defaults(
|
||||
team_id=team_id,
|
||||
team_alias=group.displayName,
|
||||
members_with_roles=members_with_roles,
|
||||
metadata={SCIM_MANAGED_TEAM_METADATA_KEY: True},
|
||||
),
|
||||
http_request=Request(scope={"type": "http", "path": "/scim/v2/Groups"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
|
|
|
|||
|
|
@ -262,6 +262,7 @@ async def add_team_callbacks(
|
|||
- langfuse_secret_key: The secret key for the Langfuse callback
|
||||
- langfuse_secret: The secret for the Langfuse callback
|
||||
- langfuse_host: The host for the Langfuse callback
|
||||
- langfuse_environment: The tracing environment for the Langfuse callback (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)
|
||||
- gcs_bucket_name: The name of the GCS bucket
|
||||
- gcs_path_service_account: The path to the GCS service account
|
||||
- langsmith_api_key: The API key for the Langsmith callback
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import json
|
|||
import math
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
|
||||
|
|
@ -494,8 +495,13 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit: int | None = None,
|
||||
team_member_tpm_limit: int | None = None,
|
||||
team_member_budget_duration: str | None = None,
|
||||
explicitly_set_fields: AbstractSet[str] = frozenset(),
|
||||
) -> dict:
|
||||
"""Create team member budget table with provided limits"""
|
||||
"""Create team member budget table with provided limits.
|
||||
|
||||
The team's own reset period is only inherited when the caller left the
|
||||
member duration out, so an explicit null means "never resets".
|
||||
"""
|
||||
from litellm.proxy._types import BudgetNewRequest
|
||||
from litellm.proxy.management_endpoints.budget_management_endpoints import (
|
||||
new_budget,
|
||||
|
|
@ -509,7 +515,11 @@ class TeamMemberBudgetHandler:
|
|||
# Create budget request with all provided limits
|
||||
budget_request: Final = BudgetNewRequest(
|
||||
budget_id=budget_id,
|
||||
budget_duration=data.budget_duration or team_member_budget_duration,
|
||||
budget_duration=(
|
||||
team_member_budget_duration
|
||||
if "team_member_budget_duration" in explicitly_set_fields
|
||||
else data.budget_duration or team_member_budget_duration
|
||||
),
|
||||
)
|
||||
|
||||
if team_member_budget is not None:
|
||||
|
|
@ -545,8 +555,13 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit: int | None = None,
|
||||
team_member_tpm_limit: int | None = None,
|
||||
team_member_budget_duration: str | None = None,
|
||||
explicitly_set_fields: AbstractSet[str] = frozenset(),
|
||||
) -> dict:
|
||||
"""Upsert team member budget table with provided limits"""
|
||||
"""Upsert team member budget table with provided limits.
|
||||
|
||||
A field the caller explicitly sent as null is written as null, so a
|
||||
team can keep a member budget while dropping its reset period.
|
||||
"""
|
||||
from litellm.proxy._types import BudgetNewRequest
|
||||
from litellm.proxy.management_endpoints.budget_management_endpoints import (
|
||||
update_budget,
|
||||
|
|
@ -560,14 +575,16 @@ class TeamMemberBudgetHandler:
|
|||
# Budget exists - create update request with only provided values
|
||||
budget_request: Final = BudgetNewRequest(budget_id=team_member_budget_id)
|
||||
|
||||
if team_member_budget is not None:
|
||||
if team_member_budget is not None or "team_member_budget" in explicitly_set_fields:
|
||||
budget_request.max_budget = team_member_budget
|
||||
if team_member_rpm_limit is not None:
|
||||
if team_member_rpm_limit is not None or "team_member_rpm_limit" in explicitly_set_fields:
|
||||
budget_request.rpm_limit = team_member_rpm_limit
|
||||
if team_member_tpm_limit is not None:
|
||||
if team_member_tpm_limit is not None or "team_member_tpm_limit" in explicitly_set_fields:
|
||||
budget_request.tpm_limit = team_member_tpm_limit
|
||||
if team_member_budget_duration is not None:
|
||||
if team_member_budget_duration is not None or "team_member_budget_duration" in explicitly_set_fields:
|
||||
budget_request.budget_duration = team_member_budget_duration
|
||||
if team_member_budget_duration is None:
|
||||
budget_request.budget_reset_at = None
|
||||
|
||||
budget_row: Final = await _as_budget_write(update_budget)(
|
||||
budget_obj=budget_request,
|
||||
|
|
@ -593,6 +610,7 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit=team_member_rpm_limit,
|
||||
team_member_tpm_limit=team_member_tpm_limit,
|
||||
team_member_budget_duration=team_member_budget_duration,
|
||||
explicitly_set_fields=explicitly_set_fields,
|
||||
)
|
||||
|
||||
# Remove team member fields from updated_kv
|
||||
|
|
@ -1479,6 +1497,7 @@ async def new_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
explicitly_set_fields=data.model_fields_set,
|
||||
)
|
||||
|
||||
## ADD TO TEAM TABLE
|
||||
|
|
@ -2184,6 +2203,7 @@ async def update_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
explicitly_set_fields=_team_member_fields_in_request,
|
||||
)
|
||||
# Backfill team_memberships for members who joined before the
|
||||
# budget was configured — they won't have a membership row yet.
|
||||
|
|
|
|||
|
|
@ -615,7 +615,7 @@ class VertexPassthroughLoggingHandler:
|
|||
response_cost: Final = litellm.completion_cost(
|
||||
completion_response=litellm_model_response,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ from litellm.types.utils import StandardPassThroughResponseObject
|
|||
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.gemini_passthrough_logging_handler import (
|
||||
GeminiPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.openai_passthrough_logging_handler import (
|
||||
OpenAIPassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -243,6 +246,26 @@ class PassThroughStreamingHandler:
|
|||
)
|
||||
standard_logging_response_object = vertex_passthrough_logging_handler_result["result"]
|
||||
kwargs = vertex_passthrough_logging_handler_result["kwargs"]
|
||||
elif endpoint_type == EndpointType.GEMINI:
|
||||
gemini_passthrough_logging_handler_result: Final = (
|
||||
GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( # pyright: ignore[reportPrivateUsage] # mirrors sibling handler dispatch
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
passthrough_success_handler_obj=passthrough_success_handler_obj,
|
||||
url_route=url_route,
|
||||
request_body=request_body,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
all_chunks=all_chunks,
|
||||
end_time=end_time,
|
||||
model=model,
|
||||
)
|
||||
)
|
||||
standard_logging_response_object = ( # rebind-ok: branch bind in shared if/elif dispatch
|
||||
gemini_passthrough_logging_handler_result["result"]
|
||||
)
|
||||
kwargs = ( # rebind-ok: branch bind in shared if/elif dispatch
|
||||
gemini_passthrough_logging_handler_result["kwargs"]
|
||||
)
|
||||
elif endpoint_type == EndpointType.OPENAI:
|
||||
openai_passthrough_logging_handler_result: Final = (
|
||||
OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks(
|
||||
|
|
|
|||
|
|
@ -323,6 +323,7 @@ def create_versioned_prompt_spec(db_prompt: _PromptRow) -> PromptSpec:
|
|||
prompt_info=prompt_info,
|
||||
created_at=row.created_at,
|
||||
updated_at=row.updated_at,
|
||||
version=row.version,
|
||||
environment=row.environment,
|
||||
created_by=row.created_by,
|
||||
)
|
||||
|
|
@ -334,6 +335,21 @@ class Prompt(BaseModel):
|
|||
prompt_info: PromptInfo | None = None
|
||||
|
||||
|
||||
AMBIGUOUS_PROMPT_DATA_ERROR: Final = (
|
||||
"litellm_params.prompt_id cannot be combined with prompt_data keyed by template name. "
|
||||
'Send a flat template, prompt_data={"content": "...", "metadata": {...}}, together with litellm_params.prompt_id, '
|
||||
'or send prompt_data={"<template_id>": {"content": "...", "metadata": {...}}} without litellm_params.prompt_id.'
|
||||
)
|
||||
|
||||
|
||||
def is_ambiguous_keyed_prompt_data(litellm_params: PromptLiteLLMParams) -> bool:
|
||||
extra_fields: Final = litellm_params.model_extra or {}
|
||||
prompt_data: Final = extra_fields.get("prompt_data")
|
||||
if not litellm_params.prompt_id or not isinstance(prompt_data, dict):
|
||||
return False
|
||||
return bool(prompt_data) and "content" not in prompt_data
|
||||
|
||||
|
||||
class PatchPromptRequest(BaseModel):
|
||||
litellm_params: PromptLiteLLMParams | None = None
|
||||
prompt_info: PromptInfo | None = None
|
||||
|
|
@ -737,11 +753,9 @@ async def create_prompt(
|
|||
-d '{
|
||||
"prompt_id": "my_prompt",
|
||||
"litellm_params": {
|
||||
"prompt_id": "json_prompt",
|
||||
"prompt_id": "my_prompt",
|
||||
"prompt_integration": "dotprompt",
|
||||
### EITHER prompt_directory OR prompt_data MUST BE PROVIDED
|
||||
"prompt_directory": "/path/to/dotprompt/folder",
|
||||
"prompt_data": {"json_prompt": {"content": "This is a prompt", "metadata": {"model": "gpt-4"}}}
|
||||
"prompt_data": {"content": "This is a prompt", "metadata": {"model": "gpt-4"}}
|
||||
},
|
||||
"prompt_info": {
|
||||
"prompt_type": "config"
|
||||
|
|
@ -763,6 +777,9 @@ async def create_prompt(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
if is_ambiguous_keyed_prompt_data(request.litellm_params):
|
||||
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
|
||||
|
||||
try:
|
||||
# Extract environment from request
|
||||
environment: Final = (
|
||||
|
|
@ -857,6 +874,9 @@ async def update_prompt(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
if is_ambiguous_keyed_prompt_data(request.litellm_params):
|
||||
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
|
||||
|
||||
try:
|
||||
# Strip version suffix from prompt_id if present (e.g., "jack_success.v1" -> "jack_success")
|
||||
base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id)
|
||||
|
|
@ -1001,19 +1021,7 @@ async def delete_prompt(
|
|||
# Delete versions from the database (scoped to environment if provided)
|
||||
await _prompt_table(prisma_client).delete_many(where=delete_where)
|
||||
|
||||
# Remove matching prompts from memory — scope to environment if provided
|
||||
if environment:
|
||||
prompts_to_delete: Final = [
|
||||
pid
|
||||
for pid, prompt in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.items()
|
||||
if get_base_prompt_id(prompt_id=pid) == base_prompt_id and prompt.environment == environment
|
||||
]
|
||||
for pid in prompts_to_delete:
|
||||
del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[pid]
|
||||
if pid in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt:
|
||||
del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[pid]
|
||||
else:
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id)
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id, environment=environment or None)
|
||||
|
||||
env_msg: Final = f" from {environment}" if environment else ""
|
||||
return {"message": f"Prompt {base_prompt_id} deleted successfully{env_msg}"}
|
||||
|
|
@ -1025,15 +1033,8 @@ async def delete_prompt(
|
|||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
def _reload_prompt_in_registry(
|
||||
registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec
|
||||
) -> PromptSpec:
|
||||
"""Remove stale entry and re-initialize the prompt in the in-memory registry."""
|
||||
if versioned_id in registry.IN_MEMORY_PROMPTS:
|
||||
del registry.IN_MEMORY_PROMPTS[versioned_id]
|
||||
if versioned_id in registry.prompt_id_to_custom_prompt:
|
||||
del registry.prompt_id_to_custom_prompt[versioned_id]
|
||||
initialized: Final = registry.initialize_prompt(prompt=updated_prompt_spec, config_file_path=None)
|
||||
def _reload_prompt_in_registry(registry: "InMemoryPromptRegistry", updated_prompt_spec: PromptSpec) -> PromptSpec:
|
||||
initialized: Final = registry.reload_prompt(prompt=updated_prompt_spec)
|
||||
if initialized is None:
|
||||
raise HTTPException(status_code=500, detail="Failed to patch prompt")
|
||||
return initialized
|
||||
|
|
@ -1086,6 +1087,9 @@ async def patch_prompt(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
if request.litellm_params is not None and is_ambiguous_keyed_prompt_data(request.litellm_params):
|
||||
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
|
||||
|
||||
try:
|
||||
# Resolve the target row: find the latest version in the given environment
|
||||
base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id)
|
||||
|
|
@ -1123,25 +1127,15 @@ async def patch_prompt(
|
|||
detail="Cannot update config prompts.",
|
||||
)
|
||||
|
||||
# Use existing prompt from memory or build from DB row for field merging
|
||||
if existing_prompt:
|
||||
current_litellm_params = existing_prompt.litellm_params
|
||||
current_prompt_info = existing_prompt.prompt_info
|
||||
else:
|
||||
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
|
||||
current_litellm_params = current_spec.litellm_params
|
||||
current_prompt_info = current_spec.prompt_info
|
||||
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
|
||||
|
||||
# Update fields if provided
|
||||
updated_litellm_params: Final = (
|
||||
request.litellm_params if request.litellm_params is not None else current_litellm_params
|
||||
request.litellm_params if request.litellm_params is not None else current_spec.litellm_params
|
||||
)
|
||||
|
||||
updated_prompt_info: Final = request.prompt_info if request.prompt_info is not None else current_prompt_info
|
||||
|
||||
# Ensure we have valid litellm_params
|
||||
if updated_litellm_params is None:
|
||||
raise HTTPException(status_code=400, detail="litellm_params cannot be None")
|
||||
updated_prompt_info: Final = (
|
||||
request.prompt_info if request.prompt_info is not None else current_spec.prompt_info
|
||||
)
|
||||
|
||||
# Build update data dict
|
||||
update_data: Final[dict[str, str]] = {
|
||||
|
|
@ -1165,7 +1159,7 @@ async def patch_prompt(
|
|||
|
||||
updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry)
|
||||
|
||||
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec)
|
||||
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, updated_prompt_spec)
|
||||
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -118,7 +118,16 @@ class InMemoryPromptRegistry:
|
|||
verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS")
|
||||
return self.IN_MEMORY_PROMPTS[prompt_id]
|
||||
|
||||
custom_prompt_callback: CustomPromptManagement | None = None
|
||||
parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt)
|
||||
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
|
||||
|
||||
# store references to the prompt in memory
|
||||
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
|
||||
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
|
||||
|
||||
return parsed_prompt
|
||||
|
||||
def _build_prompt_callback(self, prompt: PromptSpec) -> tuple[PromptSpec, CustomPromptManagement]:
|
||||
litellm_params_data: Final = prompt.litellm_params
|
||||
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
|
||||
|
||||
|
|
@ -132,29 +141,48 @@ class InMemoryPromptRegistry:
|
|||
raise ValueError("prompt_integration is required")
|
||||
|
||||
initializer: Final = prompt_initializer_registry.get(prompt_integration)
|
||||
|
||||
if initializer:
|
||||
custom_prompt_callback = initializer(litellm_params, prompt)
|
||||
if not isinstance(custom_prompt_callback, CustomPromptManagement):
|
||||
raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}")
|
||||
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
|
||||
else:
|
||||
if initializer is None:
|
||||
raise ValueError(f"Unsupported prompt: {prompt_integration}")
|
||||
|
||||
custom_prompt_callback: Final = initializer(litellm_params, prompt)
|
||||
if not isinstance(custom_prompt_callback, CustomPromptManagement):
|
||||
raise ValueError( # noqa: TRY004 # prompt endpoints map ValueError to HTTP 400; keep the existing contract
|
||||
f"CustomPromptManagement is required, got {type(custom_prompt_callback)}"
|
||||
)
|
||||
|
||||
parsed_prompt: Final = PromptSpec(
|
||||
prompt_id=prompt_id,
|
||||
prompt_id=prompt.prompt_id,
|
||||
litellm_params=litellm_params,
|
||||
prompt_info=prompt.prompt_info or PromptInfo(prompt_type="config"),
|
||||
created_at=prompt.created_at,
|
||||
updated_at=prompt.updated_at,
|
||||
version=prompt.version,
|
||||
environment=prompt.environment,
|
||||
created_by=prompt.created_by,
|
||||
)
|
||||
return parsed_prompt, custom_prompt_callback
|
||||
|
||||
# store references to the prompt in memory
|
||||
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
|
||||
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
|
||||
def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None:
|
||||
import litellm
|
||||
|
||||
parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt)
|
||||
stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None)
|
||||
self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None)
|
||||
if stale_callback is not None:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback)
|
||||
litellm.logging_callback_manager.add_litellm_callback(new_callback)
|
||||
self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt
|
||||
self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback
|
||||
return parsed_prompt
|
||||
|
||||
def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None:
|
||||
existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id)
|
||||
if existing is None:
|
||||
return self.initialize_prompt(prompt=prompt)
|
||||
if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info:
|
||||
return existing
|
||||
return self.reload_prompt(prompt=prompt)
|
||||
|
||||
def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None:
|
||||
"""
|
||||
Get a prompt by its ID from memory
|
||||
|
|
@ -167,12 +195,22 @@ class InMemoryPromptRegistry:
|
|||
"""
|
||||
return self.prompt_id_to_custom_prompt.get(prompt_id)
|
||||
|
||||
def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]:
|
||||
def remove_prompt(self, prompt_id: str) -> None:
|
||||
import litellm
|
||||
|
||||
self.IN_MEMORY_PROMPTS.pop(prompt_id, None)
|
||||
stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt_id, None)
|
||||
if stale_callback is not None:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback)
|
||||
|
||||
def delete_prompts_by_base_id(self, base_prompt_id: str, environment: str | None = None) -> list[str]:
|
||||
"""
|
||||
Delete all prompts matching the given base prompt ID from memory.
|
||||
Delete all prompts matching the given base prompt ID from memory, along with their
|
||||
registered callbacks; scoped to one environment when given.
|
||||
|
||||
Args:
|
||||
base_prompt_id: The base prompt ID (without version suffix)
|
||||
environment: When set, only delete prompts deployed to this environment
|
||||
|
||||
Returns:
|
||||
List of prompt IDs that were deleted
|
||||
|
|
@ -180,13 +218,14 @@ class InMemoryPromptRegistry:
|
|||
from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id
|
||||
|
||||
prompts_to_delete: Final = [
|
||||
pid for pid in self.IN_MEMORY_PROMPTS if get_base_prompt_id(prompt_id=pid) == base_prompt_id
|
||||
pid
|
||||
for pid, prompt in self.IN_MEMORY_PROMPTS.items()
|
||||
if get_base_prompt_id(prompt_id=pid) == base_prompt_id
|
||||
and (environment is None or prompt.environment == environment)
|
||||
]
|
||||
|
||||
for pid in prompts_to_delete:
|
||||
del self.IN_MEMORY_PROMPTS[pid]
|
||||
if pid in self.prompt_id_to_custom_prompt:
|
||||
del self.prompt_id_to_custom_prompt[pid]
|
||||
self.remove_prompt(prompt_id=pid)
|
||||
|
||||
return prompts_to_delete
|
||||
|
||||
|
|
|
|||
|
|
@ -248,7 +248,6 @@ from litellm.constants import (
|
|||
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD,
|
||||
WEEKLY_SPEND_REPORT_JOB_ID,
|
||||
)
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
|
|
@ -329,6 +328,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.common_utils.healthy_model_filter import (
|
||||
get_hidden_unhealthy_model_names,
|
||||
is_healthy_only_listing_default,
|
||||
)
|
||||
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body,
|
||||
|
|
@ -1502,9 +1505,9 @@ def get_openapi_schema():
|
|||
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
|
||||
|
||||
# Stub unloaded lazy features so they appear as Swagger sections.
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs, loaded_lazy_modules
|
||||
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema)
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema, loaded_lazy_modules(app))
|
||||
openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema)
|
||||
|
||||
# Fix Swagger UI execute path error when server_root_path is set
|
||||
|
|
@ -1534,9 +1537,9 @@ def custom_openapi():
|
|||
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
|
||||
|
||||
# Stub unloaded lazy features so they appear as Swagger sections.
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs, loaded_lazy_modules
|
||||
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema)
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema, loaded_lazy_modules(app))
|
||||
openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema)
|
||||
|
||||
# Fix Swagger UI execute path error when server_root_path is set
|
||||
|
|
@ -3393,9 +3396,7 @@ def _rss_mb_for_log() -> str:
|
|||
return f"{rss_mb:.2f}"
|
||||
|
||||
|
||||
def _is_unexpected_keyword_argument_type_error(exc: BaseException) -> bool:
|
||||
"""True when ``exc`` is a TypeError from passing a kwarg the callee does not accept."""
|
||||
return isinstance(exc, TypeError) and ("unexpected keyword argument" in str(exc).lower())
|
||||
_UNEXPECTED_KWARG: Final = re.compile(r"unexpected keyword argument '(?P<name>[^']+)'")
|
||||
|
||||
|
||||
async def _run_direct_health_check_with_instrumentation(
|
||||
|
|
@ -3404,31 +3405,33 @@ async def _run_direct_health_check_with_instrumentation(
|
|||
max_concurrency: int | None,
|
||||
instrumentation_context: dict,
|
||||
):
|
||||
"""Call ``perform_health_check``, retrying with fewer kwargs on unexpected-kw TypeErrors."""
|
||||
_hc_filter: Final = health_check_filter_kwargs_from_general_settings(general_settings)
|
||||
last_type_error: TypeError | None = None
|
||||
for extra_kwargs in (
|
||||
"""Call ``perform_health_check``, dropping exactly the optional kwarg each TypeError names.
|
||||
|
||||
A callee that predates an argument rejects it by name, so only that one is dropped. A
|
||||
hand-written ladder of combinations would drop working options alongside it, and would
|
||||
need a new rung every time an argument is added.
|
||||
"""
|
||||
optional: Mapping[str, object] = MappingProxyType( # rebind-ok: loses the kwarg the callee rejected
|
||||
{
|
||||
"router": llm_router,
|
||||
"instrumentation_context": instrumentation_context,
|
||||
**_hc_filter,
|
||||
},
|
||||
{"instrumentation_context": instrumentation_context},
|
||||
dict(_hc_filter),
|
||||
{},
|
||||
):
|
||||
**health_check_filter_kwargs_from_general_settings(general_settings),
|
||||
}
|
||||
)
|
||||
for _ in range(len(optional) + 1):
|
||||
try:
|
||||
return await perform_health_check(
|
||||
model_list=model_list,
|
||||
details=details,
|
||||
max_concurrency=max_concurrency,
|
||||
**extra_kwargs,
|
||||
**optional,
|
||||
)
|
||||
except TypeError as e:
|
||||
if not _is_unexpected_keyword_argument_type_error(e):
|
||||
rejected = _UNEXPECTED_KWARG.search(str(e))
|
||||
if rejected is None or rejected["name"] not in optional:
|
||||
raise
|
||||
last_type_error = e
|
||||
assert last_type_error is not None
|
||||
raise last_type_error
|
||||
optional = MappingProxyType({k: v for k, v in optional.items() if k != rejected["name"]})
|
||||
raise AssertionError("perform_health_check rejected every optional argument")
|
||||
|
||||
|
||||
def _schedule_background_health_check_db_save(
|
||||
|
|
@ -3483,6 +3486,13 @@ def _write_health_state_to_router_cache(
|
|||
"""
|
||||
Write deployment health states to the router's health state cache
|
||||
for health-check-driven routing. No-op if the feature is disabled.
|
||||
|
||||
`model_list_healthy_only` reads the same cache to hide unhealthy models from
|
||||
the listing endpoints, so it also keeps the cache populated. That is a pure
|
||||
write: every routing-time reader is itself gated on
|
||||
`enable_health_check_routing`, and the cooldown/failure bookkeeping below
|
||||
stays behind that flag, so routing is untouched when only the listing filter
|
||||
is on.
|
||||
"""
|
||||
from litellm.proxy.health_check import build_deployment_health_states
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
|
|
@ -3493,7 +3503,10 @@ def _write_health_state_to_router_cache(
|
|||
_exceptions: Final[dict] = exceptions_by_model_id or {}
|
||||
|
||||
try:
|
||||
if llm_router is None or not llm_router.enable_health_check_routing:
|
||||
if llm_router is None:
|
||||
return
|
||||
health_check_routing_enabled: Final = llm_router.enable_health_check_routing
|
||||
if not health_check_routing_enabled and not is_healthy_only_listing_default(general_settings):
|
||||
return
|
||||
|
||||
# When health_check_ignore_transient_errors is set, treat 429/408
|
||||
|
|
@ -3516,6 +3529,9 @@ def _write_health_state_to_router_cache(
|
|||
sum(1 for s in states.values() if not s.get("is_healthy")),
|
||||
)
|
||||
|
||||
if not health_check_routing_enabled:
|
||||
return
|
||||
|
||||
for endpoint in unhealthy_endpoints:
|
||||
model_id = endpoint.get("model_id")
|
||||
if not model_id:
|
||||
|
|
@ -3683,6 +3699,7 @@ async def _run_background_health_check():
|
|||
model_list=_llm_model_list,
|
||||
details=details_bool,
|
||||
max_concurrency=health_check_concurrency,
|
||||
router=llm_router,
|
||||
**_hc_filter,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -6268,7 +6285,14 @@ class ProxyConfig:
|
|||
):
|
||||
from litellm.utils import _update_dictionary
|
||||
|
||||
combined_router_settings = _update_dictionary(config_router_settings, db_router_settings.param_value)
|
||||
db_overlay_deferring_empty_lists_to_config: Final = {
|
||||
k: v
|
||||
for k, v in db_router_settings.param_value.items()
|
||||
if not (k in config_router_settings and isinstance(v, list) and len(v) == 0)
|
||||
}
|
||||
combined_router_settings = _update_dictionary(
|
||||
config_router_settings, db_overlay_deferring_empty_lists_to_config
|
||||
)
|
||||
elif config_router_settings is not None and isinstance(config_router_settings, dict):
|
||||
combined_router_settings = config_router_settings
|
||||
elif db_router_settings is not None and isinstance(db_router_settings.param_value, dict):
|
||||
|
|
@ -7232,12 +7256,53 @@ class ProxyConfig:
|
|||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
|
||||
def parse_row(db_prompt: object) -> PromptSpec | None:
|
||||
try:
|
||||
return self._get_prompt_spec_for_db_prompt(db_prompt=db_prompt)
|
||||
except Exception as row_error: # noqa: BLE001 # a malformed row must not block syncing the remaining prompts
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to parse prompt row %s: %s",
|
||||
getattr(db_prompt, "prompt_id", None),
|
||||
row_error,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
prompt_ids_loaded_before_db_read: Final = frozenset(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS)
|
||||
prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many()
|
||||
for prompt in prompts_in_db:
|
||||
# Convert DB object to dict and create versioned prompt_id
|
||||
prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt)
|
||||
IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec)
|
||||
parsed_specs: Final[tuple[PromptSpec, ...]] = tuple(
|
||||
spec for row in prompts_in_db if (spec := parse_row(row)) is not None
|
||||
)
|
||||
newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType(
|
||||
{
|
||||
spec.prompt_id: spec
|
||||
for spec in sorted(
|
||||
parsed_specs,
|
||||
key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"),
|
||||
)
|
||||
}
|
||||
)
|
||||
for prompt_spec in newest_spec_per_id.values():
|
||||
try:
|
||||
IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec)
|
||||
except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to sync prompt %s: %s",
|
||||
prompt_spec.prompt_id,
|
||||
prompt_sync_error,
|
||||
)
|
||||
# An unparsable row still exists in the DB, so skip the sweep rather than unload its in-memory copy
|
||||
every_row_parsed: Final = len(parsed_specs) == len(prompts_in_db)
|
||||
if every_row_parsed:
|
||||
deleted_db_prompt_ids: Final = tuple(
|
||||
prompt_id
|
||||
for prompt_id in prompt_ids_loaded_before_db_read
|
||||
if (loaded_spec := IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.get(prompt_id)) is not None
|
||||
and loaded_spec.prompt_info.prompt_type == "db"
|
||||
and prompt_id not in newest_spec_per_id
|
||||
)
|
||||
for deleted_prompt_id in deleted_db_prompt_ids:
|
||||
IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id=deleted_prompt_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e)
|
||||
|
||||
|
|
@ -7472,11 +7537,9 @@ class ProxyConfig:
|
|||
len(db_search_tools),
|
||||
)
|
||||
|
||||
if llm_router is not None and search_tools:
|
||||
if llm_router is not None:
|
||||
await SearchAPIRouter.update_router_search_tools(router_instance=llm_router, search_tools=search_tools)
|
||||
verbose_proxy_logger.info("Successfully loaded %s search tool(s) into router", len(search_tools))
|
||||
elif llm_router is not None:
|
||||
verbose_proxy_logger.debug("No search tools found in config or database, skipping router update")
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Router not initialized yet, search tools will be added when router is created"
|
||||
|
|
@ -7487,6 +7550,26 @@ class ProxyConfig:
|
|||
"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - %s", e
|
||||
)
|
||||
|
||||
async def reload_search_tools_from_db(self) -> None:
|
||||
"""Refresh this worker's router from the search tools table.
|
||||
|
||||
Driven by the management endpoints so the worker that served the write is correct
|
||||
immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same
|
||||
way as startup, so an admin who excluded search_tools from supported_db_objects opts out.
|
||||
|
||||
Serialized by MODEL_RECONCILE_LOCK for the reason add_deployment documents: the body is a
|
||||
read-modify-write of the shared ``llm_router`` global, so two of them interleaving lets the
|
||||
older snapshot's wholesale assignment land last and restore a tool the newer one deleted.
|
||||
The lock belongs here rather than in _init_search_tools_in_db, which _init_non_llm_objects_in_db
|
||||
already calls while holding it.
|
||||
"""
|
||||
if not self._should_load_db_object(object_type="search_tools"):
|
||||
return
|
||||
if prisma_client is None:
|
||||
return
|
||||
async with MODEL_RECONCILE_LOCK:
|
||||
await self._init_search_tools_in_db(prisma_client=prisma_client)
|
||||
|
||||
@staticmethod
|
||||
def _merge_config_and_db_search_tools(
|
||||
config_search_tools: list[SearchToolTypedDict],
|
||||
|
|
@ -7973,10 +8056,6 @@ def _fast_serialize_simple_model_response_stream(
|
|||
for top_level_key in ("id", "object", "created"):
|
||||
if payload[top_level_key] is None:
|
||||
payload.pop(top_level_key)
|
||||
|
||||
router_model_name: Final = getattr(chunk, ROUTER_MODEL_NAME_RESPONSE_FIELD, None)
|
||||
if router_model_name is not None:
|
||||
payload[ROUTER_MODEL_NAME_RESPONSE_FIELD] = router_model_name
|
||||
return orjson.dumps(payload)
|
||||
|
||||
|
||||
|
|
@ -8274,9 +8353,6 @@ async def async_data_generator(
|
|||
model_mismatch_logged = False
|
||||
fallback_metadata_event_sent = False
|
||||
include_fallback_errors: Final = _should_include_fallback_errors(request_data)
|
||||
# Fallbacks resolve on the first ``__anext__``, so the selected group is read
|
||||
# per chunk off this object rather than snapshotted here.
|
||||
router_logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
# Use a running string instead of list + join to avoid O(n^2) overhead.
|
||||
# Previously "".join(str_so_far_parts) was called every chunk, re-joining
|
||||
# the entire accumulated response. String += is O(n) amortized total.
|
||||
|
|
@ -8366,10 +8442,6 @@ async def async_data_generator(
|
|||
fallback_was_attempted=fallback_was_attempted,
|
||||
fallback_model_from_metadata=fallback_model_from_metadata,
|
||||
)
|
||||
ProxyBaseLLMRequestProcessing.set_router_selected_model_field(
|
||||
response_obj=chunk,
|
||||
router_model_name=ProxyBaseLLMRequestProcessing.get_router_selected_model_name(router_logging_obj),
|
||||
)
|
||||
|
||||
if strip_stream_usage and _is_injected_stream_usage_artifact(chunk):
|
||||
if pending_fallback_event:
|
||||
|
|
@ -9106,7 +9178,18 @@ class ProxyStartupEvent:
|
|||
|
||||
if store_model_in_db is not True:
|
||||
await proxy_config.init_mcp_servers_from_db()
|
||||
# Without this branch's own refresh, a UI-created search tool never reaches the router:
|
||||
# the add_deployment job that carries it in store_model_in_db=True mode is not scheduled.
|
||||
await proxy_config.reload_search_tools_from_db()
|
||||
if prisma_client is not None:
|
||||
scheduler.add_job(
|
||||
proxy_config.reload_search_tools_from_db,
|
||||
"interval",
|
||||
seconds=config_reload_interval_seconds,
|
||||
id="reload_search_tools_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
# DB-backed MCP servers are live objects in every mode, so the registry refresh that
|
||||
# store_model_in_db=True deployments get via the add_deployment job must run here
|
||||
# too; without it, a server whose OAuth discovery failed at startup is rebuilt only
|
||||
|
|
@ -9762,9 +9845,13 @@ async def model_list(
|
|||
When scope=expand is passed, proxy admins, team admins, and org admins
|
||||
will receive all proxy models as if they are a proxy admin.
|
||||
- healthy_only: When true, hide models whose backing deployments are all marked
|
||||
unhealthy by background health checks. Requires
|
||||
`background_health_checks: true` in general_settings; without
|
||||
health state the listing is returned unfiltered (fail open).
|
||||
unhealthy by background health checks. Set
|
||||
`general_settings.model_list_healthy_only: true` to apply this
|
||||
to every caller without the query parameter. Requires
|
||||
`background_health_checks: true` in general_settings, plus
|
||||
either `model_list_healthy_only` or `enable_health_check_routing`
|
||||
to keep deployment health state cached; without health state
|
||||
the listing is returned unfiltered (fail open).
|
||||
Models expanded from wildcard routes (e.g. `openai/*`) are not
|
||||
filtered, and nothing is hidden when `allowed_fails_policy` is
|
||||
configured (cooldown remains the sole exclusion mechanism).
|
||||
|
|
@ -9814,14 +9901,11 @@ async def model_list(
|
|||
|
||||
# Opt-in: also hide models whose deployments are all unhealthy per background
|
||||
# health checks. Empty when health state is unavailable or stale (fail open).
|
||||
unhealthy_names: set[str] = set()
|
||||
if healthy_only and llm_router is not None:
|
||||
unhealthy_names = await llm_router.async_get_fully_unhealthy_model_names()
|
||||
if not unhealthy_names:
|
||||
verbose_proxy_logger.debug(
|
||||
"healthy_only=true but no unhealthy deployment state is available "
|
||||
"(requires background_health_checks); returning unfiltered model list"
|
||||
)
|
||||
unhealthy_names: Final = await get_hidden_unhealthy_model_names(
|
||||
healthy_only=healthy_only,
|
||||
general_settings=settings,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
hidden_names: Final = blocked_names | unhealthy_names
|
||||
|
||||
|
|
@ -9984,9 +10068,11 @@ async def model_info(
|
|||
# Mirror /v1/models' visibility filter so first-occurrence resolution
|
||||
# cannot land on a deployment the listing had hidden.
|
||||
blocked_names: Final = llm_router.get_fully_blocked_model_names() if llm_router is not None else set()
|
||||
unhealthy_names: set[str] = set()
|
||||
if healthy_only and llm_router is not None:
|
||||
unhealthy_names = await llm_router.async_get_fully_unhealthy_model_names()
|
||||
unhealthy_names: Final = await get_hidden_unhealthy_model_names(
|
||||
healthy_only=healthy_only,
|
||||
general_settings=settings,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
hidden_names: Final = blocked_names | unhealthy_names
|
||||
if hidden_names:
|
||||
all_models = [m for m in all_models if m not in hidden_names]
|
||||
|
|
@ -14011,6 +14097,7 @@ async def model_info_v1(
|
|||
None,
|
||||
description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids",
|
||||
),
|
||||
healthy_only: bool | None = False,
|
||||
):
|
||||
"""
|
||||
Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)
|
||||
|
|
@ -14022,6 +14109,15 @@ async def model_info_v1(
|
|||
- When litellm_model_id is not passed, it will return the info for all models
|
||||
- include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info).
|
||||
- teamId: Filter to models accessible by the given team.
|
||||
- healthy_only: When true, hide models whose backing deployments are all marked
|
||||
unhealthy by background health checks, matching `/v1/models?healthy_only=true`.
|
||||
Set `general_settings.model_list_healthy_only: true` to apply this to every
|
||||
caller without the query parameter. Requires `background_health_checks: true`,
|
||||
plus either `model_list_healthy_only` or `enable_health_check_routing` to keep
|
||||
deployment health state cached; without health state the listing is returned
|
||||
unfiltered (fail open). Ignored when `litellm_model_id` is passed, since that
|
||||
is a direct lookup of one deployment rather than a listing. Hiding is
|
||||
presentation-only: a hidden model can still be called directly.
|
||||
|
||||
Each model in the list response includes `model_info.access_via_team_ids` and
|
||||
`model_info.direct_access` when the proxy database is connected.
|
||||
|
|
@ -14186,8 +14282,15 @@ async def model_info_v1(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("all_models: %s", all_models)
|
||||
return {"data": all_models}
|
||||
hidden_names: Final = await get_hidden_unhealthy_model_names(
|
||||
healthy_only=healthy_only,
|
||||
general_settings=general_settings,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
visible_models: Final = [model for model in all_models if model.get("model_name") not in hidden_names]
|
||||
|
||||
verbose_proxy_logger.debug("all_models: %s", visible_models)
|
||||
return {"data": visible_models}
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
|
|||
|
|
@ -51,6 +51,20 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
|
|||
TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]]
|
||||
|
||||
|
||||
async def _refresh_router_search_tools() -> None:
|
||||
"""Push the search tools table into this worker's router.
|
||||
|
||||
Best-effort: the row is already committed, so a refresh failure must not surface as a 500 and
|
||||
push the caller into a retry that creates duplicates.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
try:
|
||||
await proxy_config.reload_search_tools_from_db()
|
||||
except Exception as e: # noqa: BLE001 # the row is committed; no refresh failure may reach the caller
|
||||
verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e)
|
||||
|
||||
|
||||
async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable:
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import (
|
||||
|
|
@ -305,8 +319,10 @@ async def create_search_tool(request: CreateSearchToolRequest):
|
|||
search_tool=request.search_tool, prisma_client=prisma_client
|
||||
)
|
||||
|
||||
await _refresh_router_search_tools()
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Successfully added search tool '%s' to database. Router will be updated by the cron job.",
|
||||
"Successfully added search tool '%s' to database.",
|
||||
result.get("search_tool_name"),
|
||||
)
|
||||
|
||||
|
|
@ -388,8 +404,10 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
await _refresh_router_search_tools()
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Successfully updated search tool '%s' in database. Router will be updated by the cron job.",
|
||||
"Successfully updated search tool '%s' in database.",
|
||||
result.get("search_tool_name"),
|
||||
)
|
||||
|
||||
|
|
@ -445,9 +463,9 @@ async def delete_search_tool(search_tool_id: str):
|
|||
search_tool_id=search_tool_id, prisma_client=prisma_client
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Successfully deleted search tool from database. Router will be updated by the cron job."
|
||||
)
|
||||
await _refresh_router_search_tools()
|
||||
|
||||
verbose_proxy_logger.debug("Successfully deleted search tool from database.")
|
||||
|
||||
return result
|
||||
except HTTPException as e:
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from typing import (
|
|||
)
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
|
|
@ -208,9 +208,18 @@ async def _find_spend_logs(
|
|||
prisma_client: PrismaClient,
|
||||
where: Mapping[str, object],
|
||||
order: Mapping[str, str],
|
||||
take: int,
|
||||
http_response: Response,
|
||||
) -> Sequence[_SupportsModelDump]:
|
||||
"""Read spend log rows as Prisma model instances."""
|
||||
return await _spend_logs_table(prisma_client).find_many(where=where, order=order)
|
||||
"""Read spend log rows as Prisma model instances, capped at ``take`` rows."""
|
||||
rows: Final = await _spend_logs_table(prisma_client).find_many(where=where, order=order, take=take)
|
||||
if len(rows) == take:
|
||||
http_response.headers["x-litellm-spend-logs-truncated"] = "true"
|
||||
verbose_proxy_logger.warning(
|
||||
"/spend/logs result truncated to the %s most recent rows; use /spend/logs/v2 for paginated access",
|
||||
take,
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None:
|
||||
|
|
@ -2851,6 +2860,7 @@ async def ui_view_request_response_for_request_id(
|
|||
},
|
||||
)
|
||||
async def view_spend_logs(
|
||||
fastapi_response: Response,
|
||||
api_key: str | None = fastapi.Query(
|
||||
default=None,
|
||||
description="Get spend logs based on api key",
|
||||
|
|
@ -2881,6 +2891,8 @@ async def view_spend_logs(
|
|||
[DEPRECATED] This endpoint is not paginated and can cause performance issues.
|
||||
Please use `/spend/logs/v2` instead for paginated access to spend logs.
|
||||
|
||||
Row results are capped at 10,000 most recent entries per response.
|
||||
|
||||
View all spend logs, if request_id is provided, only logs for that request_id will be returned
|
||||
|
||||
When start_date and end_date are provided:
|
||||
|
|
@ -2931,7 +2943,6 @@ async def view_spend_logs(
|
|||
raise Exception(
|
||||
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||||
)
|
||||
spend_logs = []
|
||||
if (
|
||||
start_date is not None
|
||||
and isinstance(start_date, str)
|
||||
|
|
@ -2970,6 +2981,8 @@ async def view_spend_logs(
|
|||
prisma_client,
|
||||
where=filter_query,
|
||||
order={"startTime": "desc"},
|
||||
take=SPEND_LOGS_PAGINATION_COUNT_CAP,
|
||||
http_response=fastapi_response,
|
||||
)
|
||||
return data
|
||||
|
||||
|
|
@ -3040,14 +3053,12 @@ async def view_spend_logs(
|
|||
if user_id is not None and isinstance(user_id, str):
|
||||
scoped_filter["user"] = user_id
|
||||
|
||||
if not scoped_filter:
|
||||
spend_logs = await prisma_client.get_data(table_name="spend", query_type="find_all")
|
||||
return spend_logs
|
||||
|
||||
data = await _find_spend_logs(
|
||||
prisma_client,
|
||||
where=scoped_filter,
|
||||
order={"startTime": "desc"},
|
||||
take=SPEND_LOGS_PAGINATION_COUNT_CAP,
|
||||
http_response=fastapi_response,
|
||||
)
|
||||
return data
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_litellm_metadata_from_kwargs,
|
||||
reconstruct_model_name,
|
||||
)
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
|
||||
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes
|
||||
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
|
||||
|
|
@ -277,7 +278,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
usage: dict = {}
|
||||
if call_type in ["ocr", "aocr"]:
|
||||
usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict)
|
||||
else:
|
||||
elif not is_unbilled_non_inference_call(call_type, metadata, response_obj_dict):
|
||||
# Use response_obj_dict instead of response_obj to avoid calling .get() on Pydantic models
|
||||
_usage: Final = response_obj_dict.get("usage", None) or {}
|
||||
if isinstance(_usage, litellm.Usage):
|
||||
|
|
|
|||
|
|
@ -1402,6 +1402,7 @@ class ProxyLogging:
|
|||
get_latest_version_prompt_id,
|
||||
)
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.utils import get_non_default_completion_params
|
||||
|
||||
if prompt_version is None:
|
||||
|
|
@ -1420,13 +1421,20 @@ class ProxyLogging:
|
|||
data.pop("prompt_id", None)
|
||||
|
||||
if custom_logger and prompt_spec is not None:
|
||||
is_responses_call: Final = call_type == "aresponses"
|
||||
original_responses_input: Final = data.get("input", "") if is_responses_call else ""
|
||||
client_messages: Final = (
|
||||
ResponsesAPIRequestUtils.responses_input_to_chat_messages(original_responses_input)
|
||||
if is_responses_call
|
||||
else data.get("messages", [])
|
||||
)
|
||||
(
|
||||
model,
|
||||
messages,
|
||||
optional_params,
|
||||
) = await litellm_logging_obj.async_get_chat_completion_prompt(
|
||||
model=data.get("model", ""),
|
||||
messages=data.get("messages", []),
|
||||
messages=client_messages,
|
||||
non_default_params=get_non_default_completion_params(kwargs=data) or {},
|
||||
prompt_id=litellm_prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
|
|
@ -1438,7 +1446,14 @@ class ProxyLogging:
|
|||
|
||||
data.update(optional_params)
|
||||
data["model"] = model
|
||||
data["messages"] = messages
|
||||
if is_responses_call:
|
||||
data["input"] = ResponsesAPIRequestUtils.merge_prompt_management_input(
|
||||
original_input=original_responses_input,
|
||||
client_input=client_messages,
|
||||
merged_input=messages,
|
||||
)
|
||||
else:
|
||||
data["messages"] = messages
|
||||
# prevent re-processing the prompt template
|
||||
data.pop("prompt_id", None)
|
||||
data.pop("prompt_variables", None)
|
||||
|
|
@ -1653,7 +1668,7 @@ class ProxyLogging:
|
|||
not guardrails_only
|
||||
and litellm_logging_obj is not None
|
||||
and prompt_id is not None
|
||||
and (call_type == "completion" or call_type == "acompletion")
|
||||
and (call_type == "completion" or call_type == "acompletion" or call_type == "aresponses")
|
||||
):
|
||||
await self._process_prompt_template(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -557,6 +557,34 @@ async def _arealtime(
|
|||
raise ValueError(f"Unsupported model: {model}")
|
||||
|
||||
|
||||
def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) -> bool:
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models
|
||||
return False
|
||||
if model_info.get("mode") == "audio_transcription":
|
||||
return True
|
||||
return "/v1/realtime/transcription_sessions" in (model_info.get("supported_endpoints") or ())
|
||||
|
||||
|
||||
_TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"}
|
||||
|
||||
|
||||
def _azure_realtime_health_protocol(
|
||||
model: str, realtime_protocol: str | None, model_params: Mapping[str, Any]
|
||||
) -> tuple[str, RealtimeQueryParams | None]:
|
||||
query_params: Final = _TRANSCRIPTION_QUERY_PARAMS if _is_transcription_only_realtime_model(model, "azure") else None
|
||||
configured_raw: Final = (
|
||||
realtime_protocol or model_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL")
|
||||
)
|
||||
configured: Final = configured_raw if isinstance(configured_raw, str) else None
|
||||
if configured is not None:
|
||||
return configured, query_params
|
||||
if query_params is not None:
|
||||
return "GA", query_params
|
||||
return "beta", None
|
||||
|
||||
|
||||
def _realtime_health_check_auth_headers(
|
||||
custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, Any]
|
||||
) -> Mapping[str, str | None]:
|
||||
|
|
@ -586,7 +614,9 @@ async def _realtime_health_check(
|
|||
api_version: Optional[str] - api version
|
||||
api_key: str - api key
|
||||
custom_llm_provider: str - custom llm provider
|
||||
realtime_protocol: Optional[str] - protocol version ("GA"/"v1" for GA path, "beta"/None for beta path)
|
||||
realtime_protocol: Optional[str] - protocol version ("GA"/"v1" for GA path, "beta" for beta path);
|
||||
None resolves it for Azure from model_params/env, with transcription-only models probing GA
|
||||
plus intent=transcription the way real calls do
|
||||
|
||||
Returns:
|
||||
bool - True if connection is successful, False otherwise
|
||||
|
|
@ -602,11 +632,17 @@ async def _realtime_health_check(
|
|||
model_params=model_params or _EMPTY_MODEL_PARAMS,
|
||||
)
|
||||
if custom_llm_provider == "azure":
|
||||
resolved_protocol, azure_query_params = _azure_realtime_health_protocol(
|
||||
model=model,
|
||||
realtime_protocol=realtime_protocol,
|
||||
model_params=model_params or _EMPTY_MODEL_PARAMS,
|
||||
)
|
||||
url = azure_realtime._construct_url(
|
||||
api_base=api_base or "",
|
||||
model=model,
|
||||
api_version=api_version or "2024-10-01-preview",
|
||||
realtime_protocol=realtime_protocol,
|
||||
realtime_protocol=resolved_protocol,
|
||||
query_params=azure_query_params,
|
||||
)
|
||||
elif custom_llm_provider == "openai":
|
||||
url = openai_realtime._construct_url(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
from collections.abc import Coroutine, Iterable, Mapping
|
||||
from collections.abc import Coroutine, Generator, Iterable, Mapping
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
|
||||
|
||||
|
|
@ -13,6 +14,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
|
|||
LiteLLMResponsesTransformationHandler,
|
||||
)
|
||||
from litellm.constants import request_timeout
|
||||
from litellm.integrations.anthropic_cache_control_hook import CARRY_UNMATCHED_MESSAGE_POINTS
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
|
|
@ -26,7 +28,6 @@ from litellm.responses.litellm_completion_transformation.handler import (
|
|||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
PromptObject,
|
||||
Reasoning,
|
||||
ResponseIncludable,
|
||||
|
|
@ -390,6 +391,60 @@ async def aresponses_api_with_mcp(
|
|||
return response
|
||||
|
||||
|
||||
def _bridges_to_chat_completions(
|
||||
responses_api_provider_config: BaseResponsesAPIConfig | None, use_chat_completions_api: bool
|
||||
) -> bool:
|
||||
"""Whether the request reaches its provider as a chat completion, not a Responses call."""
|
||||
return responses_api_provider_config is None or use_chat_completions_api is True
|
||||
|
||||
|
||||
def _will_bridge_to_chat_completions(
|
||||
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool
|
||||
) -> bool:
|
||||
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
|
||||
|
||||
Resolving the config is a pure lookup, so this asks the same question the dispatch
|
||||
asks rather than restating its condition. Both callers resolve the provider before
|
||||
this runs, so the only way to be wrong is a prompt manager that moves the model
|
||||
across the bridge boundary, which would leave the deferred points to a pass that
|
||||
never comes.
|
||||
"""
|
||||
normalized_model: Final = _normalize_openai_chat_completions_responses_model(model)
|
||||
if custom_llm_provider is None:
|
||||
return True
|
||||
return _bridges_to_chat_completions(
|
||||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=normalized_model[0], provider=custom_llm_provider
|
||||
),
|
||||
use_chat_completions_api or normalized_model[1],
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs: dict[str, Any], # mutable-ok: the signal is read and popped out of the caller's own kwargs
|
||||
bridged: bool,
|
||||
) -> Generator[None, None]:
|
||||
"""Tell the cache-control hook that this layer's messages are not the ones sent upstream.
|
||||
|
||||
A Responses request keeps its system prompt in ``instructions``, which only becomes a
|
||||
system message when the chat-completion bridge builds one, so a role-targeted point
|
||||
is placed by the bridge's pass rather than this one.
|
||||
|
||||
Only raised for a request that will be bridged. A provider serving Responses natively
|
||||
gets no second pass, so this layer is the last one that can place anything and handing
|
||||
a point forward there drops it.
|
||||
"""
|
||||
if not bridged:
|
||||
yield
|
||||
return
|
||||
kwargs[CARRY_UNMATCHED_MESSAGE_POINTS] = True
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
kwargs.pop(CARRY_UNMATCHED_MESSAGE_POINTS, None)
|
||||
|
||||
|
||||
@client
|
||||
async def aresponses(
|
||||
input: str | ResponseInputParam,
|
||||
|
|
@ -463,23 +518,26 @@ async def aresponses(
|
|||
if isinstance(
|
||||
litellm_logging_obj, LiteLLMLoggingObj
|
||||
) and litellm_logging_obj.should_run_prompt_management_hooks(prompt_id=prompt_id, non_default_params=kwargs):
|
||||
if isinstance(input, str):
|
||||
client_input: list[AllMessageValues] = [{"role": "user", "content": input}]
|
||||
else:
|
||||
client_input = [item for item in input if isinstance(item, dict) and "role" in item]
|
||||
(
|
||||
model,
|
||||
merged_input,
|
||||
merged_optional_params,
|
||||
) = await litellm_logging_obj.async_get_chat_completion_prompt(
|
||||
model=model,
|
||||
messages=client_input,
|
||||
non_default_params=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
prompt_label=kwargs.get("prompt_label", None),
|
||||
prompt_version=kwargs.get("prompt_version", None),
|
||||
)
|
||||
client_input: Final = ResponsesAPIRequestUtils.responses_input_to_chat_messages(input)
|
||||
with _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs,
|
||||
bridged=_will_bridge_to_chat_completions(
|
||||
model, custom_llm_provider, bool(kwargs.get("use_chat_completions_api"))
|
||||
),
|
||||
):
|
||||
(
|
||||
model,
|
||||
merged_input,
|
||||
merged_optional_params,
|
||||
) = await litellm_logging_obj.async_get_chat_completion_prompt(
|
||||
model=model,
|
||||
messages=client_input,
|
||||
non_default_params=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
prompt_label=kwargs.get("prompt_label", None),
|
||||
prompt_version=kwargs.get("prompt_version", None),
|
||||
)
|
||||
input = cast(
|
||||
str | ResponseInputParam,
|
||||
ResponsesAPIRequestUtils.merge_prompt_management_input(
|
||||
|
|
@ -489,7 +547,13 @@ async def aresponses(
|
|||
),
|
||||
)
|
||||
if model != original_model:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
|
||||
custom_llm_provider = _resolve_prompt_swapped_provider(
|
||||
original_model=original_model,
|
||||
swapped_model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
kwargs.pop("prompt_id", None)
|
||||
kwargs["_async_prompt_merged_params"] = merged_optional_params
|
||||
|
||||
|
|
@ -559,6 +623,35 @@ async def aresponses(
|
|||
)
|
||||
|
||||
|
||||
def _resolve_prompt_swapped_provider(
|
||||
original_model: str,
|
||||
swapped_model: str,
|
||||
custom_llm_provider: str | None,
|
||||
kwargs: Mapping[str, object],
|
||||
prompt_id: str | None,
|
||||
) -> str:
|
||||
swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1]
|
||||
if kwargs.get("api_key") is None and kwargs.get("api_base") is None:
|
||||
return swapped_provider
|
||||
try:
|
||||
original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1]
|
||||
except litellm.BadRequestError:
|
||||
return swapped_provider
|
||||
if swapped_provider == original_provider:
|
||||
return swapped_provider
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
f"prompt_id '{prompt_id}' swaps model '{original_model}' -> '{swapped_model}', which changes the "
|
||||
f"provider from '{original_provider}' to '{swapped_provider}' after credentials for "
|
||||
f"'{original_provider}' were already resolved. Refusing to send them to '{swapped_provider}'. "
|
||||
"Point the request at a model whose provider matches the prompt's metadata.model, or set "
|
||||
"ignore_prompt_manager_model on the prompt to keep the requested model."
|
||||
),
|
||||
model=swapped_model,
|
||||
llm_provider=swapped_provider,
|
||||
)
|
||||
|
||||
|
||||
def _apply_prompt_management_to_responses_call(
|
||||
input: str | ResponseInputParam,
|
||||
model: str,
|
||||
|
|
@ -566,6 +659,7 @@ def _apply_prompt_management_to_responses_call(
|
|||
litellm_logging_obj: LiteLLMLoggingObj | None,
|
||||
kwargs: dict[str, Any],
|
||||
local_vars: dict[str, object],
|
||||
use_chat_completions_api: bool,
|
||||
) -> tuple[str | ResponseInputParam, str, str | None]:
|
||||
async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None)
|
||||
if async_merged is not None:
|
||||
|
|
@ -577,27 +671,28 @@ def _apply_prompt_management_to_responses_call(
|
|||
prompt_variables: Final = cast(dict | None, kwargs.get("prompt_variables", None))
|
||||
original_model: Final = model
|
||||
|
||||
if isinstance(input, str):
|
||||
client_input: list[AllMessageValues] = [{"role": "user", "content": input}]
|
||||
else:
|
||||
client_input = [item for item in input if isinstance(item, dict) and "role" in item]
|
||||
client_input: Final = ResponsesAPIRequestUtils.responses_input_to_chat_messages(input)
|
||||
|
||||
if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and litellm_logging_obj.should_run_prompt_management_hooks(
|
||||
prompt_id=prompt_id, non_default_params=kwargs
|
||||
):
|
||||
(
|
||||
model,
|
||||
merged_input,
|
||||
merged_optional_params,
|
||||
) = litellm_logging_obj.get_chat_completion_prompt(
|
||||
model=model,
|
||||
messages=client_input,
|
||||
non_default_params=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
prompt_label=kwargs.get("prompt_label", None),
|
||||
prompt_version=kwargs.get("prompt_version", None),
|
||||
)
|
||||
with _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs,
|
||||
bridged=_will_bridge_to_chat_completions(model, custom_llm_provider, use_chat_completions_api),
|
||||
):
|
||||
(
|
||||
model,
|
||||
merged_input,
|
||||
merged_optional_params,
|
||||
) = litellm_logging_obj.get_chat_completion_prompt(
|
||||
model=model,
|
||||
messages=client_input,
|
||||
non_default_params=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
prompt_label=kwargs.get("prompt_label", None),
|
||||
prompt_version=kwargs.get("prompt_version", None),
|
||||
)
|
||||
input = cast(
|
||||
str | ResponseInputParam,
|
||||
ResponsesAPIRequestUtils.merge_prompt_management_input(
|
||||
|
|
@ -609,7 +704,13 @@ def _apply_prompt_management_to_responses_call(
|
|||
local_vars["input"] = input
|
||||
local_vars["model"] = model
|
||||
if model != original_model:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
|
||||
custom_llm_provider = _resolve_prompt_swapped_provider(
|
||||
original_model=original_model,
|
||||
swapped_model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
local_vars["custom_llm_provider"] = custom_llm_provider
|
||||
for key, value in merged_optional_params.items():
|
||||
local_vars[key] = value
|
||||
|
|
@ -927,6 +1028,33 @@ def responses(
|
|||
# Update local_vars to include the converted text parameter
|
||||
local_vars["text"] = text
|
||||
|
||||
#########################################################
|
||||
# PROMPT MANAGEMENT
|
||||
# If aresponses() already ran the async hook, it pops prompt_id and
|
||||
# passes the result via _async_prompt_merged_params — apply those
|
||||
# directly and skip the sync hook to avoid double-merging.
|
||||
#########################################################
|
||||
_stripped_model, _from_chat_completions_prefix = _normalize_openai_chat_completions_responses_model(model)
|
||||
model = _stripped_model
|
||||
local_vars["model"] = model
|
||||
use_chat_completions_api = use_chat_completions_api or _from_chat_completions_prefix
|
||||
|
||||
if custom_llm_provider is None:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model=model, api_base=local_vars.get("base_url", None)
|
||||
)
|
||||
local_vars["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
input, model, custom_llm_provider = _apply_prompt_management_to_responses_call(
|
||||
input=input,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
kwargs=kwargs,
|
||||
local_vars=local_vars,
|
||||
use_chat_completions_api=use_chat_completions_api,
|
||||
)
|
||||
|
||||
# get llm provider logic
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
|
||||
|
|
@ -936,11 +1064,6 @@ def responses(
|
|||
if litellm_params.mock_response and isinstance(litellm_params.mock_response, str):
|
||||
return mock_responses_api_response(mock_response=litellm_params.mock_response)
|
||||
|
||||
_stripped_model, _from_chat_completions_prefix = _normalize_openai_chat_completions_responses_model(model)
|
||||
model = _stripped_model
|
||||
local_vars["model"] = model
|
||||
use_chat_completions_api = use_chat_completions_api or _from_chat_completions_prefix
|
||||
|
||||
model, custom_llm_provider = _resolve_model_provider_for_responses(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
@ -948,21 +1071,6 @@ def responses(
|
|||
local_vars=local_vars,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# PROMPT MANAGEMENT
|
||||
# If aresponses() already ran the async hook, it pops prompt_id and
|
||||
# passes the result via _async_prompt_merged_params — apply those
|
||||
# directly and skip the sync hook to avoid double-merging.
|
||||
#########################################################
|
||||
input, model, custom_llm_provider = _apply_prompt_management_to_responses_call(
|
||||
input=input,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
kwargs=kwargs,
|
||||
local_vars=local_vars,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Update input and tools with provider-specific file IDs if managed files are used
|
||||
#########################################################
|
||||
|
|
@ -1063,7 +1171,7 @@ def responses(
|
|||
if _file_search_dispatch is not None:
|
||||
return _file_search_dispatch
|
||||
|
||||
if responses_api_provider_config is None or use_chat_completions_api is True:
|
||||
if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api):
|
||||
return litellm_completion_transformation_handler.response_api_handler(
|
||||
model=model,
|
||||
input=input,
|
||||
|
|
|
|||
|
|
@ -573,13 +573,16 @@ class BaseResponsesAPIStreamingIterator:
|
|||
):
|
||||
return
|
||||
|
||||
if litellm.cache is None:
|
||||
cache: Final = litellm.cache
|
||||
if cache is None:
|
||||
return
|
||||
|
||||
cached_response: Final = response_obj.model_dump_json()
|
||||
if is_async:
|
||||
cache_write_task: Final = asyncio.create_task(
|
||||
litellm.cache.async_add_cache(
|
||||
from litellm.caching.caching_handler import create_cache_write_task
|
||||
|
||||
cache_write_task: Final = create_cache_write_task(
|
||||
lambda: cache.async_add_cache(
|
||||
cached_response,
|
||||
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
|
||||
**request_kwargs,
|
||||
|
|
@ -592,7 +595,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
)
|
||||
)
|
||||
else:
|
||||
litellm.cache.add_cache(
|
||||
cache.add_cache(
|
||||
cached_response,
|
||||
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
|
||||
**request_kwargs,
|
||||
|
|
|
|||
|
|
@ -72,6 +72,16 @@ class ResponsesAPIRequestUtils:
|
|||
shaped_content: Final = [_as_input_text_part(part) for part in content] # mutable-ok: Responses-shaped copy
|
||||
return {**message, "content": shaped_content} # mutable-ok: copy, the hook's message stays untouched
|
||||
|
||||
@staticmethod
|
||||
def responses_input_to_chat_messages(
|
||||
input: str | ResponseInputParam | None,
|
||||
) -> list[AllMessageValues]:
|
||||
if input is None:
|
||||
return []
|
||||
if isinstance(input, str):
|
||||
return [{"role": "user", "content": input}]
|
||||
return [item for item in input if isinstance(item, dict) and "role" in item]
|
||||
|
||||
@staticmethod
|
||||
def merge_prompt_management_input(
|
||||
original_input: str | ResponseInputParam,
|
||||
|
|
|
|||
|
|
@ -45,7 +45,6 @@ from litellm.caching.caching import (
|
|||
RedisClusterCache,
|
||||
)
|
||||
from litellm.constants import (
|
||||
AUTO_ROUTED_REQUEST_METADATA_KEY,
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY,
|
||||
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
|
|
@ -205,6 +204,7 @@ from litellm.types.router import (
|
|||
PreRoutingStrategy,
|
||||
RetryPolicy,
|
||||
RouterCacheEnum,
|
||||
RouterErrors,
|
||||
RouterGeneralSettings,
|
||||
RouterModelGroupAliasItem,
|
||||
RouterRateLimitError,
|
||||
|
|
@ -11520,10 +11520,8 @@ class Router:
|
|||
|
||||
# If still no deployments after checking for fallbacks, raise an error
|
||||
if len(healthy_deployments) == 0:
|
||||
message: Final = f"You passed in model={model}. There are no healthy deployments for this model"
|
||||
|
||||
raise litellm.BadRequestError(
|
||||
message=message,
|
||||
message=f"You passed in model={model}. {RouterErrors.no_healthy_deployments.value}",
|
||||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
|
|
@ -11534,11 +11532,18 @@ class Router:
|
|||
] # update the model to the actual value if an alias has been passed in
|
||||
|
||||
marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments)
|
||||
if all(marker_flags) or not any(marker_flags):
|
||||
if not any(marker_flags):
|
||||
return model, healthy_deployments
|
||||
return model, [ # mutable-ok: matches this function's list contract expected by downstream filters
|
||||
selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters
|
||||
d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker
|
||||
]
|
||||
if not selectable:
|
||||
raise litellm.BadRequestError(
|
||||
message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}",
|
||||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
return model, selectable
|
||||
|
||||
def _filter_deployments_by_model_access_groups(
|
||||
self,
|
||||
|
|
@ -12127,7 +12132,15 @@ class Router:
|
|||
This hook is called before the routing decision is made.
|
||||
|
||||
Used for the litellm auto-router to modify the request before the routing decision is made.
|
||||
|
||||
`model` is whatever the caller asked for, which may be a `model_group_alias` key, while the
|
||||
strategy registries and the marker deployment are keyed by the marker's own `model_name`, so
|
||||
every lookup below resolves the alias first. Only the lookups: the caller-facing name stays
|
||||
the alias, since spend metadata is stamped before routing and the response carries the tier
|
||||
group the strategy picked.
|
||||
"""
|
||||
registered_model_name: Final = self._get_model_from_alias(model=model) or model
|
||||
|
||||
#########################################################
|
||||
# Run the routing-plugin pipeline, if any plugins are configured.
|
||||
# Plugins narrow the candidate deployment pool (consumed later by
|
||||
|
|
@ -12135,9 +12148,13 @@ class Router:
|
|||
# downstream strategies (auto-router, complexity-router, ...) to read.
|
||||
#########################################################
|
||||
if self.routing_plugins:
|
||||
await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages)
|
||||
await self._run_routing_plugins(
|
||||
model=registered_model_name, request_kwargs=request_kwargs, messages=messages
|
||||
)
|
||||
|
||||
selected_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs)
|
||||
selected_strategy: Final = self._select_pre_routing_strategy(
|
||||
model=registered_model_name, request_kwargs=request_kwargs
|
||||
)
|
||||
if selected_strategy is None:
|
||||
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
|
||||
self._stamp_or_clear_metadata_key(
|
||||
|
|
@ -12146,13 +12163,10 @@ class Router:
|
|||
self._stamp_or_clear_metadata_key(
|
||||
request_kwargs=request_kwargs, key=CONSUMED_REQUEST_TAGS_METADATA_KEY, value=None
|
||||
)
|
||||
self._stamp_or_clear_metadata_key(
|
||||
request_kwargs=request_kwargs, key=AUTO_ROUTED_REQUEST_METADATA_KEY, value=None
|
||||
)
|
||||
return None
|
||||
|
||||
pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook(
|
||||
model=model,
|
||||
model=registered_model_name,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
|
|
@ -12176,13 +12190,6 @@ class Router:
|
|||
request_tags=_get_tags_from_request_kwargs(request_kwargs),
|
||||
),
|
||||
)
|
||||
# Gates the proxy's `router_model_name` response field; the body `model` is
|
||||
# always restamped back to the alias the client sent.
|
||||
self._stamp_or_clear_metadata_key(
|
||||
request_kwargs=request_kwargs,
|
||||
key=AUTO_ROUTED_REQUEST_METADATA_KEY,
|
||||
value=(True if pre_routing_hook_response is not None else None),
|
||||
)
|
||||
|
||||
# `model` (the alias, e.g. "smart-router") is never the deployment actually
|
||||
# called - apply the router marker's own litellm_params to the request,
|
||||
|
|
@ -12204,7 +12211,7 @@ class Router:
|
|||
# Per-tier `litellm_params` on the hook response are deliberate overrides
|
||||
# the caller applies on top, so those keys are never forwarded here.
|
||||
marker_params: Final = (
|
||||
self._forwardable_alias_marker_params(model=model, strategy_tags=selected_strategy.tags)
|
||||
self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags)
|
||||
if pre_routing_hook_response is not None
|
||||
else ()
|
||||
)
|
||||
|
|
|
|||
|
|
@ -178,6 +178,50 @@ response = litellm.completion(
|
|||
|
||||
## Special Behaviors
|
||||
|
||||
### Heuristic-first chaining
|
||||
|
||||
`classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM
|
||||
classifier for the ones the scorer could not place cheaply. It takes the same classifier settings as
|
||||
`classifier_type: llm`, plus `heuristic_first_max_tier`:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
classifier_type: heuristic_first
|
||||
heuristic_first_max_tier: SIMPLE
|
||||
classifier_llm_config:
|
||||
model: gpt-4o-mini
|
||||
tiers:
|
||||
SIMPLE: gpt-4o-mini
|
||||
MEDIUM: gpt-4o
|
||||
COMPLEX: claude-sonnet-4
|
||||
REASONING: o1-preview
|
||||
```
|
||||
|
||||
A request short-circuits, meaning it routes on the scorer's own tier with no classifier call, when
|
||||
two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least
|
||||
one signal. Everything else goes to the classifier, which then decides as it normally would.
|
||||
|
||||
The signal requirement is what keeps this from quietly routing everything to your cheapest model.
|
||||
A prompt where no dimension fires scores exactly 0.0, which is below `simple_medium`, so the score
|
||||
to tier mapping calls it SIMPLE by default rather than by evidence. Around half of general traffic
|
||||
scores that way. Those requests reach the classifier instead, which is the whole reason to configure
|
||||
one. Note the converse too: the score is not a confidence, and a prompt that fires a single weak
|
||||
signal and still lands under the boundary does short-circuit, so a lower threshold buys accuracy and
|
||||
a higher one buys savings.
|
||||
|
||||
`heuristic_first_max_tier` names a built-in tier and may not name the highest one, since that would
|
||||
short-circuit everything and leave the classifier unreachable. Operator-defined tier sets
|
||||
(`tier_definitions`) are not supported here, because the scorer only produces the built-in tiers.
|
||||
When the classifier call fails, the fallback works exactly as it does under `classifier_type: llm`,
|
||||
except that the heuristic outcome is the one already computed rather than a second scoring pass.
|
||||
|
||||
Spend logs record `routing_decision.cause` as `heuristic_first_short_circuit` when the classifier
|
||||
was skipped, and `llm_classifier` when it ran, so the two are told apart per request.
|
||||
|
||||
### Reasoning Override
|
||||
|
||||
If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone.
|
||||
|
|
|
|||
|
|
@ -719,6 +719,7 @@ class ClassificationOutcome(NamedTuple):
|
|||
"heuristic_scorer",
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
"heuristic_first_short_circuit",
|
||||
"classifier_plugin",
|
||||
"classifier_fallback",
|
||||
"default_model_fallback",
|
||||
|
|
@ -859,7 +860,7 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
# Both are pure functions of the config, so building them per classifier call would
|
||||
# re-run create_model and the schema conversion on every request for the same result.
|
||||
llm_classifier_configured: Final = self.config.classifier_type == "llm" and (
|
||||
llm_classifier_configured: Final = self.config.uses_llm_classifier and (
|
||||
self.config.classifier_llm_config is not None
|
||||
)
|
||||
self._classifier_system_prompt: str | None = (
|
||||
|
|
@ -1237,17 +1238,63 @@ class ComplexityRouter(CustomLogger):
|
|||
"""
|
||||
Classify a prompt by complexity, using the LLM classifier when configured.
|
||||
|
||||
Falls back to the local heuristic scorer if classifier_type is "heuristic". If the LLM call
|
||||
or the classifier plugin fails, times out, or produces no usable tier, the configured
|
||||
fallback_tier wins on a custom tier set, and classifier_fallback otherwise decides between
|
||||
the heuristic scorer and default_model. The outcome's `cause` reports which path actually ran.
|
||||
Falls back to the local heuristic scorer if classifier_type is "heuristic". Under
|
||||
"heuristic_first" the scorer runs first and the classifier is called only for requests it
|
||||
could not place at or below heuristic_first_max_tier. If the LLM call or the classifier
|
||||
plugin fails, times out, or produces no usable tier, the configured fallback_tier wins on a
|
||||
custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and
|
||||
default_model. The outcome's `cause` reports which path actually ran.
|
||||
"""
|
||||
if self.config.classifier_type == "custom":
|
||||
return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages)
|
||||
if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None:
|
||||
return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages)
|
||||
if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None:
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
|
||||
|
||||
async def _classify_heuristic_first(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
) -> ClassificationOutcome:
|
||||
"""Score locally, and only pay for the classifier call when the scorer did not confidently
|
||||
place the request at or below heuristic_first_max_tier.
|
||||
|
||||
Confidence is `signals`, not `score`. A prompt where no dimension fired scores exactly 0.0,
|
||||
which is below simple_medium and so lands SIMPLE by default rather than by evidence, and a
|
||||
threshold check alone would hand that traffic to the cheapest model without ever consulting
|
||||
the classifier. Scores also go negative when simple indicators fire, so a score threshold
|
||||
would reject exactly the trivial prompts this path exists to serve.
|
||||
"""
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
threshold: Final = self.config.heuristic_first_max_tier
|
||||
decided_cheaply: Final = (
|
||||
threshold is not None
|
||||
and bool(signals)
|
||||
and self._active_tier_severity(tier) <= self._active_tier_severity(threshold)
|
||||
)
|
||||
if decided_cheaply:
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit")
|
||||
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored)
|
||||
|
||||
async def _llm_classifier_outcome(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
scored: ClassificationOutcome | None = None,
|
||||
) -> ClassificationOutcome:
|
||||
"""Call the LLM classifier and turn its verdict, or its failure, into an outcome.
|
||||
|
||||
`scored` is the heuristic outcome the caller already computed, which only "heuristic_first"
|
||||
has. It is handed to the failure path so a classifier error does not re-run the scorer.
|
||||
"""
|
||||
try:
|
||||
tier, classifier_cost = await self._classify_with_llm(prompt, system_prompt, request_kwargs, messages)
|
||||
return ClassificationOutcome(
|
||||
|
|
@ -1258,11 +1305,20 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_cost=classifier_cost,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path
|
||||
return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt)
|
||||
return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored)
|
||||
|
||||
def _classifier_failure_outcome(self, reason: str, prompt: str, system_prompt: str | None) -> ClassificationOutcome:
|
||||
def _classifier_failure_outcome(
|
||||
self,
|
||||
reason: str,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
scored: ClassificationOutcome | None = None,
|
||||
) -> ClassificationOutcome:
|
||||
"""The outcome when the LLM classifier or classifier plugin produced no usable tier:
|
||||
fallback_tier on a custom tier set, classifier_fallback otherwise."""
|
||||
fallback_tier on a custom tier set, classifier_fallback otherwise.
|
||||
|
||||
A caller that already scored the prompt passes `scored` so the heuristic arm returns that
|
||||
verdict instead of running the same scan again on the request path."""
|
||||
fallback_tier: Final = self.config.fallback_tier
|
||||
if fallback_tier is not None:
|
||||
verbose_router_logger.warning("ComplexityRouter: %s, routing to fallback_tier %s", reason, fallback_tier)
|
||||
|
|
@ -1277,6 +1333,8 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
if self.config.classifier_fallback == "default_model":
|
||||
return self._default_model_fallback_outcome()
|
||||
if scored is not None:
|
||||
return scored
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,11 @@ class ClassificationRubric(str, Enum):
|
|||
# routers get the calibrated rubric without changing what is already running.
|
||||
DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubric.LEGACY
|
||||
|
||||
# The classifier_type values that can call classifier_llm_config.model. Every consumer asking
|
||||
# "is the classifier model a real dependency of this router" resolves it here, including the ones
|
||||
# that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier.
|
||||
LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first"})
|
||||
|
||||
|
||||
TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
|
||||
ComplexityTier.SIMPLE,
|
||||
|
|
@ -591,13 +596,30 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
|
||||
# Classifier strategy
|
||||
classifier_type: Literal["heuristic", "llm", "custom"] = Field(
|
||||
classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field(
|
||||
default="heuristic",
|
||||
description="Classification strategy: local regex/keyword scoring, an LLM call, or a custom classifier plugin",
|
||||
description=(
|
||||
"Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier "
|
||||
"plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier "
|
||||
"when the local scorer does not confidently land a cheap tier"
|
||||
),
|
||||
)
|
||||
classifier_llm_config: ClassifierLLMConfig | None = Field(
|
||||
default=None,
|
||||
description="Configuration for the LLM classifier; required when classifier_type is 'llm'",
|
||||
description="Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first'",
|
||||
)
|
||||
heuristic_first_max_tier: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The highest tier the local scorer may decide on its own; required when classifier_type is "
|
||||
"'heuristic_first' and rejected otherwise. A request whose heuristic tier is at or below this "
|
||||
"one skips the LLM classifier and routes straight to that heuristic tier, so the classifier "
|
||||
"call is only paid for on traffic the scorer could not place cheaply. The scorer must also "
|
||||
"have produced at least one signal: a prompt where no dimension fired scores 0.0 and would "
|
||||
"otherwise land SIMPLE by default rather than by evidence, which is how a chained router "
|
||||
"would silently send unclassified traffic to the cheapest model. Names a built-in tier, and "
|
||||
"may not name the highest one, since that would make the LLM classifier unreachable."
|
||||
),
|
||||
)
|
||||
classifier_plugin: ClassifierPlugin | None = Field(
|
||||
default=None,
|
||||
|
|
@ -626,7 +648,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"which is what a classifier on some other taxonomy wants: a prompt that grades data "
|
||||
"sensitivity has no use for a complexity score, and scoring one produces a tier unrelated to "
|
||||
"what the operator configured. Requires default_model when set to 'default_model'. Only "
|
||||
"applies when classifier_type is 'llm' or 'custom'."
|
||||
"applies when classifier_type is 'llm', 'custom', or 'heuristic_first'."
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -936,8 +958,8 @@ class ComplexityRouterConfig(BaseModel):
|
|||
|
||||
@model_validator(mode="after")
|
||||
def _validate_classifier_config(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type == "llm" and self.classifier_llm_config is None:
|
||||
raise ValueError("classifier_llm_config is required when classifier_type is 'llm'")
|
||||
if self.uses_llm_classifier and self.classifier_llm_config is None:
|
||||
raise ValueError(f"classifier_llm_config is required when classifier_type is {self.classifier_type!r}")
|
||||
if self.classifier_type == "custom" and self.classifier_plugin is None:
|
||||
raise ValueError("classifier_plugin is required when classifier_type is 'custom'")
|
||||
if self.classifier_plugin is not None and self.classifier_type != "custom":
|
||||
|
|
@ -947,6 +969,49 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
return self
|
||||
|
||||
@field_validator("heuristic_first_max_tier", mode="before")
|
||||
@classmethod
|
||||
def _coerce_heuristic_first_max_tier(cls, value: object) -> object:
|
||||
if isinstance(value, ComplexityTier):
|
||||
return value.value
|
||||
if isinstance(value, str):
|
||||
return value.strip()
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_heuristic_first_max_tier(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type != "heuristic_first":
|
||||
if self.heuristic_first_max_tier is not None:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier is set but classifier_type is {self.classifier_type!r}; "
|
||||
"the local scorer would never gate the classifier. Set classifier_type "
|
||||
"'heuristic_first' or remove heuristic_first_max_tier"
|
||||
)
|
||||
return self
|
||||
threshold: Final = self.heuristic_first_max_tier
|
||||
if threshold is None:
|
||||
raise ValueError(
|
||||
"heuristic_first_max_tier is required when classifier_type is 'heuristic_first': without a "
|
||||
"threshold there is nothing to decide whether a request escalates to the LLM classifier"
|
||||
)
|
||||
names: Final = self.tier_names()
|
||||
if threshold not in names:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold!r} is not an active tier: it must name one of {', '.join(names)}"
|
||||
)
|
||||
if threshold == names[-1]:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold} is the highest tier, so every request would short-circuit "
|
||||
"and the LLM classifier would never run; name a lower tier or use classifier_type 'heuristic'"
|
||||
)
|
||||
if threshold not in self.tiers:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold} has no model configured in tiers; a threshold pointing at "
|
||||
"an unconfigured tier would route short-circuited requests to the default fallback instead of the "
|
||||
"pool the operator intended"
|
||||
)
|
||||
return self
|
||||
|
||||
@field_validator("fallback_tier", "classification_prompt")
|
||||
@classmethod
|
||||
def _reject_blank_optional_text(cls, value: str | None) -> str | None:
|
||||
|
|
@ -969,6 +1034,14 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"""True when the operator replaced the built-in tier set via tier_definitions."""
|
||||
return self.tier_definitions is not None
|
||||
|
||||
@property
|
||||
def uses_llm_classifier(self) -> bool:
|
||||
"""True when this router can call classifier_llm_config.model, so the model is a real
|
||||
dependency: authorized against the caller's key, counted in the health graph, and given a
|
||||
prebuilt rubric. 'heuristic_first' only calls it for traffic the local scorer escalates,
|
||||
which still makes it a dependency on every one of those requests."""
|
||||
return self.classifier_type in LLM_CLASSIFIER_TYPES
|
||||
|
||||
def tier_names(self) -> tuple[str, ...]:
|
||||
"""The active tier names: the defined names, or the built-in set in severity order."""
|
||||
if self.tier_definitions is not None:
|
||||
|
|
@ -1063,7 +1136,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
if duplicated:
|
||||
raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}")
|
||||
if self.classifier_type == "heuristic":
|
||||
if self.classifier_type in ("heuristic", "heuristic_first"):
|
||||
raise ValueError(
|
||||
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
|
||||
"produces the built-in tiers"
|
||||
|
|
|
|||
|
|
@ -10,13 +10,28 @@ the router silently dropping the deployment at load time under
|
|||
``ignore_invalid_deployments``.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from litellm.router_strategy.complexity_router.config import LLM_CLASSIFIER_TYPES
|
||||
|
||||
AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/"
|
||||
|
||||
StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"]
|
||||
|
||||
StrategyRouterDependencyRole: TypeAlias = Literal["tier", "default", "classifier", "embedding"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StrategyRouterDependency:
|
||||
"""A model name a strategy router must be able to reach to do its job."""
|
||||
|
||||
model_name: str
|
||||
role: StrategyRouterDependencyRole
|
||||
|
||||
|
||||
STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"auto_router_config",
|
||||
|
|
@ -63,6 +78,88 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None:
|
|||
return "semantic"
|
||||
|
||||
|
||||
def _named(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]:
|
||||
"""One dependency from a scalar field, or none when it is absent or not a name."""
|
||||
return (StrategyRouterDependency(value, role),) if isinstance(value, str) and value else ()
|
||||
|
||||
|
||||
def _pool(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]:
|
||||
"""Dependencies from a field holding either a single name or a pool of them."""
|
||||
if isinstance(value, str):
|
||||
return _named(value, role)
|
||||
if isinstance(value, Sequence):
|
||||
return tuple(dep for entry in value for dep in _named(entry, role))
|
||||
return ()
|
||||
|
||||
|
||||
_NO_CONFIG: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _mapping(value: object) -> Mapping[str, object]:
|
||||
return value if isinstance(value, Mapping) else _NO_CONFIG
|
||||
|
||||
|
||||
def strategy_router_dependencies(
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[StrategyRouterDependency, ...]:
|
||||
"""The model names a strategy-router deployment must reach, in no particular order.
|
||||
|
||||
A field is a dependency only under the condition the runtime itself reads it: the
|
||||
classifier model needs `classifier_type: llm`, and the complexity embedding model needs
|
||||
`semantic_keyword_matching`. Listing one the router never calls reds a working deployment.
|
||||
|
||||
The two default-model spellings are not symmetric. A quality router falls back to its
|
||||
config's `default_model`, so both are read. A complexity router ignores that field and
|
||||
derives its default from the tiers instead (`fallback_tier`, then MEDIUM, then SIMPLE),
|
||||
overwriting the config value at init, so only the `litellm_params` spelling is a
|
||||
dependency here; the derived one is already covered as a tier.
|
||||
|
||||
Returns empty for a regular deployment, and for any name this module cannot reach from
|
||||
the deployment dict alone: a semantic router's routes live in an `auto_router_config`
|
||||
JSON string or an `auto_router_config_path` file, so only its default and embedding
|
||||
models are enumerable here. Every field is read defensively, since a caller may hold a
|
||||
config the router itself would refuse, and a health check must not raise on one.
|
||||
"""
|
||||
kind: Final = classify_strategy_router_model(str(litellm_params.get("model", "")))
|
||||
if kind is None:
|
||||
return ()
|
||||
if kind == "semantic":
|
||||
return _named(litellm_params.get("auto_router_default_model"), "default") + _named(
|
||||
litellm_params.get("auto_router_embedding_model"), "embedding"
|
||||
)
|
||||
if kind == "adaptive":
|
||||
return _pool(_mapping(litellm_params.get("adaptive_router_config")).get("available_models"), "tier")
|
||||
if kind == "quality":
|
||||
quality: Final = _mapping(litellm_params.get("quality_router_config"))
|
||||
return tuple(
|
||||
dict.fromkeys(
|
||||
_pool(quality.get("available_models"), "tier")
|
||||
+ _named(
|
||||
litellm_params.get("quality_router_default_model") or quality.get("default_model"),
|
||||
"default",
|
||||
)
|
||||
)
|
||||
)
|
||||
complexity: Final = _mapping(litellm_params.get("complexity_router_config"))
|
||||
classifier: Final = _mapping(complexity.get("classifier_llm_config"))
|
||||
return tuple(
|
||||
dict.fromkeys(
|
||||
tuple(dep for tier in _mapping(complexity.get("tiers")).values() for dep in _pool(tier, "tier"))
|
||||
+ _named(litellm_params.get("complexity_router_default_model"), "default")
|
||||
+ (
|
||||
_named(classifier.get("model"), "classifier")
|
||||
if complexity.get("classifier_type") in LLM_CLASSIFIER_TYPES
|
||||
else ()
|
||||
)
|
||||
+ (
|
||||
_named(complexity.get("embedding_model"), "embedding")
|
||||
if complexity.get("semantic_keyword_matching")
|
||||
else ()
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None:
|
||||
"""Reject a complexity config the router would refuse to build a deployment from.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
||||
class LangfuseLoggingConfig(TypedDict):
|
||||
langfuse_secret: str | None
|
||||
langfuse_public_key: str | None
|
||||
langfuse_host: str | None
|
||||
langfuse_environment: ReadOnly[str | None]
|
||||
|
||||
|
||||
class LangfuseUsageDetails(TypedDict):
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue