mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge remote-tracking branch 'origin/main' into litellm-providers/price-sync-databricks
This commit is contained in:
commit
24b80738a2
209 changed files with 34741 additions and 20515 deletions
15
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
15
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -4,9 +4,20 @@
|
|||
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
|
||||
"## Affected release" with nothing under it must not appear in the final description -->
|
||||
|
||||
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
|
||||
## Handwritten 1-3 sentence description
|
||||
|
||||
<!-- The PR author rewrites this in their own words before asking for review. Being able to sum up the PR in a few
|
||||
plain sentences proves the author understands what it does, and human-written technical English is far easier
|
||||
for a reviewer to parse than AI-written English
|
||||
An AI agent may draft a starting point here, but always wraps its draft in the make-me-human-readable tags below,
|
||||
backticks included so the tags stay visible on GitHub, so nobody mistakes it for the author's words. The author rewrites it in their own words and deletes the tags.
|
||||
Tags still present at review time mean the summary is AI-written and has not been rewritten yet
|
||||
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
|
||||
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
|
||||
|
||||
`<make-me-human-readable>`
|
||||
<blah>
|
||||
`</make-me-human-readable>`
|
||||
|
||||
<details>
|
||||
<summary>Details</summary>
|
||||
|
|
|
|||
2
.github/workflows/create-rc-branch.yml
vendored
2
.github/workflows/create-rc-branch.yml
vendored
|
|
@ -2,7 +2,7 @@ name: Create RC Branch
|
|||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 3 * * 5"
|
||||
- cron: "0 3 * * 1"
|
||||
timezone: "America/Los_Angeles"
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
|
|||
|
|
@ -41,6 +41,8 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
|
||||
|
||||
Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --num_workers 2 --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`, and swap its `--config` for a copy of `litellm/proxy/dev_config.yaml` that points both `router_settings` (`redis_host`, `redis_port`) and `general_settings.coordination_redis` (`host`, `port`) at one dedicated Redis instance, with credentials when required, so the workers share router state and spend, budget, and rate-limit counters. Check that two worker PIDs are serving, and use the same worker count and Redis setup for the before and after proof. When the PR is about `--run_gunicorn`, `--run_hypercorn`, or `--run_granian`, which `--reload` does not support, drop `--reload`. When it is about behavior without Redis, drop Redis. When the proxy prints `"workers" flag is ignored when reloading is enabled`, that checkout or its uvicorn reloads only one worker, so drop the two-worker check for it. Say in the proof which part you dropped and why
|
||||
|
||||
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
|
||||
|
||||
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
|
||||
|
|
|
|||
|
|
@ -696,7 +696,7 @@ For companies that need better security, user management and professional suppor
|
|||
[Talk to founders](https://enterprise.litellm.ai/demo)
|
||||
|
||||
This covers:
|
||||
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/proxy/enterprise):**
|
||||
- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/enterprise):**
|
||||
- ✅ **Feature Prioritization**
|
||||
- ✅ **Custom Integrations**
|
||||
- ✅ **Professional Support - Dedicated discord + slack**
|
||||
|
|
|
|||
|
|
@ -1647,6 +1647,7 @@ def set_global_gitlab_config(config: Mapping[str, object]) -> None:
|
|||
# Lazy loading system for heavy modules to reduce initial import time and memory usage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from . import guardrails as guardrails
|
||||
from litellm.types.utils import ModelInfo as _ModelInfoType
|
||||
from litellm.types.utils import PriorityReservationSettings
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
|
@ -2499,10 +2500,10 @@ def __getattr__(name: str) -> Any:
|
|||
handler_func: Final = registry[name]
|
||||
return handler_func(name)
|
||||
|
||||
if name == "proxy":
|
||||
if name in ("proxy", "guardrails"):
|
||||
import importlib
|
||||
|
||||
return importlib.import_module("litellm.proxy")
|
||||
return importlib.import_module(f"litellm.{name}")
|
||||
|
||||
if name == "harness" or name in _AGENT_EXPORTS:
|
||||
import importlib
|
||||
|
|
|
|||
|
|
@ -432,6 +432,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
|||
DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
|
||||
NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1"
|
||||
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
|
||||
DEEPGRAM_SPEAK_DEFAULT_MODEL: Final = "aura-asteria-en"
|
||||
|
||||
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
|
|
|
|||
66
litellm/guardrails.py
Normal file
66
litellm/guardrails.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
"""Invoke configured gateway guardrails without generating a completion."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from enum import Enum
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared factory has untyped optional parameters
|
||||
)
|
||||
from litellm.types.guardrails import NativeGuardrailResponse
|
||||
|
||||
_NATIVE_RESPONSE: Final[TypeAdapter[NativeGuardrailResponse]] = TypeAdapter(NativeGuardrailResponse)
|
||||
|
||||
|
||||
class _InputType(Enum):
|
||||
REQUEST = "request"
|
||||
RESPONSE = "response"
|
||||
|
||||
|
||||
async def apply(
|
||||
guardrail_name: str,
|
||||
*,
|
||||
input: str | None = None,
|
||||
response: str | None = None,
|
||||
messages: Sequence[Mapping[str, object]] | None = None,
|
||||
metadata: Mapping[str, object] | None = None,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> NativeGuardrailResponse:
|
||||
"""Return a configured guardrail's native structured findings.
|
||||
|
||||
Supply exactly one of ``input`` or ``response``. Grounding guardrails can
|
||||
use ``messages`` as evidence. Connection settings default to
|
||||
``litellm.api_base`` and ``litellm.api_key``.
|
||||
"""
|
||||
if (input is None) == (response is None):
|
||||
raise ValueError("Supply exactly one of input or response.")
|
||||
resolved_base: Final = api_base if api_base is not None else litellm.api_base
|
||||
if not resolved_base:
|
||||
raise ValueError("Set api_base or litellm.api_base to your LiteLLM gateway URL.")
|
||||
resolved_key: Final = api_key if api_key is not None else litellm.api_key
|
||||
input_type: Final = _InputType.RESPONSE if response is not None else _InputType.REQUEST
|
||||
payload: Final = dict[str, object](
|
||||
guardrail_name=guardrail_name,
|
||||
text=response if response is not None else input,
|
||||
input_type=input_type.value,
|
||||
response_format="native",
|
||||
)
|
||||
if messages is not None:
|
||||
payload["messages"] = [dict(message) for message in messages]
|
||||
if metadata is not None:
|
||||
payload["metadata"] = dict(metadata)
|
||||
gateway_root: Final = resolved_base.rstrip("/").removesuffix("/v1")
|
||||
client: Final = get_async_httpx_client(llm_provider="litellm_guardrails")
|
||||
result: Final = await client.post( # pyright: ignore[reportUnknownMemberType] # shared handler has untyped optional parameters
|
||||
url=f"{gateway_root}/guardrails/apply_guardrail",
|
||||
json=payload,
|
||||
headers={"Authorization": f"Bearer {resolved_key}"} if resolved_key else {},
|
||||
timeout=timeout,
|
||||
)
|
||||
return _NATIVE_RESPONSE.validate_json(result.content)
|
||||
|
|
@ -13,7 +13,7 @@ import copy
|
|||
import itertools
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Container, Iterable, Mapping, Sequence
|
||||
from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -32,6 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
tool_call_ids_answered_by_client,
|
||||
tool_call_ids_answered_in,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import string_field
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
is_claude_code_one_shot_subagent_request,
|
||||
supports_anthropic_cache_control,
|
||||
|
|
@ -128,7 +129,7 @@ def _message_mark_reaches_provider(fields: Mapping[str, object]) -> bool:
|
|||
return False
|
||||
return (
|
||||
not isinstance(fields.get("content"), list)
|
||||
or fields.get("role") in _ROLES_SENDING_MESSAGE_MARK_WITH_LIST_CONTENT
|
||||
or string_field(fields, "role") in _ROLES_SENDING_MESSAGE_MARK_WITH_LIST_CONTENT
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -445,22 +446,41 @@ def _as_object_list(value: object | None) -> list[object] | None:
|
|||
return _validated_object_list(value)
|
||||
|
||||
|
||||
def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool:
|
||||
def _server_tool_call_id(tool_call: object) -> str | None:
|
||||
tool_call_id: Final = _attribute_or_key(tool_call, "id")
|
||||
if isinstance(tool_call_id, str) and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX):
|
||||
return tool_call_id
|
||||
return None
|
||||
|
||||
|
||||
def _tool_call_carries_cache_breakpoint(tool_call: object, client_tool_call_ids: Container[str]) -> bool:
|
||||
if _attribute_or_key(tool_call, "cache_control") is None:
|
||||
return False
|
||||
server_tool_call_id: Final = _server_tool_call_id(tool_call)
|
||||
return server_tool_call_id is None or server_tool_call_id in client_tool_call_ids
|
||||
|
||||
return not tool_call_is_rebuilt_as_server_tool_use(
|
||||
_attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids
|
||||
|
||||
def _server_tool_call_ids_without_a_kept_result(message: object) -> Iterator[str]:
|
||||
provider_specific_fields: Final = _attribute_or_key(message, "provider_specific_fields")
|
||||
tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls")) or ()
|
||||
return (
|
||||
tool_call_id
|
||||
for tool_call in tool_calls
|
||||
if (tool_call_id := _server_tool_call_id(tool_call)) is not None
|
||||
and not tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields)
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
|
||||
def _client_tool_call_ids_when_a_server_tool_call_is_marked(
|
||||
messages: Sequence[object], model: str | None = None, custom_llm_provider: str | None = None
|
||||
) -> frozenset[str]:
|
||||
if not _marks_a_server_tool_call(messages):
|
||||
return frozenset()
|
||||
if model is not None and not _builds_anthropic_messages(model, custom_llm_provider):
|
||||
return tool_call_ids_answered_in(messages)
|
||||
server_calls_sent_as_client_calls: Final = itertools.chain.from_iterable(
|
||||
map(_server_tool_call_ids_without_a_kept_result, messages)
|
||||
)
|
||||
return tool_call_ids_answered_in(messages) | frozenset(server_calls_sent_as_client_calls)
|
||||
return tool_call_ids_answered_by_client(messages, _keeps_mid_conversation_system(model, custom_llm_provider))
|
||||
|
||||
|
||||
|
|
@ -517,12 +537,7 @@ def _model_info_flags_mid_conversation_system(model: str, custom_llm_provider: s
|
|||
|
||||
|
||||
def _is_marked_server_tool_call(tool_call: object) -> bool:
|
||||
tool_call_id: Final = _attribute_or_key(tool_call, "id")
|
||||
return (
|
||||
isinstance(tool_call_id, str)
|
||||
and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
|
||||
and _attribute_or_key(tool_call, "cache_control") is not None
|
||||
)
|
||||
return _server_tool_call_id(tool_call) is not None and _attribute_or_key(tool_call, "cache_control") is not None
|
||||
|
||||
|
||||
def _tool_carries_cache_breakpoint(tool: object) -> bool:
|
||||
|
|
@ -668,7 +683,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
|
||||
remaining_points, external_breakpoints, openai_dialect
|
||||
)
|
||||
answered_tool_ids: Final = _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
|
||||
answered_tool_ids: Final = _client_tool_call_ids_when_a_server_tool_call_is_marked(
|
||||
processed_messages, model, custom_llm_provider
|
||||
)
|
||||
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(
|
||||
|
|
@ -757,7 +772,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
)
|
||||
message_list: Final = tuple(messages)
|
||||
client_answers: Final = (
|
||||
_tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(message_list)
|
||||
_client_tool_call_ids_when_a_server_tool_call_is_marked(message_list)
|
||||
if answered_tool_ids is None
|
||||
else answered_tool_ids
|
||||
)
|
||||
|
|
@ -849,7 +864,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
so earlier points win when slots are scarce.
|
||||
"""
|
||||
client_answers: Final = (
|
||||
_tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(messages)
|
||||
_client_tool_call_ids_when_a_server_tool_call_is_marked(messages)
|
||||
if answered_tool_ids is None
|
||||
else answered_tool_ids
|
||||
)
|
||||
|
|
@ -939,11 +954,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0
|
||||
tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls"))
|
||||
tool_call_count: Final = (
|
||||
sum(
|
||||
1
|
||||
for tool_call in tool_calls
|
||||
if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids)
|
||||
)
|
||||
sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, answered_tool_ids))
|
||||
if tool_calls
|
||||
else 0
|
||||
)
|
||||
|
|
|
|||
|
|
@ -315,6 +315,8 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = {
|
|||
"model_context_window_exceeded": "length",
|
||||
# Cohere
|
||||
"COMPLETE": "stop",
|
||||
"STOP_SEQUENCE": "stop",
|
||||
"TOOL_CALL": "tool_calls",
|
||||
"ERROR_TOXIC": "content_filter",
|
||||
"ERROR": "stop",
|
||||
# HuggingFace / Together AI
|
||||
|
|
|
|||
|
|
@ -3272,7 +3272,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
start_time = self.start_time
|
||||
if end_time is None:
|
||||
end_time = datetime.datetime.now()
|
||||
if self.completion_start_time is None:
|
||||
# A request re-dispatched to a fallback reuses this logging object but runs in
|
||||
# its own window, so a first-token time from the previous attempt cannot belong
|
||||
# to it. Keeping it makes the attempt report a negative TTFT.
|
||||
first_token_seconds: Final = _timestamp_seconds(self.completion_start_time)
|
||||
attempt_start_seconds: Final = _timestamp_seconds(start_time)
|
||||
if first_token_seconds is None or (
|
||||
attempt_start_seconds is not None and first_token_seconds < attempt_start_seconds
|
||||
):
|
||||
self.completion_start_time = end_time
|
||||
self.model_call_details["completion_start_time"] = self.completion_start_time
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import io
|
|||
import json
|
||||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from graphlib import CycleError, TopologicalSorter
|
||||
from itertools import chain, groupby, islice
|
||||
|
|
@ -3235,3 +3235,15 @@ def text_completion_prompt_to_messages(prompt: object) -> tuple[AllMessageValues
|
|||
if string_entries and len(string_entries) == len(entries):
|
||||
return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in string_entries)
|
||||
raise ValueError(f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {prompt_type_name}.")
|
||||
|
||||
|
||||
def is_present_non_string(value: object) -> bool:
|
||||
return value is not None and not isinstance(value, str)
|
||||
|
||||
|
||||
def invalid_type_message(field_path: str, value: object) -> str:
|
||||
return f"Invalid type for '{field_path}': expected a string, got {type(value).__name__}"
|
||||
|
||||
|
||||
def is_unhashable_value_behind(error: TypeError, value: object) -> bool:
|
||||
return not isinstance(value, Hashable) and f"unhashable type: '{type(value).__name__}'" in str(error)
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
|
|||
parts_of,
|
||||
place_mid_conversation_system,
|
||||
split_leading_system_run,
|
||||
string_field,
|
||||
)
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value
|
||||
|
|
@ -2023,7 +2024,7 @@ def _server_tool_answers_replaced_after(messages: Sequence[object], index: int)
|
|||
following: Final = _messages_until_the_next_assistant(messages, index)
|
||||
a_user_turn_remains: Final = any(
|
||||
message_field(message, "role") in ("user", "tool", "function")
|
||||
and message_field(message, "tool_call_id") not in kept_server_result_ids
|
||||
and string_field(message, "tool_call_id") not in kept_server_result_ids
|
||||
for message in following
|
||||
)
|
||||
return _tool_call_ids_in_tool_messages(following) & kept_server_result_ids if a_user_turn_remains else frozenset()
|
||||
|
|
|
|||
|
|
@ -85,6 +85,11 @@ def message_field(message: object, key: str) -> object:
|
|||
return mapping.get(key) if mapping is not None else getattr(message, key, None)
|
||||
|
||||
|
||||
def string_field(message: object, key: str) -> str | None:
|
||||
value: Final = message_field(message, key)
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def is_system_message(message: object) -> bool:
|
||||
return message_field(message, "role") == "system"
|
||||
|
||||
|
|
|
|||
|
|
@ -905,13 +905,21 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def _served_openrouter_model(self) -> str | None:
|
||||
if self.custom_llm_provider == LlmProviders.OPENROUTER.value or (
|
||||
self.custom_llm_provider == "cached_response"
|
||||
and self.logging_obj.model_call_details.get("custom_llm_provider") == LlmProviders.OPENROUTER.value
|
||||
):
|
||||
return self._provider_response_model
|
||||
return None
|
||||
|
||||
def model_response_creator(
|
||||
self,
|
||||
chunk: dict | None = None,
|
||||
hidden_params: Mapping[str, object] | None = None,
|
||||
delta: Delta | None = None,
|
||||
) -> ModelResponseStream:
|
||||
_model: Final = self._cached_model_name
|
||||
_model: Final = self._served_openrouter_model() or self._cached_model_name
|
||||
_logging_obj_llm_provider: Final = self._cached_logging_llm_provider
|
||||
|
||||
if chunk is None:
|
||||
|
|
@ -1699,7 +1707,7 @@ class CustomStreamWrapper:
|
|||
return dispatch_result.value
|
||||
response_obj = dispatch_result.response_obj
|
||||
|
||||
model_response.model = self.model
|
||||
model_response.model = self._served_openrouter_model() or self.model
|
||||
## FUNCTION CALL PARSING
|
||||
original_chunk: Final[_UpstreamChunkLike | None] = (
|
||||
response_obj.get("original_chunk") if response_obj is not None else None
|
||||
|
|
@ -2561,7 +2569,7 @@ class CustomStreamWrapper:
|
|||
if partial is None or not isinstance(usage, Usage):
|
||||
return None
|
||||
if self.model:
|
||||
partial.model = self.model
|
||||
partial.model = self._served_openrouter_model() or self.model
|
||||
backfill_missing_cache_usage_fields(usage)
|
||||
return partial
|
||||
|
||||
|
|
|
|||
|
|
@ -160,8 +160,11 @@ async def resolve_a2a_hop_auth_header(
|
|||
litellm_params: Mapping[str, object],
|
||||
custom_llm_provider: object,
|
||||
resolve_entra_header: AgentAuthHeaderResolver = resolve_azure_ai_agent_auth_header,
|
||||
agent_extra_headers: Mapping[str, str] | None = None,
|
||||
) -> Mapping[str, str] | None:
|
||||
"""Entra credentials authenticate the A2A hop only; a completion-bridge agent hands them to the model provider it bridges to."""
|
||||
if not a2a_hop_uses_entra(litellm_params, custom_llm_provider):
|
||||
if a2a_hop_uses_entra(litellm_params, custom_llm_provider):
|
||||
return await resolve_entra_header(litellm_params)
|
||||
if custom_llm_provider or any(name.lower() == "authorization" for name in (agent_extra_headers or {})):
|
||||
return None
|
||||
return await resolve_entra_header(litellm_params)
|
||||
api_key: Final = litellm_params.get("api_key")
|
||||
return {"Authorization": f"Bearer {api_key}"} if isinstance(api_key, str) and api_key else None
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@ import copy
|
|||
import json
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from functools import partial
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast
|
||||
|
||||
|
|
@ -27,7 +29,10 @@ from litellm.constants import (
|
|||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
DEFS_MAX_INLINED_BYTES,
|
||||
invalid_type_message,
|
||||
is_object_schema,
|
||||
is_present_non_string,
|
||||
is_unhashable_value_behind,
|
||||
sanitize_input_schema_for_anthropic,
|
||||
schema_fragment_ref,
|
||||
schema_grammar_branch_count,
|
||||
|
|
@ -45,6 +50,8 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
|||
inline_remote_image_urls,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
|
||||
message_field,
|
||||
parts_of,
|
||||
place_mid_conversation_system,
|
||||
split_leading_system_run,
|
||||
)
|
||||
|
|
@ -179,6 +186,76 @@ NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS: Final = 24
|
|||
NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES: Final = 2048
|
||||
|
||||
|
||||
def _tool_call_ids_failing(
|
||||
index: int, tool_calls: object, fails: Callable[[object], bool]
|
||||
) -> Iterator[tuple[str, object]]:
|
||||
ids: Final = ((position, message_field(tool_call, "id")) for position, tool_call in enumerate(parts_of(tool_calls)))
|
||||
return (
|
||||
(f"messages[{index}].tool_calls[{position}].id", tool_call_id)
|
||||
for position, tool_call_id in ids
|
||||
if fails(tool_call_id)
|
||||
)
|
||||
|
||||
|
||||
def _roles_failing(messages: Sequence[object], fails: Callable[[object], bool]) -> Iterator[tuple[str, object]]:
|
||||
roles: Final = ((index, message_field(message, "role")) for index, message in enumerate(messages))
|
||||
return ((f"messages[{index}].role", role) for index, role in roles if fails(role))
|
||||
|
||||
|
||||
def _is_non_empty_non_string(value: object) -> bool:
|
||||
return bool(value) and is_present_non_string(value)
|
||||
|
||||
|
||||
def _assistant_tool_call_ids_failing(
|
||||
messages: Sequence[object], fails: Callable[[object], bool]
|
||||
) -> Iterator[tuple[str, object]]:
|
||||
assistant_turns: Final = (
|
||||
(index, message) for index, message in enumerate(messages) if message_field(message, "role") == "assistant"
|
||||
)
|
||||
return chain.from_iterable(
|
||||
_tool_call_ids_failing(index, message_field(message, "tool_calls"), fails) for index, message in assistant_turns
|
||||
)
|
||||
|
||||
|
||||
def _tool_answer_ids_failing(
|
||||
messages: Sequence[object], fails: Callable[[object], bool]
|
||||
) -> Iterator[tuple[str, object]]:
|
||||
answers: Final = (
|
||||
(index, message_field(message, "tool_call_id"))
|
||||
for index, message in enumerate(messages)
|
||||
if message_field(message, "role") in ("tool", "function")
|
||||
)
|
||||
return ((f"messages[{index}].tool_call_id", tool_call_id) for index, tool_call_id in answers if fails(tool_call_id))
|
||||
|
||||
|
||||
def _is_non_string_value_behind(error: AttributeError, value: object) -> bool:
|
||||
return is_present_non_string(value) and value is error.obj
|
||||
|
||||
|
||||
def _fields_behind(messages: Sequence[object], error: Exception) -> Iterable[tuple[str, object]]:
|
||||
match error:
|
||||
case TypeError():
|
||||
caused_error: Final = partial(is_unhashable_value_behind, error)
|
||||
return chain(
|
||||
_assistant_tool_call_ids_failing(messages, caused_error),
|
||||
_tool_answer_ids_failing(messages, caused_error),
|
||||
_roles_failing(messages, caused_error),
|
||||
)
|
||||
case AttributeError():
|
||||
return _assistant_tool_call_ids_failing(messages, partial(_is_non_string_value_behind, error))
|
||||
case _:
|
||||
return ()
|
||||
|
||||
|
||||
def non_string_message_field_error(messages: Sequence[object], error: Exception) -> AnthropicError | None:
|
||||
malformed_field: Final = next(
|
||||
chain(_roles_failing(messages, _is_non_empty_non_string), _fields_behind(messages, error)), None
|
||||
)
|
||||
if malformed_field is None:
|
||||
return None
|
||||
return AnthropicError(status_code=400, message=invalid_type_message(*malformed_field))
|
||||
|
||||
|
||||
def _tool_type(tool: object) -> str | None:
|
||||
match tool:
|
||||
case {"type": str() as tool_type}:
|
||||
|
|
@ -2370,12 +2447,18 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
# Handling anthropic API Prompt Caching
|
||||
if len(anthropic_system_message_list) > 0:
|
||||
optional_params["system"] = anthropic_system_message_list
|
||||
conversation: Final = place_mid_conversation_system(
|
||||
later_messages,
|
||||
supports_mid_conversation_system=supports_mid_conversation_system(
|
||||
model=model, custom_llm_provider=self.custom_llm_provider
|
||||
),
|
||||
)
|
||||
try:
|
||||
conversation: Final = place_mid_conversation_system(
|
||||
later_messages,
|
||||
supports_mid_conversation_system=supports_mid_conversation_system(
|
||||
model=model, custom_llm_provider=self.custom_llm_provider
|
||||
),
|
||||
)
|
||||
except TypeError as error:
|
||||
field_error: Final = non_string_message_field_error(messages, error)
|
||||
if field_error is None:
|
||||
raise
|
||||
raise field_error from error
|
||||
# Format rest of message according to anthropic guidelines
|
||||
try:
|
||||
anthropic_messages = anthropic_messages_pt(
|
||||
|
|
@ -2385,7 +2468,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
drop_params=drop_params,
|
||||
)
|
||||
except Exception as e:
|
||||
raise AnthropicError(
|
||||
conversion_field_error: Final = non_string_message_field_error(messages, e)
|
||||
raise conversion_field_error or AnthropicError(
|
||||
status_code=400,
|
||||
message=f"{e}\nReceived Messages={request_messages}",
|
||||
) # don't use verbose_logger.exception, if exception is raised
|
||||
|
|
|
|||
|
|
@ -118,9 +118,10 @@ def _extract_proxy_litellm_metadata(
|
|||
|
||||
def _requested_router(kwargs: Mapping[str, object]) -> "Router | None":
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
requested: Final = kwargs.get("litellm_router")
|
||||
return requested if isinstance(requested, Router) else None
|
||||
return requested if isinstance(requested, (Router, PythonCore)) else None # pyright: ignore[reportReturnType] # the router passes its core, which carries Router's members
|
||||
|
||||
|
||||
def request_safeguards_evaluator(
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ class _ClearToolUsesFields(BaseModel):
|
|||
clear_tool_inputs: bool | tuple[str, ...] | None = None
|
||||
|
||||
|
||||
_BLOCK: Final = TypeAdapter(Mapping[str, object])
|
||||
_BLOCKS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
# api.anthropic.com, measured 2026-10-19 and 2026-10-28: its wording for the shapes pydantic words differently
|
||||
_API_MESSAGES: Final = MappingProxyType(
|
||||
|
|
@ -317,12 +318,10 @@ def _clear_tool_results(
|
|||
new_blocks: list[object] = []
|
||||
mutated = False
|
||||
for block in content:
|
||||
if (
|
||||
isinstance(block, dict)
|
||||
and block.get("type") == "tool_result"
|
||||
and block.get("tool_use_id") in ids_to_clear
|
||||
):
|
||||
new_blocks.append(_cleared_result(block))
|
||||
if (parsed := _parse(_BLOCK, block)) is not None and _block_id(
|
||||
parsed, "tool_result", "tool_use_id"
|
||||
) in ids_to_clear:
|
||||
new_blocks.append(_cleared_result(parsed))
|
||||
mutated = True
|
||||
cleared += 1
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ def responses_batch_body_to_chat_body(
|
|||
openai_request_body: Mapping[str, object],
|
||||
custom_llm_provider: str | None = None,
|
||||
drop_params: bool | None = None,
|
||||
*,
|
||||
caller_input: str | ResponseInputParam | None = None,
|
||||
) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict
|
||||
"""
|
||||
Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body.
|
||||
|
|
@ -52,5 +54,6 @@ def responses_batch_body_to_chat_body(
|
|||
responses_api_request=responses_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
drop_params=drop_params,
|
||||
caller_input=caller_input,
|
||||
metadata=openai_request_body.get("metadata"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
extract_file_data,
|
||||
text_completion_prompt_to_messages,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import AnthropicError
|
||||
from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body
|
||||
|
|
@ -347,7 +348,14 @@ class _BatchInputLine:
|
|||
record: _OpenAIBatchRecord
|
||||
|
||||
|
||||
_BATCH_RECORD_CONVERSION_ERRORS: Final = (AttributeError, LookupError, TypeError, ValueError, BadRequestError)
|
||||
_BATCH_RECORD_CONVERSION_ERRORS: Final = (
|
||||
AttributeError,
|
||||
LookupError,
|
||||
TypeError,
|
||||
ValueError,
|
||||
BadRequestError,
|
||||
AnthropicError,
|
||||
)
|
||||
|
||||
|
||||
def _batch_line_error(line_number: int, problem: str) -> BedrockError:
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import Final, Literal
|
|||
|
||||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.cohere import CohereV2StreamedToolCall
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionToolCallChunk,
|
||||
|
|
@ -222,6 +223,7 @@ class CohereV2ModelResponseIterator:
|
|||
self.content_blocks: list = []
|
||||
self.tool_index = -1
|
||||
self.json_mode = json_mode
|
||||
self.tool_call_indexes_with_arguments: frozenset[int] = frozenset()
|
||||
|
||||
def _parse_content_delta(self, chunk: dict) -> str:
|
||||
"""Parse content-delta chunks to extract text."""
|
||||
|
|
@ -234,20 +236,23 @@ class CohereV2ModelResponseIterator:
|
|||
return content
|
||||
return ""
|
||||
|
||||
def _parse_tool_call_delta(self, chunk: dict) -> ChatCompletionToolCallChunk | None:
|
||||
"""Parse tool-call-delta chunks to extract tool calls."""
|
||||
delta: Final = chunk.get("delta", {})
|
||||
tool_calls: Final = delta.get("tool_calls", [])
|
||||
if tool_calls:
|
||||
return {
|
||||
"id": tool_calls[0].get("id", ""),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_calls[0].get("name", ""),
|
||||
"arguments": tool_calls[0].get("arguments", ""),
|
||||
},
|
||||
}
|
||||
return None
|
||||
def _parse_tool_call_delta(self, chunk: dict, index: int) -> ChatCompletionToolCallChunk:
|
||||
"""Parse tool-call-start and tool-call-delta chunks to extract tool calls."""
|
||||
tool_call: Final = CohereV2StreamedToolCall.model_validate(chunk)
|
||||
if tool_call.arguments:
|
||||
self.tool_call_indexes_with_arguments = self.tool_call_indexes_with_arguments | {index}
|
||||
return {
|
||||
"id": tool_call.id,
|
||||
"type": "function",
|
||||
"function": {"name": tool_call.name, "arguments": tool_call.arguments},
|
||||
"index": index,
|
||||
}
|
||||
|
||||
def _parse_tool_call_end(self, index: int) -> ChatCompletionToolCallChunk | None:
|
||||
"""Close a tool call that streamed no arguments with an empty JSON object, as OpenAI streams one."""
|
||||
if index in self.tool_call_indexes_with_arguments:
|
||||
return None
|
||||
return {"id": None, "type": "function", "function": {"name": None, "arguments": "{}"}, "index": index}
|
||||
|
||||
def _parse_tool_plan_delta(self, chunk: dict) -> dict | None:
|
||||
"""Parse tool-plan-delta events to extract tool plan."""
|
||||
|
|
@ -301,7 +306,8 @@ class CohereV2ModelResponseIterator:
|
|||
|
||||
v2 format:
|
||||
- Content: chunk.type == "content-delta" -> chunk.delta.message.content.text
|
||||
- Tool calls: chunk.type == "tool-call-delta" -> chunk.delta.tool_calls
|
||||
- Tool calls: chunk.type in ("tool-call-start", "tool-call-delta") -> chunk.delta.message.tool_calls
|
||||
- Tool call end: chunk.type == "tool-call-end" -> "{}" arguments when the call streamed none
|
||||
- Tool plan: chunk.event == "tool-plan-delta" -> chunk.data.delta.message.tool_plan
|
||||
- Citations: chunk.event == "citation-start" -> chunk.data.delta.message.citations
|
||||
- Finish: chunk.event == "message-end" -> chunk.data.delta.finish_reason
|
||||
|
|
@ -321,8 +327,12 @@ class CohereV2ModelResponseIterator:
|
|||
# Handle different chunk types
|
||||
if chunk_type == "content-delta":
|
||||
text = self._parse_content_delta(chunk)
|
||||
elif chunk_type == "tool-call-delta":
|
||||
tool_use = self._parse_tool_call_delta(chunk)
|
||||
elif chunk_type in ("tool-call-start", "tool-call-delta", "tool-call-end"):
|
||||
tool_use = (
|
||||
self._parse_tool_call_end(index)
|
||||
if chunk_type == "tool-call-end"
|
||||
else self._parse_tool_call_delta(chunk, index)
|
||||
)
|
||||
elif event_type == "tool-plan-delta":
|
||||
provider_specific_fields = self._parse_tool_plan_delta(chunk)
|
||||
elif event_type == "citation-start":
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig
|
|||
from litellm.llms.base_llm.image_variations.transformation import (
|
||||
BaseImageVariationConfig,
|
||||
)
|
||||
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
|
||||
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -90,7 +91,7 @@ class BaseLLMAIOHTTPHandler:
|
|||
if self.transport is None:
|
||||
connector: Final = self._get_connector()
|
||||
if connector:
|
||||
return aiohttp.ClientSession(connector=connector)
|
||||
return aiohttp.ClientSession(connector=connector, trace_configs=[build_pool_wait_trace_config()])
|
||||
|
||||
transport: Final = self.transport or self._get_or_create_transport()
|
||||
if transport is not None and hasattr(transport, "_get_valid_client_session"):
|
||||
|
|
@ -99,7 +100,7 @@ class BaseLLMAIOHTTPHandler:
|
|||
except RuntimeError:
|
||||
pass
|
||||
|
||||
return aiohttp.ClientSession()
|
||||
return aiohttp.ClientSession(trace_configs=[build_pool_wait_trace_config()])
|
||||
|
||||
def _get_async_client_session(self, dynamic_client_session: ClientSession | None = None) -> ClientSession:
|
||||
if dynamic_client_session:
|
||||
|
|
|
|||
130
litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py
Normal file
130
litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
from aiohttp import (
|
||||
ClientSession,
|
||||
TraceConfig,
|
||||
TraceConnectionQueuedEndParams,
|
||||
TraceConnectionQueuedStartParams,
|
||||
TraceRequestExceptionParams,
|
||||
TraceRequestRedirectParams,
|
||||
TraceRequestStartParams,
|
||||
)
|
||||
from yarl import URL
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
_POOL_WAIT_WARNING_SECONDS: Final = 1.0
|
||||
|
||||
|
||||
class _PoolWaitContext(SimpleNamespace):
|
||||
def __init__(self, trace_request_ctx: object = None) -> None:
|
||||
super().__init__()
|
||||
self.host: str | None = None
|
||||
self.port: int | None = None
|
||||
self.started_at: float | None = None
|
||||
self.waited: float = 0.0
|
||||
self.warning: asyncio.TimerHandle | None = None
|
||||
self.warned: bool = False
|
||||
|
||||
def set_destination(self, params: TraceRequestStartParams) -> None:
|
||||
self.host = params.url.host
|
||||
self.port = params.url.port
|
||||
|
||||
def redirect(self, params: TraceRequestRedirectParams) -> None:
|
||||
location: Final = params.response.headers.get("Location") or params.response.headers.get("URI")
|
||||
if location is None:
|
||||
return
|
||||
try:
|
||||
destination: Final = params.url.join(URL(location))
|
||||
host: Final = destination.host
|
||||
port: Final = destination.port
|
||||
except ValueError:
|
||||
return
|
||||
self.host = host
|
||||
self.port = port
|
||||
|
||||
def start(self, session: ClientSession) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
self.started_at = loop.time()
|
||||
if not self.warned:
|
||||
self.warning = loop.call_later(max(0.0, _POOL_WAIT_WARNING_SECONDS - self.waited), self.warn, session)
|
||||
|
||||
def elapsed(self) -> float:
|
||||
return self.waited + (
|
||||
asyncio.get_running_loop().time() - self.started_at if self.started_at is not None else 0.0
|
||||
)
|
||||
|
||||
def warn(self, session: ClientSession) -> None:
|
||||
if self.warned or self.started_at is None:
|
||||
return
|
||||
self.warned = True
|
||||
connector: Final = session.connector
|
||||
if connector is None:
|
||||
return
|
||||
verbose_logger.warning(
|
||||
"aiohttp connection pool saturated: host=%s port=%s waited=%.3fs for a connection slot "
|
||||
"(acquired_total=%s, limit=%s, limit_per_host=%s). "
|
||||
"Wait occurred before sending the request upstream.",
|
||||
self.host,
|
||||
self.port,
|
||||
self.elapsed(),
|
||||
len(connector._acquired), # pyright: ignore[reportPrivateUsage] # aiohttp exposes no public acquired count
|
||||
connector.limit,
|
||||
connector.limit_per_host,
|
||||
)
|
||||
|
||||
def finish(self, session: ClientSession, failed: bool = False) -> None:
|
||||
if self.warning is not None:
|
||||
self.warning.cancel()
|
||||
self.warning = None
|
||||
if self.started_at is None:
|
||||
return
|
||||
if failed or self.elapsed() >= _POOL_WAIT_WARNING_SECONDS:
|
||||
self.warn(session)
|
||||
self.waited = self.elapsed()
|
||||
self.started_at = None
|
||||
|
||||
|
||||
async def _on_request_start(session: ClientSession, context: SimpleNamespace, params: TraceRequestStartParams) -> None:
|
||||
if isinstance(context, _PoolWaitContext):
|
||||
context.set_destination(params)
|
||||
|
||||
|
||||
async def _on_queued_start(
|
||||
session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedStartParams
|
||||
) -> None:
|
||||
if isinstance(context, _PoolWaitContext):
|
||||
context.start(session)
|
||||
|
||||
|
||||
async def _on_queued_end(
|
||||
session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedEndParams
|
||||
) -> None:
|
||||
if isinstance(context, _PoolWaitContext):
|
||||
context.finish(session)
|
||||
|
||||
|
||||
async def _on_request_exception(
|
||||
session: ClientSession, context: SimpleNamespace, params: TraceRequestExceptionParams
|
||||
) -> None:
|
||||
if isinstance(context, _PoolWaitContext):
|
||||
context.finish(session, failed=True)
|
||||
|
||||
|
||||
async def _on_request_redirect(
|
||||
session: ClientSession, context: SimpleNamespace, params: TraceRequestRedirectParams
|
||||
) -> None:
|
||||
if isinstance(context, _PoolWaitContext):
|
||||
context.redirect(params)
|
||||
|
||||
|
||||
def build_pool_wait_trace_config() -> TraceConfig:
|
||||
trace: Final = TraceConfig(trace_config_ctx_factory=_PoolWaitContext)
|
||||
trace.on_request_start.append(_on_request_start)
|
||||
trace.on_connection_queued_start.append(_on_queued_start)
|
||||
trace.on_connection_queued_end.append(_on_queued_end)
|
||||
trace.on_request_exception.append(_on_request_exception)
|
||||
trace.on_request_redirect.append(_on_request_redirect)
|
||||
return trace
|
||||
|
|
@ -44,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
|||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
|
||||
from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread
|
||||
from litellm.types.llms.custom_http import *
|
||||
|
||||
|
|
@ -1456,6 +1457,7 @@ class AsyncHTTPHandler:
|
|||
connector=TCPConnector(**transport_connector_kwargs),
|
||||
cookie_jar=DummyCookieJar(),
|
||||
trust_env=trust_env,
|
||||
trace_configs=[build_pool_wait_trace_config()],
|
||||
)
|
||||
|
||||
# Use shared session if provided and valid
|
||||
|
|
|
|||
|
|
@ -1,16 +1,21 @@
|
|||
import json
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from urllib.parse import parse_qs, urlparse, urlsplit, urlunsplit
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL, DEEPGRAM_SPEAK_DEFAULT_MODEL
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionUsage
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration
|
||||
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"})
|
||||
|
|
@ -199,3 +204,71 @@ def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]
|
|||
if frame.get("type") == "Results"
|
||||
if (transcript := _final_transcript(frame)) is not None
|
||||
)
|
||||
|
||||
|
||||
_AURA_GENERATION: Final = re.compile(r"aura(?:-\d+)?(?=-[a-z])")
|
||||
|
||||
|
||||
def deepgram_speak_websocket_target(api_base: str | None, query_string: str) -> str:
|
||||
speak_url: Final = urlsplit(join_api_base(api_base or DEEPGRAM_DEFAULT_API_BASE, "speak"))
|
||||
scheme: Final = _WEBSOCKET_SCHEMES.get(speak_url.scheme, speak_url.scheme)
|
||||
params: Final = _first_occurrences(query_string)
|
||||
query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_SPEAK_DEFAULT_MODEL)
|
||||
return urlunsplit((scheme, speak_url.netloc, speak_url.path, str(query), ""))
|
||||
|
||||
|
||||
def deepgram_speak_model(upstream_url: str) -> str:
|
||||
models: Final = parse_qs(urlparse(upstream_url).query).get("model")
|
||||
return models[0] if models else DEEPGRAM_SPEAK_DEFAULT_MODEL
|
||||
|
||||
|
||||
def deepgram_speak_requested_model(query_string: str) -> str:
|
||||
"""Read off the upstream target, so the voice a key is authorized for is the voice Deepgram receives"""
|
||||
return deepgram_speak_model(deepgram_speak_websocket_target(None, query_string))
|
||||
|
||||
|
||||
def deepgram_speak_registry_key(upstream_url: str) -> str:
|
||||
"""Deepgram bills every voice of an Aura generation at one rate, so ``aura-2-thalia-en`` is priced at the
|
||||
``deepgram/aura-2`` row"""
|
||||
model: Final = deepgram_speak_model(upstream_url)
|
||||
generation: Final = _AURA_GENERATION.match(model)
|
||||
return f"{LlmProviders.DEEPGRAM.value}/{generation.group() if generation else model}"
|
||||
|
||||
|
||||
def deepgram_speak_character_rate(upstream_url: str) -> float | None:
|
||||
"""Only an exact registry hit counts, so a voice is never billed at a neighbouring row's rate"""
|
||||
registry_key: Final = deepgram_speak_registry_key(upstream_url)
|
||||
model_info: Final = lookup_model_info_for_registration(registry_key, LlmProviders.DEEPGRAM.value)
|
||||
if isinstance(model_info, ModelInfoLookupMiss) or model_info["key"] != registry_key:
|
||||
return None
|
||||
return model_info.get("input_cost_per_character")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeepgramSpeakMeter:
|
||||
"""Holds the rate checked when the socket opened, so a cost-map reload mid-session can't change or drop the price"""
|
||||
|
||||
upstream_url: str
|
||||
character_rate: float
|
||||
|
||||
def measure(self, client_text: str) -> int:
|
||||
"""Counts a Speak message's text as it is sent. Deepgram also speaks, and bills, unflushed text on its own (a
|
||||
long buffer, Close, a close frame), so counting only flushed text would let a client get audio free, while
|
||||
counting on send can only over-bill text Deepgram discards, such as a Clear or an oversized Speak. Never
|
||||
raises, so a malformed frame can't end the session"""
|
||||
try:
|
||||
message: Final[object] = json.loads(client_text)
|
||||
except (ValueError, RecursionError):
|
||||
return 0
|
||||
match message:
|
||||
case {"type": "Speak", "text": str() as text}:
|
||||
return len(text)
|
||||
case _:
|
||||
return 0
|
||||
|
||||
def price(self, total: int) -> WebSocketSessionUsage:
|
||||
return WebSocketSessionUsage(
|
||||
model=deepgram_speak_model(self.upstream_url),
|
||||
custom_llm_provider=LlmProviders.DEEPGRAM.value,
|
||||
response_cost=total * self.character_rate,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ async def search_skills(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> SkillSearchOutcome:
|
||||
if embedding_model is None:
|
||||
if embedding_model is None or not embedding_model.strip():
|
||||
return SkillSearchNotConfigured(
|
||||
reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -437,6 +437,8 @@ _NAMED_UNKNOWN_VOICE_KEYS: Final = 10
|
|||
|
||||
_NAMED_VOICE_KEY_CHARS: Final = 100
|
||||
|
||||
_NAMED_AUDIO_FORMAT_CHARS: Final = 100
|
||||
|
||||
|
||||
def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
camel_config: Final = as_str_mapping(value.get("speechConfig"))
|
||||
|
|
@ -473,6 +475,31 @@ def reject_unknown_gemini_tts_voice_keys(
|
|||
)
|
||||
|
||||
|
||||
def _named_audio_format(audio_format: object) -> str:
|
||||
if isinstance(audio_format, str):
|
||||
return repr(audio_format[:_NAMED_AUDIO_FORMAT_CHARS])
|
||||
return repr(audio_format)[:_NAMED_AUDIO_FORMAT_CHARS]
|
||||
|
||||
|
||||
def unsupported_gemini_tts_audio_format(
|
||||
model: str, audio: Mapping[str, object], drop_params: bool, llm_provider: str | None
|
||||
) -> UnsupportedParamsError | None:
|
||||
nested_config: Final = _gemini_speech_config_wrapper(audio) or {}
|
||||
audio_format: Final = next(
|
||||
(item for item in (audio.get("format"), nested_config.get("format")) if item not in (None, "pcm16")), None
|
||||
)
|
||||
if audio_format is None or drops_unsupported_params(drop_params):
|
||||
return None
|
||||
return UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not take the audio format {_named_audio_format(audio_format)}. Set format "
|
||||
"to 'pcm16' or leave it out, or set drop_params to drop it."
|
||||
),
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
|
||||
def normalize_gemini_speech_config(
|
||||
value: Mapping[str, object], drop_unknown_fields: bool = False
|
||||
) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict
|
||||
|
|
@ -1555,15 +1582,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
languageCode: "en-US",
|
||||
}
|
||||
"""
|
||||
# Validate audio format - Gemini TTS only supports pcm16
|
||||
audio_format: Final = value.get("format")
|
||||
if audio_format is not None and audio_format != "pcm16":
|
||||
raise ValueError(
|
||||
f"Unsupported audio format for Gemini TTS models: {audio_format}. "
|
||||
f"Gemini TTS models only support 'pcm16' format as they return audio data in L16 PCM format. "
|
||||
f"Please set audio format to 'pcm16'."
|
||||
)
|
||||
|
||||
audio_format_error: Final = unsupported_gemini_tts_audio_format(
|
||||
"Gemini TTS", value, drop_unknown_fields, self.custom_llm_provider
|
||||
)
|
||||
if audio_format_error is not None:
|
||||
raise audio_format_error
|
||||
camel_speech_config: Final = value.get("speechConfig")
|
||||
snake_speech_config: Final = value.get("speech_config")
|
||||
speech_config_value: Final = (
|
||||
|
|
@ -1601,6 +1624,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
def _map_chat_audio_params(
|
||||
self, model: str, value: Mapping[str, object], drop_params: bool
|
||||
) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary
|
||||
audio_format_error: Final = unsupported_gemini_tts_audio_format(
|
||||
model, value, drop_params, self.custom_llm_provider
|
||||
)
|
||||
if audio_format_error is not None:
|
||||
raise audio_format_error
|
||||
if _gemini_speech_config_wrapper(value) is not None:
|
||||
reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider)
|
||||
return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params))
|
||||
|
|
|
|||
|
|
@ -15555,18 +15555,20 @@
|
|||
"input_cost_per_token": 2.65e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 8000,
|
||||
"max_output_tokens": 2048,
|
||||
"max_output_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.5e-06,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
|
||||
"supports_pdf_input": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 8000,
|
||||
"max_output_tokens": 2048,
|
||||
"max_output_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
|
||||
"supports_pdf_input": true
|
||||
},
|
||||
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
|
||||
|
|
@ -48970,6 +48972,7 @@
|
|||
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
|
||||
"deprecation_date": "2026-11-04",
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"deprecation_date": "2026-11-04",
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1010000,
|
||||
|
|
@ -53970,23 +53973,23 @@
|
|||
"supports_vision": true
|
||||
},
|
||||
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
|
||||
"max_input_tokens": 131000,
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"output_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V3.1": {
|
||||
"supports_reasoning": true,
|
||||
"max_input_tokens": 161000,
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"input_cost_per_token": 5.5e-07,
|
||||
"output_cost_per_token": 1.65e-06,
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
|
||||
},
|
||||
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
|
||||
"max_input_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -105,10 +105,10 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
LazyFeature(
|
||||
name="a2a",
|
||||
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
|
||||
# ``/v1/a2a/{agent_id}/message/send`` is caught via the suffix so the
|
||||
# ``/v1/a2a/{agent_id}/message/send`` and ``/message/stream`` are caught via suffixes so the
|
||||
# ``/v1/a2a`` prefix doesn't subsume the discover prefix below.
|
||||
path_prefixes=("/a2a",),
|
||||
path_suffixes=("/message/send",),
|
||||
path_suffixes=("/message/send", "/message/stream"),
|
||||
),
|
||||
LazyFeature(
|
||||
name="a2a_registration",
|
||||
|
|
|
|||
|
|
@ -54,6 +54,57 @@
|
|||
],
|
||||
"title": "ValidationError",
|
||||
"type": "object"
|
||||
},
|
||||
"_A2AJsonRpcResponse": {
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
"error": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Error"
|
||||
},
|
||||
"id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Id"
|
||||
},
|
||||
"jsonrpc": {
|
||||
"const": "2.0",
|
||||
"title": "Jsonrpc",
|
||||
"type": "string"
|
||||
},
|
||||
"result": {
|
||||
"anyOf": [
|
||||
{},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Result"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"jsonrpc",
|
||||
"id"
|
||||
],
|
||||
"title": "_A2AJsonRpcResponse",
|
||||
"type": "object"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
@ -242,6 +293,119 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/a2a/{agent_id}/message/stream": {
|
||||
"post": {
|
||||
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
|
||||
"operationId": "invoke_agent_a2a_a2a__agent_id__message_stream_post",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "agent_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Agent Id",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"example": {
|
||||
"id": "stream-1",
|
||||
"jsonrpc": "2.0",
|
||||
"method": "message/stream",
|
||||
"params": {
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"messageId": "message-1",
|
||||
"parts": [
|
||||
{
|
||||
"kind": "text",
|
||||
"text": "Hello"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
}
|
||||
}
|
||||
},
|
||||
"schema": {
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
"id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Id"
|
||||
},
|
||||
"jsonrpc": {
|
||||
"const": "2.0",
|
||||
"title": "Jsonrpc",
|
||||
"type": "string"
|
||||
},
|
||||
"method": {
|
||||
"title": "Method",
|
||||
"type": "string"
|
||||
},
|
||||
"params": {
|
||||
"additionalProperties": true,
|
||||
"title": "Params",
|
||||
"type": "object"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"jsonrpc",
|
||||
"method",
|
||||
"params"
|
||||
],
|
||||
"title": "_A2AJsonRpcRequest",
|
||||
"type": "object"
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": true
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Invoke Agent A2A",
|
||||
"tags": [
|
||||
"a2a"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/a2a/{agent_id}/message/send": {
|
||||
"post": {
|
||||
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
|
||||
|
|
@ -287,6 +451,119 @@
|
|||
"a2a"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/a2a/{agent_id}/message/stream": {
|
||||
"post": {
|
||||
"description": "Invoke an agent using the A2A protocol (JSON-RPC 2.0).\n\nSupported methods:\n- message/send: Send a message and get a response\n- message/stream: Send a message and stream the response",
|
||||
"operationId": "invoke_agent_a2a_v1_a2a__agent_id__message_stream_post",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "agent_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Agent Id",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"example": {
|
||||
"id": "stream-1",
|
||||
"jsonrpc": "2.0",
|
||||
"method": "message/stream",
|
||||
"params": {
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"messageId": "message-1",
|
||||
"parts": [
|
||||
{
|
||||
"kind": "text",
|
||||
"text": "Hello"
|
||||
}
|
||||
],
|
||||
"role": "user"
|
||||
}
|
||||
}
|
||||
},
|
||||
"schema": {
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
"id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Id"
|
||||
},
|
||||
"jsonrpc": {
|
||||
"const": "2.0",
|
||||
"title": "Jsonrpc",
|
||||
"type": "string"
|
||||
},
|
||||
"method": {
|
||||
"title": "Method",
|
||||
"type": "string"
|
||||
},
|
||||
"params": {
|
||||
"additionalProperties": true,
|
||||
"title": "Params",
|
||||
"type": "object"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"jsonrpc",
|
||||
"method",
|
||||
"params"
|
||||
],
|
||||
"title": "_A2AJsonRpcRequest",
|
||||
"type": "object"
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": true
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/_A2AJsonRpcResponse"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Invoke Agent A2A",
|
||||
"tags": [
|
||||
"a2a"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
@ -15326,6 +15603,10 @@
|
|||
],
|
||||
"title": "Metadata"
|
||||
},
|
||||
"response_format": {
|
||||
"$ref": "#/components/schemas/GuardrailResponseFormat",
|
||||
"default": "standard"
|
||||
},
|
||||
"text": {
|
||||
"title": "Text",
|
||||
"type": "string"
|
||||
|
|
@ -16203,6 +16484,14 @@
|
|||
"title": "BedrockChecksSensitiveInformationModel",
|
||||
"type": "object"
|
||||
},
|
||||
"BedrockGuardrailOutputScope": {
|
||||
"enum": [
|
||||
"INTERVENTIONS",
|
||||
"FULL"
|
||||
],
|
||||
"title": "BedrockGuardrailOutputScope",
|
||||
"type": "string"
|
||||
},
|
||||
"BlockedWord": {
|
||||
"description": "Represents a blocked word with its action and optional description",
|
||||
"properties": {
|
||||
|
|
@ -16746,6 +17035,14 @@
|
|||
"title": "GuardrailInfoResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"GuardrailResponseFormat": {
|
||||
"enum": [
|
||||
"standard",
|
||||
"native"
|
||||
],
|
||||
"title": "GuardrailResponseFormat",
|
||||
"type": "string"
|
||||
},
|
||||
"GuardrailSubmissionItem": {
|
||||
"properties": {
|
||||
"created_at": {
|
||||
|
|
@ -16998,6 +17295,7 @@
|
|||
"title": "HTTPValidationError",
|
||||
"type": "object"
|
||||
},
|
||||
"JsonValue": {},
|
||||
"LakeraCategoryThresholds": {
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
|
|
@ -17924,6 +18222,18 @@
|
|||
"description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.",
|
||||
"title": "Fail On Error"
|
||||
},
|
||||
"file_checkpoint_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).",
|
||||
"title": "File Checkpoint Id"
|
||||
},
|
||||
"file_guardrail_timeout": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -18426,6 +18736,17 @@
|
|||
],
|
||||
"description": "Optional parameters for the guardrail"
|
||||
},
|
||||
"outputScope": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/BedrockGuardrailOutputScope"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "ApplyGuardrail response scope. FULL includes detected and non-detected findings. Omit to retain Bedrock's default of returning detected findings only."
|
||||
},
|
||||
"output_parse_pii": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -20308,7 +20629,24 @@
|
|||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ApplyGuardrailResponse"
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/ApplyGuardrailResponse"
|
||||
},
|
||||
{
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/JsonValue"
|
||||
},
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/JsonValue"
|
||||
},
|
||||
"type": "array"
|
||||
}
|
||||
],
|
||||
"title": "Response Apply Guardrail Apply Guardrail Post"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
@ -20400,7 +20738,24 @@
|
|||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ApplyGuardrailResponse"
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/ApplyGuardrailResponse"
|
||||
},
|
||||
{
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/JsonValue"
|
||||
},
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/JsonValue"
|
||||
},
|
||||
"type": "array"
|
||||
}
|
||||
],
|
||||
"title": "Response Apply Guardrail Guardrails Apply Guardrail Post"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -501,6 +501,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/a2a/{agent_id}/message/send",
|
||||
"/v1/a2a/{agent_id}/message/send",
|
||||
"/a2a/{agent_id}/message/stream",
|
||||
"/v1/a2a/{agent_id}/message/stream",
|
||||
"/a2a/{agent_id}/.well-known/agent-card.json",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ from collections.abc import Mapping
|
|||
from copy import deepcopy
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
|
||||
SupportedA2AVersion = Literal["0.3", "1.0"]
|
||||
|
||||
# Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent;
|
||||
|
|
@ -21,7 +23,6 @@ SUPPORTED_A2A_PROTOCOL_VERSIONS: Final[tuple[SupportedA2AVersion, ...]] = ("0.3"
|
|||
# Default served version when the agent card does not pin one.
|
||||
LITELLM_A2A_PROTOCOL_VERSION: Final = "1.0"
|
||||
|
||||
|
||||
_PROTOCOL_VERSION_PATTERN: Final = re.compile(
|
||||
r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$"
|
||||
)
|
||||
|
|
@ -50,6 +51,24 @@ def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str:
|
|||
return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION
|
||||
|
||||
|
||||
def point_agent_card_at_proxy(
|
||||
agent_card: Mapping[str, object],
|
||||
proxy_url: str,
|
||||
) -> dict[str, object]: # mutable-ok: the public helper contract requires a dictionary
|
||||
interfaces: Final = agent_card.get("supportedInterfaces")
|
||||
interface_items: Final = (
|
||||
tuple(iter_items(agent_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
|
||||
)
|
||||
first: Final[object] = interface_items[0] if interface_items else None
|
||||
first_mapping: Final = as_str_mapping(first)
|
||||
pointed_interfaces: Final = (
|
||||
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}, *interface_items[1:]]}
|
||||
if isinstance(first, dict) and first_mapping is not None
|
||||
else {}
|
||||
)
|
||||
return deepcopy({**agent_card, "url": proxy_url, **pointed_interfaces})
|
||||
|
||||
|
||||
# Security scheme exposed by the LiteLLM-fronted agent card. Always replaces
|
||||
# whatever upstream advertised — the client must authenticate to the proxy,
|
||||
# not the upstream agent.
|
||||
|
|
|
|||
|
|
@ -11,15 +11,15 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from copy import deepcopy
|
||||
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from fastapi.routing import APIRoute
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -27,7 +27,8 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
|||
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
||||
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
|
||||
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy
|
||||
from litellm.proxy.a2a.version_convert import (
|
||||
A2AVersion,
|
||||
normalize_agent_card,
|
||||
|
|
@ -42,6 +43,8 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
|
|||
)
|
||||
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body
|
||||
from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception
|
||||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
SSE_COMMENT_PING,
|
||||
coerce_keepalive_interval,
|
||||
|
|
@ -57,6 +60,42 @@ if TYPE_CHECKING:
|
|||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
class _A2AJsonRpcRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
|
||||
jsonrpc: Literal["2.0"]
|
||||
id: str | int | None = None
|
||||
method: str
|
||||
params: Mapping[str, object]
|
||||
|
||||
|
||||
class _A2AJsonRpcResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
|
||||
jsonrpc: Literal["2.0"]
|
||||
id: str | int | None
|
||||
result: object | None = None
|
||||
error: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
_A2A_MESSAGE_STREAM_REQUEST_BODY: Final = inline_request_body(
|
||||
_A2AJsonRpcRequest,
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "stream-1",
|
||||
"method": "message/stream",
|
||||
"params": {
|
||||
"message": {
|
||||
"messageId": "message-1",
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hello"}],
|
||||
"kind": "message",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
# Mirrors the native seam's own headers: a reverse proxy that batches the whole
|
||||
# stream would swallow the keepalives this route sends to defeat idle timeouts.
|
||||
_SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
|
|
@ -172,13 +211,16 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str,
|
|||
async def _resolve_backend_auth(
|
||||
litellm_params: dict[str, object],
|
||||
custom_llm_provider: object,
|
||||
agent_extra_headers: Mapping[str, str] | None = None,
|
||||
) -> AgentBackendAuth:
|
||||
databricks_auth: Final = await resolve_databricks_backend_auth(litellm_params)
|
||||
if has_databricks_oauth(litellm_params):
|
||||
return databricks_auth
|
||||
return AgentBackendAuth(
|
||||
litellm_params=databricks_auth.litellm_params,
|
||||
header=await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider),
|
||||
header=await resolve_a2a_hop_auth_header(
|
||||
litellm_params, custom_llm_provider, agent_extra_headers=agent_extra_headers
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -201,6 +243,35 @@ def _forwarding_headers(
|
|||
return merged or None
|
||||
|
||||
|
||||
async def _log_agent_rejection(
|
||||
*,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_id: object,
|
||||
method: str | None,
|
||||
rejection: HTTPException,
|
||||
target_agent_id: str,
|
||||
model: str,
|
||||
) -> None:
|
||||
request_data: Final[dict[str, object]] = { # mutable-ok: the failure hook stamps its logging state onto it
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": method,
|
||||
"model": model,
|
||||
"agent_id": target_agent_id,
|
||||
"metadata": {"agent_id": target_agent_id},
|
||||
}
|
||||
try:
|
||||
await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig
|
||||
request_data=request_data,
|
||||
original_exception=rejection,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route="/a2a/{agent_id}",
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a logging failure must not change the rejection the caller gets
|
||||
verbose_proxy_logger.exception("Failed to log rejected A2A request")
|
||||
|
||||
|
||||
def _jsonrpc_error(
|
||||
request_id: object,
|
||||
code: int,
|
||||
|
|
@ -218,6 +289,57 @@ def _jsonrpc_error(
|
|||
)
|
||||
|
||||
|
||||
class _JsonRpcRequestId(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str | int | None = None
|
||||
|
||||
|
||||
async def _jsonrpc_request_id(request: Request) -> str | int | None:
|
||||
try:
|
||||
return _JsonRpcRequestId.model_validate_json(await request.body()).id
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def _jsonrpc_refusal(request: Request, refusal: ProxyException | HTTPException) -> JSONResponse:
|
||||
from litellm.proxy.proxy_server import (
|
||||
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does
|
||||
)
|
||||
|
||||
error: Final = refusal if isinstance(refusal, ProxyException) else status_typed_proxy_exception(refusal)
|
||||
_close_dangling_otel_server_span(request, int(error.code), exc=refusal)
|
||||
return JSONResponse(
|
||||
content={
|
||||
"jsonrpc": "2.0",
|
||||
"id": await _jsonrpc_request_id(request),
|
||||
"error": {
|
||||
"code": -32000,
|
||||
"message": error.message,
|
||||
"data": {"type": error.type, "param": error.param, "code": error.code},
|
||||
},
|
||||
},
|
||||
status_code=int(error.code),
|
||||
headers=error.headers if isinstance(refusal, ProxyException) else refusal.headers,
|
||||
)
|
||||
|
||||
|
||||
class _JsonRpcRefusalRoute(APIRoute):
|
||||
def get_route_handler(self) -> Callable[[Request], Coroutine[object, object, Response]]:
|
||||
handler: Final = super().get_route_handler()
|
||||
|
||||
async def answer_refusals_as_jsonrpc(request: Request) -> Response:
|
||||
try:
|
||||
return await handler(request)
|
||||
except (ProxyException, HTTPException) as refusal:
|
||||
return await _jsonrpc_refusal(request, refusal)
|
||||
|
||||
return answer_refusals_as_jsonrpc
|
||||
|
||||
|
||||
jsonrpc_router: Final = APIRouter(route_class=_JsonRpcRefusalRoute)
|
||||
|
||||
|
||||
async def _get_agent(agent_id: str) -> "AgentResponse | None":
|
||||
"""Look up an agent by ID or name. Returns None if not found."""
|
||||
from litellm.proxy.common_utils.registry_read_through import (
|
||||
|
|
@ -628,16 +750,12 @@ async def get_agent_card(
|
|||
)
|
||||
|
||||
proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}")
|
||||
agent_card = deepcopy(agent.agent_card_params)
|
||||
agent_card["url"] = proxy_url
|
||||
interfaces: Final = agent_card.get("supportedInterfaces")
|
||||
if isinstance(interfaces, list) and interfaces:
|
||||
interfaces[0]["url"] = proxy_url
|
||||
agent_card: Final = point_agent_card_at_proxy(agent.agent_card_params, proxy_url)
|
||||
served_version: Final = _served_version(agent, request)
|
||||
agent_card = normalize_agent_card(agent_card, served_version)
|
||||
normalized_agent_card: Final = normalize_agent_card(agent_card, served_version)
|
||||
|
||||
verbose_proxy_logger.debug("Returning agent card for '%s' with proxy URL: %s", agent_id, proxy_url)
|
||||
return JSONResponse(content=agent_card)
|
||||
return JSONResponse(content=normalized_agent_card)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
|
|
@ -646,21 +764,35 @@ async def get_agent_card(
|
|||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post(
|
||||
@jsonrpc_router.post(
|
||||
"/a2a/{agent_id}",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.post(
|
||||
@jsonrpc_router.post(
|
||||
"/a2a/{agent_id}/message/send",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.post(
|
||||
@jsonrpc_router.post(
|
||||
"/v1/a2a/{agent_id}/message/send",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.post(
|
||||
"/a2a/{agent_id}/message/stream",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
|
||||
responses={200: {"model": _A2AJsonRpcResponse}},
|
||||
)
|
||||
@router.post(
|
||||
"/v1/a2a/{agent_id}/message/stream",
|
||||
tags=["[beta] A2A Agents"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
openapi_extra={"requestBody": _A2A_MESSAGE_STREAM_REQUEST_BODY},
|
||||
responses={200: {"model": _A2AJsonRpcResponse}},
|
||||
)
|
||||
async def invoke_agent_a2a(
|
||||
agent_id: str,
|
||||
request: Request,
|
||||
|
|
@ -722,6 +854,15 @@ async def invoke_agent_a2a(
|
|||
# Find the agent
|
||||
agent: Final = await _get_agent(agent_id)
|
||||
if agent is None:
|
||||
await _log_agent_rejection(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_id=request_id,
|
||||
method=method,
|
||||
rejection=HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found"),
|
||||
target_agent_id=agent_id,
|
||||
model=f"a2a_agent/{agent_id}",
|
||||
)
|
||||
return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404)
|
||||
|
||||
served_version: Final = _served_version(agent, request, original_method)
|
||||
|
|
@ -731,10 +872,20 @@ async def invoke_agent_a2a(
|
|||
user_api_key_auth=user_api_key_dict,
|
||||
)
|
||||
if not is_allowed:
|
||||
raise HTTPException(
|
||||
denial: Final = HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
|
||||
)
|
||||
await _log_agent_rejection(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_id=request_id,
|
||||
method=method,
|
||||
rejection=denial,
|
||||
target_agent_id=agent.agent_id,
|
||||
model=f"a2a_agent/{(as_str_mapping(agent.agent_card_params or {}) or {}).get('name', agent_id)}",
|
||||
)
|
||||
raise denial
|
||||
|
||||
user_api_key_dict.invoked_agent_id = agent.agent_id
|
||||
strict_policy: Final = agent.identity_managed or bool(user_api_key_dict.requires_fresh_policy)
|
||||
|
|
@ -834,14 +985,15 @@ async def invoke_agent_a2a(
|
|||
if header_name:
|
||||
dynamic_headers[header_name] = val
|
||||
|
||||
backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider)
|
||||
merged_agent_headers: Final = merge_agent_headers(
|
||||
dynamic_headers=dynamic_headers or None,
|
||||
static_headers=static_headers or None,
|
||||
)
|
||||
backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider, merged_agent_headers)
|
||||
agent_extra_headers: Final = _forwarding_headers(
|
||||
caller_identity=caller_identity,
|
||||
request_data=data,
|
||||
agent_extra_headers=merge_agent_headers(
|
||||
dynamic_headers=dynamic_headers or None,
|
||||
static_headers=static_headers or None,
|
||||
),
|
||||
agent_extra_headers=merged_agent_headers,
|
||||
backend_auth_header=backend_auth.header,
|
||||
)
|
||||
forwarded_litellm_params: Final = backend_auth.litellm_params
|
||||
|
|
@ -1051,3 +1203,6 @@ async def invoke_agent_a2a(
|
|||
if isinstance(e, litellm.BadRequestError):
|
||||
return _jsonrpc_error(body.get("id"), -32602, e.message, 400)
|
||||
return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500)
|
||||
|
||||
|
||||
router.include_router(jsonrpc_router)
|
||||
|
|
|
|||
|
|
@ -140,7 +140,7 @@ async def search_agents(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> AgentSearchOutcome:
|
||||
if embedding_model is None:
|
||||
if embedding_model is None or not embedding_model.strip():
|
||||
return AgentSearchNotConfigured(
|
||||
reason="agent search needs litellm_settings.agent_search_embedding_model set to an embedding model from model_list"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from typing_extensions import ReadOnly, Required, assert_never
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, AGENT_HEALTH_CHECK_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
|
|
@ -33,6 +34,7 @@ from litellm.proxy.a2a.agent_card import (
|
|||
SUPPORTED_A2A_PROTOCOL_VERSIONS,
|
||||
merge_agent_card,
|
||||
normalize_protocol_version,
|
||||
point_agent_card_at_proxy,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
AgentIdWhere,
|
||||
|
|
@ -180,10 +182,36 @@ def _redact_agent_litellm_params_dict(
|
|||
return dict(parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params)))
|
||||
|
||||
|
||||
def _non_admin_agent_card(
|
||||
card: Mapping[str, object], proxy_url: str
|
||||
) -> dict[str, object]: # mutable-ok: AgentResponse.agent_card_params requires a mutable dictionary
|
||||
pointed_card: Final = point_agent_card_at_proxy(card, proxy_url)
|
||||
interfaces: Final = pointed_card.get("supportedInterfaces")
|
||||
interface_items: Final = (
|
||||
tuple(iter_items(pointed_card.get("supportedInterfaces"))) if isinstance(interfaces, list) else ()
|
||||
)
|
||||
first_interface: Final[object] = interface_items[0] if interface_items else None
|
||||
first_mapping: Final = as_str_mapping(first_interface)
|
||||
proxy_interface: Final = (
|
||||
{"supportedInterfaces": [{**first_mapping, "url": proxy_url}]}
|
||||
if isinstance(first_interface, dict) and first_mapping is not None and first_mapping.get("url") == proxy_url
|
||||
else {}
|
||||
)
|
||||
return {
|
||||
**{
|
||||
key: value
|
||||
for key, value in pointed_card.items()
|
||||
if key not in ("additionalInterfaces", "supportedInterfaces")
|
||||
},
|
||||
**proxy_interface,
|
||||
}
|
||||
|
||||
|
||||
def _redact_sensitive_agent_fields(
|
||||
agents: Sequence[AgentResponse],
|
||||
*,
|
||||
is_admin: bool,
|
||||
request_base_url: str,
|
||||
) -> list[AgentResponse]:
|
||||
"""
|
||||
Return copies of the given agents with credential-bearing litellm_params
|
||||
|
|
@ -207,6 +235,11 @@ def _redact_sensitive_agent_fields(
|
|||
copy.keys = None
|
||||
copy.kill_switch = None
|
||||
copy.identity = None
|
||||
if copy.agent_card_params:
|
||||
copy.agent_card_params = _non_admin_agent_card(
|
||||
copy.agent_card_params,
|
||||
get_custom_url(request_base_url, route=f"a2a/{copy.agent_id}"),
|
||||
)
|
||||
if copy.litellm_params:
|
||||
copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
|
||||
copy.kill_switch = redact_kill_switch(copy.kill_switch)
|
||||
|
|
@ -461,14 +494,6 @@ async def get_agents(
|
|||
global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
|
||||
)
|
||||
|
||||
# litellm_params secrets are always redacted; keys/headers stay
|
||||
# admin-only.
|
||||
is_admin: Final = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin)
|
||||
|
||||
if health_check:
|
||||
agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")]
|
||||
agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")]
|
||||
|
|
@ -493,6 +518,18 @@ async def get_agents(
|
|||
healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
|
||||
returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url
|
||||
|
||||
# litellm_params secrets are always redacted; keys/headers stay
|
||||
# admin-only.
|
||||
is_admin: Final = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
returned_agents = _redact_sensitive_agent_fields(
|
||||
returned_agents,
|
||||
is_admin=is_admin,
|
||||
request_base_url=str(request.base_url),
|
||||
)
|
||||
|
||||
if query is None:
|
||||
return returned_agents
|
||||
return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict)
|
||||
|
|
@ -691,7 +728,11 @@ async def create_agent(
|
|||
|
||||
# The caller is a proxy admin (enforced above); litellm_params
|
||||
# secrets are still never echoed back in the response.
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
|
|
@ -708,6 +749,7 @@ async def create_agent(
|
|||
)
|
||||
async def get_agent_by_id(
|
||||
agent_id: str,
|
||||
request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
|
|
@ -770,7 +812,11 @@ async def get_agent_by_id(
|
|||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0]
|
||||
agent = _redact_sensitive_agent_fields(
|
||||
(agent,),
|
||||
is_admin=is_admin,
|
||||
request_base_url=str(request.base_url),
|
||||
)[0]
|
||||
|
||||
return agent
|
||||
except HTTPException:
|
||||
|
|
@ -879,7 +925,11 @@ async def update_agent(
|
|||
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
|
||||
)
|
||||
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
@ -987,7 +1037,11 @@ async def patch_agent(
|
|||
"Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
|
||||
)
|
||||
|
||||
return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
|
||||
return _redact_sensitive_agent_fields(
|
||||
(result,),
|
||||
is_admin=True,
|
||||
request_base_url=str(http_request.base_url),
|
||||
)[0]
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
|
|||
172
litellm/proxy/dev_reload.py
Normal file
172
litellm/proxy/dev_reload.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
import importlib
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import partial
|
||||
from inspect import signature
|
||||
from multiprocessing.process import BaseProcess
|
||||
from socket import socket
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
from uvicorn import Config, Server
|
||||
from uvicorn import config as uvicorn_config
|
||||
from uvicorn.supervisors import ChangeReload, Multiprocess
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class ServerOptions(TypedDict, total=False):
|
||||
host: ReadOnly[str]
|
||||
port: ReadOnly[int]
|
||||
server_header: ReadOnly[bool]
|
||||
timeout_graceful_shutdown: ReadOnly[float]
|
||||
log_config: ReadOnly[ # mutable-ok: uvicorn.Config requires a concrete logging dict
|
||||
str | dict[str, JsonValue] | None
|
||||
]
|
||||
log_level: ReadOnly[str | int]
|
||||
timeout_keep_alive: ReadOnly[int]
|
||||
timeout_worker_healthcheck: ReadOnly[int]
|
||||
limit_max_requests: ReadOnly[int]
|
||||
limit_max_requests_jitter: ReadOnly[int]
|
||||
limit_concurrency: ReadOnly[int]
|
||||
ssl_keyfile: ReadOnly[str]
|
||||
ssl_certfile: ReadOnly[str]
|
||||
loop: ReadOnly[Literal["none", "auto", "asyncio", "uvloop"]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class ReloadOptions(ServerOptions):
|
||||
app: Required[ReadOnly[str]]
|
||||
reload: ReadOnly[bool]
|
||||
reload_dirs: ReadOnly[tuple[str, ...]]
|
||||
reload_includes: ReadOnly[tuple[str, ...]]
|
||||
|
||||
|
||||
RELOAD_OPTIONS: Final = TypeAdapter(ReloadOptions)
|
||||
SERVER_OPTIONS: Final = TypeAdapter(ServerOptions)
|
||||
ASGI_APP: Final = TypeAdapter(ImportString[Callable[..., object]])
|
||||
LOGGER: Final = logging.getLogger("uvicorn.error")
|
||||
RELOAD_WORKER_HEALTHCHECK_SECONDS: Final = 30
|
||||
RELOADER_ONLY_KEYS: Final = frozenset(ReloadOptions.__annotations__) - frozenset(ServerOptions.__annotations__)
|
||||
|
||||
|
||||
def multi_worker_reload_supported(platform: str = sys.platform) -> bool:
|
||||
return platform != "win32" and hasattr(uvicorn_config, "STARTUP_FAILURE")
|
||||
|
||||
|
||||
def server_options(options: ReloadOptions) -> ServerOptions:
|
||||
return SERVER_OPTIONS.validate_python(
|
||||
{key: value for key, value in options.items() if key not in RELOADER_ONLY_KEYS}
|
||||
)
|
||||
|
||||
|
||||
def load_app_or_exit(import_path: str, factory: bool = False) -> Callable[..., object]:
|
||||
module_name: Final = import_path.partition(":")[0]
|
||||
try:
|
||||
importlib.import_module(module_name)
|
||||
app: Final = ASGI_APP.validate_python(import_path)
|
||||
return ASGI_APP.validate_python(app()) if factory else app
|
||||
except Exception:
|
||||
LOGGER.exception("Error loading ASGI app %s", import_path)
|
||||
sys.exit(uvicorn_config.STARTUP_FAILURE)
|
||||
|
||||
|
||||
def stop_worker_when_pool_exits(pool_parent: BaseProcess) -> None:
|
||||
pool_parent.join()
|
||||
LOGGER.info("Worker pool process [%s] exited, stopping worker [%s].", pool_parent.pid, os.getpid())
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[..., object]:
|
||||
pool_parent: Final = multiprocessing.parent_process()
|
||||
if pool_parent is not None:
|
||||
threading.Thread(target=stop_worker_when_pool_exits, args=(pool_parent,), daemon=True).start()
|
||||
return load_app_or_exit(import_path, worker_app_factory)
|
||||
|
||||
|
||||
def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config:
|
||||
overrides: Final = worker_overrides or {}
|
||||
worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"]))
|
||||
worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False))
|
||||
return Config(
|
||||
partial(load_worker_app, worker_app, worker_app_factory),
|
||||
factory=True,
|
||||
workers=workers,
|
||||
**SERVER_OPTIONS.validate_python(
|
||||
{
|
||||
"timeout_worker_healthcheck": RELOAD_WORKER_HEALTHCHECK_SECONDS,
|
||||
**server_options(options),
|
||||
**{key: value for key, value in overrides.items() if key not in ("app", "factory")},
|
||||
}
|
||||
), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
|
||||
)
|
||||
|
||||
|
||||
def stop_pool_when_reloader_exits(reloader: BaseProcess, pool_should_exit: threading.Event) -> None:
|
||||
reloader.join()
|
||||
LOGGER.info("Reloader process [%s] exited, stopping the worker pool.", reloader.pid)
|
||||
pool_should_exit.set()
|
||||
|
||||
|
||||
class ReloadWorkerPool(Multiprocess):
|
||||
quit_received: bool = False
|
||||
|
||||
def handle_quit(self) -> None:
|
||||
LOGGER.info("Received SIGQUIT, exiting.")
|
||||
self.quit_received = True
|
||||
self.should_exit.set()
|
||||
|
||||
def terminate_all(self) -> None:
|
||||
quit_signal: Final = getattr(signal, "SIGQUIT", None)
|
||||
if quit_signal is not None and quit_signal in self.signal_queue:
|
||||
self.handle_quit()
|
||||
if not self.quit_received:
|
||||
super().terminate_all()
|
||||
return
|
||||
for process in self.processes:
|
||||
process.kill()
|
||||
|
||||
|
||||
def _run_worker_pool(config: Config, sockets: Sequence[socket] | None) -> None:
|
||||
shared_sockets: Final = list(sockets or ())
|
||||
pool: Final = (
|
||||
ReloadWorkerPool(
|
||||
config=config,
|
||||
target=Server(config=config).run, # pyright: ignore[reportCallIssue] # Older uvicorn versions require target
|
||||
sockets=shared_sockets,
|
||||
)
|
||||
if "target" in signature(Multiprocess).parameters
|
||||
else ReloadWorkerPool(config=config, sockets=shared_sockets)
|
||||
)
|
||||
reloader: Final = multiprocessing.parent_process()
|
||||
if reloader is not None:
|
||||
threading.Thread(target=stop_pool_when_reloader_exits, args=(reloader, pool.should_exit), daemon=True).start()
|
||||
pool.run()
|
||||
|
||||
|
||||
def run_multi_worker_reload(
|
||||
options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None
|
||||
) -> None:
|
||||
worker_config: Final = pool_config(options, workers, worker_overrides)
|
||||
reload_config: Final = Config(
|
||||
options["app"],
|
||||
workers=1,
|
||||
reload=options["reload"],
|
||||
reload_dirs=list(options["reload_dirs"]),
|
||||
reload_includes=list(options["reload_includes"]),
|
||||
**server_options(options), # pyright: ignore[reportArgumentType] # uvicorn types the float graceful shutdown timeout as int
|
||||
)
|
||||
shared_socket: Final = worker_config.bind_socket()
|
||||
try:
|
||||
ChangeReload(
|
||||
config=reload_config,
|
||||
target=partial(_run_worker_pool, worker_config),
|
||||
sockets=[shared_socket],
|
||||
).run()
|
||||
finally:
|
||||
shared_socket.close()
|
||||
|
|
@ -9,7 +9,7 @@ import os
|
|||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, Union, cast, get_args, get_origin
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
|
|
@ -43,6 +43,11 @@ from litellm.proxy.guardrails.guardrail_registry import (
|
|||
encrypt_guardrail_litellm_params,
|
||||
parse_tolerant_litellm_params,
|
||||
)
|
||||
from litellm.proxy.guardrails.response_handlers import (
|
||||
GUARDRAIL_RESPONSE_HANDLERS,
|
||||
GuardrailEndpointResponse,
|
||||
StandardGuardrailResponseHandler,
|
||||
)
|
||||
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
|
||||
from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports
|
||||
_user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -53,7 +58,6 @@ from litellm.repositories.table_repositories import GuardrailsRepository
|
|||
from litellm.types.guardrails import (
|
||||
PII_ENTITY_CATEGORIES_MAP,
|
||||
ApplyGuardrailRequest,
|
||||
ApplyGuardrailResponse,
|
||||
BaseLitellmParams,
|
||||
BedrockGuardrailConfigModel,
|
||||
BedrockGuardrailStreamingParams,
|
||||
|
|
@ -2350,13 +2354,7 @@ async def _run_test_guardrail(execute_guardrail: Callable[[], object], timeout:
|
|||
return await await_with_timeout(raw_result, remaining, label="test")
|
||||
|
||||
|
||||
def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: str) -> Literal["request", "response"]:
|
||||
"""Return the effective input_type, auto-upgrading to 'response' for post_call guardrails."""
|
||||
if input_type == "request":
|
||||
hook: Final = getattr(active_guardrail, "event_hook", None)
|
||||
if hook == GuardrailEventHooks.post_call or hook == "post_call":
|
||||
return "response"
|
||||
return "response" if input_type == "response" else "request"
|
||||
_RESPONSE: Final[TypeAdapter[GuardrailEndpointResponse]] = TypeAdapter(GuardrailEndpointResponse)
|
||||
|
||||
|
||||
class _GuardrailLoggingObj(Protocol):
|
||||
|
|
@ -2384,35 +2382,43 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj,
|
|||
)
|
||||
|
||||
|
||||
async def _emit_guardrail_success_logs(
|
||||
async def _guardrail_response_after_success_hook(
|
||||
response_handler: StandardGuardrailResponseHandler,
|
||||
proxy_logging_obj: _GuardrailProxyLogging,
|
||||
litellm_logging_obj: _GuardrailLoggingObj | None,
|
||||
data: dict,
|
||||
data: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: ApplyGuardrailResponse,
|
||||
start_time: datetime,
|
||||
) -> ApplyGuardrailResponse:
|
||||
"""Fire proxy and LiteLLM success hooks after a successful guardrail run.
|
||||
|
||||
Each hook is wrapped defensively so a callback failure never prevents the
|
||||
caller from receiving the guardrail response. Returns the (possibly
|
||||
hook-modified) response.
|
||||
"""
|
||||
response: GuardrailEndpointResponse,
|
||||
) -> GuardrailEndpointResponse:
|
||||
try:
|
||||
modified: Final = await proxy_logging_obj.post_call_success_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
)
|
||||
if isinstance(modified, ApplyGuardrailResponse):
|
||||
response = modified
|
||||
return response_handler.accept_hook_result(response, modified)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception("apply_guardrail: post_call_success_hook failed")
|
||||
return response
|
||||
|
||||
|
||||
async def _emit_guardrail_success_logs(
|
||||
response_handler: StandardGuardrailResponseHandler,
|
||||
proxy_logging_obj: _GuardrailProxyLogging,
|
||||
litellm_logging_obj: _GuardrailLoggingObj | None,
|
||||
data: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: GuardrailEndpointResponse,
|
||||
start_time: datetime,
|
||||
) -> GuardrailEndpointResponse:
|
||||
"""Apply the success hook and log the response returned to the caller."""
|
||||
updated_response: Final = await _guardrail_response_after_success_hook(
|
||||
response_handler, proxy_logging_obj, data, user_api_key_dict, response
|
||||
)
|
||||
|
||||
# Build the logging payload after post_call_success_hook so that logged
|
||||
# data matches what the caller actually receives if the hook modified
|
||||
# the response.
|
||||
response_for_logging: Final = {"response": response.model_dump(exclude_none=True)}
|
||||
response_for_logging: Final = {"response": _RESPONSE.dump_python(updated_response, mode="json")}
|
||||
|
||||
if litellm_logging_obj is not None:
|
||||
end_time: Final = datetime.now(timezone.utc)
|
||||
|
|
@ -2426,11 +2432,11 @@ async def _emit_guardrail_success_logs(
|
|||
except Exception:
|
||||
verbose_proxy_logger.exception("apply_guardrail: async_success_handler failed")
|
||||
|
||||
return response
|
||||
return updated_response
|
||||
|
||||
|
||||
@router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse)
|
||||
@router.post("/apply_guardrail", response_model=ApplyGuardrailResponse)
|
||||
@router.post("/guardrails/apply_guardrail", response_model=GuardrailEndpointResponse)
|
||||
@router.post("/apply_guardrail", response_model=GuardrailEndpointResponse)
|
||||
async def apply_guardrail(
|
||||
fastapi_request: Request,
|
||||
request: ApplyGuardrailRequest,
|
||||
|
|
@ -2455,7 +2461,7 @@ async def apply_guardrail(
|
|||
)
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
|
||||
data: dict = {
|
||||
data: dict[str, object] = {
|
||||
"guardrail_name": request.guardrail_name,
|
||||
"input": [request.text],
|
||||
"messages": request.messages or [],
|
||||
|
|
@ -2495,18 +2501,15 @@ async def apply_guardrail(
|
|||
if litellm_logging_obj is not None:
|
||||
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
|
||||
|
||||
request_data: Final[dict] = {
|
||||
**({"messages": request.messages} if request.messages is not None else {}),
|
||||
**({"metadata": request.metadata} if request.metadata is not None else {}),
|
||||
}
|
||||
_input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
|
||||
response_handler: Final = GUARDRAIL_RESPONSE_HANDLERS[request.response_format]
|
||||
request_data: Final = dict(response_handler.prepare_request_data(request))
|
||||
input_type: Final = response_handler.resolve_input_type(request, active_guardrail)
|
||||
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
|
||||
inputs={"texts": [request.text]},
|
||||
request_data=request_data,
|
||||
input_type=_input_type,
|
||||
input_type=input_type,
|
||||
)
|
||||
response_text: Final = guardrailed_inputs.get("texts", [])
|
||||
response = ApplyGuardrailResponse(response_text=response_text[0] if response_text else request.text)
|
||||
response: Final = response_handler.build_response(request, guardrailed_inputs, request_data)
|
||||
except Exception as e:
|
||||
if litellm_logging_obj is not None and not isinstance(e, HTTPException):
|
||||
try:
|
||||
|
|
@ -2537,7 +2540,8 @@ async def apply_guardrail(
|
|||
raise handle_exception_on_proxy(status_typed_proxy_exception(e) if isinstance(e, HTTPException) else e)
|
||||
|
||||
# Success logging outside except so a hook error never triggers failure handlers.
|
||||
response = await _emit_guardrail_success_logs(
|
||||
logged_response: Final = await _emit_guardrail_success_logs(
|
||||
response_handler=response_handler,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
data=data,
|
||||
|
|
@ -2545,7 +2549,7 @@ async def apply_guardrail(
|
|||
response=response,
|
||||
start_time=start_time,
|
||||
)
|
||||
return response
|
||||
return logged_response
|
||||
|
||||
|
||||
# Usage (dashboard) endpoints: overview, detail, logs
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ from litellm.proxy.guardrails.anthropic_sse import (
|
|||
)
|
||||
from litellm.types.guardrails import (
|
||||
BedrockChecksConfigModel,
|
||||
BedrockGuardrailOutputScope,
|
||||
BedrockGuardrailStreamingParams,
|
||||
GuardrailEventHooks,
|
||||
LitellmParams,
|
||||
|
|
@ -289,6 +290,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
streaming_buffer_release_on_scan: bool | None = None,
|
||||
streaming_scan_in_background: bool | None = None,
|
||||
async_handler: AsyncHTTPHandler | None = None,
|
||||
outputScope: BedrockGuardrailOutputScope | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = async_handler or get_async_httpx_client(
|
||||
|
|
@ -309,6 +311,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
self.guardrailIdentifier = guardrailIdentifier
|
||||
self.guardrailVersion = guardrailVersion
|
||||
self.outputScope = outputScope
|
||||
self.guardrail_provider = "bedrock"
|
||||
self.chunk_budget_chars = chunk_budget_chars
|
||||
self.contextual_grounding_from_messages = contextual_grounding_from_messages
|
||||
|
|
@ -1280,10 +1283,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
AWS billed them to ``completed_chunk_usages``, and the attempt log sums those
|
||||
with the blocking call's own usage.
|
||||
"""
|
||||
bedrock_request_data: Final = {
|
||||
**base_request_data,
|
||||
"content": content,
|
||||
}
|
||||
bedrock_request_data: Final = dict(base_request_data)
|
||||
bedrock_request_data["content"] = content
|
||||
if self.outputScope is not None and "outputScope" not in bedrock_request_data:
|
||||
bedrock_request_data["outputScope"] = self.outputScope.value
|
||||
prepared_request: Final = await run_aws_signing(
|
||||
self._prepare_request,
|
||||
credentials=credentials,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
application_id: Final = getattr(litellm_params, "application_id", None)
|
||||
pre_checkpoint_id: Final = getattr(litellm_params, "pre_checkpoint_id", None)
|
||||
post_checkpoint_id: Final = getattr(litellm_params, "post_checkpoint_id", None)
|
||||
file_checkpoint_id: Final = getattr(litellm_params, "file_checkpoint_id", None)
|
||||
|
||||
_ovalix_callback: Final = OvalixGuardrail(
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
|
|
@ -27,6 +28,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
application_id=application_id,
|
||||
pre_checkpoint_id=pre_checkpoint_id,
|
||||
post_checkpoint_id=post_checkpoint_id,
|
||||
file_checkpoint_id=file_checkpoint_id,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
timeout=litellm_params.timeout,
|
||||
|
|
|
|||
|
|
@ -4,13 +4,19 @@ Use Ovalix Guardrails for your LLM calls. Supports pre_call (user input) and
|
|||
post_call (model output) checkpoints with optional correction/blocking.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import datetime
|
||||
import gzip
|
||||
import hashlib
|
||||
import mimetypes
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict, Unpack
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
|
|
@ -18,10 +24,23 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix_extraction import (
|
||||
FilePart,
|
||||
extract_file_parts_from_images,
|
||||
extract_file_parts_from_messages,
|
||||
extract_tool_calls_from_messages,
|
||||
extract_tool_results,
|
||||
make_tool_data,
|
||||
str_or_none,
|
||||
tool_call_to_tool_data,
|
||||
tool_data_key,
|
||||
tool_result_text_indices,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -32,13 +51,34 @@ if TYPE_CHECKING:
|
|||
|
||||
BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix"
|
||||
BLOCKED_ACTION_TYPE: Final = "block"
|
||||
_MODIFY_ACTION_TYPES: Final = ("anonymize", "sanitize")
|
||||
_DEFAULT_FILE_SIZE_LIMIT: Final = 64 * 1024 * 1024
|
||||
_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_FILE_BLOCK_ESCALATION_REASON: Final = (
|
||||
"This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM"
|
||||
)
|
||||
_UNDECODABLE_FILE_BLOCK_REASON: Final = (
|
||||
"This message was blocked by Ovalix because an attached file couldn't be decoded for inspection"
|
||||
)
|
||||
_TOOL_BLOCK_ESCALATION_REASON: Final = (
|
||||
"This message was blocked by Ovalix because tool call anonymization isn't possible via LiteLLM"
|
||||
)
|
||||
_TOOL_RESULT_BLOCK_ESCALATION_REASON: Final = (
|
||||
"This message was blocked by Ovalix because tool result anonymization isn't possible via LiteLLM"
|
||||
)
|
||||
|
||||
|
||||
def _encode_file_wire_format(raw: bytes) -> str:
|
||||
return base64.b64encode(gzip.compress(raw)).decode()
|
||||
|
||||
|
||||
def _mapping_or_empty(value: object) -> Mapping[str, object]:
|
||||
return as_str_mapping(value) or _NO_METADATA
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
"""Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
|
||||
|
||||
supported_event_hooks: ReadOnly[list[GuardrailEventHooks]]
|
||||
|
||||
|
||||
class OvalixGuardrailMissingSecrets(Exception):
|
||||
"""Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing."""
|
||||
|
|
@ -87,6 +127,8 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
application_id: str | None = None,
|
||||
pre_checkpoint_id: str | None = None,
|
||||
post_checkpoint_id: str | None = None,
|
||||
file_checkpoint_id: str | None = None,
|
||||
supported_event_hooks: list[GuardrailEventHooks] | None = None,
|
||||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
):
|
||||
self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE")
|
||||
|
|
@ -94,23 +136,26 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID")
|
||||
self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID")
|
||||
self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID")
|
||||
self._file_checkpoint_id = file_checkpoint_id or os.environ.get("OVALIX_FILE_CHECKPOINT_ID")
|
||||
|
||||
supported_event_hooks: Final = kwargs.get("supported_event_hooks", [])
|
||||
event_hooks: Final = self._validated_event_hooks(supported_event_hooks or ())
|
||||
|
||||
self._validate_config(supported_event_hooks)
|
||||
|
||||
self._tracker_headers = httpx.Headers(
|
||||
{
|
||||
"Authorization": f"Bearer {self._tracker_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
encoding="utf-8",
|
||||
self._tracker_headers = dict(
|
||||
httpx.Headers(
|
||||
MappingProxyType(
|
||||
{
|
||||
"Authorization": f"Bearer {self._tracker_api_key}",
|
||||
"x-api-key": self._tracker_api_key or "",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
)
|
||||
|
||||
self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
||||
forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks}
|
||||
super().__init__(**forwarded)
|
||||
super().__init__(supported_event_hooks=event_hooks, **kwargs)
|
||||
verbose_proxy_logger.debug(
|
||||
"Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s",
|
||||
self._tracker_api_base,
|
||||
|
|
@ -119,44 +164,53 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
self._post_checkpoint_id,
|
||||
)
|
||||
|
||||
def _validate_config(self, supported_event_hooks: list[GuardrailEventHooks]) -> None:
|
||||
"""Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present."""
|
||||
errors: Final[list[str]] = []
|
||||
|
||||
if not self._tracker_api_base:
|
||||
errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base")
|
||||
if not self._tracker_api_key:
|
||||
errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key")
|
||||
if not self._application_id:
|
||||
errors.append("Application ID, set OVALIX_APPLICATION_ID or pass application_id")
|
||||
if not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in supported_event_hooks:
|
||||
errors.append("Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id")
|
||||
if not self._post_checkpoint_id and GuardrailEventHooks.post_call in supported_event_hooks:
|
||||
errors.append("Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id")
|
||||
if not self._pre_checkpoint_id and not self._post_checkpoint_id:
|
||||
errors.append(
|
||||
"Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id"
|
||||
def _validated_event_hooks(self, requested: Sequence[GuardrailEventHooks]) -> list[GuardrailEventHooks]:
|
||||
"""Ensure required secrets and checkpoint IDs are set; return the requested hooks plus one per configured checkpoint."""
|
||||
errors: Final = tuple(
|
||||
message
|
||||
for missing, message in (
|
||||
(not self._tracker_api_base, "Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base"),
|
||||
(not self._tracker_api_key, "Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key"),
|
||||
(not self._application_id, "Application ID, set OVALIX_APPLICATION_ID or pass application_id"),
|
||||
(
|
||||
not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in requested,
|
||||
"Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id",
|
||||
),
|
||||
(
|
||||
not self._post_checkpoint_id and GuardrailEventHooks.post_call in requested,
|
||||
"Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id",
|
||||
),
|
||||
(
|
||||
not self._pre_checkpoint_id and not self._post_checkpoint_id,
|
||||
"Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id",
|
||||
),
|
||||
)
|
||||
if missing
|
||||
)
|
||||
|
||||
if errors:
|
||||
raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors))
|
||||
|
||||
# auto-add hooks when checkpoint IDs are present
|
||||
if self._pre_checkpoint_id and GuardrailEventHooks.pre_call not in supported_event_hooks:
|
||||
supported_event_hooks.append(GuardrailEventHooks.pre_call)
|
||||
if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks:
|
||||
supported_event_hooks.append(GuardrailEventHooks.post_call)
|
||||
auto_added: Final = tuple(
|
||||
hook
|
||||
for configured, hook in (
|
||||
(self._pre_checkpoint_id, GuardrailEventHooks.pre_call),
|
||||
(self._post_checkpoint_id, GuardrailEventHooks.post_call),
|
||||
)
|
||||
if configured and hook not in requested
|
||||
)
|
||||
return [*requested, *auto_added]
|
||||
|
||||
def _get_actor(self, data: dict) -> str:
|
||||
def _get_actor(self, data: Mapping[str, object]) -> str:
|
||||
"""Return a stable actor identifier from request metadata (e.g. user email or id)."""
|
||||
metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {}
|
||||
if metadata.get("user_api_key_user_email"):
|
||||
return metadata["user_api_key_user_email"]
|
||||
if metadata.get("user_api_key_user_id"):
|
||||
return metadata["user_api_key_user_id"]
|
||||
return "unknown"
|
||||
metadata: Final = _mapping_or_empty(data.get("metadata") or data.get("litellm_metadata"))
|
||||
return (
|
||||
str_or_none(metadata.get("user_api_key_user_email"))
|
||||
or str_or_none(metadata.get("user_api_key_user_id"))
|
||||
or "unknown"
|
||||
)
|
||||
|
||||
def _get_tracker_actor_id(self, data: dict) -> str:
|
||||
def _get_tracker_actor_id(self, data: Mapping[str, object]) -> str:
|
||||
"""Normalize the actor string into a short, stable id for Tracker API payloads."""
|
||||
# NOTE: this hash is purely for normalization — it collapses an arbitrary actor
|
||||
# string (email, user id, or "unknown") into a compact, fixed-length, consistent
|
||||
|
|
@ -166,7 +220,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
normalized_actor_id: Final = hashlib.sha256(actor_id).hexdigest()[:8]
|
||||
return normalized_actor_id
|
||||
|
||||
def _get_session_id(self, data: dict) -> str:
|
||||
def _get_session_id(self, data: Mapping[str, object]) -> str:
|
||||
"""Return a unique identifier for the chat/session (actor + date + application_id)."""
|
||||
actor_hash: Final = self._get_tracker_actor_id(data)
|
||||
today: Final = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
|
||||
|
|
@ -174,106 +228,254 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
|
||||
async def _call_checkpoint(
|
||||
self,
|
||||
content: str,
|
||||
data_type: str,
|
||||
data: Mapping[str, object],
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Call the Ovalix Tracker checkpoint API and return the JSON response."""
|
||||
oversize: bool = False,
|
||||
) -> Mapping[str, object]:
|
||||
"""Call the Ovalix Tracker checkpoint API and return the JSON response.
|
||||
|
||||
Both routes live on the tracker's /beta litellm router, which accepts the api key this
|
||||
guardrail already sends.
|
||||
"""
|
||||
application_id: Final = self._application_id
|
||||
if not application_id or not checkpoint_id:
|
||||
raise ValueError("Ovalix: application_id or checkpoint_id not resolved")
|
||||
|
||||
url: Final = f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
|
||||
headers: Final = dict(self._tracker_headers)
|
||||
route: Final = "file_checkpoint" if data_type == "FILE" else "checkpoint"
|
||||
payload: Final = {
|
||||
"application_id": application_id,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
"actor": actor,
|
||||
"session_id": session_id,
|
||||
"data_type": "TEXT",
|
||||
"data": {"content": content},
|
||||
"data_type": data_type,
|
||||
"data": data,
|
||||
"tool": "LiteLLM",
|
||||
**(MappingProxyType({"oversize": True}) if oversize else _NO_METADATA),
|
||||
}
|
||||
response: Final = await self._async_handler.post(url, headers=headers, json=payload, timeout=self.timeout)
|
||||
response: Final = await self._async_handler.post(
|
||||
f"{self._tracker_api_base}/tracking/beta/{route}",
|
||||
headers=self._tracker_headers,
|
||||
json=payload,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
def _verdict(self, resp: Mapping[str, object]) -> tuple[str, str | None]:
|
||||
return (str_or_none(resp.get("action_type")) or "").lower(), self._get_trackers_corrected_message(resp)
|
||||
|
||||
async def _block_reason_for_item(
|
||||
self,
|
||||
data_type: str,
|
||||
data: Mapping[str, object],
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
escalation_reason: str,
|
||||
oversize: bool = False,
|
||||
) -> str | None:
|
||||
try:
|
||||
resp: Final = await self._call_checkpoint(
|
||||
data_type, data, checkpoint_id, actor, session_id, oversize=oversize
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: {e}",
|
||||
should_wrap_with_default_message=False,
|
||||
) from e
|
||||
action, corrected = self._verdict(resp)
|
||||
if action == BLOCKED_ACTION_TYPE:
|
||||
return corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
|
||||
if action in _MODIFY_ACTION_TYPES:
|
||||
return escalation_reason
|
||||
return None
|
||||
|
||||
async def _check_items_block_only(
|
||||
self,
|
||||
items: Sequence[tuple[str, Mapping[str, object]]],
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
escalation_reason: str,
|
||||
) -> str | None:
|
||||
for data_type, data in items:
|
||||
reason = await self._block_reason_for_item(
|
||||
data_type, data, checkpoint_id, actor, session_id, escalation_reason
|
||||
)
|
||||
if reason is not None:
|
||||
return reason
|
||||
return None
|
||||
|
||||
async def _check_files_for_block(
|
||||
self,
|
||||
file_parts: Sequence[FilePart],
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
) -> str | None:
|
||||
if any(part.undecodable for part in file_parts):
|
||||
return _UNDECODABLE_FILE_BLOCK_REASON
|
||||
for part in sorted(file_parts, key=lambda p: p.message_index, reverse=True):
|
||||
data = await self._file_part_to_data(part)
|
||||
reason = await self._block_reason_for_item(
|
||||
"FILE",
|
||||
data,
|
||||
checkpoint_id,
|
||||
actor,
|
||||
session_id,
|
||||
_FILE_BLOCK_ESCALATION_REASON,
|
||||
oversize=part.oversize,
|
||||
)
|
||||
if reason is not None:
|
||||
return reason
|
||||
return None
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
request_data: Mapping[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""
|
||||
Apply Ovalix guardrail to the given inputs (request or response text).
|
||||
|
||||
Used by the unified guardrail flow and the /apply_guardrail API.
|
||||
For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint.
|
||||
|
||||
Args:
|
||||
inputs: Guardrail API inputs (e.g. texts to check).
|
||||
request_data: Full request payload (messages, metadata, response).
|
||||
input_type: "request" (pre_call) or "response" (post_call).
|
||||
logging_obj: Optional logging context.
|
||||
|
||||
Returns:
|
||||
Updated inputs (e.g. with replaced/corrected texts, or unchanged).
|
||||
"""
|
||||
if not self._pre_checkpoint_id and not self._post_checkpoint_id:
|
||||
return inputs
|
||||
|
||||
tracker_actor_id: Final = self._get_tracker_actor_id(request_data)
|
||||
actor: Final = self._get_actor(request_data)
|
||||
session_id: Final = self._get_session_id(request_data)
|
||||
texts: Final = inputs.get("texts") or []
|
||||
if not texts or not isinstance(texts, list):
|
||||
is_response: Final = input_type == "response"
|
||||
|
||||
prompt_checkpoint: Final = self._post_checkpoint_id if is_response else self._pre_checkpoint_id
|
||||
file_checkpoint: Final = self._file_checkpoint_id or prompt_checkpoint
|
||||
if not file_checkpoint:
|
||||
verbose_proxy_logger.debug(
|
||||
"Ovalix guardrail: application %s has no %s checkpoint, leaving this direction uninspected",
|
||||
self._application_id,
|
||||
input_type,
|
||||
)
|
||||
return inputs
|
||||
|
||||
if input_type == "response":
|
||||
if not self._post_checkpoint_id:
|
||||
return inputs
|
||||
corrected_llm_responses: Final = await self._generate_post_guardrail_llm_texts(
|
||||
texts, tracker_actor_id, session_id, self._post_checkpoint_id
|
||||
structured_messages: Final = inputs.get("structured_messages") or ()
|
||||
file_parts: Final = (
|
||||
extract_file_parts_from_images(inputs.get("images"), size_limit=_DEFAULT_FILE_SIZE_LIMIT)
|
||||
if is_response
|
||||
else extract_file_parts_from_messages(structured_messages, size_limit=_DEFAULT_FILE_SIZE_LIMIT)
|
||||
)
|
||||
file_block: Final = await self._check_files_for_block(file_parts, file_checkpoint, actor, session_id)
|
||||
if file_block is not None:
|
||||
self._block_current_message(file_block)
|
||||
|
||||
if not prompt_checkpoint:
|
||||
verbose_proxy_logger.debug(
|
||||
"Ovalix guardrail: application %s has only a %s file checkpoint, skipping text and tool inspection",
|
||||
self._application_id,
|
||||
input_type,
|
||||
)
|
||||
return {**inputs, "texts": corrected_llm_responses}
|
||||
return inputs
|
||||
|
||||
if self._pre_checkpoint_id:
|
||||
post_guardrail_texts: Final = await self._generate_post_guardrail_llm_texts(
|
||||
texts, tracker_actor_id, session_id, self._pre_checkpoint_id
|
||||
)
|
||||
return {**inputs, "texts": post_guardrail_texts}
|
||||
return inputs
|
||||
tool_calls: Final = (
|
||||
*(inputs.get("tool_calls") or ()),
|
||||
*extract_tool_calls_from_messages(structured_messages),
|
||||
)
|
||||
unique_tool_data: Final = MappingProxyType(
|
||||
{tool_data_key(data): data for data in (tool_call_to_tool_data(tc) for tc in tool_calls) if data}
|
||||
)
|
||||
tool_call_items: Final = tuple(("TOOL", data) for data in unique_tool_data.values())
|
||||
tool_block: Final = await self._check_items_block_only(
|
||||
tool_call_items,
|
||||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
_TOOL_BLOCK_ESCALATION_REASON,
|
||||
)
|
||||
if tool_block is not None:
|
||||
self._block_current_message(tool_block)
|
||||
|
||||
async def _generate_post_guardrail_llm_texts(
|
||||
self, texts: list[str], actor: str, session_id: str, checkpoint_id: str
|
||||
) -> list[str]:
|
||||
"""Generate post-guardrail LLM responses for the given LLM responses."""
|
||||
post_guardrail_texts: Final[list[str]] = []
|
||||
tool_results: Final = extract_tool_results(structured_messages)
|
||||
tool_result_items: Final = tuple(("TOOL", make_tool_data(name, content)) for name, content, _ in tool_results)
|
||||
tool_result_block: Final = await self._check_items_block_only(
|
||||
tool_result_items,
|
||||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
_TOOL_RESULT_BLOCK_ESCALATION_REASON,
|
||||
)
|
||||
if tool_result_block is not None:
|
||||
self._block_current_message(tool_result_block)
|
||||
|
||||
is_first_response = True
|
||||
for llm_response in reversed(texts):
|
||||
try:
|
||||
resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: {e}",
|
||||
should_wrap_with_default_message=False,
|
||||
) from e
|
||||
texts: Final = inputs.get("texts") or ()
|
||||
if not texts:
|
||||
return inputs
|
||||
output_texts: Final = await self._check_texts(
|
||||
texts,
|
||||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
tool_result_text_indices(structured_messages, texts),
|
||||
)
|
||||
if output_texts is None:
|
||||
return inputs
|
||||
return {**inputs, "texts": output_texts}
|
||||
|
||||
action_type = (resp.get("action_type") or "").lower()
|
||||
blocking_message = self._get_trackers_corrected_message(resp) or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
|
||||
if action_type == BLOCKED_ACTION_TYPE and is_first_response:
|
||||
self._block_current_message(blocking_message)
|
||||
elif action_type == BLOCKED_ACTION_TYPE:
|
||||
post_guardrail_texts.insert(0, blocking_message)
|
||||
else:
|
||||
corrected_text = self._get_trackers_corrected_message(resp) or llm_response
|
||||
post_guardrail_texts.insert(0, corrected_text)
|
||||
is_first_response = False
|
||||
return post_guardrail_texts
|
||||
async def _file_part_to_data(self, part: FilePart) -> Mapping[str, object]:
|
||||
extension: Final = mimetypes.guess_extension(part.mime_hint) if part.mime_hint else None
|
||||
name: Final = part.name or (f"file{extension}" if extension else "file")
|
||||
content: Final = (
|
||||
await asyncio.get_event_loop().run_in_executor(None, _encode_file_wire_format, part.data)
|
||||
if part.data
|
||||
else None
|
||||
)
|
||||
return {"name": name, "content": content}
|
||||
|
||||
async def _check_texts(
|
||||
self,
|
||||
texts: Sequence[str],
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
skip_indices: frozenset[int],
|
||||
) -> list[str] | None:
|
||||
newest_first: Final = tuple(index for index in reversed(range(len(texts))) if index not in skip_indices)
|
||||
checked: Final = MappingProxyType(
|
||||
{
|
||||
index: await self._check_text(
|
||||
texts[index], checkpoint_id, actor, session_id, is_newest=index == len(texts) - 1
|
||||
)
|
||||
for index in newest_first
|
||||
}
|
||||
)
|
||||
output: Final = [checked.get(index, text) for index, text in enumerate(texts)]
|
||||
return output if output != list(texts) else None
|
||||
|
||||
async def _check_text(
|
||||
self,
|
||||
content: str,
|
||||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
is_newest: bool,
|
||||
) -> str:
|
||||
"""Return the text that replaces `content` after Ovalix checks it; a blocked newest text raises."""
|
||||
try:
|
||||
resp: Final = await self._call_checkpoint("TEXT", {"content": content}, checkpoint_id, actor, session_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: {e}",
|
||||
should_wrap_with_default_message=False,
|
||||
) from e
|
||||
action, corrected = self._verdict(resp)
|
||||
if action == BLOCKED_ACTION_TYPE:
|
||||
block_message: Final = corrected or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
|
||||
if is_newest:
|
||||
self._block_current_message(block_message)
|
||||
return block_message
|
||||
if action in _MODIFY_ACTION_TYPES and corrected is not None:
|
||||
return corrected
|
||||
return content
|
||||
|
||||
def _block_current_message(self, blocking_message: str) -> None:
|
||||
"""Raise OvalixGuardrailBlockedException with the given message (no default wrapper)."""
|
||||
|
|
@ -283,12 +485,13 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
should_wrap_with_default_message=False,
|
||||
)
|
||||
|
||||
def _get_trackers_corrected_message(self, resp: dict) -> str | None:
|
||||
def _get_trackers_corrected_message(self, resp: Mapping[str, object]) -> str | None:
|
||||
"""Extract corrected/blocking message content from Tracker checkpoint response."""
|
||||
modified: Final = resp.get("modified_data")
|
||||
if isinstance(modified, dict) and "content" in modified:
|
||||
return modified["content"]
|
||||
return None
|
||||
modified: Final = as_str_mapping(resp.get("modified_data"))
|
||||
if modified is None:
|
||||
return None
|
||||
content: Final = modified.get("content")
|
||||
return content if isinstance(content, str) else None
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type["GuardrailConfigModel"] | None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,396 @@
|
|||
import base64
|
||||
import json
|
||||
import posixpath
|
||||
import re
|
||||
from bisect import bisect_left
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from itertools import chain, groupby
|
||||
from operator import attrgetter
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NamedTuple
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
|
||||
_TOOL_NAME_MAX_LENGTH: Final = 100
|
||||
_DEFAULT_TOOL_RESULT_NAME: Final = "tool_result"
|
||||
_NO_TOOL_INPUT: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
_DATA_URL_RE: Final = re.compile(r"^data:(?P<mime>[^;,]+)?(?P<params>(?:;[^;,]+)*?)(?P<b64>;base64)?,", re.IGNORECASE)
|
||||
_URLSAFE_TO_STANDARD_B64: Final = str.maketrans("-_", "+/")
|
||||
_OBJECT_LIST: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
def str_or_none(value: object) -> str | None:
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _list_items(value: object) -> Sequence[object]:
|
||||
try:
|
||||
return _OBJECT_LIST.validate_python(value, strict=True)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
class FilePart(NamedTuple):
|
||||
name: str | None
|
||||
data: bytes | None
|
||||
mime_hint: str | None
|
||||
inline: bool
|
||||
oversize: bool
|
||||
message_index: int
|
||||
undecodable: bool = False
|
||||
|
||||
|
||||
def _split_data_url(value: str) -> tuple[str | None, str | None]:
|
||||
match: Final = _DATA_URL_RE.match(value)
|
||||
if not match:
|
||||
if value.lower().startswith("data:"):
|
||||
return None, None
|
||||
return None, value
|
||||
mime: Final = match.group("mime") or None
|
||||
if not match.group("b64"):
|
||||
return mime, None
|
||||
return mime, value[match.end() :]
|
||||
|
||||
|
||||
def _b64decode_or_none(payload: str) -> bytes | None:
|
||||
try:
|
||||
return base64.b64decode(payload, validate=True)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _decoded_base64(cleaned: str) -> bytes | None:
|
||||
strict: Final = _b64decode_or_none(cleaned)
|
||||
if strict is not None or ("-" not in cleaned and "_" not in cleaned):
|
||||
return strict
|
||||
return _b64decode_or_none(cleaned.translate(_URLSAFE_TO_STANDARD_B64))
|
||||
|
||||
|
||||
def _inline_part(name: str | None, value: str, size_limit: int | None, message_index: int) -> FilePart | None:
|
||||
mime_hint, payload = _split_data_url(value)
|
||||
if payload is None:
|
||||
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
|
||||
cleaned: Final = "".join(payload.split())
|
||||
if not cleaned:
|
||||
return None
|
||||
if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit:
|
||||
return FilePart(name, None, mime_hint, True, True, message_index)
|
||||
data: Final = _decoded_base64(cleaned)
|
||||
if data is None:
|
||||
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
|
||||
if size_limit is not None and len(data) > size_limit:
|
||||
return FilePart(name, None, mime_hint, True, True, message_index)
|
||||
return FilePart(name, data, mime_hint, True, False, message_index)
|
||||
|
||||
|
||||
def _name_from_url(url: str) -> str | None:
|
||||
try:
|
||||
return unquote(posixpath.basename(urlparse(url).path)) or None
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _part_from_file_block(block: Mapping[str, object], size_limit: int | None, message_index: int) -> FilePart | None:
|
||||
file_obj: Final = as_str_mapping(block.get("file"))
|
||||
if file_obj is None:
|
||||
return None
|
||||
name: Final = str_or_none(file_obj.get("filename")) or str_or_none(file_obj.get("file_id"))
|
||||
file_data: Final = str_or_none(file_obj.get("file_data"))
|
||||
inline: Final = _inline_part(name, file_data, size_limit, message_index) if file_data else None
|
||||
return inline or FilePart(name, None, None, False, False, message_index)
|
||||
|
||||
|
||||
def _part_from_image_url_block(
|
||||
block: Mapping[str, object], size_limit: int | None, message_index: int
|
||||
) -> FilePart | None:
|
||||
image_url: Final = block.get("image_url")
|
||||
image_url_fields: Final = as_str_mapping(image_url)
|
||||
url: Final = image_url_fields.get("url") if image_url_fields is not None else image_url
|
||||
if not isinstance(url, str) or not url:
|
||||
return None
|
||||
if url.startswith("data:"):
|
||||
return _inline_part(None, url, size_limit, message_index)
|
||||
return FilePart(_name_from_url(url), None, None, False, False, message_index)
|
||||
|
||||
|
||||
def _part_from_input_file_block(
|
||||
block: Mapping[str, object], size_limit: int | None, message_index: int
|
||||
) -> FilePart | None:
|
||||
declared_name: Final = str_or_none(block.get("filename")) or str_or_none(block.get("file_id"))
|
||||
file_data: Final = str_or_none(block.get("file_data"))
|
||||
inline: Final = _inline_part(declared_name, file_data, size_limit, message_index) if file_data else None
|
||||
if inline is not None:
|
||||
return inline
|
||||
file_url: Final = block.get("file_url")
|
||||
name: Final = (
|
||||
_name_from_url(file_url) if isinstance(file_url, str) and file_url and not declared_name else declared_name
|
||||
)
|
||||
return FilePart(name, None, None, False, False, message_index)
|
||||
|
||||
|
||||
def _part_from_input_audio_block(
|
||||
block: Mapping[str, object], size_limit: int | None, message_index: int
|
||||
) -> FilePart | None:
|
||||
audio: Final = as_str_mapping(block.get("input_audio"))
|
||||
if audio is None:
|
||||
return None
|
||||
data_b64: Final = audio.get("data")
|
||||
if not isinstance(data_b64, str) or not data_b64:
|
||||
return None
|
||||
name: Final = f"audio.{audio.get('format') or 'bin'}"
|
||||
return _inline_part(name, data_b64, size_limit, message_index) or FilePart(
|
||||
name, None, None, False, False, message_index
|
||||
)
|
||||
|
||||
|
||||
_BLOCK_PARSERS: Final[Mapping[str, Callable[[Mapping[str, object], int | None, int], FilePart | None]]] = (
|
||||
MappingProxyType(
|
||||
{
|
||||
"file": _part_from_file_block,
|
||||
"image_url": _part_from_image_url_block,
|
||||
"input_image": _part_from_image_url_block,
|
||||
"input_file": _part_from_input_file_block,
|
||||
"input_audio": _part_from_input_audio_block,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _file_parts_of_message(
|
||||
message: Mapping[str, object], size_limit: int | None, message_index: int
|
||||
) -> Iterator[FilePart]:
|
||||
for block in map(as_str_mapping, _list_items(message.get("content"))):
|
||||
if block is None:
|
||||
continue
|
||||
block_type = block.get("type")
|
||||
if not isinstance(block_type, str):
|
||||
continue
|
||||
parser = _BLOCK_PARSERS.get(block_type)
|
||||
if parser is None:
|
||||
continue
|
||||
try:
|
||||
part = parser(block, size_limit, message_index)
|
||||
except (TypeError, ValueError, AttributeError, KeyError):
|
||||
continue
|
||||
if part is not None and (part.inline or part.name):
|
||||
yield part
|
||||
|
||||
|
||||
def _file_parts_of_messages(structured_messages: Sequence[object], size_limit: int | None) -> Iterator[FilePart]:
|
||||
for message_index, message in enumerate(map(as_str_mapping, structured_messages)):
|
||||
if message is not None:
|
||||
yield from _file_parts_of_message(message, size_limit, message_index)
|
||||
|
||||
|
||||
def extract_file_parts_from_messages(
|
||||
structured_messages: Sequence[object] | None, size_limit: int | None = None
|
||||
) -> tuple[FilePart, ...]:
|
||||
return tuple(_file_parts_of_messages(structured_messages or (), size_limit))
|
||||
|
||||
|
||||
def _file_part_of_image(value: str, size_limit: int | None, index: int) -> FilePart | None:
|
||||
if value.startswith(("http://", "https://")):
|
||||
name: Final = _name_from_url(value)
|
||||
return FilePart(name, None, None, False, False, index) if name else None
|
||||
return _inline_part(None, value, size_limit, index)
|
||||
|
||||
|
||||
def extract_file_parts_from_images(
|
||||
images: Sequence[object] | None, size_limit: int | None = None
|
||||
) -> tuple[FilePart, ...]:
|
||||
candidates: Final = (
|
||||
_file_part_of_image(value, size_limit, index)
|
||||
for index, value in enumerate(images or ())
|
||||
if isinstance(value, str) and value
|
||||
)
|
||||
return tuple(part for part in candidates if part is not None)
|
||||
|
||||
|
||||
def make_tool_data(
|
||||
name: str, content: str | None, tool_input: Mapping[str, object] | None = None
|
||||
) -> Mapping[str, object]:
|
||||
action_name: Final = str(name) if str(name).strip() else _DEFAULT_TOOL_RESULT_NAME
|
||||
truncated: Final = action_name[:_TOOL_NAME_MAX_LENGTH]
|
||||
tool_name: Final = truncated if truncated.strip() else _DEFAULT_TOOL_RESULT_NAME
|
||||
return {
|
||||
"content": content,
|
||||
"tool_name": tool_name,
|
||||
"action_name": action_name,
|
||||
"tool_input": dict(tool_input or ()),
|
||||
}
|
||||
|
||||
|
||||
def _tool_call_field(tool_call: object, key: str) -> object:
|
||||
fields: Final = as_str_mapping(tool_call)
|
||||
return fields.get(key) if fields is not None else getattr(tool_call, key, None)
|
||||
|
||||
|
||||
def _json_or_str(value: object) -> str:
|
||||
try:
|
||||
return json.dumps(value)
|
||||
except (TypeError, ValueError):
|
||||
return str(value)
|
||||
|
||||
|
||||
def _parsed_tool_input(raw_arguments: str) -> Mapping[str, object]:
|
||||
if not raw_arguments:
|
||||
return _NO_TOOL_INPUT
|
||||
try:
|
||||
parsed: Final = json.loads(raw_arguments)
|
||||
except (ValueError, TypeError):
|
||||
return _NO_TOOL_INPUT
|
||||
return as_str_mapping(parsed) or _NO_TOOL_INPUT
|
||||
|
||||
|
||||
def _tool_content_and_input(raw_arguments: object) -> tuple[str, Mapping[str, object]]:
|
||||
if isinstance(raw_arguments, str):
|
||||
return raw_arguments, _parsed_tool_input(raw_arguments)
|
||||
if raw_arguments is None:
|
||||
return "", _NO_TOOL_INPUT
|
||||
arguments: Final = as_str_mapping(raw_arguments)
|
||||
return _json_or_str(raw_arguments), arguments if arguments is not None else _NO_TOOL_INPUT
|
||||
|
||||
|
||||
def tool_call_to_tool_data(tool_call: object) -> Mapping[str, object] | None:
|
||||
function: Final = _tool_call_field(tool_call, "function")
|
||||
name: Final = _tool_call_field(function, "name")
|
||||
if not name or not str(name).strip():
|
||||
return None
|
||||
content, tool_input = _tool_content_and_input(_tool_call_field(function, "arguments"))
|
||||
return make_tool_data(str(name), content, tool_input)
|
||||
|
||||
|
||||
def _message_tool_calls(message: object) -> Sequence[object]:
|
||||
fields: Final = as_str_mapping(message)
|
||||
return _list_items(fields.get("tool_calls")) if fields is not None else ()
|
||||
|
||||
|
||||
def extract_tool_calls_from_messages(structured_messages: Sequence[object] | None) -> tuple[object, ...]:
|
||||
"""Tool calls declared on the messages themselves.
|
||||
|
||||
Surfaces such as the Anthropic request path populate ``structured_messages`` but leave the
|
||||
top-level ``tool_calls`` input empty, so calls made in prior assistant turns are only visible here.
|
||||
"""
|
||||
return tuple(chain.from_iterable(_message_tool_calls(message) for message in structured_messages or ()))
|
||||
|
||||
|
||||
def tool_data_key(tool_data: Mapping[str, object]) -> str:
|
||||
"""Stable identity for a tool payload, so a call reached from two sources is only scanned once."""
|
||||
return json.dumps(tool_data, sort_keys=True, default=str)
|
||||
|
||||
|
||||
def _tool_content_block_text(block: object) -> str | None:
|
||||
if isinstance(block, str):
|
||||
return block
|
||||
fields: Final = as_str_mapping(block)
|
||||
return str_or_none(fields.get("text")) if fields is not None else None
|
||||
|
||||
|
||||
def _tool_content_text(content: object) -> str:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
fields: Final = as_str_mapping(content)
|
||||
if fields is not None:
|
||||
return _json_or_str(fields)
|
||||
texts: Final = (_tool_content_block_text(block) for block in _list_items(content))
|
||||
return "\n".join(text for text in texts if text is not None)
|
||||
|
||||
|
||||
def _extract_tool_content(content: object) -> str | None:
|
||||
text: Final = _tool_content_text(content)
|
||||
return text if text.strip() else None
|
||||
|
||||
|
||||
class _Declaration(NamedTuple):
|
||||
call_id: str
|
||||
message_index: int
|
||||
name: str
|
||||
|
||||
|
||||
def _declarations_in(message_index: int, message: object) -> Iterator[_Declaration]:
|
||||
fields: Final = as_str_mapping(message)
|
||||
if fields is None or fields.get("role") != "assistant":
|
||||
return
|
||||
for tool_call in map(as_str_mapping, _list_items(fields.get("tool_calls"))):
|
||||
if tool_call is None:
|
||||
continue
|
||||
call_id = tool_call.get("id")
|
||||
function = as_str_mapping(tool_call.get("function"))
|
||||
name = str_or_none(function.get("name")) if function is not None else None
|
||||
if isinstance(call_id, str) and call_id and name and name.strip():
|
||||
yield _Declaration(call_id, message_index, name)
|
||||
|
||||
|
||||
def _declarations_by_call_id(messages: Sequence[object]) -> Mapping[str, tuple[_Declaration, ...]]:
|
||||
"""Every named declaration of each tool call id, in message order, built in one pass."""
|
||||
declarations: Final = chain.from_iterable(
|
||||
_declarations_in(index, message) for index, message in enumerate(messages)
|
||||
)
|
||||
ordered: Final = sorted(declarations, key=attrgetter("call_id", "message_index"))
|
||||
return MappingProxyType({call_id: tuple(group) for call_id, group in groupby(ordered, key=attrgetter("call_id"))})
|
||||
|
||||
|
||||
def _resolve_tool_name(
|
||||
declarations: Mapping[str, tuple[_Declaration, ...]], tool_index: int, tool_call_id: object
|
||||
) -> str:
|
||||
"""Name from the nearest declaration of ``tool_call_id`` before ``tool_index``."""
|
||||
if not isinstance(tool_call_id, str) or not tool_call_id:
|
||||
return _DEFAULT_TOOL_RESULT_NAME
|
||||
declared: Final = declarations.get(tool_call_id, ())
|
||||
earlier: Final = bisect_left(declared, tool_index, key=attrgetter("message_index"))
|
||||
return declared[earlier - 1].name if earlier else _DEFAULT_TOOL_RESULT_NAME
|
||||
|
||||
|
||||
def extract_tool_results(
|
||||
structured_messages: Sequence[object] | None,
|
||||
) -> tuple[tuple[str, str, object], ...]:
|
||||
messages: Final = tuple(structured_messages or ())
|
||||
declarations: Final = _declarations_by_call_id(messages)
|
||||
|
||||
def _results() -> Iterator[tuple[str, str, object]]:
|
||||
for index, message in enumerate(map(as_str_mapping, messages)):
|
||||
if message is None or message.get("role") != "tool":
|
||||
continue
|
||||
content = _extract_tool_content(message.get("content"))
|
||||
if content is None:
|
||||
continue
|
||||
tool_call_id = message.get("tool_call_id")
|
||||
yield _resolve_tool_name(declarations, index, tool_call_id), content, tool_call_id
|
||||
|
||||
return tuple(_results())
|
||||
|
||||
|
||||
def _block_texts(content: object) -> Iterator[object]:
|
||||
blocks: Final = (as_str_mapping(block) for block in _list_items(content))
|
||||
return (block.get("text") for block in blocks if block is not None and block.get("text") is not None)
|
||||
|
||||
|
||||
def _message_text_origins(structured_messages: Sequence[object] | None) -> Iterator[tuple[object, bool]]:
|
||||
for message in map(as_str_mapping, structured_messages or ()):
|
||||
if message is None:
|
||||
continue
|
||||
content = message.get("content")
|
||||
from_tool_result = message.get("role") == "tool" and _extract_tool_content(content) is not None
|
||||
if isinstance(content, str):
|
||||
yield content, from_tool_result
|
||||
else:
|
||||
yield from ((text, from_tool_result) for text in _block_texts(content))
|
||||
|
||||
|
||||
def tool_result_text_indices(structured_messages: Sequence[object] | None, texts: Sequence[str]) -> frozenset[int]:
|
||||
"""Positions in ``texts`` that hold content already submitted under the TOOL policy.
|
||||
|
||||
The chat-completions guardrail flow builds ``texts`` and ``structured_messages`` from the
|
||||
same message list, so tool-role content lands in both and would otherwise be checked twice.
|
||||
Other surfaces (e.g. Anthropic messages) build ``texts`` from a differently shaped payload,
|
||||
so the mapping is only trusted when replaying it reproduces ``texts`` exactly; anything else
|
||||
falls back to checking every text.
|
||||
"""
|
||||
origins: Final = tuple(_message_text_origins(structured_messages))
|
||||
if tuple(text for text, _ in origins) != tuple(texts):
|
||||
return frozenset()
|
||||
return frozenset(index for index, (_, from_tool_result) in enumerate(origins) if from_tool_result)
|
||||
|
|
@ -18,6 +18,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
|
|||
event_hook=litellm_params.mode,
|
||||
guardrailIdentifier=litellm_params.guardrailIdentifier,
|
||||
guardrailVersion=litellm_params.guardrailVersion,
|
||||
outputScope=litellm_params.outputScope,
|
||||
checks=litellm_params.checks if isinstance(litellm_params.checks, BedrockChecksConfigModel) else None,
|
||||
content_filter_threshold=litellm_params.content_filter_threshold,
|
||||
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
|
||||
|
|
|
|||
117
litellm/proxy/guardrails/response_handlers.py
Normal file
117
litellm/proxy/guardrails/response_handlers.py
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
"""Response policies for the standalone guardrail endpoint."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.types.guardrails import (
|
||||
ApplyGuardrailRequest,
|
||||
ApplyGuardrailResponse,
|
||||
GuardrailEventHooks,
|
||||
GuardrailResponseFormat,
|
||||
NativeGuardrailResponse,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
GuardrailEndpointResponse: TypeAlias = ApplyGuardrailResponse | NativeGuardrailResponse
|
||||
_NATIVE_RESPONSE: Final[TypeAdapter[NativeGuardrailResponse]] = TypeAdapter(NativeGuardrailResponse)
|
||||
|
||||
|
||||
class StandardGuardrailResponseHandler:
|
||||
"""Preserve the endpoint's existing text response and hook behavior."""
|
||||
|
||||
def prepare_request_data(self, request: ApplyGuardrailRequest) -> Mapping[str, object]:
|
||||
return request.model_dump(include={"messages", "metadata"}, exclude_none=True)
|
||||
|
||||
def resolve_input_type(
|
||||
self, request: ApplyGuardrailRequest, guardrail: CustomGuardrail
|
||||
) -> Literal["request", "response"]:
|
||||
if request.input_type == "request" and guardrail.event_hook in (GuardrailEventHooks.post_call, "post_call"):
|
||||
return "response"
|
||||
return "response" if request.input_type == "response" else "request"
|
||||
|
||||
def build_response(
|
||||
self,
|
||||
request: ApplyGuardrailRequest,
|
||||
processed_inputs: GenericGuardrailAPIInputs,
|
||||
request_data: Mapping[str, object],
|
||||
) -> GuardrailEndpointResponse:
|
||||
texts: Final = processed_inputs.get("texts", [])
|
||||
return ApplyGuardrailResponse(response_text=texts[0] if texts else request.text)
|
||||
|
||||
def accept_hook_result(self, original: GuardrailEndpointResponse, modified: object) -> GuardrailEndpointResponse:
|
||||
return modified if isinstance(modified, ApplyGuardrailResponse) else original
|
||||
|
||||
|
||||
class _NativeGuardrailRecord(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
guardrail_name: str | None = None
|
||||
guardrail_response: JsonValue = None
|
||||
|
||||
|
||||
class _NativeGuardrailMetadata(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
standard_logging_guardrail_information: tuple[_NativeGuardrailRecord, ...] = ()
|
||||
|
||||
|
||||
class NativeGuardrailResponseHandler(StandardGuardrailResponseHandler):
|
||||
"""Return structured findings recorded by this guardrail invocation."""
|
||||
|
||||
def prepare_request_data(self, request: ApplyGuardrailRequest) -> Mapping[str, object]:
|
||||
return request.model_dump(
|
||||
include={"messages", "metadata"},
|
||||
exclude={"metadata": {"standard_logging_guardrail_information"}},
|
||||
exclude_none=True,
|
||||
)
|
||||
|
||||
def resolve_input_type(
|
||||
self, request: ApplyGuardrailRequest, guardrail: CustomGuardrail
|
||||
) -> Literal["request", "response"]:
|
||||
if "input_type" in request.model_fields_set:
|
||||
return "response" if request.input_type == "response" else "request"
|
||||
return super().resolve_input_type(request, guardrail)
|
||||
|
||||
def build_response(
|
||||
self,
|
||||
request: ApplyGuardrailRequest,
|
||||
processed_inputs: GenericGuardrailAPIInputs,
|
||||
request_data: Mapping[str, object],
|
||||
) -> GuardrailEndpointResponse:
|
||||
metadata: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data))
|
||||
recorded: Final = _NativeGuardrailMetadata.model_validate_json(
|
||||
safe_dumps(metadata if metadata is not None else {})
|
||||
)
|
||||
findings: Final = tuple(
|
||||
record.guardrail_response
|
||||
for record in recorded.standard_logging_guardrail_information
|
||||
if record.guardrail_name == request.guardrail_name and isinstance(record.guardrail_response, (dict, list))
|
||||
)
|
||||
if not findings:
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail=f"Guardrail '{request.guardrail_name}' did not report structured native findings.",
|
||||
)
|
||||
return findings[0] if len(findings) == 1 else findings
|
||||
|
||||
def accept_hook_result(self, original: GuardrailEndpointResponse, modified: object) -> GuardrailEndpointResponse:
|
||||
if modified is None:
|
||||
return original
|
||||
return _NATIVE_RESPONSE.validate_python(modified)
|
||||
|
||||
|
||||
GUARDRAIL_RESPONSE_HANDLERS: Final[Mapping[GuardrailResponseFormat, StandardGuardrailResponseHandler]] = (
|
||||
MappingProxyType(
|
||||
{
|
||||
GuardrailResponseFormat.STANDARD: StandardGuardrailResponseHandler(),
|
||||
GuardrailResponseFormat.NATIVE: NativeGuardrailResponseHandler(),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
|
@ -118,7 +118,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
|
|||
elif priority is not None and litellm.priority_reservation is not None:
|
||||
if os.getenv("LITELLM_LICENSE", None) is None:
|
||||
verbose_proxy_logger.error(
|
||||
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
|
||||
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/enterprise."
|
||||
)
|
||||
else:
|
||||
value: Final = litellm.priority_reservation[priority]
|
||||
|
|
|
|||
|
|
@ -133,7 +133,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
|
|||
elif priority is not None and litellm.priority_reservation is not None:
|
||||
if os.getenv("LITELLM_LICENSE", None) is None:
|
||||
verbose_proxy_logger.error(
|
||||
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
|
||||
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/enterprise."
|
||||
)
|
||||
else:
|
||||
value: Final = litellm.priority_reservation[priority]
|
||||
|
|
|
|||
|
|
@ -2405,7 +2405,7 @@ async def generate_key_fn(
|
|||
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
|
||||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
|
||||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||||
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
|
||||
- allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
|
||||
|
|
@ -2613,7 +2613,7 @@ async def generate_service_account_key_fn(
|
|||
- tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit.
|
||||
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
|
||||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
|
||||
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
|
||||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
Examples:
|
||||
|
|
@ -4045,7 +4045,7 @@ async def update_key_fn(
|
|||
- models: Optional[list] - Model_name's a user is allowed to call
|
||||
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/enterprise#governance-and-cost)
|
||||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
|
|
|
|||
|
|
@ -64,11 +64,16 @@ from litellm.llms.azure.passthrough.transformation import (
|
|||
from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
DeepgramSpeakMeter,
|
||||
deepgram_listen_callback_params,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_requested_model,
|
||||
deepgram_listen_websocket_target,
|
||||
deepgram_speak_character_rate,
|
||||
deepgram_speak_registry_key,
|
||||
deepgram_speak_requested_model,
|
||||
deepgram_speak_websocket_target,
|
||||
)
|
||||
from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
|
|
@ -129,6 +134,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
YieldsToConfiguredPassThroughRoute,
|
||||
create_pass_through_route,
|
||||
create_websocket_passthrough_route,
|
||||
truncated_close_reason,
|
||||
websocket_passthrough_request,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
|
||||
|
|
@ -145,6 +151,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
WebSocketSessionMeter,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.tinyfish import (
|
||||
TINYFISH_AUTHENTICATED_RUN_FIELDS,
|
||||
|
|
@ -3460,6 +3467,7 @@ class _WebsocketRelay(Protocol):
|
|||
forward_headers: bool,
|
||||
endpoint: str,
|
||||
accept_websocket: bool,
|
||||
meter: WebSocketSessionMeter | None = None,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
|
|
@ -3606,7 +3614,9 @@ async def deepgram_listen_websocket_route(
|
|||
if not deepgram_listen_is_priced(target):
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_listen_registry_key(target)),
|
||||
reason=truncated_close_reason(
|
||||
_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_listen_registry_key(target))
|
||||
),
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -3621,6 +3631,54 @@ async def deepgram_listen_websocket_route(
|
|||
)
|
||||
|
||||
|
||||
async def deepgram_speak_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(
|
||||
websocket, model=deepgram_speak_requested_model(websocket.url.query)
|
||||
)
|
||||
|
||||
|
||||
@router.websocket("/deepgram/v1/speak")
|
||||
@router.websocket("/deepgram/speak")
|
||||
async def deepgram_speak_websocket_route(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_speak_user_api_key_auth)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
) -> None:
|
||||
deepgram_api_key: Final = passthrough_endpoint_router.get_credentials(
|
||||
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
|
||||
region_name=None,
|
||||
)
|
||||
if deepgram_api_key is None:
|
||||
await websocket.close(code=1011, reason=_DEEPGRAM_WS_MISSING_KEY_REASON)
|
||||
return
|
||||
|
||||
await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket))
|
||||
target: Final = deepgram_speak_websocket_target(
|
||||
api_base=get_secret_str("DEEPGRAM_API_BASE"),
|
||||
query_string=websocket.url.query,
|
||||
)
|
||||
character_rate: Final = deepgram_speak_character_rate(target)
|
||||
if character_rate is None:
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=truncated_close_reason(
|
||||
_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_speak_registry_key(target))
|
||||
),
|
||||
)
|
||||
return
|
||||
|
||||
await relay(
|
||||
websocket=websocket,
|
||||
target=target,
|
||||
custom_headers={"Authorization": f"Token {deepgram_api_key}"},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
forward_headers=False,
|
||||
endpoint=websocket.url.path,
|
||||
accept_websocket=False,
|
||||
meter=DeepgramSpeakMeter(target, character_rate),
|
||||
)
|
||||
|
||||
|
||||
_OPENAI_BETA_HEADER: Final = "openai-beta"
|
||||
_CLIENT_OPENAI_BETA_RESOURCES: Final = frozenset({"agents", "vaults", "chatkit"})
|
||||
|
||||
|
|
|
|||
|
|
@ -2584,7 +2584,7 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
|
|||
return decoded if isinstance(decoded, dict) else None
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
def truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
"""
|
||||
|
|
@ -2958,7 +2958,7 @@ async def websocket_passthrough_request(
|
|||
if upstream_close is not None and _client_socket_is_open(websocket):
|
||||
await websocket.close(
|
||||
code=upstream_close.code,
|
||||
reason=_truncated_close_reason(upstream_close.reason),
|
||||
reason=truncated_close_reason(upstream_close.reason),
|
||||
)
|
||||
|
||||
end_time: Final = datetime.now()
|
||||
|
|
|
|||
|
|
@ -1193,7 +1193,7 @@ class ProxyInitializationHelpers:
|
|||
"--reload",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
|
||||
help="Enable uvicorn hot reload (dev only), including every --num_workers worker on uvicorn>=0.50.0 outside Windows. Also reloads when the --config YAML or .env changes. Incompatible with --run_gunicorn, --run_hypercorn, and --run_granian.",
|
||||
)
|
||||
@click.option(
|
||||
"--prometheus_metrics_port",
|
||||
|
|
@ -1785,9 +1785,28 @@ def run_server(
|
|||
if num_workers > 1:
|
||||
start_query_engine_reaper()
|
||||
if num_workers > 1 or reload:
|
||||
uvicorn_args.update(
|
||||
ProxyInitializationHelpers.child_process_uvicorn_overrides(uvicorn_args.get("log_config"))
|
||||
worker_overrides: Final = ProxyInitializationHelpers.child_process_uvicorn_overrides(
|
||||
uvicorn_args.get("log_config")
|
||||
)
|
||||
if reload and num_workers > 1:
|
||||
from litellm.proxy.dev_reload import (
|
||||
RELOAD_OPTIONS,
|
||||
multi_worker_reload_supported,
|
||||
run_multi_worker_reload,
|
||||
)
|
||||
|
||||
if multi_worker_reload_supported():
|
||||
run_multi_worker_reload(
|
||||
RELOAD_OPTIONS.validate_python(uvicorn_args),
|
||||
workers=TypeAdapter(int).validate_python(num_workers),
|
||||
worker_overrides=worker_overrides,
|
||||
)
|
||||
return
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: --reload with --num_workers>1 requires uvicorn>=0.50.0 "
|
||||
"and an OS other than Windows, so only a single worker reloads.\033[0m"
|
||||
)
|
||||
uvicorn_args.update(worker_overrides)
|
||||
uvicorn.run(
|
||||
**uvicorn_args,
|
||||
workers=num_workers,
|
||||
|
|
|
|||
|
|
@ -392,6 +392,7 @@ from litellm.llms.anthropic.messages.convertible_blocks import (
|
|||
drop_non_string_tool_result_texts,
|
||||
malformed_tool_result_count_error,
|
||||
)
|
||||
from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.openai.live import parse_live_session_start
|
||||
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
|
||||
|
|
@ -846,6 +847,8 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware
|
|||
from litellm.proxy.telemetry.runtime import TelemetryRuntime
|
||||
from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors
|
||||
from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings
|
||||
from litellm.router import router_arguments
|
||||
from litellm.router_config import build_router_config
|
||||
|
||||
try:
|
||||
from litellm.proxy.enterprise_billing.billing_metrics import (
|
||||
|
|
@ -1207,7 +1210,7 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai
|
|||
ui_message += f"\n\n🔎 [```LiteLLM Model Hub```]({model_hub_link}). See available models on the proxy. [**Docs**](https://docs.litellm.ai/docs/proxy/ai_hub)"
|
||||
|
||||
custom_swagger_message: Final = (
|
||||
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)"
|
||||
"[**Customize Swagger Docs**](https://docs.litellm.ai/docs/enterprise#operations--branding)"
|
||||
)
|
||||
|
||||
### CUSTOM BRANDING [ENTERPRISE FEATURE] ###
|
||||
|
|
@ -1458,7 +1461,11 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None":
|
|||
connector_kwargs["socket_factory"] = socket_factory
|
||||
|
||||
connector: Final = TCPConnector(**connector_kwargs)
|
||||
session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar())
|
||||
session: Final = ClientSession(
|
||||
connector=connector,
|
||||
cookie_jar=DummyCookieJar(),
|
||||
trace_configs=[build_pool_wait_trace_config()],
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)",
|
||||
|
|
@ -3068,6 +3075,29 @@ redis_usage_cache: RedisCache | None = None # redis cache used for tracking spe
|
|||
polling_via_cache_enabled: Literal["all"] | list[str] | bool = False
|
||||
native_background_mode: list[str] = [] # Models that should use native provider background mode instead of polling
|
||||
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
|
||||
|
||||
|
||||
class _ResponseBackgroundSettings(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
polling_via_cache: Literal["all"] | tuple[str, ...] | bool = False
|
||||
native_background_mode: tuple[str, ...] = ()
|
||||
ttl: int = 3600
|
||||
|
||||
|
||||
def _apply_background_response_settings(value: SettingsJsonValue) -> None:
|
||||
global polling_via_cache_enabled, native_background_mode, polling_cache_ttl
|
||||
|
||||
responses: Final = _as_settings_mapping(value)
|
||||
settings: Final = _ResponseBackgroundSettings.model_validate(responses.get("background_mode") or {})
|
||||
polling: Final = settings.polling_via_cache
|
||||
polling_via_cache_enabled, native_background_mode, polling_cache_ttl = ( # rebind-ok: Refresh runtime config.
|
||||
list(polling) if isinstance(polling, tuple) else polling,
|
||||
list(settings.native_background_mode),
|
||||
settings.ttl,
|
||||
)
|
||||
|
||||
|
||||
user_custom_auth = None
|
||||
user_custom_key_generate = None
|
||||
# Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload.
|
||||
|
|
@ -6930,20 +6960,7 @@ class ProxyConfig:
|
|||
# see usage here: https://docs.litellm.ai/docs/proxy/caching
|
||||
pass
|
||||
elif key == "responses":
|
||||
# Initialize global polling via cache settings
|
||||
global polling_via_cache_enabled, native_background_mode, polling_cache_ttl
|
||||
background_mode = value.get("background_mode", {})
|
||||
polling_via_cache_enabled = background_mode.get("polling_via_cache", False)
|
||||
native_background_mode = background_mode.get("native_background_mode", [])
|
||||
polling_cache_ttl = background_mode.get("ttl", 3600)
|
||||
verbose_proxy_logger.debug(
|
||||
"%s Initialized polling via cache: enabled=%s, native_background_mode=%s, ttl=%s%s",
|
||||
blue_color_code,
|
||||
polling_via_cache_enabled,
|
||||
native_background_mode,
|
||||
polling_cache_ttl,
|
||||
reset_color_code,
|
||||
)
|
||||
_apply_background_response_settings(value)
|
||||
elif key == "max_ui_session_budget":
|
||||
litellm.max_ui_session_budget = float(value) if value is not None else None
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -7407,20 +7424,24 @@ class ProxyConfig:
|
|||
verbose_proxy_logger.warning(
|
||||
"Key '%s' is not a valid argument for Router.__init__(). Ignoring this key.", k
|
||||
)
|
||||
router = litellm.Router(
|
||||
**router_params,
|
||||
assistants_config=assistants_config,
|
||||
search_tools=search_tools,
|
||||
router_general_settings=RouterGeneralSettings(
|
||||
async_only_mode=True # only init async clients
|
||||
),
|
||||
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
router = litellm.Router.from_config(
|
||||
build_router_config(
|
||||
router_arguments(
|
||||
**router_params,
|
||||
assistants_config=assistants_config,
|
||||
search_tools=search_tools,
|
||||
router_general_settings=RouterGeneralSettings(
|
||||
async_only_mode=True # only init async clients
|
||||
),
|
||||
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
if redis_usage_cache is not None and router.health().store.redis_cache is None:
|
||||
|
|
@ -7910,19 +7931,23 @@ class ProxyConfig:
|
|||
# Router can function with model_list=[] if search_tools are configured
|
||||
if len(_model_list) > 0 or search_tools:
|
||||
verbose_proxy_logger.debug("_model_list: %s", _model_list)
|
||||
llm_router = litellm.Router(
|
||||
model_list=_model_list,
|
||||
router_general_settings=RouterGeneralSettings(
|
||||
async_only_mode=True # only init async clients
|
||||
),
|
||||
search_tools=search_tools,
|
||||
ignore_invalid_deployments=True,
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
llm_router = litellm.Router.from_config(
|
||||
build_router_config(
|
||||
router_arguments(
|
||||
model_list=_model_list,
|
||||
router_general_settings=RouterGeneralSettings(
|
||||
async_only_mode=True # only init async clients
|
||||
),
|
||||
search_tools=search_tools,
|
||||
ignore_invalid_deployments=True,
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
)
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
|
||||
else:
|
||||
|
|
@ -8643,6 +8668,7 @@ class ProxyConfig:
|
|||
|
||||
def _apply_litellm_settings_db_values(self, db_values: Mapping[str, SettingsJsonValue]) -> None:
|
||||
self.litellm_settings.apply_db_row("litellm_settings", db_values)
|
||||
_apply_background_response_settings(self.litellm_settings.get("responses", {}))
|
||||
for key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES:
|
||||
if key in db_values and (value := self.litellm_settings.get(key)) is not None:
|
||||
setattr(litellm, key, value)
|
||||
|
|
@ -18100,6 +18126,8 @@ async def model_deprecations(
|
|||
def get_model_group_info(
|
||||
llm_router: Router, all_models_str: Sequence[str], model_group: str | None, team_id: str | None = None
|
||||
) -> list[ModelGroupInfoProxy]:
|
||||
from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request
|
||||
|
||||
model_groups: Final[list[ModelGroupInfoProxy]] = []
|
||||
|
||||
unique_models: Final = tuple(dict.fromkeys(all_models_str))
|
||||
|
|
@ -18109,6 +18137,14 @@ def get_model_group_info(
|
|||
if model_group is not None and model_group != model:
|
||||
continue
|
||||
|
||||
supports_background_chat = should_use_polling_for_request(
|
||||
background_mode=True,
|
||||
polling_via_cache_enabled=polling_via_cache_enabled,
|
||||
redis_cache=redis_usage_cache,
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
native_background_mode=native_background_mode,
|
||||
)
|
||||
_model_group_info = registry.group(model, team_id=team_id)
|
||||
|
||||
if _model_group_info is not None:
|
||||
|
|
@ -18116,6 +18152,7 @@ def get_model_group_info(
|
|||
ModelGroupInfoProxy.model_validate(
|
||||
{
|
||||
**_model_group_info.model_dump(),
|
||||
"supports_background_chat": supports_background_chat,
|
||||
**model_group_custom_fields(
|
||||
registry.deployments(
|
||||
resolve_model_group_alias(registry.aliases(), model) or model,
|
||||
|
|
@ -18130,6 +18167,7 @@ def get_model_group_info(
|
|||
model_group_info = ModelGroupInfoProxy(
|
||||
model_group=model,
|
||||
providers=[],
|
||||
supports_background_chat=supports_background_chat,
|
||||
)
|
||||
model_groups.append(model_group_info)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ https://platform.openai.com/docs/api-reference/responses-streaming
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterable, Callable, Mapping, Sequence
|
||||
from time import perf_counter
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable
|
||||
|
||||
from fastapi import Request, Response
|
||||
|
|
@ -128,6 +129,16 @@ async def background_streaming_task(
|
|||
https://platform.openai.com/docs/api-reference/responses/object
|
||||
"""
|
||||
|
||||
started_at: Final = perf_counter()
|
||||
first_token_at: float | None = None # rebind-ok: first content event closes the TTFT interval
|
||||
finished_at: float | None = None # rebind-ok: terminal event closes the execution interval
|
||||
|
||||
def timing_snapshot() -> Mapping[str, float]:
|
||||
return {
|
||||
**({"time_to_first_token_ms": (first_token_at - started_at) * 1000} if first_token_at is not None else {}),
|
||||
**({"total_latency_ms": (finished_at - started_at) * 1000} if finished_at is not None else {}),
|
||||
}
|
||||
|
||||
try:
|
||||
verbose_proxy_logger.info("Starting background streaming for %s", polling_id)
|
||||
|
||||
|
|
@ -191,7 +202,7 @@ async def background_streaming_task(
|
|||
incomplete_details_data = None
|
||||
|
||||
state_dirty = False # Track if state needs to be synced
|
||||
last_update_time = asyncio.get_event_loop().time()
|
||||
last_update_time = perf_counter()
|
||||
UPDATE_INTERVAL: Final = 0.150 # 150ms batching interval
|
||||
|
||||
# Track the terminal event from the stream (may not be "completed")
|
||||
|
|
@ -210,13 +221,14 @@ async def background_streaming_task(
|
|||
"""Flush accumulated state to Redis if interval elapsed or forced"""
|
||||
nonlocal state_dirty, last_update_time
|
||||
|
||||
current_time: Final = asyncio.get_event_loop().time()
|
||||
current_time: Final = perf_counter()
|
||||
if state_dirty and (force or (current_time - last_update_time) >= UPDATE_INTERVAL):
|
||||
# Convert output_items dict to list for update
|
||||
output_list: Final = list(output_items.values())
|
||||
await polling_handler.update_state(
|
||||
polling_id=polling_id,
|
||||
output=output_list,
|
||||
timing=timing_snapshot(),
|
||||
)
|
||||
state_dirty = False
|
||||
last_update_time = current_time
|
||||
|
|
@ -242,6 +254,20 @@ async def background_streaming_task(
|
|||
event: _StreamEvent = _StreamEventParser.parse(chunk_data)
|
||||
event_type = event.get("type", "")
|
||||
|
||||
if (
|
||||
first_token_at is None
|
||||
and event_type
|
||||
in (
|
||||
"response.output_text.delta",
|
||||
"response.reasoning_text.delta",
|
||||
"response.reasoning_summary_text.delta",
|
||||
"response.reasoning.delta",
|
||||
)
|
||||
and event.get("delta")
|
||||
):
|
||||
first_token_at = perf_counter()
|
||||
state_dirty = True
|
||||
|
||||
# Process different event types based on OpenAI streaming spec
|
||||
if event_type == "response.output_item.added":
|
||||
# New output item added
|
||||
|
|
@ -334,6 +360,7 @@ async def background_streaming_task(
|
|||
):
|
||||
# Terminal event - extract all ResponsesAPIResponse fields
|
||||
# https://platform.openai.com/docs/api-reference/responses-streaming
|
||||
finished_at = perf_counter()
|
||||
response_data = event.get("response", {})
|
||||
terminal_status = response_data.get(
|
||||
"status",
|
||||
|
|
@ -384,10 +411,13 @@ async def background_streaming_task(
|
|||
|
||||
# Use the terminal status from the stream, default to "completed"
|
||||
final_status: Final = terminal_status or "completed"
|
||||
if finished_at is None:
|
||||
finished_at = perf_counter() # rebind-ok: close the interval when the stream has no terminal event
|
||||
|
||||
await polling_handler.update_state(
|
||||
polling_id=polling_id,
|
||||
status=final_status,
|
||||
timing=timing_snapshot(),
|
||||
usage=usage_data,
|
||||
error=terminal_error,
|
||||
reasoning=reasoning_data,
|
||||
|
|
@ -417,6 +447,7 @@ async def background_streaming_task(
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
finished_at = perf_counter() # rebind-ok: capture execution time at failure
|
||||
verbose_proxy_logger.error("Error in background streaming task for %s: %s", polling_id, e)
|
||||
import traceback
|
||||
|
||||
|
|
@ -425,6 +456,7 @@ async def background_streaming_task(
|
|||
await polling_handler.update_state(
|
||||
polling_id=polling_id,
|
||||
status="failed",
|
||||
timing=timing_snapshot(),
|
||||
error={
|
||||
"type": "internal_error",
|
||||
"message": str(e),
|
||||
|
|
|
|||
|
|
@ -3,9 +3,9 @@ Response Polling Handler for Background Responses with Cache
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
|
|
@ -15,6 +15,9 @@ from litellm._uuid import uuid4
|
|||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
_RESPONSE_POLLING_TARGET: Final = "response_polling"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_JSON_LIST: Final = TypeAdapter(list[object])
|
||||
|
|
@ -116,6 +119,7 @@ class ResponsePollingHandler:
|
|||
parallel_tool_calls: bool | None = None,
|
||||
user: str | None = None,
|
||||
store: bool | None = None,
|
||||
timing: Mapping[str, float] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Update the polling state in Redis
|
||||
|
|
@ -144,6 +148,7 @@ class ResponsePollingHandler:
|
|||
parallel_tool_calls: Whether parallel tool calls are enabled
|
||||
user: User identifier
|
||||
store: Whether to store the response
|
||||
timing: Server execution timings in milliseconds
|
||||
"""
|
||||
if not self.redis_cache:
|
||||
return
|
||||
|
|
@ -167,6 +172,9 @@ class ResponsePollingHandler:
|
|||
if output is not None:
|
||||
state["output"] = output
|
||||
|
||||
if timing is not None:
|
||||
state["litellm_response_timing"] = timing
|
||||
|
||||
# Update usage
|
||||
if usage:
|
||||
state["usage"] = usage
|
||||
|
|
@ -264,10 +272,10 @@ class ResponsePollingHandler:
|
|||
|
||||
def should_use_polling_for_request(
|
||||
background_mode: bool,
|
||||
polling_via_cache_enabled, # Can be False, "all", or List[str]
|
||||
redis_cache, # RedisCache or None
|
||||
polling_via_cache_enabled: bool | str | Sequence[str],
|
||||
redis_cache: RedisCache | None,
|
||||
model: str,
|
||||
llm_router, # Router instance or None
|
||||
llm_router: "Router | None",
|
||||
native_background_mode: list[str] | None = None, # List of models that should use native background mode
|
||||
) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -18,7 +18,9 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.router import router_arguments
|
||||
from litellm.router_api import router_init_args
|
||||
from litellm.router_config import build_router_config
|
||||
|
||||
# Client-supplied params that make the router or the call path fabricate a
|
||||
# failure or a delay instead of calling the provider. The ``mock_testing_*``
|
||||
|
|
@ -98,7 +100,9 @@ async def _route_user_config_request(data: dict, route_type: str):
|
|||
filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args}
|
||||
|
||||
logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type)
|
||||
cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config))
|
||||
cached_router: Final = user_config_routers.get(
|
||||
filtered_config, lambda: litellm.Router.from_config(build_router_config(router_arguments(**filtered_config)))
|
||||
)
|
||||
cached_router.observe(logging_obj)
|
||||
return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj})
|
||||
|
||||
|
|
|
|||
|
|
@ -216,6 +216,20 @@ class _SettingsWithSchema(LiteLLMBaseModel):
|
|||
field_schema: dict[str, object]
|
||||
|
||||
|
||||
class BackgroundChatSettings(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
enabled: bool = Field(description="Enable cache-backed background Responses for Chat and API requests.")
|
||||
|
||||
|
||||
class BackgroundChatSettingsResponse(BackgroundChatSettings):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
redis_configured: bool = Field(description="Whether this proxy has coordination Redis configured.")
|
||||
managed_by_config: bool = Field(description="Whether the responses configuration is owned by the config file.")
|
||||
can_save: bool = Field(description="Whether database-backed configuration updates are available.")
|
||||
|
||||
|
||||
class SSOSettingsResponse(SettingsResponse):
|
||||
"""Response model for SSO settings"""
|
||||
|
||||
|
|
@ -1815,6 +1829,87 @@ async def sync_ui_settings_to_general_settings(
|
|||
return apply_runtime_general_settings_flags(parsed)
|
||||
|
||||
|
||||
def _background_chat_settings_response(config: Mapping[str, object]) -> BackgroundChatSettingsResponse:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
settings: Final = _as_settings_section(config.get("litellm_settings"))
|
||||
responses: Final = _as_settings_section(settings.get("responses"))
|
||||
background: Final = _as_settings_section(responses.get("background_mode"))
|
||||
polling: Final = background.get("polling_via_cache", False)
|
||||
return BackgroundChatSettingsResponse(
|
||||
enabled=bool(polling),
|
||||
redis_configured=proxy_server.redis_usage_cache is not None,
|
||||
managed_by_config=proxy_server.proxy_config.litellm_settings.owned_by_config("responses"),
|
||||
can_save=proxy_server.prisma_client is not None and proxy_server.store_model_in_db is True,
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/background_chat_settings",
|
||||
tags=["Settings"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=BackgroundChatSettingsResponse,
|
||||
)
|
||||
async def get_background_chat_settings() -> BackgroundChatSettingsResponse:
|
||||
"""Get saved background-response settings and the prerequisites for changing them."""
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
config: Final = _as_settings_section(_DECODED_JSON.validate_python(await proxy_config.get_config()))
|
||||
return _background_chat_settings_response(config)
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/update/background_chat_settings",
|
||||
tags=["Settings"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=BackgroundChatSettingsResponse,
|
||||
)
|
||||
async def update_background_chat_settings(
|
||||
settings: BackgroundChatSettings,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> BackgroundChatSettingsResponse:
|
||||
"""Save background-response polling settings for the proxy configuration refresh."""
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_config_audit_log,
|
||||
prisma_client,
|
||||
proxy_config,
|
||||
redis_usage_cache,
|
||||
store_model_in_db,
|
||||
)
|
||||
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(status_code=403, detail="Only proxy admins can update background chat settings.")
|
||||
if prisma_client is None or store_model_in_db is not True:
|
||||
raise HTTPException(status_code=400, detail="Enable database-backed configuration to save this setting.")
|
||||
if settings.enabled and redis_usage_cache is None:
|
||||
raise HTTPException(status_code=400, detail="Configure coordination Redis and restart the proxy first.")
|
||||
|
||||
raw_config: Final = await proxy_config.get_config()
|
||||
config: Final = _as_settings_section(_DECODED_JSON.validate_python(raw_config))
|
||||
litellm_settings: Final = _as_settings_section(config.get("litellm_settings"))
|
||||
responses: Final = _UI_SETTINGS_OBJECT.validate_python(_as_settings_section(litellm_settings.get("responses")))
|
||||
background: Final = _UI_SETTINGS_OBJECT.validate_python(_as_settings_section(responses.get("background_mode")))
|
||||
updated_responses: Final[Mapping[str, JsonValue]] = {
|
||||
**responses,
|
||||
"background_mode": {**background, "polling_via_cache": "all" if settings.enabled else False},
|
||||
}
|
||||
proxy_config.reject_config_owned_writes(
|
||||
section_name="litellm_settings", changed_keys={"responses": dict(updated_responses)}
|
||||
)
|
||||
raw_config["litellm_settings"] = {**litellm_settings, "responses": dict(updated_responses)}
|
||||
await proxy_config.save_config(new_config=config)
|
||||
asyncio.create_task(
|
||||
create_config_audit_log(
|
||||
param_name="responses.background_mode.polling_via_cache",
|
||||
action="updated",
|
||||
before_value=background.get("polling_via_cache", False),
|
||||
after_value="all" if settings.enabled else False,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
)
|
||||
return _background_chat_settings_response(config)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/ui_settings",
|
||||
tags=["UI Settings"],
|
||||
|
|
|
|||
|
|
@ -128,6 +128,15 @@ def _with_hoisted_additional_tools(responses_request: Mapping[str, object], drop
|
|||
return {**responses_request, "input": hoisted.input, "tools": list(hoisted.tools)}
|
||||
|
||||
|
||||
def _caller_input(responses_request: Mapping[str, object]) -> ResponseInputParam | None:
|
||||
responses_input: Final = responses_request.get("input")
|
||||
return (
|
||||
cast(ResponseInputParam, responses_input) # cast-ok: forwarded as sent, as inference does
|
||||
if isinstance(responses_input, list)
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
tools: Final = responses_request.get("tools")
|
||||
if tools is None:
|
||||
|
|
@ -200,6 +209,7 @@ async def chat_count_request(
|
|||
),
|
||||
custom_llm_provider,
|
||||
policy.drop_params,
|
||||
caller_input=_caller_input(responses_request),
|
||||
),
|
||||
model,
|
||||
custom_llm_provider,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Handler for transforming responses api requests to litellm.completion requests
|
|||
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -32,6 +33,13 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_BRIDGE_TRANSFORM_KEYWORDS: Final = frozenset({"caller_input", "names_non_string_fields"})
|
||||
|
||||
|
||||
def _without_bridge_transform_keywords(request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{key: value for key, value in request_kwargs.items() if key not in _BRIDGE_TRANSFORM_KEYWORDS}
|
||||
)
|
||||
|
||||
|
||||
def _chat_verbosity(responses_api_request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -95,6 +103,7 @@ class LiteLLMCompletionTransformationHandler:
|
|||
previous_response_id,
|
||||
hoisted,
|
||||
model=model,
|
||||
caller_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
|
|
@ -108,11 +117,13 @@ class LiteLLMCompletionTransformationHandler:
|
|||
**LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
|
||||
model=model,
|
||||
input=bridged_input,
|
||||
caller_input=input,
|
||||
names_non_string_fields=True,
|
||||
responses_api_request=bridged_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
extra_headers=extra_headers,
|
||||
**kwargs,
|
||||
**_without_bridge_transform_keywords(_REQUEST_KWARGS_ADAPTER.validate_python(kwargs)),
|
||||
),
|
||||
**_chat_verbosity(bridged_request),
|
||||
}
|
||||
|
|
@ -163,6 +174,7 @@ class LiteLLMCompletionTransformationHandler:
|
|||
/,
|
||||
*,
|
||||
model: str,
|
||||
caller_input: str | ResponseInputParam,
|
||||
responses_api_request: ResponsesAPIOptionalRequestParams,
|
||||
custom_llm_provider: str | None,
|
||||
stream: bool | None,
|
||||
|
|
@ -186,11 +198,13 @@ class LiteLLMCompletionTransformationHandler:
|
|||
**LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
|
||||
model=model,
|
||||
input=bridged.input,
|
||||
caller_input=caller_input,
|
||||
names_non_string_fields=True,
|
||||
responses_api_request=bridged.responses_api_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
extra_headers=extra_headers,
|
||||
**request_kwargs,
|
||||
**_without_bridge_transform_keywords(request_kwargs),
|
||||
),
|
||||
**_chat_verbosity(bridged.responses_api_request),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ from litellm.constants import (
|
|||
REDACTED_BY_LITELLM,
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER,
|
||||
)
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.integrations.anthropic_cache_control_hook import (
|
||||
AnthropicCacheControlHook,
|
||||
block_with_cache_control,
|
||||
|
|
@ -65,6 +66,12 @@ from litellm.litellm_core_utils.hidden_params import (
|
|||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
invalid_type_message,
|
||||
is_present_non_string,
|
||||
is_unhashable_value_behind,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of
|
||||
from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control
|
||||
from litellm.llms.anthropic.messages.native.utils import anthropic_stop_reason, openai_chat_refusal_text
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
|
@ -268,6 +275,7 @@ _IMAGE_URL_ADAPTER: Final = TypeAdapter(ChatCompletionImageUrlObject)
|
|||
_FILE_PART_ADAPTER: Final = TypeAdapter(_ResponsesFilePart)
|
||||
_MessageT = TypeVar("_MessageT")
|
||||
_MediaT: Final = TypeVar("_MediaT")
|
||||
_InputItemT: Final = TypeVar("_InputItemT")
|
||||
_CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent)
|
||||
_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
|
||||
_TEXT_ADAPTER: Final = TypeAdapter(str)
|
||||
|
|
@ -445,6 +453,86 @@ TOOL_CALL_OUTPUT_ITEM_TYPES: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
def _read_positions(
|
||||
input_items: Sequence[_InputItemT], read_input: str | ResponseInputParam
|
||||
) -> Iterator[tuple[int, _InputItemT]]:
|
||||
read_ids: Final = frozenset(map(id, parts_of(read_input)))
|
||||
return ((index, item) for index, item in enumerate(input_items) if id(item) in read_ids)
|
||||
|
||||
|
||||
def _first_invalid_input_type(fields: Iterator[tuple[str, object]]) -> BadRequestError | None:
|
||||
first: Final = next(fields, None)
|
||||
return None if first is None else _invalid_input_type(*first)
|
||||
|
||||
|
||||
def _is_tool_output_item(item: object) -> bool:
|
||||
item_type: Final = message_field(item, "type")
|
||||
return isinstance(item_type, str) and item_type in TOOL_CALL_OUTPUT_ITEM_TYPES
|
||||
|
||||
|
||||
def _converts_before(
|
||||
convert_input: Callable[[str | ResponseInputParam], object], sent_input: str | ResponseInputParam, position: int
|
||||
) -> bool:
|
||||
try:
|
||||
convert_input(sent_input[:position])
|
||||
except TypeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _unhashable_tool_output_call_id_error(
|
||||
input: str | ResponseInputParam,
|
||||
read_input: str | ResponseInputParam,
|
||||
error: TypeError,
|
||||
converts_before: Callable[[int], bool],
|
||||
) -> BadRequestError | None:
|
||||
call_ids: Final = (
|
||||
(f"input[{index}].call_id", message_field(item, "call_id"), item)
|
||||
for index, item in _read_positions(parts_of(input), read_input)
|
||||
if _is_tool_output_item(item)
|
||||
)
|
||||
first: Final = next((candidate for candidate in call_ids if is_unhashable_value_behind(error, candidate[1])), None)
|
||||
if first is None:
|
||||
return None
|
||||
path, call_id, item = first
|
||||
read_position: Final = next(
|
||||
position for position, read_item in enumerate(parts_of(read_input)) if read_item is item
|
||||
)
|
||||
return _invalid_input_type(path, call_id) if converts_before(read_position) else None
|
||||
|
||||
|
||||
def _carries_non_string_field(message: object) -> bool:
|
||||
tool_call_ids: Final = (
|
||||
message_field(tool_call, "id") for tool_call in parts_of(message_field(message, "tool_calls"))
|
||||
)
|
||||
return is_present_non_string(message_field(message, "role")) or any(map(is_present_non_string, tool_call_ids))
|
||||
|
||||
|
||||
def _non_string_item_fields(
|
||||
item: Mapping[str, object], convert_item: Callable[[Mapping[str, object]], Sequence[object]]
|
||||
) -> Iterator[tuple[str, object]]:
|
||||
messages: Final = convert_item(item)
|
||||
tool_call_id_field: Final = "call_id" if item.get("call_id") else "id"
|
||||
roles: Final = (("role", message_field(message, "role")) for message in messages)
|
||||
tool_calls: Final = chain.from_iterable(parts_of(message_field(message, "tool_calls")) for message in messages)
|
||||
tool_call_ids: Final = ((tool_call_id_field, message_field(tool_call, "id")) for tool_call in tool_calls)
|
||||
return ((field, value) for field, value in chain(roles, tool_call_ids) if is_present_non_string(value))
|
||||
|
||||
|
||||
def _indexed_non_string_item_fields(
|
||||
input: ResponseInputParam,
|
||||
read_input: str | ResponseInputParam,
|
||||
convert_item: Callable[[Mapping[str, object]], Sequence[object]],
|
||||
) -> Iterator[tuple[str, object]]:
|
||||
for index, item in _read_positions(input, read_input):
|
||||
for field, value in _non_string_item_fields(item, convert_item):
|
||||
yield f"input[{index}].{field}", value
|
||||
|
||||
|
||||
def _invalid_input_type(field_path: str, value: object) -> BadRequestError:
|
||||
return BadRequestError(message=invalid_type_message(field_path, value), model="", llm_provider="")
|
||||
|
||||
|
||||
def _input_item_role(input_item: Mapping[str, object]) -> str:
|
||||
return cast(str, input_item.get("role") or "user") # cast-ok: client-supplied role forwarded verbatim, unvalidated
|
||||
|
||||
|
|
@ -1101,6 +1189,9 @@ class LiteLLMCompletionResponsesConfig:
|
|||
stream: bool | None = None,
|
||||
extra_headers: Mapping[str, object] | None = None,
|
||||
drop_params: object = None,
|
||||
*,
|
||||
caller_input: str | ResponseInputParam | None = None,
|
||||
names_non_string_fields: bool = False,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
"""
|
||||
|
|
@ -1149,6 +1240,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
native_compaction: Final = _replays_compaction_block(
|
||||
request_input, context_management
|
||||
) and provider_takes_compaction_blocks(model, custom_llm_provider)
|
||||
numbered_input: Final = request_input if caller_input is None else caller_input
|
||||
bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
|
|
@ -1157,7 +1249,28 @@ class LiteLLMCompletionResponsesConfig:
|
|||
tool_call_provider_fields=not embeds_thought_signature_in_tool_call_id(request_kwargs),
|
||||
native_compaction=native_compaction,
|
||||
only_since_compaction=not native_compaction,
|
||||
caller_input=numbered_input,
|
||||
)
|
||||
malformed_field_error: Final = (
|
||||
_first_invalid_input_type(
|
||||
_indexed_non_string_item_fields(
|
||||
numbered_input,
|
||||
_input_to_read(request_input, not native_compaction, native_compaction),
|
||||
partial(
|
||||
LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message,
|
||||
replay_reasoning=True,
|
||||
native_compaction=native_compaction,
|
||||
),
|
||||
)
|
||||
)
|
||||
if names_non_string_fields
|
||||
and isinstance(numbered_input, list)
|
||||
and any(map(_carries_non_string_field, bridged_messages))
|
||||
and provider_builds_anthropic_request(model, custom_llm_provider)
|
||||
else None
|
||||
)
|
||||
if malformed_field_error is not None:
|
||||
raise malformed_field_error
|
||||
|
||||
litellm_completion_request: dict = {
|
||||
"messages": (
|
||||
|
|
@ -1221,6 +1334,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
tool_call_provider_fields: bool = False,
|
||||
native_compaction: bool = False,
|
||||
only_since_compaction: bool = False,
|
||||
caller_input: str | ResponseInputParam | None = None,
|
||||
) -> list[
|
||||
AllMessageValues
|
||||
| GenericChatCompletionMessage
|
||||
|
|
@ -1274,13 +1388,27 @@ class LiteLLMCompletionResponsesConfig:
|
|||
|
||||
dumped_input: Final = with_compaction_items_as_params(input)
|
||||
read_input: Final = _input_to_read(dumped_input, only_since_compaction, native_compaction)
|
||||
messages.extend(
|
||||
LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message(
|
||||
input=_without_item_keys(read_input, _dropped_item_keys(item_cache_control, tool_call_provider_fields)),
|
||||
replay_reasoning=replay_reasoning,
|
||||
native_compaction=native_compaction,
|
||||
)
|
||||
sent_input: Final = _without_item_keys(
|
||||
read_input, _dropped_item_keys(item_cache_control, tool_call_provider_fields)
|
||||
)
|
||||
convert_input: Final = partial(
|
||||
LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message,
|
||||
replay_reasoning=replay_reasoning,
|
||||
native_compaction=native_compaction,
|
||||
)
|
||||
try:
|
||||
converted: Final = convert_input(sent_input)
|
||||
except TypeError as error:
|
||||
call_id_error: Final = _unhashable_tool_output_call_id_error(
|
||||
dumped_input if caller_input is None else caller_input,
|
||||
read_input,
|
||||
error,
|
||||
partial(_converts_before, convert_input, sent_input),
|
||||
)
|
||||
if call_id_error is None:
|
||||
raise
|
||||
raise call_id_error from error
|
||||
messages.extend(converted)
|
||||
|
||||
return messages
|
||||
|
||||
|
|
|
|||
|
|
@ -233,6 +233,11 @@ def spawn_logging_task(
|
|||
|
||||
|
||||
async def drain_logging_tasks(timeout: float = 5.0) -> None:
|
||||
# A task that finished as its loop closed never ran its discard callback, and a
|
||||
# task from another loop cannot be awaited here: forget both, wait on the rest.
|
||||
running_loop: Final = asyncio.get_running_loop()
|
||||
stale: Final = {task for task in _PENDING_LOGGING_TASKS if task.done() or task.get_loop() is not running_loop}
|
||||
_PENDING_LOGGING_TASKS.difference_update(stale)
|
||||
if not _PENDING_LOGGING_TASKS:
|
||||
return
|
||||
_, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout)
|
||||
|
|
|
|||
17811
litellm/router.py
17811
litellm/router.py
File diff suppressed because it is too large
Load diff
|
|
@ -42,6 +42,8 @@ from dataclasses import dataclass, field
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, overload
|
||||
|
||||
from litellm.router_config import detached
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any # noqa: TID251 # only to spell the cores' untyped (Unknown) containers and one Any type
|
||||
|
||||
|
|
@ -176,9 +178,9 @@ def deployment_usable_by_team(deployment: "Mapping[str, object] | Deployment", t
|
|||
Whether ``team_id`` may use ``deployment``: a team-scoped deployment (``model_info.team_id``
|
||||
set) only by that team, any other deployment by everyone. Needs no router instance.
|
||||
"""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return Router._core_deployment_usable_by_team(deployment, team_id)
|
||||
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
|
||||
|
||||
|
||||
def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str:
|
||||
|
|
@ -200,9 +202,9 @@ def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str
|
|||
|
||||
def router_init_args() -> tuple[str, ...]:
|
||||
"""The arguments ``Router.__init__`` accepts (was ``Router.get_valid_args``). Needs no router instance."""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return tuple(Router.get_valid_args())
|
||||
return tuple(PythonCore.get_valid_args())
|
||||
|
||||
|
||||
def mark_advisor_orchestration_failure(exception: BaseException) -> None:
|
||||
|
|
@ -266,9 +268,9 @@ def deployment_metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, ob
|
|||
|
||||
|
||||
def _json_default_stable_id(obj: object) -> str:
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return Router._json_default_stable_id(obj)
|
||||
return PythonCore._json_default_stable_id(obj)
|
||||
|
||||
|
||||
class RouterRegistry:
|
||||
|
|
@ -529,10 +531,10 @@ class RouterRegistry:
|
|||
``should_include_deployment``), or, without ``model_name``, whether the team may use it at
|
||||
all (was ``deployment_usable_by_team``).
|
||||
"""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
if model_name is None:
|
||||
return Router._core_deployment_usable_by_team(deployment, team_id)
|
||||
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
|
||||
return self._router._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id)
|
||||
|
||||
@overload
|
||||
|
|
@ -1053,6 +1055,12 @@ SETTING_FIELDS: Final[tuple[str, ...]] = (
|
|||
"cache_responses",
|
||||
)
|
||||
|
||||
# Runtime settings kept as plain router attributes rather than ``settings()`` fields. A write
|
||||
# to one (``router.num_retries = 2``) is a changeset too: ``ChangeSet(assign={...})``.
|
||||
PLAIN_SETTING_ATTRIBUTES: Final[frozenset[str]] = frozenset({"num_retries", "max_fallbacks", "retry_after", "timeout"})
|
||||
# Former attributes whose value lives under another name, assignable through a changeset.
|
||||
ASSIGNABLE_ALIASES: Final[Mapping[str, str]] = MappingProxyType({"model_group_alias": "_core_model_group_alias"})
|
||||
|
||||
_MISSING: Final = object()
|
||||
|
||||
|
||||
|
|
@ -1168,16 +1176,16 @@ class RouterHealth:
|
|||
@staticmethod
|
||||
def is_blocked(deployment: "Deployment") -> bool:
|
||||
"""Whether an admin paused this deployment (was ``is_deployment_blocked``)."""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return Router._core_is_deployment_blocked(deployment)
|
||||
return PythonCore._core_is_deployment_blocked(deployment)
|
||||
|
||||
@staticmethod
|
||||
def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool:
|
||||
"""Whether every deployment in the list is paused (was ``are_all_deployments_blocked``)."""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return Router._core_are_all_deployments_blocked(deployments)
|
||||
return PythonCore._core_are_all_deployments_blocked(deployments)
|
||||
|
||||
def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set
|
||||
"""Groups whose every deployment is paused (was ``get_fully_blocked_model_names``)."""
|
||||
|
|
@ -1296,13 +1304,17 @@ class RoutingObservers:
|
|||
def apply_changes(router: "Router", changes: ChangeSet) -> ApplyResult:
|
||||
if changes.refresh_model_info:
|
||||
raise TypeError("ChangeSet(refresh_model_info=True) needs aapply()")
|
||||
unknown: Final = sorted(set(changes.assign) - set(SETTING_FIELDS))
|
||||
unknown: Final = sorted(set(changes.assign) - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES})
|
||||
if unknown:
|
||||
raise ValueError(f"not assignable router settings: {unknown}")
|
||||
for name, value in changes.assign.items():
|
||||
setattr(router, "_core_" + name, value)
|
||||
setattr(
|
||||
router,
|
||||
name if name in PLAIN_SETTING_ATTRIBUTES else ASSIGNABLE_ALIASES.get(name, "_core_" + name),
|
||||
detached(value),
|
||||
)
|
||||
if changes.settings:
|
||||
router._core_update_settings(**changes.settings)
|
||||
router._core_update_settings(**{name: detached(value) for name, value in changes.settings.items()})
|
||||
upserted: Final = [router._core_upsert_deployment(deployment=d) for d in changes.upsert]
|
||||
deleted: Final = [router._core_delete_deployment(id=i) for i in changes.delete]
|
||||
if changes.redis_cache is not None:
|
||||
|
|
@ -1388,7 +1400,7 @@ class RoutingContext:
|
|||
return self._router.observe()
|
||||
|
||||
def belongs_to(self, router: "Router") -> bool:
|
||||
return self._router is router
|
||||
return self._router is router or router.__dict__.get("_core") is self._router
|
||||
|
||||
@property
|
||||
def explicit_allowed_fails(self) -> int | None:
|
||||
|
|
|
|||
157
litellm/router_config.py
Normal file
157
litellm/router_config.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""
|
||||
The router's configuration, built outside the router.
|
||||
|
||||
``RouterConfig`` is static: the constructor arguments, frozen for the router's life.
|
||||
``RouterDynamicConfig`` is a snapshot of the process settings the router reads while it serves
|
||||
(``litellm.num_retries``, ``litellm.model_cost``, the fallback lists, ...). A call reads one snapshot
|
||||
at its start. This module is the only router code that reads ``litellm.*`` settings.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RouterConfig:
|
||||
"""The arguments the router was built with, exactly as passed (unset ones are absent)."""
|
||||
|
||||
arguments: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
def get(self, name: str, default: object = None) -> object:
|
||||
return self.arguments.get(name, default)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RouterDynamicConfig:
|
||||
"""Process settings the router reads per call. Values are the live objects, not copies."""
|
||||
|
||||
version: int
|
||||
num_retries: int | None
|
||||
max_fallbacks: int | None
|
||||
router_max_fallbacks: int
|
||||
allowed_fails: int
|
||||
request_timeout: float
|
||||
fallbacks: Sequence[object] | None
|
||||
default_fallbacks: Sequence[str] | None
|
||||
context_window_fallbacks: Sequence[object] | None
|
||||
content_policy_fallbacks: Sequence[object] | None
|
||||
drop_params: bool | None
|
||||
provider_list: Sequence[str]
|
||||
expose_router_debug_in_errors: bool
|
||||
redact_user_api_key_info: bool | None
|
||||
model_alias_map: Mapping[str, str]
|
||||
model_cost: Mapping[str, object]
|
||||
known_custom_logger_compatible_callbacks: Sequence[str]
|
||||
exception_types: Sequence[type[BaseException]]
|
||||
callbacks: Sequence[object]
|
||||
success_callback: Sequence[object]
|
||||
failure_callback: Sequence[object]
|
||||
input_callback: Sequence[object]
|
||||
callback_manager: "LoggingCallbackManager"
|
||||
cache: "Cache | None"
|
||||
|
||||
|
||||
_versions: Final = itertools.count(1)
|
||||
|
||||
|
||||
def detached(value: object) -> object:
|
||||
"""``value`` with its plain containers (dict, list, tuple, set) copied, all the way down, so the
|
||||
caller's own objects can't change the router afterwards. Other objects (clients, caches,
|
||||
callbacks, pydantic models) are kept as they are."""
|
||||
if isinstance(value, dict):
|
||||
return {key: detached(item) for key, item in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [detached(item) for item in value]
|
||||
if isinstance(value, tuple) and type(value) is tuple:
|
||||
return tuple(detached(item) for item in value)
|
||||
if isinstance(value, set):
|
||||
return {detached(item) for item in value}
|
||||
return value
|
||||
|
||||
|
||||
def build_router_config(arguments: Mapping[str, object]) -> RouterConfig:
|
||||
"""The constructor arguments, detached from the caller's objects: settings are static."""
|
||||
return RouterConfig(arguments=MappingProxyType({name: detached(value) for name, value in arguments.items()}))
|
||||
|
||||
|
||||
def current_dynamic_config() -> RouterDynamicConfig:
|
||||
"""A snapshot of the process settings as they are right now."""
|
||||
return RouterDynamicConfig(
|
||||
version=next(_versions),
|
||||
num_retries=litellm.num_retries,
|
||||
max_fallbacks=litellm.max_fallbacks,
|
||||
router_max_fallbacks=litellm.ROUTER_MAX_FALLBACKS,
|
||||
allowed_fails=litellm.allowed_fails,
|
||||
request_timeout=litellm.request_timeout,
|
||||
fallbacks=litellm.fallbacks,
|
||||
default_fallbacks=litellm.default_fallbacks,
|
||||
context_window_fallbacks=litellm.context_window_fallbacks,
|
||||
content_policy_fallbacks=litellm.content_policy_fallbacks,
|
||||
drop_params=litellm.drop_params,
|
||||
provider_list=litellm.provider_list,
|
||||
expose_router_debug_in_errors=litellm.expose_router_debug_in_errors,
|
||||
redact_user_api_key_info=litellm.redact_user_api_key_info,
|
||||
model_alias_map=litellm.model_alias_map,
|
||||
model_cost=litellm.model_cost,
|
||||
known_custom_logger_compatible_callbacks=litellm._known_custom_logger_compatible_callbacks,
|
||||
exception_types=litellm.LITELLM_EXCEPTION_TYPES,
|
||||
callbacks=litellm.callbacks,
|
||||
success_callback=litellm.success_callback,
|
||||
failure_callback=litellm.failure_callback,
|
||||
input_callback=litellm.input_callback,
|
||||
callback_manager=litellm.logging_callback_manager,
|
||||
cache=litellm.cache,
|
||||
)
|
||||
|
||||
|
||||
_LIVE_SOURCES: Final = MappingProxyType(
|
||||
{
|
||||
"num_retries": "num_retries",
|
||||
"max_fallbacks": "max_fallbacks",
|
||||
"router_max_fallbacks": "ROUTER_MAX_FALLBACKS",
|
||||
"allowed_fails": "allowed_fails",
|
||||
"request_timeout": "request_timeout",
|
||||
"fallbacks": "fallbacks",
|
||||
"default_fallbacks": "default_fallbacks",
|
||||
"context_window_fallbacks": "context_window_fallbacks",
|
||||
"content_policy_fallbacks": "content_policy_fallbacks",
|
||||
"drop_params": "drop_params",
|
||||
"provider_list": "provider_list",
|
||||
"expose_router_debug_in_errors": "expose_router_debug_in_errors",
|
||||
"redact_user_api_key_info": "redact_user_api_key_info",
|
||||
"model_alias_map": "model_alias_map",
|
||||
"model_cost": "model_cost",
|
||||
"known_custom_logger_compatible_callbacks": "_known_custom_logger_compatible_callbacks",
|
||||
"exception_types": "LITELLM_EXCEPTION_TYPES",
|
||||
"callbacks": "callbacks",
|
||||
"success_callback": "success_callback",
|
||||
"failure_callback": "failure_callback",
|
||||
"input_callback": "input_callback",
|
||||
"callback_manager": "logging_callback_manager",
|
||||
"cache": "cache",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class LiveDynamicConfig:
|
||||
"""``RouterDynamicConfig``'s fields read from the process settings at access time, with no snapshot."""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __getattr__(self, name: str) -> object:
|
||||
source: Final = _LIVE_SOURCES.get(name)
|
||||
if source is None:
|
||||
raise AttributeError(name)
|
||||
return getattr(litellm, source)
|
||||
|
||||
|
||||
live_dynamic_config: Final = LiveDynamicConfig()
|
||||
56
litellm/router_core.py
Normal file
56
litellm/router_core.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
"""
|
||||
``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one.
|
||||
|
||||
``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core
|
||||
(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when
|
||||
the router is built, from its ``RouterConfig``.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Protocol, runtime_checkable
|
||||
|
||||
from litellm.router_config import RouterConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_api import (
|
||||
ApplyResult,
|
||||
ChangeSet,
|
||||
RouterHealth,
|
||||
RouterRegistry,
|
||||
RouterSettings,
|
||||
RoutingObservers,
|
||||
)
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.types.router import CustomRoutingStrategyBase
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RouterCore(Protocol):
|
||||
async def call(self, op: str, request: Mapping[str, object]) -> object: ...
|
||||
|
||||
def registry(self) -> "RouterRegistry": ...
|
||||
|
||||
def settings(self) -> "RouterSettings": ...
|
||||
|
||||
def health(self) -> "RouterHealth": ...
|
||||
|
||||
async def select(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict
|
||||
) -> object: ...
|
||||
|
||||
def apply(self, changes: "ChangeSet") -> "ApplyResult": ...
|
||||
|
||||
def observe(self) -> "RoutingObservers": ...
|
||||
|
||||
def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: ...
|
||||
|
||||
def retire(self) -> None: ...
|
||||
|
||||
|
||||
def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore:
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
|
||||
|
|
@ -31,7 +31,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload
|
||||
|
||||
from litellm.router_api import SETTING_FIELDS, ChangeSet
|
||||
from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged
|
||||
|
|
@ -65,7 +65,7 @@ _INTERNAL_PREFIX: Final = "litellm.router"
|
|||
|
||||
|
||||
# Router's own attribute hooks: a write or delete through an old name passes through them first.
|
||||
_ATTRIBUTE_HOOKS: Final = frozenset({"__setattr__", "__delattr__"})
|
||||
_ATTRIBUTE_HOOKS: Final = frozenset({"__getattribute__", "__setattr__", "__delattr__"})
|
||||
|
||||
|
||||
def _note(name: str, replacement: str) -> None:
|
||||
|
|
@ -88,9 +88,10 @@ def _note(name: str, replacement: str) -> None:
|
|||
def _as_router(obj: "DeprecatedRouterMembers") -> "Router":
|
||||
"""The Router this mixin is part of (it is mixed into ``litellm.Router`` only)."""
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
if isinstance(obj, Router):
|
||||
return obj
|
||||
if isinstance(obj, (Router, PythonCore)):
|
||||
return obj # pyright: ignore[reportReturnType] # a PythonCore carries every member the mixin reads
|
||||
raise TypeError(f"{type(obj).__name__} is not a litellm.Router; DeprecatedRouterMembers is mixed into Router only")
|
||||
|
||||
|
||||
|
|
@ -134,7 +135,7 @@ class _DeprecatedAttribute(Generic[_T]):
|
|||
|
||||
def __set__(self, obj: "DeprecatedRouterMembers", value: _T) -> None:
|
||||
_note(self.name, self.replacement)
|
||||
if self.name in SETTING_FIELDS:
|
||||
if self.name in SETTING_FIELDS or self.name in ASSIGNABLE_ALIASES:
|
||||
_as_router(obj).apply(ChangeSet(assign={self.name: value}))
|
||||
else:
|
||||
setattr(obj, "_core_" + self.name, value)
|
||||
|
|
@ -169,6 +170,13 @@ _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] =
|
|||
)
|
||||
|
||||
|
||||
def _override_key(router: object, name: str) -> tuple[int, str]:
|
||||
"""The running-override key for ``router``: a Router and its core share one, so an old member
|
||||
entered on the Router is seen by the override routing that runs on the core."""
|
||||
core: Final = getattr(router, "__dict__", {}).get("_core")
|
||||
return (id(router if core is None else core), name)
|
||||
|
||||
|
||||
def _enter_override(name: str, override: Callable[..., object]) -> Callable[..., object]:
|
||||
"""
|
||||
A direct call of a subclass override (``router.<name>()``) marks it running, so its
|
||||
|
|
@ -184,7 +192,7 @@ def _enter_override(name: str, override: Callable[..., object]) -> Callable[...,
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {(id(router), name)})
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
|
||||
try:
|
||||
return await async_override(router, *args, **kwargs)
|
||||
finally:
|
||||
|
|
@ -198,7 +206,7 @@ def _enter_override(name: str, override: Callable[..., object]) -> Callable[...,
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {(id(router), name)})
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
|
||||
try:
|
||||
return override(router, *args, **kwargs)
|
||||
finally:
|
||||
|
|
@ -236,7 +244,7 @@ def _route_override(cls: type, name: str, override: Callable[..., object]) -> Ca
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
key: Final = (id(router), name)
|
||||
key: Final = _override_key(router, name)
|
||||
running: Final = _OVERRIDES_RUNNING.get()
|
||||
if key in running:
|
||||
parent_async: Final[Callable[..., Awaitable[object]]] = getattr(super(cls, router), core)
|
||||
|
|
@ -255,7 +263,7 @@ def _route_override(cls: type, name: str, override: Callable[..., object]) -> Ca
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
key: Final = (id(router), name)
|
||||
key: Final = _override_key(router, name)
|
||||
running: Final = _OVERRIDES_RUNNING.get()
|
||||
if key in running:
|
||||
return parent(router)(*args, **kwargs)
|
||||
|
|
@ -502,9 +510,9 @@ class DeprecatedRouterMembers:
|
|||
@staticmethod
|
||||
def deployment_usable_by_team(model: "Mapping[str, object] | Deployment", team_id: str | None) -> bool:
|
||||
_note("deployment_usable_by_team", "registry().includes(deployment, team_id=...)")
|
||||
from litellm.router import Router
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return Router._core_deployment_usable_by_team(model, team_id)
|
||||
return PythonCore._core_deployment_usable_by_team(model, team_id)
|
||||
|
||||
_deployment_usable_by_team = deployment_usable_by_team
|
||||
|
||||
|
|
|
|||
17669
litellm/router_python_core.py
Normal file
17669
litellm/router_python_core.py
Normal file
File diff suppressed because it is too large
Load diff
135
litellm/router_python_host.py
Normal file
135
litellm/router_python_host.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
"""
|
||||
``PythonHost``: what the router core reaches outside itself, served from the litellm package.
|
||||
|
||||
The core never touches ``litellm.*`` directly. Inference calls (``invoke``), backoff sleeps,
|
||||
callback registration and the process-wide cache all go through the host, and the process settings
|
||||
come from ``dynamic()`` (``litellm/router_config.py``). A Rust core gets the same operations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Literal,
|
||||
TypeAlias,
|
||||
cast, # noqa: TID251 # untyped litellm globals: the live config view and the callback lists
|
||||
)
|
||||
|
||||
from typing_extensions import overload
|
||||
|
||||
import litellm
|
||||
from litellm.router_config import RouterDynamicConfig, live_dynamic_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIFileObject
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
AdapterCompletionStreamWrapper,
|
||||
EmbeddingResponse,
|
||||
ImageResponse,
|
||||
LiteLLMBatch,
|
||||
ModelInfo,
|
||||
ModelResponse,
|
||||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
|
||||
CallbackList: TypeAlias = Literal[
|
||||
"success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback", "input_callback"
|
||||
]
|
||||
|
||||
Endpoint: TypeAlias = Callable[..., object]
|
||||
|
||||
|
||||
class PythonHost:
|
||||
__slots__ = ()
|
||||
|
||||
# The endpoints whose results the router core reads, typed as litellm declares them; any other
|
||||
# op gets the untyped Endpoint.
|
||||
@overload
|
||||
def invoke(self, op: Literal["acompletion"]) -> Callable[..., Awaitable["ModelResponse | CustomStreamWrapper"]]: ...
|
||||
@overload
|
||||
def invoke(
|
||||
self, op: Literal["atext_completion"]
|
||||
) -> Callable[..., Awaitable["TextCompletionResponse | TextCompletionStreamWrapper"]]: ...
|
||||
@overload
|
||||
def invoke(
|
||||
self, op: Literal["aadapter_completion"]
|
||||
) -> Callable[..., Awaitable["BaseModel | AdapterCompletionStreamWrapper | None"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["aembedding"]) -> Callable[..., Awaitable["EmbeddingResponse"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["aimage_generation"]) -> Callable[..., Awaitable["ImageResponse"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["atranscription"]) -> Callable[..., Awaitable["TranscriptionResponse"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["aspeech"]) -> Callable[..., Awaitable["HttpxBinaryResponseContent"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["arerank"]) -> Callable[..., Awaitable["RerankResponse"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["acreate_file"]) -> Callable[..., Awaitable["OpenAIFileObject"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["acreate_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["aretrieve_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["acancel_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["alist_batches"]) -> Callable[..., Awaitable[object]]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["get_llm_provider"]) -> Callable[..., "tuple[str, str, str | None, str | None]"]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["get_model_info"]) -> Callable[..., "ModelInfo"]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["get_supported_openai_params"]) -> Callable[..., "Sequence[str] | None"]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["token_counter"]) -> Callable[..., int]: ...
|
||||
@overload
|
||||
def invoke(self, op: Literal["calculate_retry_after"]) -> Callable[..., float]: ...
|
||||
@overload
|
||||
def invoke(self, op: str) -> Endpoint: ...
|
||||
def invoke(self, op: str) -> Endpoint:
|
||||
"""The litellm function named ``op`` (``acompletion``, ``get_model_info``, ...), looked up at call time."""
|
||||
endpoint: Final[object] = getattr(litellm, op, None)
|
||||
if endpoint is None:
|
||||
raise AttributeError(f"litellm has no endpoint {op!r}")
|
||||
if not callable(endpoint):
|
||||
raise TypeError(f"litellm.{op} is not an endpoint")
|
||||
return endpoint
|
||||
|
||||
def sleep(self, seconds: float) -> Awaitable[None]:
|
||||
return asyncio.sleep(seconds)
|
||||
|
||||
def dynamic(self) -> RouterDynamicConfig:
|
||||
"""The process settings, read live on each access (``current_dynamic_config()`` takes a snapshot)."""
|
||||
return cast(RouterDynamicConfig, live_dynamic_config) # cast-ok: the live view reads the same fields
|
||||
|
||||
@property
|
||||
def callbacks(self) -> "LoggingCallbackManager":
|
||||
return litellm.logging_callback_manager
|
||||
|
||||
def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself
|
||||
value: Final = getattr(litellm, name)
|
||||
return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global
|
||||
|
||||
def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given
|
||||
setattr(litellm, name, value)
|
||||
|
||||
def global_cache(self) -> "Cache | None":
|
||||
return litellm.cache
|
||||
|
||||
def set_global_cache(self, cache: "Cache | None") -> None:
|
||||
litellm.cache = cache
|
||||
|
||||
def suppress_debug_info(self) -> None:
|
||||
litellm.suppress_debug_info = True
|
||||
|
||||
|
||||
python_host: Final = PythonHost()
|
||||
|
|
@ -14,13 +14,13 @@ else:
|
|||
RoutingContext = Any
|
||||
|
||||
|
||||
def litellm_to_list(embeds: litellm.EmbeddingResponse) -> list[list[float]]:
|
||||
def litellm_to_list(embeds: object) -> list[list[float]]:
|
||||
"""Convert a LiteLLM embedding response to a list of embeddings.
|
||||
|
||||
:param embeds: The LiteLLM embedding response.
|
||||
:return: A list of embeddings.
|
||||
"""
|
||||
if not embeds or not embeds.data:
|
||||
if not embeds or not isinstance(embeds, litellm.EmbeddingResponse) or not embeds.data:
|
||||
raise ValueError("No embeddings found in LiteLLM embedding response.")
|
||||
return [x["embedding"] for x in embeds.data]
|
||||
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ class JevBudgetedRequest:
|
|||
telemetry: StandardLoggingRoutingDecision
|
||||
|
||||
|
||||
@renamed_router_param("router", "context")
|
||||
def jev_input_targets(context: "RoutingContext", config: "OpenSourceClassifierConfig") -> tuple[JevInputTarget, ...]:
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
|
|
|
|||
|
|
@ -392,9 +392,11 @@ def _resolve_constraint_only_pool(
|
|||
|
||||
def _all_deployments_or_fallback(
|
||||
ctx: "RoutingContext",
|
||||
model: str,
|
||||
model: str | None,
|
||||
fallback: _DeploymentPool,
|
||||
) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]:
|
||||
if model is None: # a group of no name has no deployments, it is not a failed lookup
|
||||
return ()
|
||||
try:
|
||||
return ctx.registry().deployments(model, source="all")
|
||||
except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ else:
|
|||
|
||||
@renamed_router_param("litellm_router_instance")
|
||||
async def send_llm_exception_alert(
|
||||
ctx: "RoutingContext",
|
||||
ctx: "RoutingContext | None",
|
||||
request_kwargs: dict,
|
||||
error_traceback_str: str,
|
||||
original_exception,
|
||||
|
|
@ -39,6 +39,9 @@ async def send_llm_exception_alert(
|
|||
Returns:
|
||||
None
|
||||
"""
|
||||
if ctx is None:
|
||||
return
|
||||
|
||||
slack_alerting_logger: Final = ctx.observers().slack_alerting_logger
|
||||
if slack_alerting_logger is None:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, TypeAlias, cast
|
||||
|
||||
from pydantic import ConfigDict, Field, field_validator, model_validator
|
||||
from pydantic import ConfigDict, Field, JsonValue, field_validator, model_validator
|
||||
from pydantic.fields import FieldInfo
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
|
|
@ -604,6 +604,11 @@ class BedrockChecksConfigModel(LiteLLMBaseModel):
|
|||
return self
|
||||
|
||||
|
||||
class BedrockGuardrailOutputScope(Enum):
|
||||
INTERVENTIONS = "INTERVENTIONS"
|
||||
FULL = "FULL"
|
||||
|
||||
|
||||
class BedrockGuardrailConfigModel(LiteLLMBaseModel):
|
||||
"""Configuration parameters for the AWS Bedrock guardrail"""
|
||||
|
||||
|
|
@ -612,6 +617,11 @@ class BedrockGuardrailConfigModel(LiteLLMBaseModel):
|
|||
default=None,
|
||||
description="The version of your Bedrock guardrail (e.g., DRAFT or version number)",
|
||||
)
|
||||
outputScope: BedrockGuardrailOutputScope | None = Field(
|
||||
default=None,
|
||||
description="ApplyGuardrail response scope. FULL includes detected and non-detected findings. "
|
||||
"Omit to retain Bedrock's default of returning detected findings only.",
|
||||
)
|
||||
disable_exception_on_block: bool | None = Field(
|
||||
default=False,
|
||||
description="If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.",
|
||||
|
|
@ -1540,6 +1550,14 @@ class PresidioPerRequestConfig(LiteLLMBaseModel):
|
|||
entities: list[PiiEntityType] | None = None
|
||||
|
||||
|
||||
NativeGuardrailResponse: TypeAlias = Mapping[str, JsonValue] | Sequence[JsonValue]
|
||||
|
||||
|
||||
class GuardrailResponseFormat(Enum):
|
||||
STANDARD = "standard"
|
||||
NATIVE = "native"
|
||||
|
||||
|
||||
class ApplyGuardrailRequest(LiteLLMBaseModel):
|
||||
guardrail_name: str
|
||||
text: str
|
||||
|
|
@ -1548,6 +1566,7 @@ class ApplyGuardrailRequest(LiteLLMBaseModel):
|
|||
input_type: str = "request"
|
||||
messages: list[dict[str, object]] | None = None
|
||||
metadata: dict[str, object] | None = None
|
||||
response_format: GuardrailResponseFormat = GuardrailResponseFormat.STANDARD
|
||||
|
||||
|
||||
class ApplyGuardrailResponse(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import Literal
|
||||
|
||||
from pydantic import AliasPath, BaseModel, ConfigDict, Field
|
||||
from typing_extensions import Required, TypedDict
|
||||
|
||||
|
||||
|
|
@ -98,3 +99,15 @@ class CohereV2ChatResponse(TypedDict):
|
|||
message: CohereV2ChatResponseMessage
|
||||
usage: CohereV2ChatResponseUsage
|
||||
logprobs: CohereV2ChatResponseLogProbs
|
||||
|
||||
|
||||
class CohereV2StreamedToolCall(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str | None = Field(default=None, validation_alias=AliasPath("delta", "message", "tool_calls", "id"))
|
||||
name: str | None = Field(
|
||||
default=None, validation_alias=AliasPath("delta", "message", "tool_calls", "function", "name")
|
||||
)
|
||||
arguments: str = Field(
|
||||
default="", validation_alias=AliasPath("delta", "message", "tool_calls", "function", "arguments")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -28,6 +28,10 @@ class OvalixGuardrailConfigModel(GuardrailConfigModel):
|
|||
default=None,
|
||||
description="Post-checkpoint ID for the Ovalix Tracker service.",
|
||||
)
|
||||
file_checkpoint_id: str | None = Field(
|
||||
default=None,
|
||||
description="File-checkpoint ID for the Ovalix Tracker service (falls back to pre/post).",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
|
|
|
|||
|
|
@ -35,6 +35,9 @@ class PatchModelRequest(updateDeployment):
|
|||
|
||||
|
||||
class ModelGroupInfoProxy(ModelGroupInfo):
|
||||
supports_background_chat: bool = Field(
|
||||
default=False, description="Whether cache-backed background Responses runs are enabled for this model group"
|
||||
)
|
||||
is_public_model_group: bool = Field(default=False)
|
||||
health_status: str | None = Field(default=None)
|
||||
health_response_time: float | None = Field(default=None)
|
||||
|
|
|
|||
|
|
@ -1151,6 +1151,8 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
|
|||
# A2A (Agent-to-Agent)
|
||||
"/a2a/{agent_id}": [CallTypes.asend_message, CallTypes.send_message],
|
||||
"/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message],
|
||||
"/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
|
||||
"/v1/a2a/{agent_id}/message/stream": [CallTypes.asend_message, CallTypes.send_message],
|
||||
# Passthrough endpoints
|
||||
"/llm_passthrough": [
|
||||
CallTypes.allm_passthrough_route,
|
||||
|
|
|
|||
|
|
@ -15555,18 +15555,20 @@
|
|||
"input_cost_per_token": 2.65e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 8000,
|
||||
"max_output_tokens": 2048,
|
||||
"max_output_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.5e-06,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-70b-instruct.md",
|
||||
"supports_pdf_input": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 8000,
|
||||
"max_output_tokens": 2048,
|
||||
"max_output_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-8b-instruct.md",
|
||||
"supports_pdf_input": true
|
||||
},
|
||||
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": {
|
||||
|
|
@ -48970,6 +48972,7 @@
|
|||
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
|
||||
"deprecation_date": "2026-11-04",
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"deprecation_date": "2026-11-04",
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1010000,
|
||||
|
|
@ -53970,23 +53973,23 @@
|
|||
"supports_vision": true
|
||||
},
|
||||
"wandb/meta-llama/Llama-3.1-8B-Instruct": {
|
||||
"max_input_tokens": 131000,
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"input_cost_per_token": 2.2e-07,
|
||||
"output_cost_per_token": 2.2e-07,
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
|
||||
},
|
||||
"wandb/deepseek-ai/DeepSeek-V3.1": {
|
||||
"supports_reasoning": true,
|
||||
"max_input_tokens": 161000,
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 128000,
|
||||
"input_cost_per_token": 5.5e-07,
|
||||
"output_cost_per_token": 1.65e-06,
|
||||
"litellm_provider": "wandb",
|
||||
"mode": "chat",
|
||||
"source": "https://wandb.ai/site/pricing/tokens/"
|
||||
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference"
|
||||
},
|
||||
"wandb/meta-llama/Llama-3.3-70B-Instruct": {
|
||||
"max_input_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -70,7 +70,9 @@ def enqueue(body: object, status: int = 200) -> None:
|
|||
|
||||
def enqueue_stream() -> None:
|
||||
chunk: Final = {**CHAT, "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": "pong"}}]}
|
||||
RESPONSES.put((200, f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode()))
|
||||
terminal: Final = {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
|
||||
stream: Final = f"data: {json.dumps(chunk)}\n\ndata: {json.dumps(terminal)}\n\ndata: [DONE]\n\n"
|
||||
RESPONSES.put((200, stream.encode()))
|
||||
|
||||
|
||||
def bedrock_event(event: str, payload: object) -> bytes:
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import os
|
|||
|
||||
IGNORE_FUNCTIONS = [
|
||||
"_json_cost", # bounded at depth 32 and consumed under byte/node limits.
|
||||
"detached", # router_config: copies only the plain containers of the router's own constructor arguments, once per build or changeset
|
||||
"_format_type",
|
||||
"remove_additional_properties",
|
||||
"remove_strict_from_schema",
|
||||
|
|
|
|||
54
tests/code_coverage_tests/router_litellm_globals_check.py
Normal file
54
tests/code_coverage_tests/router_litellm_globals_check.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
"""
|
||||
The router and its cores never read or write ``litellm.*``.
|
||||
|
||||
Process settings reach the router through ``RouterDynamicConfig`` (litellm/router_config.py) and
|
||||
everything else outside it (inference calls, callbacks, the global cache) through ``PythonHost``
|
||||
(litellm/router_python_host.py). Those two modules are the only router code allowed to touch the
|
||||
``litellm`` package namespace. This check fails on ``import litellm`` (or ``import litellm.<x>``,
|
||||
which binds the same name) and on any ``litellm.<attr>`` access in the files below.
|
||||
|
||||
Run: python tests/code_coverage_tests/router_litellm_globals_check.py
|
||||
"""
|
||||
|
||||
import ast
|
||||
import os
|
||||
import sys
|
||||
|
||||
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
CHECKED = (
|
||||
"litellm/router.py",
|
||||
"litellm/router_api.py",
|
||||
"litellm/router_core.py",
|
||||
"litellm/router_deprecated.py",
|
||||
"litellm/router_python_core.py",
|
||||
)
|
||||
|
||||
|
||||
def violations(rel: str) -> list[str]:
|
||||
tree = ast.parse(open(os.path.join(ROOT, rel)).read(), rel)
|
||||
found = []
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
found.extend(
|
||||
f"{rel}:{node.lineno}: import {alias.name}"
|
||||
for alias in node.names
|
||||
if alias.asname is None and (alias.name == "litellm" or alias.name.startswith("litellm."))
|
||||
)
|
||||
elif isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "litellm":
|
||||
found.append(f"{rel}:{node.lineno}: litellm.{node.attr}")
|
||||
return found
|
||||
|
||||
|
||||
def main() -> int:
|
||||
found = [line for rel in CHECKED for line in violations(rel)]
|
||||
for line in found:
|
||||
print(line)
|
||||
if found:
|
||||
print(f"{len(found)} litellm.* uses in router modules; read settings from RouterDynamicConfig and use PythonHost")
|
||||
return 1
|
||||
print("router modules use no litellm.* globals")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -38,11 +38,15 @@ ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.js
|
|||
ROUTER_FILES = {
|
||||
"litellm/router.py",
|
||||
"litellm/router_api.py",
|
||||
"litellm/router_config.py",
|
||||
"litellm/router_core.py",
|
||||
"litellm/router_python_core.py",
|
||||
"litellm/router_python_host.py",
|
||||
"litellm/router_deprecated.py",
|
||||
"litellm/router_host.py",
|
||||
}
|
||||
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_deprecated.py"]
|
||||
ROUTER_CLASS_NAMES = {"Router", "DeprecatedRouterMembers"}
|
||||
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"]
|
||||
ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"}
|
||||
|
||||
SCAN_ROOTS = ["litellm", "enterprise"]
|
||||
SKIP_DIRS = {
|
||||
|
|
|
|||
|
|
@ -90,18 +90,7 @@ The suites run against a live proxy, so bring one up first by running the litell
|
|||
|
||||
They also need a proxy whose bundled UI contains the change under test, so run the proxy from your branch (an editable install serves the UI your checkout builds)
|
||||
|
||||
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE`
|
||||
|
||||
The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
store_prompts_in_spend_logs: true
|
||||
```
|
||||
|
||||
Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config <your-e2e-config>.yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts
|
||||
|
||||
A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
|
||||
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and start the proxy with `LITELLM_OTEL_V2=true`, the `otel` callback, and `OTEL_EXPORTER_OTLP_ENDPOINT` pointing at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy.
|
||||
|
||||
### Secret manager lanes
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from dataclasses import dataclass
|
|||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from e2e_config import settle_propagation
|
||||
from e2e_http import NoBody, Result, Success, get_external, is_ok
|
||||
from e2e_http import NoBody, Result, Success, is_ok
|
||||
from e2e_metadata import STEP_FRAMES, step
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
|
@ -36,11 +36,6 @@ class A2ASkill(BaseModel):
|
|||
examples: list[str] | None = None
|
||||
|
||||
|
||||
class A2AProvider(BaseModel):
|
||||
organization: str
|
||||
url: str
|
||||
|
||||
|
||||
class AgentCardParams(BaseModel):
|
||||
"""The upstream agent card an admin registers. `protocolVersion` is the field the
|
||||
proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration."""
|
||||
|
|
@ -57,29 +52,6 @@ class AgentCardParams(BaseModel):
|
|||
preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport")
|
||||
|
||||
|
||||
class UpstreamAgentCard(BaseModel):
|
||||
"""A real published agent card parsed from a public /.well-known endpoint. Keys on
|
||||
the A2A wire aliases so `model_validate_json` reads the served JSON and
|
||||
`model_dump(by_alias=True)` re-emits it unchanged for verbatim registration; it is
|
||||
only ever fetched-and-validated, never hand-constructed, so aliasing on the wire
|
||||
names does not affect any call site."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
protocol_version: str = Field(alias="protocolVersion")
|
||||
name: str
|
||||
description: str
|
||||
version: str
|
||||
url: str
|
||||
provider: A2AProvider | None = None
|
||||
documentation_url: str | None = Field(default=None, alias="documentationUrl")
|
||||
capabilities: A2ACapabilities = A2ACapabilities()
|
||||
skills: list[A2ASkill]
|
||||
default_input_modes: list[str] = Field(default=["text"], alias="defaultInputModes")
|
||||
default_output_modes: list[str] = Field(default=["text"], alias="defaultOutputModes")
|
||||
preferred_transport: str | None = Field(default=None, alias="preferredTransport")
|
||||
|
||||
|
||||
class A2ABridgeParams(BaseModel):
|
||||
"""litellm_params that route the agent through the completion bridge: an A2A
|
||||
message/send is transformed into a litellm.acompletion against this provider."""
|
||||
|
|
@ -93,7 +65,7 @@ class A2ABridgeParams(BaseModel):
|
|||
|
||||
class AgentRegisterBody(BaseModel):
|
||||
agent_name: str
|
||||
agent_card_params: AgentCardParams | UpstreamAgentCard
|
||||
agent_card_params: AgentCardParams
|
||||
litellm_params: A2ABridgeParams | None = None
|
||||
|
||||
|
||||
|
|
@ -136,32 +108,9 @@ class A2ATextPart(BaseModel):
|
|||
text: str
|
||||
|
||||
|
||||
class A2ASearchPropertiesParams(BaseModel):
|
||||
"""The strict param schema of the published property agent's `search_properties`
|
||||
skill (unknown keys are rejected upstream), so a natural-language query like
|
||||
"properties for sale in SF under $2M" is expressed as typed fields."""
|
||||
|
||||
un_locode: str | None = None
|
||||
service_type: str | None = None
|
||||
property_type: str | None = None
|
||||
bedrooms_min: int | None = None
|
||||
asking_price_max: float | None = None
|
||||
limit: int | None = None
|
||||
|
||||
|
||||
class A2ASkillInvocation(BaseModel):
|
||||
skill: str
|
||||
params: A2ASearchPropertiesParams
|
||||
|
||||
|
||||
class A2ADataPart(BaseModel):
|
||||
kind: str = "data"
|
||||
data: A2ASkillInvocation
|
||||
|
||||
|
||||
class A2AOutboundMessage(BaseModel):
|
||||
role: str = "user"
|
||||
parts: list[A2ATextPart | A2ADataPart]
|
||||
parts: list[A2ATextPart]
|
||||
message_id: str = Field(serialization_alias="messageId")
|
||||
|
||||
|
||||
|
|
@ -189,57 +138,10 @@ class A2AResponseMessage(BaseModel):
|
|||
parts: list[A2AResponsePart] = []
|
||||
|
||||
|
||||
class A2ATaskStatus(BaseModel):
|
||||
state: str | None = None
|
||||
message: A2AResponseMessage | None = None
|
||||
|
||||
|
||||
class A2AListingLocation(BaseModel):
|
||||
"""Only the location fields a test reads back off a returned listing."""
|
||||
|
||||
un_locode: str | None = None
|
||||
|
||||
|
||||
class A2AListing(BaseModel):
|
||||
"""A single property card from the agent's `search_results` artifact; only the
|
||||
identity/location fields a test asserts on are modelled."""
|
||||
|
||||
raia_id: str
|
||||
property_type: str | None = None
|
||||
service_type: str | None = None
|
||||
location: A2AListingLocation = A2AListingLocation()
|
||||
|
||||
|
||||
class A2ASearchResults(BaseModel):
|
||||
"""The DataPart payload the property agent's `search_properties` skill returns:
|
||||
the run count plus the listing cards themselves. Proof the tool actually ran and
|
||||
matched, not just that the task completed with some text."""
|
||||
|
||||
total: int
|
||||
count: int
|
||||
listings: list[A2AListing] = []
|
||||
|
||||
|
||||
class A2AArtifactPart(BaseModel):
|
||||
kind: str | None = None
|
||||
data: A2ASearchResults | None = None
|
||||
|
||||
|
||||
class A2AArtifact(BaseModel):
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
artifact_id: str | None = Field(default=None, alias="artifactId")
|
||||
name: str | None = None
|
||||
parts: list[A2AArtifactPart] = []
|
||||
|
||||
|
||||
class A2AResult(BaseModel):
|
||||
"""A message/send result. In 0.3 the message fields sit directly on the result
|
||||
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that
|
||||
runs a task replies with a `task` whose agent text lives on `status.message` and
|
||||
whose tool output lives on `artifacts`. `text` reads the agent's reply from
|
||||
whichever shape the served version produced; `search_results` reads the tool's
|
||||
structured output when the agent ran a skill."""
|
||||
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the
|
||||
agent's reply from whichever shape the served version produced."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
|
@ -248,8 +150,6 @@ class A2AResult(BaseModel):
|
|||
message_id: str | None = Field(default=None, alias="messageId")
|
||||
parts: list[A2AResponsePart] = []
|
||||
message: A2AResponseMessage | None = None
|
||||
status: A2ATaskStatus | None = None
|
||||
artifacts: list[A2AArtifact] = []
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
|
|
@ -257,8 +157,6 @@ class A2AResult(BaseModel):
|
|||
parts = self.message.parts
|
||||
elif self.parts:
|
||||
parts = self.parts
|
||||
elif self.status is not None and self.status.message is not None:
|
||||
parts = self.status.message.parts
|
||||
else:
|
||||
parts = []
|
||||
return "".join(part.text or "" for part in parts)
|
||||
|
|
@ -267,14 +165,6 @@ class A2AResult(BaseModel):
|
|||
def is_nested_v1_shape(self) -> bool:
|
||||
return self.message is not None
|
||||
|
||||
@property
|
||||
def search_results(self) -> A2ASearchResults | None:
|
||||
for artifact in self.artifacts:
|
||||
for part in artifact.parts:
|
||||
if part.data is not None:
|
||||
return part.data
|
||||
return None
|
||||
|
||||
|
||||
class A2AError(BaseModel):
|
||||
code: int
|
||||
|
|
@ -380,11 +270,3 @@ class A2AClient:
|
|||
|
||||
def build_a2a_client(proxy: ProxyClient) -> A2AClient:
|
||||
return A2AClient(proxy=proxy)
|
||||
|
||||
|
||||
@step("Fetch a published A2A agent card from its /.well-known endpoint")
|
||||
def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]:
|
||||
"""Fetch a live A2A agent card from its /.well-known endpoint and parse it into the
|
||||
registration model, so a test can register a real published card verbatim rather
|
||||
than a hand-rolled one."""
|
||||
return get_external(url, response_type=UpstreamAgentCard, timeout=timeout)
|
||||
|
|
|
|||
|
|
@ -17,18 +17,14 @@ import pytest
|
|||
from a2a_client import (
|
||||
A2ABridgeParams,
|
||||
A2AClient,
|
||||
A2ADataPart,
|
||||
A2AJsonRpcRequest,
|
||||
A2AMessageSendParams,
|
||||
A2AOutboundMessage,
|
||||
A2ASearchPropertiesParams,
|
||||
A2ASkill,
|
||||
A2ASkillInvocation,
|
||||
A2ATextPart,
|
||||
AgentCardParams,
|
||||
AgentRegisterBody,
|
||||
AgentResponse,
|
||||
fetch_agent_card,
|
||||
)
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import Result, UnknownApiError, unwrap
|
||||
|
|
@ -50,9 +46,6 @@ BRIDGE = A2ABridgeParams(
|
|||
model=BRIDGE_MODEL,
|
||||
)
|
||||
|
||||
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
|
||||
MOVEHOME_ORIGIN = "https://movehome.org"
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
|
|
@ -135,48 +128,6 @@ class TestA2AAgentLifecycle:
|
|||
assert result is not None
|
||||
assert result.text != ""
|
||||
|
||||
@pytest.mark.covers("other.a2a.message_send.real_world_agent_replies")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.AGENTS_API,
|
||||
route=Route.A2A,
|
||||
)
|
||||
)
|
||||
def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
|
||||
upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN})
|
||||
assert upstream.protocol_version == "0.3.0"
|
||||
marker = unique_marker()
|
||||
body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream)
|
||||
agent = unwrap(client.register_agent(body))
|
||||
resources.defer(lambda: client.delete_agent(agent.agent_id))
|
||||
assert agent.agent_card_params.protocol_version == "0.3"
|
||||
location = "GBLON"
|
||||
request = A2AJsonRpcRequest(
|
||||
id=f"e2e-{unique_marker()}",
|
||||
params=A2AMessageSendParams(
|
||||
message=A2AOutboundMessage(
|
||||
parts=[
|
||||
A2ADataPart(
|
||||
data=A2ASkillInvocation(
|
||||
skill="search_properties",
|
||||
params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3),
|
||||
)
|
||||
)
|
||||
],
|
||||
message_id=unique_marker(),
|
||||
)
|
||||
),
|
||||
)
|
||||
response = unwrap(client.send_message(agent.agent_id, scoped_key, request))
|
||||
assert response.error is None
|
||||
assert response.result is not None
|
||||
results = response.result.search_results
|
||||
assert results is not None, "agent returned no search_results artifact; skill did not run"
|
||||
assert results.total > 0
|
||||
assert results.listings, "search_properties matched nothing; agent returned no property cards"
|
||||
assert all(listing.raia_id for listing in results.listings)
|
||||
assert all(listing.location.un_locode == location for listing in results.listings)
|
||||
|
||||
@pytest.mark.covers("other.a2a.discovery.proxy_fronted_card")
|
||||
@meta(
|
||||
Subject(
|
||||
|
|
|
|||
|
|
@ -1,11 +1,6 @@
|
|||
# Guardrail enforcement (behavior features). Grounded in litellm/proxy/guardrails/guardrail_hooks/.
|
||||
# Rolls up into the "Logging & Guardrails" dashboard module together with logging.*
|
||||
- {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
|
||||
- {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
|
||||
- {id: guardrail.presidio.post_call.masks_generated_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, anthropic_messages_stream], source: "guardrail_hooks/presidio.py", rationale: "Mask model-generated credit-card output with the UI default scope"}
|
||||
- {id: guardrail.presidio.post_call.spend_log_stores_masked_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, messages, anthropic_messages_stream, responses], source: "guardrail_hooks/presidio.py", fail_before_fix: proven, rationale: "When an output guardrail masks the response, the spend log stores the masked text the caller received rather than the raw model output, on every endpoint and both stream modes (LIT-8325)"}
|
||||
- {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
|
||||
- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
|
||||
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
|
||||
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}
|
||||
|
|
|
|||
|
|
@ -139,11 +139,6 @@
|
|||
- {id: llm.responses.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses streams text deltas, usage and a terminal event"}
|
||||
- {id: llm.responses.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses returns one addressable get_weather call"}
|
||||
- {id: llm.responses.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses tool result round trip reaches the model"}
|
||||
- {id: llm.chat_completions.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "service_tier balanced and auto map to Sail completion windows and bill the matching price columns; Sail serves flex only to background responses and Batch"}
|
||||
- {id: llm.chat_completions.sail.service_tier.nonstream.rejects_unknown_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [rejects_unknown_tier], source: "llm_translation/test_sail_e2e.py", rationale: "A service_tier Sail has no completion window for is a 400 without drop_params"}
|
||||
- {id: llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [drops_unknown_tier_and_bills_asap], source: "llm_translation/test_sail_e2e.py", rationale: "An unknown service_tier under drop_params is dropped and billed at asap in both the cost header and spend log"}
|
||||
- {id: llm.responses.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "A caller metadata.completion_window of balanced on /v1/responses bills Sail balanced rates"}
|
||||
- {id: llm.messages.sail.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: sail, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_sail_e2e.py", rationale: "Sail over /v1/messages"}
|
||||
- {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"}
|
||||
- {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"}
|
||||
- {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"}
|
||||
|
|
|
|||
|
|
@ -55,7 +55,6 @@
|
|||
- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"}
|
||||
- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"}
|
||||
- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"}
|
||||
- {id: other.a2a.message_send.real_world_agent_replies, module: other, tier: P1, area: a2a, assertions: [real_world_agent_replies], source: "agent_endpoints/a2a_endpoints.py asend_message", rationale: "A real published a2a agent fetched live from a public /.well-known endpoint (pinning the full semver 0.3.0 the a2a-sdk emits) registers, serves the canonical 0.3, and a message/send skill invocation proxies to the live upstream and returns the agent's reply"}
|
||||
- {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"}
|
||||
- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"}
|
||||
- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"}
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ litellm_settings:
|
|||
- host: 127.0.0.1
|
||||
port: 6379
|
||||
ssl: true
|
||||
callbacks: ["arize_phoenix", "datadog", "smtp_email", "prometheus", "otel", "websearch_interception"]
|
||||
callbacks: ["datadog", "smtp_email", "prometheus", "otel", "websearch_interception"]
|
||||
websearch_interception_params:
|
||||
enabled_providers: ["bedrock"]
|
||||
search_tool_name: e2e-search
|
||||
|
|
@ -97,9 +97,3 @@ files_settings:
|
|||
finetune_settings:
|
||||
- custom_llm_provider: openai
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
mcp_servers:
|
||||
devin:
|
||||
url: "https://mcp.devin.ai/mcp"
|
||||
auth_type: api_key
|
||||
auth_value: os.environ/DEVIN_API_KEY
|
||||
|
|
|
|||
|
|
@ -41,8 +41,6 @@ from pydantic import BaseModel, Field
|
|||
GUARDRAIL_BACKEND: Final = "gemini/gemini-2.5-flash"
|
||||
|
||||
GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"]
|
||||
PiiEntity = Literal["EMAIL_ADDRESS", "PHONE_NUMBER", "PERSON", "CREDIT_CARD", "US_SSN"]
|
||||
PiiAction = Literal["MASK", "BLOCK"]
|
||||
BlockedWordAction = Literal["BLOCK", "MASK"]
|
||||
|
||||
|
||||
|
|
@ -80,22 +78,6 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase):
|
|||
guardrail: Literal["block_code_execution"] = "block_code_execution"
|
||||
|
||||
|
||||
class PresidioParamsBody(GuardrailParamsBase):
|
||||
"""Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the
|
||||
registration to a single callback on the configured mode; the default
|
||||
("both") also registers a second post_call output-masking callback, which a
|
||||
pre_call- or logging_only-scoped test must not drag in. `output_parse_pii`
|
||||
stays unset/False: True would unmask the response back to the caller."""
|
||||
|
||||
guardrail: Literal["presidio"] = "presidio"
|
||||
presidio_analyzer_api_base: str
|
||||
presidio_anonymizer_api_base: str
|
||||
presidio_filter_scope: Literal["input", "output", "both"] | None = None
|
||||
presidio_language: str | None = None
|
||||
output_parse_pii: bool | None = None
|
||||
pii_entities_config: dict[PiiEntity, PiiAction] | None = None
|
||||
|
||||
|
||||
class ToolPermissionRuleBody(BaseModel):
|
||||
"""One tool_permission rule: a decision for the tool named by `tool_name`."""
|
||||
|
||||
|
|
@ -130,7 +112,6 @@ GuardrailParamsBody = (
|
|||
| BedrockGuardrailParamsBody
|
||||
| OpenAIModerationParamsBody
|
||||
| BlockCodeExecutionParamsBody
|
||||
| PresidioParamsBody
|
||||
| ToolPermissionParamsBody
|
||||
| CustomCodeParamsBody
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,857 +0,0 @@
|
|||
"""Live e2e: the Presidio PII guardrail masks, per its configured hook point.
|
||||
|
||||
pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request
|
||||
messages BEFORE the model runs, so the model only ever sees placeholders like
|
||||
<EMAIL_ADDRESS>. A prompt asking the model to repeat a fake email + phone back
|
||||
must come back with the placeholders echoed and the raw PII absent, on
|
||||
/chat/completions and on /v1/messages (Anthropic format).
|
||||
|
||||
post_call: the mirror hook. The request reaches the model unmasked and the
|
||||
MODEL OUTPUT is what gets anonymized, so the caller never receives raw PII the
|
||||
model repeated back. The two hooks are told apart behaviorally rather than by
|
||||
configuration: the post_call prompt asks for a value derived from the raw email
|
||||
(its local part, which is not itself an entity Presidio masks) alongside the
|
||||
address itself, so the answer proves the model saw the raw address while the
|
||||
address in the same response comes back as <EMAIL_ADDRESS>.
|
||||
|
||||
The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE /
|
||||
PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
|
||||
Each guardrail registers with an explicit presidio_filter_scope so only the
|
||||
configured hook's callback exists (the default "both" registers input masking
|
||||
AND a post_call output masker), and is deleted on teardown.
|
||||
|
||||
The spend-log audit test requires general_settings.store_prompts_in_spend_logs:
|
||||
true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy
|
||||
process environment before startup. Setting it only on pytest has no effect.
|
||||
Without this opt-in, redacting guardrail_response is expected proxy behavior;
|
||||
this suite deliberately requires the detected-entity details to remain visible.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import Result, StreamingResponse, Success
|
||||
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
|
||||
from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
AnthropicMessagesResponse,
|
||||
ChatResponse,
|
||||
GuardrailEntityMatch,
|
||||
GuardrailRunRecord,
|
||||
SpendLogRow,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MODEL = "gemini-2.5-flash"
|
||||
|
||||
# A guardrail created via POST /guardrails reaches the worker that served the
|
||||
# create immediately, but every other worker only picks it up on its next
|
||||
# periodic DB sync (~30s), so the first requests can be served unguarded.
|
||||
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
|
||||
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
|
||||
|
||||
# Presidio's anonymizer replaces a detected entity with its unnumbered type
|
||||
# placeholder, e.g. <EMAIL_ADDRESS>. The pre_call assertions match on the bare
|
||||
# token because the model is echoing the masked prompt and may not preserve the
|
||||
# angle brackets; the logged payload keeps the placeholder verbatim.
|
||||
MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS"
|
||||
MASKED_PHONE_TOKEN = "PHONE_NUMBER"
|
||||
MASKED_CREDIT_CARD_TOKEN = "CREDIT_CARD"
|
||||
|
||||
# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects.
|
||||
FAKE_PHONE = "+1 415-555-0134"
|
||||
FAKE_VISA_TEST_CARD = "4111 1111 1111 1111"
|
||||
|
||||
_CARD_DIGIT_RUN: Final = re.compile(r"(?:\d[ -]?){13,19}")
|
||||
_CONTENT_KEYS: Final = frozenset({"content", "text"})
|
||||
|
||||
|
||||
def _presidio_bases() -> tuple[str, str]:
|
||||
analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip()
|
||||
anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip()
|
||||
if not analyzer or not anonymizer:
|
||||
pytest.fail(
|
||||
"Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE "
|
||||
"(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip"
|
||||
)
|
||||
return analyzer, anonymizer
|
||||
|
||||
|
||||
def _register_presidio(
|
||||
client: GuardrailsClient,
|
||||
resources: ResourceManager,
|
||||
*,
|
||||
name: str,
|
||||
mode: GuardrailMode | list[GuardrailMode] = "pre_call",
|
||||
filter_scope: Literal["input", "output", "both"] | None = "input",
|
||||
entities: dict[PiiEntity, PiiAction] | None = None,
|
||||
) -> None:
|
||||
analyzer, anonymizer = _presidio_bases()
|
||||
guardrail_id = client.register(
|
||||
name,
|
||||
PresidioParamsBody(
|
||||
mode=mode,
|
||||
default_on=False,
|
||||
presidio_analyzer_api_base=analyzer,
|
||||
presidio_anonymizer_api_base=anonymizer,
|
||||
presidio_filter_scope=filter_scope,
|
||||
pii_entities_config=entities,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
|
||||
def _fake_email() -> str:
|
||||
return f"jane.doe.{unique_marker()}@example.com"
|
||||
|
||||
|
||||
def _pii_prompt(marker: str, email: str) -> str:
|
||||
return (
|
||||
f"{marker} Repeat this sentence back to me exactly, word for word: "
|
||||
f"My email address is {email} and my phone number is {FAKE_PHONE}."
|
||||
)
|
||||
|
||||
|
||||
def _first_content(response: ChatResponse) -> str:
|
||||
if not response.choices:
|
||||
return ""
|
||||
message = response.choices[0].message
|
||||
return (message.content if message else None) or ""
|
||||
|
||||
|
||||
class _StreamDelta(BaseModel):
|
||||
content: str | None = None
|
||||
|
||||
|
||||
class _StreamChoice(BaseModel):
|
||||
delta: _StreamDelta
|
||||
|
||||
|
||||
class _StreamChunk(BaseModel):
|
||||
choices: tuple[_StreamChoice, ...] = ()
|
||||
|
||||
|
||||
class _AnthropicStreamDelta(BaseModel):
|
||||
type: str | None = None
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _AnthropicStreamEvent(BaseModel):
|
||||
type: str
|
||||
delta: _AnthropicStreamDelta | None = None
|
||||
|
||||
|
||||
def _credit_card_prompt(marker: str) -> str:
|
||||
return (
|
||||
f"{marker} Reply with only the well known Visa sandbox test card number that starts with 4111, "
|
||||
"the 16 digits grouped in fours separated by spaces, and nothing else."
|
||||
)
|
||||
|
||||
|
||||
def _passes_luhn(digits: str) -> bool:
|
||||
checksum = sum(
|
||||
digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2)
|
||||
for position, digit in enumerate(int(char) for char in reversed(digits))
|
||||
)
|
||||
return checksum % 10 == 0
|
||||
|
||||
|
||||
def _contains_card_number(text: str) -> bool:
|
||||
"""Presidio's CREDIT_CARD recognizer only reports Luhn-valid digit runs, so a
|
||||
Luhn-invalid number the model hallucinates is not something masking can catch."""
|
||||
return any(
|
||||
13 <= len(digits) <= 19 and _passes_luhn(digits)
|
||||
for digits in (re.sub(r"[ -]", "", match.group()) for match in _CARD_DIGIT_RUN.finditer(text))
|
||||
)
|
||||
|
||||
|
||||
def _stream_content(result: StreamingResponse) -> str:
|
||||
return "".join(
|
||||
choice.delta.content
|
||||
for event in result.stream_events
|
||||
if event != "[DONE]"
|
||||
for choice in _StreamChunk.model_validate_json(event).choices[:1]
|
||||
if choice.delta.content
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_stream_content(result: StreamingResponse) -> str:
|
||||
return "".join(
|
||||
event.delta.text
|
||||
for payload in result.stream_events
|
||||
for event in [_AnthropicStreamEvent.model_validate_json(payload)]
|
||||
if event.type == "content_block_delta"
|
||||
and event.delta is not None
|
||||
and event.delta.type == "text_delta"
|
||||
and event.delta.text
|
||||
)
|
||||
|
||||
|
||||
def _messages_text(response: AnthropicMessagesResponse) -> str:
|
||||
"""The text of a /v1/messages answer, whichever shape the proxy produced
|
||||
(Anthropic-native content blocks or OpenAI-normalized choices)."""
|
||||
parts: list[str] = []
|
||||
for block in response.content or []:
|
||||
if block.text:
|
||||
parts.append(block.text)
|
||||
for choice in response.choices or []:
|
||||
if choice.message and choice.message.content:
|
||||
parts.append(choice.message.content)
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def _assert_eventually_masked[R: BaseModel](
|
||||
fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str
|
||||
) -> None:
|
||||
"""Retry the call until the response comes back masked, to the propagation
|
||||
deadline. An unmasked early response is in-flight guardrail propagation, not
|
||||
a failure, and neither is a transient non-Success (a replica that has not
|
||||
reloaded the guardrail answers 404, the live model can rate-limit) - only a
|
||||
response that still carries the raw PII at the deadline is."""
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last: str = "<no successful response yet>"
|
||||
while True:
|
||||
result = fetch()
|
||||
match result:
|
||||
case Success(data=data):
|
||||
content = extract(data)
|
||||
last = content
|
||||
masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content
|
||||
if masked:
|
||||
assert FAKE_PHONE not in content, (
|
||||
f"the raw phone number must be masked before the model sees it, but the "
|
||||
f"response echoed it: {content[:300]!r}"
|
||||
)
|
||||
return
|
||||
case _:
|
||||
last = f"<non-Success result: {result}>"
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
f"presidio pre_call guardrail never masked the PII within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
class TestPresidioPreCallMasking:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.pre_call.masks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_pre_call_masks_pii_on_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-pre-chat-{unique_marker()}"
|
||||
_register_presidio(client, resources, name=name)
|
||||
|
||||
email = _fake_email()
|
||||
prompt = _pii_prompt(unique_marker(), email)
|
||||
|
||||
_assert_eventually_masked(
|
||||
lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
|
||||
_first_content,
|
||||
email=email,
|
||||
)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.pre_call.masks",
|
||||
exercised_on=["messages"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.MESSAGES,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_pre_call_masks_pii_on_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-pre-msg-{unique_marker()}"
|
||||
_register_presidio(client, resources, name=name)
|
||||
|
||||
email = _fake_email()
|
||||
prompt = _pii_prompt(unique_marker(), email)
|
||||
|
||||
_assert_eventually_masked(
|
||||
lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
|
||||
_messages_text,
|
||||
email=email,
|
||||
)
|
||||
|
||||
|
||||
#: Room for the model's reasoning tokens plus the three-line answer; a lower cap
|
||||
#: truncates the response before the address it is supposed to mask.
|
||||
_POST_CALL_MAX_TOKENS = 512
|
||||
|
||||
#: The post_call scenario masks these two entities and nothing else. Left
|
||||
#: unscoped, Presidio's broader recognizers claim the local part too (a random
|
||||
#: marker reads as an NRP), which would erase the very token that tells output
|
||||
#: masking apart from input masking.
|
||||
_POST_CALL_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
|
||||
|
||||
|
||||
def _post_call_prompt(marker: str, local_part: str) -> str:
|
||||
"""Ask for the local part and the full address in one answer. Presidio masks
|
||||
an EMAIL_ADDRESS entity and a bare local part is not one, so the two land
|
||||
differently in the same response and pin the hook point behaviorally."""
|
||||
return (
|
||||
f"{marker} My email address is {local_part}@example.com and my phone number is {FAKE_PHONE}. "
|
||||
"Reply with exactly three lines and nothing else. "
|
||||
"Line 1: the part of the email address before the @ sign. "
|
||||
"Line 2: the full email address. "
|
||||
"Line 3: the phone number."
|
||||
)
|
||||
|
||||
|
||||
class TestPresidioPostCallMasking:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.post_call.masks",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_post_call_masks_pii_in_model_output(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
"""A guardrail scoped to the output must anonymize the PII the model
|
||||
repeats back, so a caller (or a downstream log of the response) never
|
||||
receives it, while the request itself reaches the model untouched.
|
||||
|
||||
Both facts are asserted from one response: the local part comes back raw,
|
||||
which is only possible if the model saw the real address, and the address
|
||||
itself comes back as <EMAIL_ADDRESS> in the same answer.
|
||||
"""
|
||||
name = f"e2e-presidio-post-chat-{unique_marker()}"
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode="post_call",
|
||||
filter_scope="output",
|
||||
entities=_POST_CALL_ENTITIES,
|
||||
)
|
||||
|
||||
local_part = f"e2euser{unique_marker()}"
|
||||
email = f"{local_part}@example.com"
|
||||
prompt = _post_call_prompt(unique_marker(), local_part)
|
||||
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last = "<no successful response yet>"
|
||||
while True:
|
||||
result = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=_POST_CALL_MAX_TOKENS)
|
||||
match result:
|
||||
case Success(data=data):
|
||||
last = _first_content(data)
|
||||
if MASKED_EMAIL_TOKEN in last and email not in last:
|
||||
assert local_part in last, (
|
||||
"the model must have seen the RAW address (it is asked for the local "
|
||||
"part, which Presidio does not mask); the local part is missing, so "
|
||||
f"this response cannot tell post_call masking from pre_call: {last[:300]!r}"
|
||||
)
|
||||
assert MASKED_PHONE_TOKEN in last and FAKE_PHONE not in last, (
|
||||
f"the phone number in the model's answer must be masked too, got: {last[:300]!r}"
|
||||
)
|
||||
return
|
||||
case _:
|
||||
last = f"<non-Success result: {result}>"
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
f"presidio post_call guardrail never masked the model's output within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
def _assert_eventually_masks_generated_card(fetch: Callable[[], str | None]) -> None:
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last: str = "<no successful response yet>"
|
||||
while True:
|
||||
content = fetch()
|
||||
if content is not None:
|
||||
last = content
|
||||
if _contains_card_number(content):
|
||||
pytest.fail(
|
||||
"the post_call output masking let a card number through: "
|
||||
f"{content[:300]!r}"
|
||||
)
|
||||
if MASKED_CREDIT_CARD_TOKEN in content:
|
||||
return
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
"presidio post_call output masking never masked the generated card within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
class TestPresidioCreditCardOutputMasking:
|
||||
"""Proves the UI-default Presidio scope masks model-generated card output."""
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.post_call.masks_generated_output",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-card-chat-{unique_marker()}"
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode=["pre_call", "post_call"],
|
||||
filter_scope=None,
|
||||
entities={"CREDIT_CARD": "MASK"},
|
||||
)
|
||||
prompt: Final = _credit_card_prompt(unique_marker())
|
||||
|
||||
def fetch() -> str | None:
|
||||
result: Final = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=512)
|
||||
match result:
|
||||
case Success(data=data):
|
||||
return _first_content(data)
|
||||
case _:
|
||||
return None
|
||||
|
||||
_assert_eventually_masks_generated_card(fetch)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.post_call.masks_generated_output",
|
||||
exercised_on=["chat_completions_stream"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.STREAM,
|
||||
)
|
||||
)
|
||||
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-card-stream-{unique_marker()}"
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode=["pre_call", "post_call"],
|
||||
filter_scope=None,
|
||||
entities={"CREDIT_CARD": "MASK"},
|
||||
)
|
||||
prompt: Final = _credit_card_prompt(unique_marker())
|
||||
|
||||
def fetch() -> str | None:
|
||||
result: Final = client.chat_stream_raw(
|
||||
scoped_key,
|
||||
MODEL,
|
||||
prompt,
|
||||
guardrails=[name],
|
||||
max_tokens=512,
|
||||
)
|
||||
if not result.ok or result.stream_error:
|
||||
return None
|
||||
return _stream_content(result)
|
||||
|
||||
_assert_eventually_masks_generated_card(fetch)
|
||||
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.post_call.masks_generated_output",
|
||||
exercised_on=["anthropic_messages_stream"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.MESSAGES,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.STREAM,
|
||||
)
|
||||
)
|
||||
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_anthropic_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-card-messages-stream-{unique_marker()}"
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode=["pre_call", "post_call"],
|
||||
filter_scope=None,
|
||||
entities={"CREDIT_CARD": "MASK"},
|
||||
)
|
||||
prompt: Final = _credit_card_prompt(unique_marker())
|
||||
|
||||
def fetch() -> str | None:
|
||||
result: Final = client.messages_stream_raw(
|
||||
scoped_key,
|
||||
MODEL,
|
||||
prompt,
|
||||
guardrails=[name],
|
||||
max_tokens=512,
|
||||
)
|
||||
if not result.ok or result.stream_error:
|
||||
return None
|
||||
return _anthropic_stream_content(result)
|
||||
|
||||
_assert_eventually_masks_generated_card(fetch)
|
||||
|
||||
|
||||
def _wire_text(outcome: StreamingResponse) -> str:
|
||||
return "\n".join(outcome.stream_events) if outcome.is_streaming else outcome.body
|
||||
|
||||
|
||||
def _poll_until_generated_card_masked(fetch: Callable[[], StreamingResponse]) -> StreamingResponse:
|
||||
"""The raw HTTP outcome once the output masker is in effect on the serving
|
||||
worker: whichever wire shape the endpoint speaks, a masked body carries the
|
||||
CREDIT_CARD placeholder and no Luhn-valid card run. A raw card is a worker
|
||||
that has not loaded the guardrail yet, so it is polled through like any
|
||||
other unmasked answer."""
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last: str = "<no successful response yet>"
|
||||
while True:
|
||||
outcome = fetch()
|
||||
if outcome.ok and not outcome.stream_error:
|
||||
wire = _wire_text(outcome)
|
||||
last = wire
|
||||
if MASKED_CREDIT_CARD_TOKEN in wire and not _contains_card_number(wire):
|
||||
return outcome
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
"presidio post_call output masking never masked the generated card within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
def _spend_log_response_text(client: GuardrailsClient, key: str, call_id: str) -> str:
|
||||
rows = client.proxy.poll_logs_for_key(
|
||||
key,
|
||||
predicate=lambda logged: any(row.litellm_call_id == call_id for row in logged),
|
||||
)
|
||||
row = next((row for row in rows if row.litellm_call_id == call_id), None)
|
||||
assert row is not None, f"no spend log row ever appeared for x-litellm-call-id {call_id}"
|
||||
return "\n".join(_stored_content(row.response))
|
||||
|
||||
|
||||
def _stored_content(node: JsonValue, key: str | None = None) -> Iterator[str]:
|
||||
match node:
|
||||
case str() if key in _CONTENT_KEYS:
|
||||
yield node
|
||||
case dict():
|
||||
for child_key, child in node.items():
|
||||
yield from _stored_content(child, child_key)
|
||||
case list():
|
||||
for item in node:
|
||||
yield from _stored_content(item, key)
|
||||
case _:
|
||||
return
|
||||
|
||||
|
||||
class TestPresidioSpendLogStoresMaskedOutput:
|
||||
"""The spend log stores the response the caller received, on every endpoint
|
||||
and both stream modes, when an output-only post_call Presidio guardrail masks
|
||||
a card number the model generated."""
|
||||
|
||||
_CELL: Final = "guardrail.presidio.post_call.spend_log_stores_masked_output"
|
||||
|
||||
def _assert_spend_log_is_masked(
|
||||
self,
|
||||
client: GuardrailsClient,
|
||||
resources: ResourceManager,
|
||||
key: str,
|
||||
*,
|
||||
name: str,
|
||||
fetch: Callable[[str, str], StreamingResponse],
|
||||
) -> None:
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode="post_call",
|
||||
filter_scope="output",
|
||||
entities={"CREDIT_CARD": "MASK"},
|
||||
)
|
||||
prompt: Final = _credit_card_prompt(unique_marker())
|
||||
outcome = _poll_until_generated_card_masked(lambda: fetch(prompt, name))
|
||||
assert outcome.call_id, f"the served response must carry x-litellm-call-id: {dict(outcome.headers)}"
|
||||
|
||||
logged = _spend_log_response_text(client, key, outcome.call_id)
|
||||
assert not _contains_card_number(logged), (
|
||||
"the caller got the masked response but the spend log stored the raw model output: "
|
||||
f"{logged[:400]!r}"
|
||||
)
|
||||
assert MASKED_CREDIT_CARD_TOKEN in logged, (
|
||||
f"the spend log response carries neither the card nor the placeholder: {logged[:400]!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["chat_completions"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_spend_log_stores_masked_output_on_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-chat-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.chat_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["chat_completions_stream"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.STREAM,
|
||||
)
|
||||
)
|
||||
def test_spend_log_stores_masked_output_on_streaming_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-chat-stream-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.chat_stream_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["messages"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.MESSAGES,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_spend_log_stores_masked_output_on_anthropic_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-messages-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.messages_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["anthropic_messages_stream"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.MESSAGES,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.STREAM,
|
||||
)
|
||||
)
|
||||
def test_spend_log_stores_masked_output_on_streaming_anthropic_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-messages-stream-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.messages_stream_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["responses"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.RESPONSES,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_spend_log_stores_masked_output_on_responses(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-responses-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.responses(scoped_key, MODEL, prompt, guardrails=[guardrail]),
|
||||
)
|
||||
|
||||
|
||||
_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
|
||||
|
||||
_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
|
||||
|
||||
|
||||
def _applied_guardrails(outcome: StreamingResponse) -> str:
|
||||
return outcome.headers.get("x-litellm-applied-guardrails", "")
|
||||
|
||||
|
||||
def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]:
|
||||
metadata = row.metadata
|
||||
if metadata is None:
|
||||
return ()
|
||||
return tuple(metadata.guardrail_information or ())
|
||||
|
||||
|
||||
def _poll_until_guardrail_applied(
|
||||
client: GuardrailsClient, key: str, guardrail_name: str, prompt: str
|
||||
) -> StreamingResponse:
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
|
||||
while time.monotonic() < deadline:
|
||||
if last.ok and guardrail_name in _applied_guardrails(last):
|
||||
return last
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
|
||||
return last
|
||||
|
||||
|
||||
class TestPresidioSpendLogRecord:
|
||||
@pytest.mark.covers(
|
||||
"guardrail.presidio.pre_call.logs_masked_entities",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_masking_run_is_recorded_on_the_spend_log(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name = f"e2e-presidio-log-{unique_marker()}"
|
||||
_register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES)
|
||||
|
||||
email = _fake_email()
|
||||
prompt = _pii_prompt(unique_marker(), email)
|
||||
|
||||
outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt)
|
||||
assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}"
|
||||
assert name in _applied_guardrails(outcome), (
|
||||
"the response must carry x-litellm-applied-guardrails naming the guardrail; without it "
|
||||
f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}"
|
||||
)
|
||||
|
||||
request_id = ChatResponse.model_validate_json(outcome.body).id
|
||||
assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}"
|
||||
|
||||
rows = client.proxy.poll_logs_for_request_id(
|
||||
request_id,
|
||||
predicate=lambda logged: bool(_guardrail_records(logged[0])),
|
||||
)
|
||||
assert rows, f"no spend log row ever appeared for request {request_id}"
|
||||
|
||||
records = _guardrail_records(rows[0])
|
||||
assert records, (
|
||||
f"the spend log for {request_id} carries no guardrail_information, so the dashboard's "
|
||||
"guardrail panel would render nothing for a request the guardrail demonstrably ran on"
|
||||
)
|
||||
|
||||
record = next(
|
||||
(entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"),
|
||||
None,
|
||||
)
|
||||
assert record is not None, (
|
||||
"guardrail_information carries no pre_call record for this guardrail, only "
|
||||
f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}"
|
||||
)
|
||||
assert record.guardrail_status == "success", (
|
||||
f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}"
|
||||
)
|
||||
assert record.guardrail_provider == "presidio", (
|
||||
f"the record must attribute the run to presidio so the dashboard picks the right "
|
||||
f"renderer, got {record.guardrail_provider!r}"
|
||||
)
|
||||
|
||||
counts = record.masked_entity_count or {}
|
||||
assert _LOGGED_ENTITIES.keys() <= counts.keys(), (
|
||||
f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}"
|
||||
)
|
||||
assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), (
|
||||
f"each masked entity must be counted at least once, got {counts}"
|
||||
)
|
||||
|
||||
assert not isinstance(record.guardrail_response, str), (
|
||||
"This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config "
|
||||
"or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup "
|
||||
"(not just the pytest environment). Default prompt redaction is valid proxy behavior, "
|
||||
f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}"
|
||||
)
|
||||
entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
|
||||
assert entities, (
|
||||
"guardrail_response must carry the detected entities; the dashboard's Detected Entities "
|
||||
"list and its per-entity scores are rendered from exactly this array"
|
||||
)
|
||||
assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), (
|
||||
f"the detected entities must cover what was masked, got "
|
||||
f"{sorted(entity.entity_type for entity in entities)}"
|
||||
)
|
||||
for entity in entities:
|
||||
assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}"
|
||||
assert 0 <= entity.start < entity.end, (
|
||||
f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}"
|
||||
)
|
||||
|
|
@ -1,256 +0,0 @@
|
|||
"""Live e2e: Sail through the gateway, where LiteLLM turns ``service_tier`` into Sail's
|
||||
``metadata.completion_window`` and bills the price columns of the window it sent.
|
||||
|
||||
The deployment carries its own base, balanced and flex rates, each distinct, so a bill at
|
||||
the wrong tier cannot pass. They are registered on the deployment instead of read from the
|
||||
proxy's cost map, because a stack that loads the published map has no ``sail/`` rows until
|
||||
this provider ships. Requires SAIL_API_KEY on the proxy; no skip gate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS, unique_marker
|
||||
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import LiteLLMParamsBody, SpendLogRow
|
||||
from openai import OpenAI
|
||||
from proxy_client import ProxyClient
|
||||
from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
BACKEND: Final = "sail/zai-org/GLM-5.3"
|
||||
PricedTier = Literal["base", "balanced", "flex"]
|
||||
PRICED_TIERS: Final[tuple[PricedTier, ...]] = ("base", "balanced", "flex")
|
||||
PROMPT: Final = "Reply with one word."
|
||||
MAX_TOKENS: Final = 512
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rates:
|
||||
input: float
|
||||
output: float
|
||||
cache_read: float
|
||||
|
||||
|
||||
RATES: Final[Mapping[PricedTier, _Rates]] = {
|
||||
"base": _Rates(input=3e-06, output=9e-06, cache_read=1e-06),
|
||||
"balanced": _Rates(input=2e-06, output=6e-06, cache_read=7e-07),
|
||||
"flex": _Rates(input=1e-06, output=3e-06, cache_read=4e-07),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Tokens:
|
||||
prompt: int
|
||||
cached: int
|
||||
completion: int
|
||||
|
||||
|
||||
def _approx_equal(actual: float, expected: float) -> bool:
|
||||
return abs(actual - expected) <= max(1e-12, abs(expected) * 1e-2)
|
||||
|
||||
|
||||
def _cost(rates: _Rates, tokens: _Tokens) -> float:
|
||||
return (
|
||||
(tokens.prompt - tokens.cached) * rates.input
|
||||
+ tokens.cached * rates.cache_read
|
||||
+ tokens.completion * rates.output
|
||||
)
|
||||
|
||||
|
||||
def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]:
|
||||
model: Final = f"e2e-sail-{unique_marker()}"
|
||||
model_id: Final = proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model=BACKEND,
|
||||
api_key="os.environ/SAIL_API_KEY",
|
||||
input_cost_per_token=RATES["base"].input,
|
||||
output_cost_per_token=RATES["base"].output,
|
||||
cache_read_input_token_cost=RATES["base"].cache_read,
|
||||
input_cost_per_token_balanced=RATES["balanced"].input,
|
||||
output_cost_per_token_balanced=RATES["balanced"].output,
|
||||
cache_read_input_token_cost_balanced=RATES["balanced"].cache_read,
|
||||
input_cost_per_token_flex=RATES["flex"].input,
|
||||
output_cost_per_token_flex=RATES["flex"].output,
|
||||
cache_read_input_token_cost_flex=RATES["flex"].cache_read,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
return model, resources.key()
|
||||
|
||||
|
||||
def _openai(sdk: SdkClients, key: str) -> OpenAI:
|
||||
return sdk.openai(key).with_options(timeout=SLOW_PROVIDER_TIMEOUT_SECONDS)
|
||||
|
||||
|
||||
def _assert_billed_at(tier: PricedTier, tokens: _Tokens, header_cost: str | None) -> float:
|
||||
assert tokens.prompt > 0 and tokens.completion > 0, f"Sail reported no usage, so no cost is real: {tokens}"
|
||||
assert header_cost is not None, "x-litellm-response-cost header missing"
|
||||
costs: Final = {priced: _cost(rates, tokens) for priced, rates in RATES.items()}
|
||||
assert not any(_approx_equal(costs[other], costs[tier]) for other in PRICED_TIERS if other != tier), (
|
||||
f"{BACKEND} tier rates too close together to tell {tier} apart at {tokens}: {costs}"
|
||||
)
|
||||
assert _approx_equal(float(header_cost), costs[tier]), (
|
||||
f"header cost {header_cost} is not the {tier} price at {tokens}: expected {costs[tier]}, all tiers {costs}"
|
||||
)
|
||||
return float(header_cost)
|
||||
|
||||
|
||||
def _assert_spend_row_matches(proxy: ProxyClient, key: str, header_cost: float) -> None:
|
||||
def priced(rows: list[SpendLogRow]) -> bool:
|
||||
return any((row.spend or 0) > 0 for row in rows)
|
||||
|
||||
rows: Final = [row for row in proxy.poll_logs_for_key(key, predicate=priced) if (row.spend or 0) > 0]
|
||||
assert rows, f"no priced spend row landed for key {key}"
|
||||
assert rows[0].custom_llm_provider == "sail", f"spend row misattributed: {rows[0]}"
|
||||
assert rows[0].spend is not None and _approx_equal(rows[0].spend, header_cost), (
|
||||
f"logged spend {rows[0].spend} disagrees with the x-litellm-response-cost header {header_cost}"
|
||||
)
|
||||
|
||||
|
||||
class TestSailChatCompletions:
|
||||
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.cost_logged")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.SAIL,),
|
||||
models=(BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("service_tier", "billed_tier"), [("balanced", "balanced"), ("auto", "base")]
|
||||
)
|
||||
def test_service_tier_bills_the_matching_completion_window(
|
||||
self,
|
||||
proxy: ProxyClient,
|
||||
resources: ResourceManager,
|
||||
sdk: SdkClients,
|
||||
service_tier: str,
|
||||
billed_tier: PricedTier,
|
||||
) -> None:
|
||||
model, key = _register(proxy, resources)
|
||||
|
||||
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
|
||||
max_completion_tokens=MAX_TOKENS,
|
||||
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier},
|
||||
)
|
||||
usage: Final = raw.parse().usage
|
||||
assert usage is not None, "chat response carries no usage"
|
||||
details: Final = usage.prompt_tokens_details
|
||||
tokens: Final = _Tokens(
|
||||
prompt=usage.prompt_tokens,
|
||||
cached=(details.cached_tokens or 0) if details else 0,
|
||||
completion=usage.completion_tokens,
|
||||
)
|
||||
|
||||
header_cost: Final = _assert_billed_at(
|
||||
billed_tier, tokens, response_header(raw.headers, "x-litellm-response-cost")
|
||||
)
|
||||
_assert_spend_row_matches(proxy, key, header_cost)
|
||||
|
||||
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.SAIL,),
|
||||
models=(BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
@pytest.mark.parametrize("service_tier", ["bogus", 5])
|
||||
def test_unknown_service_tier_is_dropped_and_billed_asap(
|
||||
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, service_tier: str | int
|
||||
) -> None:
|
||||
model, key = _register(proxy, resources)
|
||||
|
||||
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
|
||||
max_completion_tokens=MAX_TOKENS,
|
||||
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier, "drop_params": True},
|
||||
)
|
||||
usage: Final = raw.parse().usage
|
||||
assert usage is not None, "chat response carries no usage"
|
||||
details: Final = usage.prompt_tokens_details
|
||||
tokens: Final = _Tokens(
|
||||
prompt=usage.prompt_tokens,
|
||||
cached=(details.cached_tokens or 0) if details else 0,
|
||||
completion=usage.completion_tokens,
|
||||
)
|
||||
header_cost: Final = _assert_billed_at("base", tokens, response_header(raw.headers, "x-litellm-response-cost"))
|
||||
_assert_spend_row_matches(proxy, key, header_cost)
|
||||
|
||||
|
||||
class TestSailResponses:
|
||||
@pytest.mark.covers("llm.responses.sail.service_tier.nonstream.cost_logged")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
route=Route.RESPONSES,
|
||||
providers=(Provider.SAIL,),
|
||||
models=(BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_caller_completion_window_bills_its_rates(
|
||||
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
|
||||
) -> None:
|
||||
model, key = _register(proxy, resources)
|
||||
|
||||
raw: Final = _openai(sdk, key).responses.with_raw_response.create(
|
||||
model=model,
|
||||
input=f"{PROMPT} {unique_marker()}",
|
||||
max_output_tokens=MAX_TOKENS,
|
||||
metadata={"completion_window": "balanced"},
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
usage: Final = raw.parse().usage
|
||||
assert usage is not None, "responses answer carries no usage"
|
||||
tokens: Final = _Tokens(
|
||||
prompt=usage.input_tokens,
|
||||
cached=usage.input_tokens_details.cached_tokens,
|
||||
completion=usage.output_tokens,
|
||||
)
|
||||
|
||||
header_cost: Final = _assert_billed_at(
|
||||
"balanced", tokens, response_header(raw.headers, "x-litellm-response-cost")
|
||||
)
|
||||
_assert_spend_row_matches(proxy, key, header_cost)
|
||||
|
||||
|
||||
class TestSailMessages:
|
||||
@pytest.mark.covers("llm.messages.sail.basic.nonstream.works")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.LLM_TRANSLATION,
|
||||
route=Route.MESSAGES,
|
||||
providers=(Provider.SAIL,),
|
||||
models=(BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_plain_call_returns_a_message(
|
||||
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
|
||||
) -> None:
|
||||
model, key = _register(proxy, resources)
|
||||
|
||||
message: Final = sdk.anthropic(key).messages.create(
|
||||
model=model,
|
||||
max_tokens=MAX_TOKENS,
|
||||
messages=[{"role": "user", "content": PROMPT}],
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
assert message.role == "assistant" and message.content, f"/v1/messages returned no content: {message}"
|
||||
assert message.usage.output_tokens > 0, f"/v1/messages reported no output usage: {message.usage}"
|
||||
|
|
@ -213,45 +213,6 @@ class LangfuseCreds:
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WeaveCreds:
|
||||
"""Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback.
|
||||
|
||||
The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment:
|
||||
the weave_otel logger is constructed from those before the per-key vars are
|
||||
applied, so a key-scoped callback on a proxy without them never initializes.
|
||||
The per-key vars are what direct THIS key's spans at this project.
|
||||
"""
|
||||
|
||||
api_key: str
|
||||
project_id: str
|
||||
|
||||
def key_logging_metadata(self) -> KeyMetadata:
|
||||
return KeyMetadata(
|
||||
logging=[
|
||||
KeyLoggingCallback(
|
||||
callback_name="weave_otel",
|
||||
callback_type="success_and_failure",
|
||||
callback_vars=KeyLoggingCallbackVars(
|
||||
wandb_api_key=self.api_key,
|
||||
weave_project_id=self.project_id,
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def load_weave_creds() -> WeaveCreds:
|
||||
api_key = os.getenv("WANDB_API_KEY")
|
||||
project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip()
|
||||
if not (api_key and project_id):
|
||||
pytest.fail(
|
||||
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
|
||||
"format <entity>/<project>); missing credentials is a hard failure, not a skip"
|
||||
)
|
||||
return WeaveCreds(api_key=api_key, project_id=project_id)
|
||||
|
||||
|
||||
def load_langfuse_creds() -> LangfuseCreds:
|
||||
public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
|
||||
secret_key = os.getenv("LANGFUSE_SECRET_KEY")
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None:
|
|||
details = _ReadinessDetails.model_validate_json(readiness_details_body(client))
|
||||
assert OTEL_V2_LOGGER_NAME in details.success_callbacks, (
|
||||
f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active "
|
||||
f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}"
|
||||
f"(LITELLM_OTEL_V2=true + the otel callback in the stack config); got: {details.success_callbacks}"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,212 +0,0 @@
|
|||
"""Live e2e: key-scoped Weave (Weights & Biases) delivery, success and failure.
|
||||
|
||||
Covers the two `logging.niche_integrations.*.logs_spend` cells with a real member
|
||||
of that cohort. A key carrying a `weave_otel` callback in its logging metadata
|
||||
must deliver its calls to the real Weave project, and each call must arrive
|
||||
exactly once, carrying the same cost the response header reported:
|
||||
|
||||
- success: one `litellm_request` call, OTEL status OK, `llm.response.cost` equal
|
||||
to `x-litellm-response-cost`, and non-zero tokens
|
||||
- failure: a provider-rejected call arrives too, as one call with OTEL status
|
||||
ERROR naming the provider exception, and with no cost - a failed call that
|
||||
silently never reaches the destination is an invisible outage, and a billed
|
||||
one is worse
|
||||
|
||||
Both halves assert the recorded state (the key's callback registration answers
|
||||
success and the destination holds the call) and the enforced behavior (the
|
||||
delivered payload's status and cost). Delivery is read back through Weave's own
|
||||
query API; nothing is mocked.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
|
||||
from e2e_http import StreamingResponse
|
||||
from e2e_metadata import Domain, Mode, Provider, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from logging_client import (
|
||||
INVALID_UPSTREAM_API_KEY,
|
||||
LoggingClient,
|
||||
WeaveCreds,
|
||||
costs_agree,
|
||||
first_ok,
|
||||
load_weave_creds,
|
||||
)
|
||||
from models import LiteLLMParamsBody
|
||||
from weave_reader import WeaveCall, WeaveReader, build_weave_reader
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
UNREACHABLE_ANTHROPIC_BACKEND: Final = "anthropic/claude-haiku-4-5"
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def weave_creds() -> WeaveCreds:
|
||||
return load_weave_creds()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def weave_reader() -> WeaveReader:
|
||||
return build_weave_reader()
|
||||
|
||||
|
||||
#: How far before the request the Weave read-back window opens, to absorb clock
|
||||
#: skew between this host and Weave. Without it a host running slightly fast
|
||||
#: would filter out its own call.
|
||||
_WINDOW_SKEW_SECONDS = 120.0
|
||||
|
||||
|
||||
def _window_start() -> float:
|
||||
return time.time() - _WINDOW_SKEW_SECONDS
|
||||
|
||||
|
||||
def _exactly_one(calls: tuple[WeaveCall, ...], *, marker: str, what: str) -> WeaveCall:
|
||||
assert calls, f"no Weave call for the {what} (marker {marker}) reached the project within the deadline"
|
||||
assert len(calls) == 1, (
|
||||
f"expected exactly ONE Weave call for the {what} (marker {marker}), got {len(calls)}: "
|
||||
f"{[call.id for call in calls]} - more than one call for one request is the "
|
||||
"duplicate-delivery bug"
|
||||
)
|
||||
return calls[0]
|
||||
|
||||
|
||||
WEAVE_STAGE_RED_REASON = (
|
||||
"stage red: product gap, key-scoped weave_otel spans are not delivered when the OTEL v2 callback is active"
|
||||
)
|
||||
|
||||
|
||||
class TestWeaveLogDelivery:
|
||||
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
|
||||
@pytest.mark.covers("logging.niche_integrations.success.logs_spend", exercised_on=["chat_completions"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.OBSERVABILITY,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(CHEAP_ANTHROPIC_MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_chat_completions_delivers_one_call_with_spend(
|
||||
self,
|
||||
client: LoggingClient,
|
||||
weave_creds: WeaveCreds,
|
||||
weave_reader: WeaveReader,
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
alias = f"weave-key-{unique_marker()}"
|
||||
key = client.key_with_alias(
|
||||
alias,
|
||||
models=[CHEAP_ANTHROPIC_MODEL],
|
||||
metadata=weave_creds.key_logging_metadata(),
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
marker = unique_marker()
|
||||
since = _window_start()
|
||||
outcome = first_ok(
|
||||
client,
|
||||
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=64),
|
||||
)
|
||||
assert outcome.response_cost is not None and outcome.response_cost > 0, (
|
||||
f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}"
|
||||
)
|
||||
|
||||
call = _exactly_one(
|
||||
weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="successful call"
|
||||
)
|
||||
|
||||
assert call.status_code == "OK", f"a successful call must land at OK span status, got {call.status_code!r}"
|
||||
cost = call.response_cost
|
||||
assert cost is not None and costs_agree(outcome.response_cost, cost), (
|
||||
f"the Weave call's llm.response.cost {cost!r} must agree with the header cost "
|
||||
f"{outcome.response_cost} - a delivered span with the wrong cost is a silent "
|
||||
"billing-attribution bug"
|
||||
)
|
||||
assert call.total_tokens is not None and call.total_tokens > 0, (
|
||||
f"the delivered call must carry token usage, got {call.total_tokens!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
|
||||
@pytest.mark.covers("logging.niche_integrations.failure.logs_spend", exercised_on=["chat_completions"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.OBSERVABILITY,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(UNREACHABLE_ANTHROPIC_BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_failed_chat_completions_delivers_one_error_call(
|
||||
self,
|
||||
client: LoggingClient,
|
||||
weave_creds: WeaveCreds,
|
||||
weave_reader: WeaveReader,
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
"""A deployment with an invalid upstream key passes proxy auth and fails
|
||||
at the provider, so exactly one provider failure exists for it. Proxy-side
|
||||
401s during key propagation never reach the provider and ship no payload,
|
||||
which is what the retry loop below relies on."""
|
||||
model_name = f"weave-err-{unique_marker()}"
|
||||
model_id = client.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(model=UNREACHABLE_ANTHROPIC_BACKEND, api_key=INVALID_UPSTREAM_API_KEY),
|
||||
)
|
||||
resources.defer(lambda: client.delete_model(model_id))
|
||||
key = client.key_with_alias(
|
||||
f"weave-err-key-{unique_marker()}",
|
||||
models=[model_name],
|
||||
metadata=weave_creds.key_logging_metadata(),
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
marker = unique_marker()
|
||||
since = _window_start()
|
||||
outcome = _provoke_provider_failure(client, key, model_name, marker)
|
||||
|
||||
call = _exactly_one(weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="failed call")
|
||||
|
||||
assert call.status_code == "ERROR", (
|
||||
f"a failed call must land at ERROR span status, got {call.status_code!r} - "
|
||||
"Weave's own summary.weave.status reads success either way, which is exactly "
|
||||
"why the span status is what this asserts on"
|
||||
)
|
||||
error = call.error
|
||||
assert error is not None and error.message is not None and "AnthropicException" in error.message, (
|
||||
f"the delivered call must carry the provider error, got {error!r}"
|
||||
)
|
||||
assert not call.response_cost, f"a failed call must not be billed, got llm.response.cost={call.response_cost!r}"
|
||||
assert outcome.status_code == 401, (
|
||||
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
|
||||
)
|
||||
|
||||
|
||||
def _provoke_provider_failure(client: LoggingClient, key: str, model_name: str, marker: str) -> StreamingResponse:
|
||||
"""Send until the provider (not the proxy) is the one rejecting the call.
|
||||
|
||||
A network failure between the test and the proxy is NOT retried: the request
|
||||
may have been served, and a retry would double-log the failure payload and
|
||||
falsely trip the exactly-one assertion.
|
||||
"""
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while True:
|
||||
outcome = client.chat_raw(key, model_name, f"trigger an upstream auth failure {marker}", max_tokens=16)
|
||||
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
|
||||
assert outcome.status_code != -1, (
|
||||
"network failure between the test and the proxy while provoking the provider failure; "
|
||||
"retrying now could double-log the failure payload and falsely trip the exactly-one "
|
||||
f"assertion - fix the rig connectivity first: {outcome.body[:200]}"
|
||||
)
|
||||
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
|
||||
break
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
assert "AnthropicException" in outcome.body, (
|
||||
"never saw the upstream provider failure before the deadline; the key may still be "
|
||||
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
|
||||
)
|
||||
return outcome
|
||||
|
|
@ -1,285 +0,0 @@
|
|||
"""Read-back for the Weave (Weights & Biases) logging tests against the real
|
||||
Weave project.
|
||||
|
||||
The proxy ships OTEL spans to https://trace.wandb.ai/otel/v1/traces with the
|
||||
``weave_otel`` callback, and the tests read the ingested calls back through
|
||||
Weave's own query API (``POST /calls/stream_query``), which answers JSON Lines:
|
||||
one JSON object per call, so the body is parsed line by line rather than as one
|
||||
document.
|
||||
|
||||
The project is shared with other traffic, so the read never relies on the target
|
||||
being among the newest N calls: the query is scoped server-side to the
|
||||
``litellm_request`` op and to calls that started after the test's own request,
|
||||
and pages with ``offset`` until the window is exhausted.
|
||||
|
||||
Weave's own ``summary.weave.status`` is a rollup that reads "success" even for a
|
||||
span the exporter marked failed, so status comes from the OTEL span itself
|
||||
(``attributes.otel_span.status.code``), and the shipped cost from
|
||||
``attributes.otel_span.attributes.llm.response.cost`` - the StandardLogging
|
||||
``response_cost``, which is what makes this a spend assertion rather than a
|
||||
delivery ping.
|
||||
|
||||
Missing configuration is a hard failure, never a skip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import count, takewhile
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
|
||||
from e2e_http import URL, AuthHeaders, send
|
||||
from e2e_metadata import step
|
||||
|
||||
_WEAVE_TRACE_API: Final = "https://trace.wandb.ai"
|
||||
|
||||
#: The op every litellm LLM call lands under. The proxy also exports a root
|
||||
#: server span ("Received Proxy Server Request") and management spans; only the
|
||||
#: LLM call carries the usage and cost this suite asserts on.
|
||||
LITELLM_REQUEST_OP: Final = "litellm_request"
|
||||
|
||||
#: How long to keep re-reading after the first matching call before trusting the
|
||||
#: exactly-one assertion. The OTEL batch exporter flushes on its own schedule, so
|
||||
#: a duplicate export can surface well after the first one, and a duplicate IS
|
||||
#: the bug being guarded against.
|
||||
WEAVE_SETTLE_SECONDS: Final = 45.0
|
||||
|
||||
#: Rows per page. The query is already scoped to this run's time window, so this
|
||||
#: only bounds one round trip, not what the read can see.
|
||||
_PAGE_SIZE: Final = 500
|
||||
|
||||
|
||||
class _WeaveSortBy(BaseModel):
|
||||
field: str
|
||||
direction: str
|
||||
|
||||
|
||||
class _WeaveOpFilter(BaseModel):
|
||||
op_names: list[str]
|
||||
|
||||
|
||||
class _WeaveGetField(BaseModel):
|
||||
get_field: str = Field(serialization_alias="$getField")
|
||||
|
||||
|
||||
class _WeaveLiteral(BaseModel):
|
||||
literal: float = Field(serialization_alias="$literal")
|
||||
|
||||
|
||||
class _WeaveGreaterThan(BaseModel):
|
||||
gt: tuple[_WeaveGetField, _WeaveLiteral] = Field(serialization_alias="$gt")
|
||||
|
||||
|
||||
class _WeaveQuery(BaseModel):
|
||||
expr: _WeaveGreaterThan = Field(serialization_alias="$expr")
|
||||
|
||||
|
||||
class _WeaveQueryBody(BaseModel):
|
||||
project_id: str
|
||||
filter: _WeaveOpFilter
|
||||
query: _WeaveQuery
|
||||
limit: int = _PAGE_SIZE
|
||||
offset: int = 0
|
||||
sort_by: list[_WeaveSortBy] = [_WeaveSortBy(field="started_at", direction="asc")]
|
||||
|
||||
|
||||
class _OtelStatus(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
code: str | None = None
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class _OtelError(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
code: str | None = None
|
||||
type: str | None = None
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class _LlmResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
cost: float | None = None
|
||||
|
||||
|
||||
class _LlmAttributes(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
response: _LlmResponse | None = None
|
||||
|
||||
|
||||
class _OtelSpanAttributes(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
llm: _LlmAttributes | None = None
|
||||
error: _OtelError | None = None
|
||||
|
||||
|
||||
class _OtelSpan(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
name: str | None = None
|
||||
status: _OtelStatus | None = None
|
||||
attributes: _OtelSpanAttributes | None = None
|
||||
|
||||
|
||||
class _CallAttributes(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
otel_span: _OtelSpan | None = None
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
total_tokens: int | None = None
|
||||
|
||||
|
||||
class _WeaveSummary(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
usage: dict[str, _Usage] = {}
|
||||
|
||||
|
||||
class WeaveCall(BaseModel):
|
||||
"""One ingested Weave call, reduced to what the scenarios assert on."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
id: str
|
||||
op_name: str
|
||||
started_at: str | None = None
|
||||
inputs: dict[str, object] = {}
|
||||
attributes: _CallAttributes | None = None
|
||||
summary: _WeaveSummary | None = Field(default=None)
|
||||
|
||||
@property
|
||||
def op(self) -> str:
|
||||
"""The bare op name out of ``weave:///<entity>/<project>/op/<op>:<digest>``."""
|
||||
return self.op_name.split("/op/")[-1].split(":")[0]
|
||||
|
||||
@property
|
||||
def status_code(self) -> str | None:
|
||||
"""The OTEL span status, not Weave's own rollup (which reads "success"
|
||||
even for a span the exporter marked ERROR)."""
|
||||
span = self.attributes.otel_span if self.attributes else None
|
||||
return span.status.code if span and span.status else None
|
||||
|
||||
@property
|
||||
def error(self) -> _OtelError | None:
|
||||
span = self.attributes.otel_span if self.attributes else None
|
||||
return span.attributes.error if span and span.attributes else None
|
||||
|
||||
@property
|
||||
def response_cost(self) -> float | None:
|
||||
span = self.attributes.otel_span if self.attributes else None
|
||||
llm = span.attributes.llm if span and span.attributes else None
|
||||
return llm.response.cost if llm and llm.response else None
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int | None:
|
||||
"""Weave keys usage by model, so the total is summed across whatever
|
||||
models the call reported."""
|
||||
if not self.summary or not self.summary.usage:
|
||||
return None
|
||||
totals = [usage.total_tokens for usage in self.summary.usage.values() if usage.total_tokens is not None]
|
||||
return sum(totals) if totals else None
|
||||
|
||||
def mentions(self, needle: str) -> bool:
|
||||
return needle in json.dumps(self.inputs, default=str)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WeaveReader:
|
||||
project_id: str
|
||||
api_key: str = field(repr=False)
|
||||
|
||||
@property
|
||||
def _headers(self) -> AuthHeaders:
|
||||
"""Weave authenticates with HTTP Basic as the fixed user ``api``."""
|
||||
token = base64.b64encode(f"api:{self.api_key}".encode()).decode()
|
||||
return AuthHeaders(authorization=f"Basic {token}")
|
||||
|
||||
def _query_body(self, *, since: float, offset: int, op: str) -> _WeaveQueryBody:
|
||||
return _WeaveQueryBody(
|
||||
project_id=self.project_id,
|
||||
filter=_WeaveOpFilter(op_names=[f"weave:///{self.project_id}/op/{op}:*"]),
|
||||
query=_WeaveQuery(
|
||||
expr=_WeaveGreaterThan(gt=(_WeaveGetField(get_field="started_at"), _WeaveLiteral(literal=since)))
|
||||
),
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
def _page(self, *, since: float, offset: int, op: str) -> tuple[WeaveCall, ...]:
|
||||
outcome = send(
|
||||
URL(f"{_WEAVE_TRACE_API}/calls/stream_query"),
|
||||
headers=self._headers,
|
||||
json=self._query_body(since=since, offset=offset, op=op),
|
||||
)
|
||||
if not outcome.ok:
|
||||
pytest.fail(
|
||||
f"Weave calls query for project {self.project_id!r} failed "
|
||||
f"({outcome.status_code}): {outcome.body[:300]}"
|
||||
)
|
||||
return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip())
|
||||
|
||||
@step("Read the Weave {op} calls carrying the marker {marker}")
|
||||
def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
|
||||
"""Every call under ``op`` started after ``since`` whose inputs carry
|
||||
``marker``, paging until the window is exhausted.
|
||||
|
||||
More than one is the duplicate-delivery bug, so this never collapses to a
|
||||
single call.
|
||||
"""
|
||||
pages = tuple(
|
||||
takewhile(
|
||||
bool,
|
||||
(self._page(since=since, offset=offset, op=op) for offset in count(0, _PAGE_SIZE)),
|
||||
)
|
||||
)
|
||||
return tuple(call for page in pages for call in page if call.mentions(marker))
|
||||
|
||||
@step("Wait for Weave to ingest a {op} call carrying the marker {marker}, then watch for duplicates")
|
||||
def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
|
||||
"""Poll until the call is readable, then keep re-reading for
|
||||
WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush
|
||||
cannot hide from the exactly-one assertion. A duplicate ends the settle
|
||||
early, because more waiting cannot clear it."""
|
||||
deadline = time.monotonic() + POLL_TIMEOUT
|
||||
while time.monotonic() < deadline:
|
||||
calls = self.calls_matching(marker, since=since, op=op)
|
||||
if calls:
|
||||
return self._settled(marker, since=since, op=op, first=calls)
|
||||
time.sleep(POLL_INTERVAL)
|
||||
return ()
|
||||
|
||||
def _settled(self, marker: str, *, since: float, op: str, first: tuple[WeaveCall, ...]) -> tuple[WeaveCall, ...]:
|
||||
"""A transiently empty re-read never downgrades what was already seen."""
|
||||
settle_deadline = time.monotonic() + WEAVE_SETTLE_SECONDS
|
||||
latest = first # rebind-ok: one settle window, re-read per poll interval
|
||||
while time.monotonic() < settle_deadline and len(latest) <= 1:
|
||||
time.sleep(POLL_INTERVAL)
|
||||
latest = self.calls_matching(marker, since=since, op=op) or latest
|
||||
return latest
|
||||
|
||||
|
||||
def build_weave_reader() -> WeaveReader:
|
||||
project_id = (os.environ.get("WEAVE_PROJECT_ID") or os.environ.get("WANDB_PROJECT_ID") or "").strip()
|
||||
api_key = os.environ.get("WANDB_API_KEY", "").strip()
|
||||
if not project_id or not api_key:
|
||||
pytest.fail(
|
||||
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
|
||||
"format <entity>/<project>): the test reads the proxy's weave_otel delivery "
|
||||
"back from the real Weave project; missing credentials is a hard failure, not a skip"
|
||||
)
|
||||
return WeaveReader(project_id=project_id, api_key=api_key)
|
||||
|
|
@ -45,8 +45,6 @@ class KeyLoggingCallbackVars(BaseModel):
|
|||
langfuse_public_key: str | None = Field(default=None, repr=False)
|
||||
langfuse_secret_key: str | None = Field(default=None, repr=False)
|
||||
langfuse_host: str | None = Field(default=None, repr=False)
|
||||
wandb_api_key: str | None = Field(default=None, repr=False)
|
||||
weave_project_id: str | None = None
|
||||
|
||||
|
||||
class KeyLoggingCallback(BaseModel):
|
||||
|
|
|
|||
145
tests/integration/_support/a2a_denial_probe.py
Normal file
145
tests/integration/_support/a2a_denial_probe.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
def _mapping(value: object) -> Mapping[str, object]:
|
||||
return value if isinstance(value, Mapping) else {}
|
||||
|
||||
|
||||
def _logged(hook: str, kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
logged: Final = _mapping(kwargs.get("standard_logging_object"))
|
||||
metadata: Final = _mapping(logged.get("metadata"))
|
||||
return {
|
||||
"hook": hook,
|
||||
"call_type": kwargs.get("call_type"),
|
||||
"litellm_call_id": kwargs.get("litellm_call_id"),
|
||||
"model": logged.get("model"),
|
||||
"status": logged.get("status"),
|
||||
"key_alias": metadata.get("user_api_key_alias"),
|
||||
"key_hash": metadata.get("user_api_key_hash"),
|
||||
"user_id": metadata.get("user_api_key_user_id"),
|
||||
"team_id": metadata.get("user_api_key_team_id"),
|
||||
}
|
||||
|
||||
|
||||
def _status(exception: Exception) -> int | None:
|
||||
status: Final = getattr(exception, "status_code", None)
|
||||
if isinstance(status, int):
|
||||
return status
|
||||
code: Final = getattr(exception, "code", None)
|
||||
return int(code) if isinstance(code, str) and code.isdigit() else None
|
||||
|
||||
|
||||
def _hooked(
|
||||
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, original_exception: Exception
|
||||
) -> Mapping[str, object]:
|
||||
return {
|
||||
"hook": "post_call_failure",
|
||||
"call_type": request_data.get("call_type"),
|
||||
"litellm_call_id": request_data.get("litellm_call_id"),
|
||||
"key_alias": user_api_key_dict.key_alias,
|
||||
"key_hash": user_api_key_dict.api_key,
|
||||
"user_id": user_api_key_dict.user_id,
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
"agent_id": request_data.get("agent_id"),
|
||||
"metadata_agent_id": _mapping(request_data.get("metadata")).get("agent_id"),
|
||||
"status_code": _status(original_exception),
|
||||
}
|
||||
|
||||
|
||||
def _send(record: Mapping[str, object]) -> None:
|
||||
with httpx.Client(trust_env=False) as client:
|
||||
client.post(os.environ["A2A_DENIAL_PROBE_URL"], content=json.dumps(record, default=str)).raise_for_status()
|
||||
|
||||
|
||||
async def _send_async(record: Mapping[str, object]) -> None:
|
||||
async with httpx.AsyncClient(trust_env=False) as client:
|
||||
response: Final = await client.post(os.environ["A2A_DENIAL_PROBE_URL"], content=json.dumps(record, default=str))
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
class A2ADenialProbe(CustomLogger):
|
||||
"""Reports the identity, target and outcome every logging and failure hook receives, one record per call."""
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
await _send_async(_logged("async_success", kwargs))
|
||||
|
||||
def log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send(_logged("sync_failure", kwargs))
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
await _send_async(_logged("async_failure", kwargs))
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
await _send_async(_hooked(request_data, user_api_key_dict, original_exception))
|
||||
|
||||
|
||||
class RaisingFailureLogger(CustomLogger):
|
||||
"""Fails inside both failure hooks."""
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
raise RuntimeError("failure logger is broken")
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
raise RuntimeError("failure hook is broken")
|
||||
|
||||
|
||||
class ReturningTeapotLogger(CustomLogger):
|
||||
"""Hands a replacement 418 back from the failure hook."""
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> HTTPException:
|
||||
return HTTPException(status_code=418, detail="transformed")
|
||||
|
||||
|
||||
class RaisingTeapotLogger(CustomLogger):
|
||||
"""Raises a replacement 418 from the failure hook."""
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
raise HTTPException(status_code=418, detail="transformed")
|
||||
|
||||
|
||||
probe: Final = A2ADenialProbe()
|
||||
raising_failure_logger: Final = RaisingFailureLogger()
|
||||
returning_teapot_logger: Final = ReturningTeapotLogger()
|
||||
raising_teapot_logger: Final = RaisingTeapotLogger()
|
||||
|
|
@ -29,6 +29,7 @@ DOCS_STREAMING_REFUSAL_URL: Final = (
|
|||
)
|
||||
|
||||
# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12.
|
||||
DOCS_REFUSAL_EXPLANATION: Final = "This request was declined because it could enable cyber harm."
|
||||
DOCS_REFUSAL: Final[Event] = {
|
||||
"id": "msg_01XFUDYJgAACzvnptvVoYEL",
|
||||
"type": "message",
|
||||
|
|
@ -39,7 +40,7 @@ DOCS_REFUSAL: Final[Event] = {
|
|||
"stop_details": {
|
||||
"type": "refusal",
|
||||
"category": "cyber",
|
||||
"explanation": "This request was declined because it could enable cyber harm.",
|
||||
"explanation": DOCS_REFUSAL_EXPLANATION,
|
||||
},
|
||||
"usage": {"input_tokens": 412, "output_tokens": 0},
|
||||
}
|
||||
|
|
@ -50,7 +51,7 @@ DOCS_PARTIAL_REFUSAL: Final[Event] = {
|
|||
"stop_details": {
|
||||
"type": "refusal",
|
||||
"category": "cyber",
|
||||
"explanation": "This request was declined because it could enable cyber harm.",
|
||||
"explanation": DOCS_REFUSAL_EXPLANATION,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from dataclasses import dataclass, field
|
|||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
from integration._support.asgi import asgi_server
|
||||
from integration._support.client import Gateway, Scenario
|
||||
|
|
@ -63,7 +64,9 @@ class Confirmation(BaseModel):
|
|||
confirmed: bool
|
||||
|
||||
|
||||
def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPServer:
|
||||
def math_service(
|
||||
name: str = "integration-math", *, rich: bool = False, on_slow: Callable[[], None] | None = None
|
||||
) -> MCPServer:
|
||||
service: Final = MCPServer(name)
|
||||
service.add_tool(add)
|
||||
service.add_tool(multiply)
|
||||
|
|
@ -77,6 +80,8 @@ def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPSe
|
|||
|
||||
@service.tool()
|
||||
async def slow(seconds: float) -> str:
|
||||
if on_slow is not None:
|
||||
on_slow()
|
||||
await asyncio.sleep(seconds)
|
||||
return "slept"
|
||||
|
||||
|
|
@ -114,6 +119,29 @@ def math_service(name: str = "integration-math", *, rich: bool = False) -> MCPSe
|
|||
return service
|
||||
|
||||
|
||||
def _cancel_on_disconnect(app: ASGIApp) -> ASGIApp:
|
||||
async def connected(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await app(scope, receive, send)
|
||||
return
|
||||
sender, reader = anyio.create_memory_object_stream[Message]()
|
||||
async with sender, reader, anyio.create_task_group() as tasks:
|
||||
|
||||
async def relay() -> None:
|
||||
while True:
|
||||
message: Final = await receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
tasks.cancel_scope.cancel()
|
||||
return
|
||||
await sender.send(message)
|
||||
|
||||
tasks.start_soon(relay)
|
||||
await app(scope, reader.receive, send)
|
||||
tasks.cancel_scope.cancel()
|
||||
|
||||
return connected
|
||||
|
||||
|
||||
def _capturing(app: Callable[[Scope, Receive, Send], object], observed: queue.Queue[dict[str, object]]):
|
||||
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
|
|
@ -247,8 +275,10 @@ def _draining_sse_watcher(app: Callable[[Scope, Receive, Send], object]):
|
|||
|
||||
|
||||
@contextmanager
|
||||
def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False) -> Iterator[McpPeer]:
|
||||
service: Final = math_service(rich=rich)
|
||||
def mcp_peer(
|
||||
transport: Literal["http", "sse"] = "http", *, rich: bool = False, on_slow: Callable[[], None] | None = None
|
||||
) -> Iterator[McpPeer]:
|
||||
service: Final = math_service(rich=rich, on_slow=on_slow)
|
||||
security: Final = TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
||||
app: Final = (
|
||||
_draining_sse_watcher(service.sse_app(transport_security=security))
|
||||
|
|
@ -256,7 +286,11 @@ def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False)
|
|||
else service.streamable_http_app(stateless_http=True, json_response=True, transport_security=security)
|
||||
)
|
||||
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
|
||||
with asgi_server(_capturing(app, observed), before_stop=_drain_sse_streams if transport == "sse" else None) as url:
|
||||
captured: Final = _capturing(app, observed)
|
||||
with asgi_server(
|
||||
_cancel_on_disconnect(captured) if transport == "http" else captured,
|
||||
before_stop=_drain_sse_streams if transport == "sse" else None,
|
||||
) as url:
|
||||
yield McpPeer(url + ("/sse" if transport == "sse" else "/mcp"), observed, transport)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
import itertools
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
|
|
@ -18,6 +19,22 @@ class _JsonRpcResult(BaseModel):
|
|||
result: dict[str, JsonValue]
|
||||
|
||||
|
||||
class _JsonRpcErrorBody(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
code: int
|
||||
message: str
|
||||
data: dict[str, JsonValue]
|
||||
|
||||
|
||||
class _JsonRpcError(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
jsonrpc: Literal["2.0"]
|
||||
id: str | int | None
|
||||
error: _JsonRpcErrorBody
|
||||
|
||||
|
||||
_STREAM_EVENTS: Final = (
|
||||
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "working"}, "final": False},
|
||||
{
|
||||
|
|
@ -49,6 +66,13 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
|||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"supportedInterfaces": [
|
||||
{
|
||||
"url": url + "/",
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "0.3",
|
||||
}
|
||||
],
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
|
|
@ -56,10 +80,15 @@ def _peer_card(url: str, name: str) -> dict[str, JsonValue]:
|
|||
}
|
||||
|
||||
|
||||
def _register_agent(gateway: Gateway, scenario: Scenario, name: str, url: str) -> str:
|
||||
created: Final = gateway.request(
|
||||
"POST", "/v1/agents", {"agent_name": name, "agent_card_params": _peer_card(url, name)}
|
||||
)
|
||||
def _register_agent(
|
||||
gateway: Gateway,
|
||||
scenario: Scenario,
|
||||
name: str,
|
||||
url: str,
|
||||
agent_card: dict[str, JsonValue] | None = None,
|
||||
) -> str:
|
||||
card: Final = _peer_card(url, name) if agent_card is None else agent_card
|
||||
created: Final = gateway.request("POST", "/v1/agents", {"agent_name": name, "agent_card_params": card})
|
||||
assert created.status_code == 200, created.text
|
||||
identity: Final = created.json()["agent_id"]
|
||||
|
||||
|
|
@ -163,6 +192,14 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
|
|||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
def jsonrpc_refusal(text: str, identity: str) -> bool:
|
||||
error: Final = _JsonRpcError.model_validate_json(text)
|
||||
return (error.id, error.error.code, error.error.message) == (
|
||||
marker,
|
||||
-32000,
|
||||
f"Agent '{identity}' is not allowed for your key/team. Contact proxy admin for access.",
|
||||
)
|
||||
|
||||
spellings: Final = (
|
||||
("send by id", lambda agent, name: f"/a2a/{agent}", send_payload, "message/send"),
|
||||
("send by name", lambda agent, name: f"/a2a/{name}", send_payload, "message/send"),
|
||||
|
|
@ -224,13 +261,13 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
|
|||
|
||||
for label, key in callers.items():
|
||||
for spelling, status, text in attempt(key, denied, denied_name):
|
||||
refused: Final = refusal(denied_name if "by name" in spelling else denied)
|
||||
assert (status, text) == (403, refused), f"{label} {spelling}: {status} {text}"
|
||||
identity: Final = denied_name if "by name" in spelling else denied
|
||||
assert status == 403 and jsonrpc_refusal(text, identity), f"{label} {spelling}: {status} {text}"
|
||||
for spelling, payload in denied_tasks:
|
||||
response: Final = gateway.client.post(
|
||||
f"/a2a/{denied}", headers={"Authorization": f"Bearer {key}"}, json=payload
|
||||
)
|
||||
assert (response.status_code, response.text) == (403, refusal(denied)), (
|
||||
assert response.status_code == 403 and jsonrpc_refusal(response.text, denied), (
|
||||
f"{label} {spelling}: {response.status_code} {response.text}"
|
||||
)
|
||||
assert card(key, denied) == (403, refusal(denied)), label
|
||||
|
|
@ -278,13 +315,72 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
|
|||
emptied: Final = gateway.request("PUT", f"/v1/access_group/{group_id}", {"access_agent_ids": []})
|
||||
assert emptied.status_code == 200, emptied.text
|
||||
emptied_outcomes: Final = attempt(callers["access group"], allowed, allowed_name)
|
||||
assert emptied_outcomes == tuple(
|
||||
(spelling, 403, refusal(allowed_name if "by name" in spelling else allowed))
|
||||
for spelling, _, _ in emptied_outcomes
|
||||
assert all(
|
||||
status == 403 and jsonrpc_refusal(text, allowed_name if "by name" in spelling else allowed)
|
||||
for spelling, status, text in emptied_outcomes
|
||||
), emptied_outcomes
|
||||
assert wire.drain() == (), emptied_outcomes
|
||||
|
||||
|
||||
def test_a2a_refusals_answer_as_jsonrpc_errors_that_echo_the_callers_request_id(gateway: Gateway) -> None:
|
||||
marker: Final = "a2arefuse" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
agent: Final = _register_agent(gateway, scenario, marker, wire.url)
|
||||
refused_callers: Final = (
|
||||
(
|
||||
"over budget",
|
||||
scenario.key(object_permission={"agents": [agent]}, max_budget=0.0001, spend=1.0),
|
||||
422,
|
||||
"Budget has been exceeded!",
|
||||
"budget_exceeded",
|
||||
),
|
||||
(
|
||||
"agent not granted",
|
||||
scenario.key(object_permission={"agents": [marker + "-other"]}),
|
||||
403,
|
||||
f"Agent '{agent}' is not allowed for your key/team. Contact proxy admin for access.",
|
||||
None,
|
||||
),
|
||||
("unknown key", "sk-" + uuid.uuid4().hex, 401, "Authentication Error", None),
|
||||
)
|
||||
message: Final = {
|
||||
"kind": "message",
|
||||
"role": "user",
|
||||
"messageId": marker + "-in",
|
||||
"parts": [{"kind": "text", "text": "ping"}],
|
||||
}
|
||||
requests: Final = tuple(
|
||||
(path, method, request_id)
|
||||
for path, (method, request_id) in itertools.product(
|
||||
(f"/a2a/{agent}", f"/a2a/{agent}/message/send", f"/v1/a2a/{agent}/message/send"),
|
||||
(("message/send", marker), ("message/stream", 7), ("SendMessage", "l10")),
|
||||
)
|
||||
)
|
||||
for label, key, status, message_start, error_type in refused_callers:
|
||||
for path, method, request_id in requests:
|
||||
response = gateway.client.post(
|
||||
path,
|
||||
headers={"Authorization": f"Bearer {key}", "Accept": "application/json, text/event-stream"},
|
||||
json={"jsonrpc": "2.0", "id": request_id, "method": method, "params": {"message": message}},
|
||||
)
|
||||
context = f"{label} {path} {method}: {response.status_code} {response.text}"
|
||||
assert response.status_code == status, context
|
||||
error = _JsonRpcError.model_validate_json(response.text)
|
||||
assert (error.id, error.error.code) == (request_id, -32000), context
|
||||
assert error.error.message.startswith(message_start), context
|
||||
assert error_type is None or error.error.data["type"] == error_type, context
|
||||
unparseable: Final = gateway.client.post(
|
||||
f"/a2a/{agent}", headers={"Authorization": "Bearer sk-" + uuid.uuid4().hex}, content=b"not json"
|
||||
)
|
||||
assert unparseable.status_code == 400, unparseable.text
|
||||
assert _JsonRpcError.model_validate_json(unparseable.text).id is None, unparseable.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: GET /a2a/{agent_id}/.well-known/agent.json returns 401 'Only proxy admin can be used' for non-admin keys because agent_inference_routes lists only agent-card.json"
|
||||
|
|
@ -332,3 +428,68 @@ def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys
|
|||
)
|
||||
assert (denied_legacy.status_code, denied_legacy.text) == (403, refusal(denied)), denied_legacy.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_agent_list_and_detail_cards_use_proxy_urls_for_scoped_keys(gateway: Gateway) -> None:
|
||||
marker: Final = "a2acardurl" + uuid.uuid4().hex[:12]
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
assert request.method == "GET", request.method
|
||||
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
backend_url: Final = wire.url + "/"
|
||||
peer_card: Final = _peer_card(wire.url, marker)
|
||||
peer_interfaces: Final = peer_card["supportedInterfaces"]
|
||||
assert isinstance(peer_interfaces, list)
|
||||
agent_card: Final = {
|
||||
**peer_card,
|
||||
"additionalInterfaces": [{"url": backend_url, "transport": "GRPC"}],
|
||||
"supportedInterfaces": [
|
||||
*peer_interfaces,
|
||||
{
|
||||
"url": backend_url,
|
||||
"protocolBinding": "GRPC",
|
||||
"protocolVersion": "0.3",
|
||||
},
|
||||
],
|
||||
}
|
||||
agent_id: Final = _register_agent(gateway, scenario, marker, wire.url, agent_card)
|
||||
user_id: Final = scenario.user(user_role="internal_user")
|
||||
scoped_key: Final = scenario.key(user_id=user_id, object_permission={"agents": [agent_id]})
|
||||
|
||||
public_card: Final = gateway.request(
|
||||
"GET",
|
||||
f"/a2a/{agent_id}/.well-known/agent-card.json",
|
||||
key=scoped_key,
|
||||
)
|
||||
listed: Final = gateway.request("GET", "/v1/agents", key=scoped_key)
|
||||
detailed: Final = gateway.request("GET", f"/v1/agents/{agent_id}", key=scoped_key)
|
||||
|
||||
assert public_card.status_code == 200, public_card.text
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert detailed.status_code == 200, detailed.text
|
||||
assert wire.url not in listed.text
|
||||
assert wire.url not in detailed.text
|
||||
assert "additionalInterfaces" not in listed.text
|
||||
assert "additionalInterfaces" not in detailed.text
|
||||
public_urls: Final = (
|
||||
public_card.json()["url"],
|
||||
public_card.json()["supportedInterfaces"][0]["url"],
|
||||
)
|
||||
listed_agent: Final = next(item for item in listed.json() if item["agent_id"] == agent_id)
|
||||
for body in (listed_agent, detailed.json()):
|
||||
card: Final = body["agent_card_params"]
|
||||
assert (card["url"], card["supportedInterfaces"][0]["url"]) == public_urls
|
||||
assert card["supportedInterfaces"] == [
|
||||
{
|
||||
"url": public_urls[1],
|
||||
"protocolBinding": "JSONRPC",
|
||||
"protocolVersion": "0.3",
|
||||
}
|
||||
]
|
||||
|
||||
admin_listed: Final = gateway.request("GET", "/v1/agents")
|
||||
assert admin_listed.status_code == 200, admin_listed.text
|
||||
admin_agent: Final = next(item for item in admin_listed.json() if item["agent_id"] == agent_id)
|
||||
assert admin_agent["agent_card_params"]["url"] == backend_url
|
||||
|
|
|
|||
809
tests/integration/authorization/test_agent_denial_logging.py
Normal file
809
tests/integration/authorization/test_agent_denial_logging.py
Normal file
|
|
@ -0,0 +1,809 @@
|
|||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
held,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
pytestmark = pytest.mark.timeout(900)
|
||||
|
||||
_PROBE: Final = "integration._support.a2a_denial_probe."
|
||||
_SETTLE: Final = 3.0
|
||||
_CONVERGE: Final = 30.0
|
||||
_NO_SETTINGS: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
_ROWS: Final = (
|
||||
'SELECT request_id, call_type, status, model, agent_id, spend, api_key, "user", team_id, metadata '
|
||||
'FROM "LiteLLM_SpendLogs" WHERE api_key = %s'
|
||||
)
|
||||
_ONE_FAILURE: Final = ("async_failure", "post_call_failure")
|
||||
_STREAM_METHODS: Final = frozenset({"message/stream", "SendStreamingMessage"})
|
||||
|
||||
|
||||
class _Event(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
hook: str
|
||||
call_type: str | None = None
|
||||
litellm_call_id: str | None = None
|
||||
model: str | None = None
|
||||
status: str | None = None
|
||||
key_alias: str | None = None
|
||||
key_hash: str | None = None
|
||||
user_id: str | None = None
|
||||
team_id: str | None = None
|
||||
agent_id: str | None = None
|
||||
metadata_agent_id: str | None = None
|
||||
status_code: int | None = None
|
||||
|
||||
|
||||
class _Events:
|
||||
def __init__(self, sink: Wire) -> None:
|
||||
self.sink: Final = sink
|
||||
self.seen: tuple[_Event, ...] = ()
|
||||
|
||||
def of(self, alias: str) -> tuple[_Event, ...]:
|
||||
self.seen = (*self.seen, *(_Event.model_validate_json(request.body) for request in self.sink.drain()))
|
||||
return tuple(event for event in self.seen if event.key_alias == alias)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Agent:
|
||||
agent_id: str
|
||||
name: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
proxy: Gateway
|
||||
events: _Events
|
||||
upstream: Wire
|
||||
allowed: _Agent
|
||||
denied: _Agent
|
||||
denied_legacy: _Agent
|
||||
bridge: _Agent
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Caller:
|
||||
key: str
|
||||
alias: str
|
||||
user_id: str
|
||||
team_id: str | None
|
||||
|
||||
@property
|
||||
def hashed(self) -> str:
|
||||
return hashlib.sha256(self.key.encode()).hexdigest()
|
||||
|
||||
|
||||
def _compact(content: Mapping[str, JsonValue]) -> bytes:
|
||||
return json.dumps(content, separators=(",", ":")).encode()
|
||||
|
||||
|
||||
def _refusal(agent: str, response: httpx.Response | None = None) -> bytes:
|
||||
message: Final = f"Agent '{agent}' is not allowed for your key/team. Contact proxy admin for access."
|
||||
if response is None:
|
||||
return _compact({"detail": message})
|
||||
return _compact(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": JSON_OBJECT.validate_json(response.request.content)["id"],
|
||||
"error": {
|
||||
"code": -32000,
|
||||
"message": message,
|
||||
"data": {"type": "permission_error", "param": None, "code": "403"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _not_found(request_id: JsonValue, agent: str) -> bytes:
|
||||
return _compact(
|
||||
{"jsonrpc": "2.0", "id": request_id, "error": {"code": -32000, "message": f"Agent '{agent}' not found"}}
|
||||
)
|
||||
|
||||
|
||||
def _card(url: str, name: str, version: str) -> dict[str, JsonValue]:
|
||||
card: Final[dict[str, JsonValue]] = {
|
||||
"protocolVersion": version,
|
||||
"name": name,
|
||||
"description": "Synthetic peer",
|
||||
"version": "1.0.0",
|
||||
"url": url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
if version == "0.3":
|
||||
return card
|
||||
return {**card, "supportedInterfaces": [{"url": url + "/", "protocolBinding": "JSONRPC", "protocolVersion": "1.0"}]}
|
||||
|
||||
|
||||
def _peer(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
return Reply(body=_compact(_card("http://" + request.headers["host"], "peer", "0.3")))
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
if body.get("method") == "message/stream":
|
||||
frame: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"kind": "status-update",
|
||||
"taskId": "t1",
|
||||
"contextId": "c1",
|
||||
"status": {"state": "completed"},
|
||||
"final": True,
|
||||
},
|
||||
}
|
||||
return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(frame)}\n\n".encode(),))
|
||||
return Reply(
|
||||
body=_compact(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": body.get("id"),
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": "peer-message",
|
||||
"parts": [{"kind": "text", "text": "granted"}],
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _register(proxy: Gateway, scenario: Scenario, body: Mapping[str, JsonValue]) -> _Agent:
|
||||
created: Final = proxy.request("POST", "/v1/agents", body)
|
||||
assert created.status_code == 200, created.text
|
||||
agent: Final = _Agent(string_value(JSON_OBJECT.validate_json(created.content)["agent_id"]), str(body["agent_name"]))
|
||||
|
||||
def cleanup() -> None:
|
||||
deleted: Final = proxy.request("DELETE", f"/v1/agents/{agent.agent_id}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
|
||||
scenario.cleanups.callback(cleanup)
|
||||
return agent
|
||||
|
||||
|
||||
def _config(directory: Path, callbacks: tuple[str, ...], general: Mapping[str, JsonValue]) -> Path:
|
||||
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = {
|
||||
**base,
|
||||
"general_settings": {**object_value(base["general_settings"]), **general},
|
||||
"litellm_settings": {
|
||||
**object_value(base["litellm_settings"]),
|
||||
"callbacks": [_PROBE + "probe", *(_PROBE + name for name in callbacks)],
|
||||
},
|
||||
}
|
||||
path: Final = directory / "a2a-denials.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _owned_rig(
|
||||
directory: Path,
|
||||
callbacks: tuple[str, ...] = (),
|
||||
general: Mapping[str, JsonValue] = _NO_SETTINGS,
|
||||
environment: Mapping[str, str] = MappingProxyType({}),
|
||||
) -> Iterator[_Rig]:
|
||||
marker: Final = "a2adeny" + uuid.uuid4().hex[:10]
|
||||
with (
|
||||
wire_server(_peer) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(
|
||||
shared,
|
||||
directory,
|
||||
{"A2A_DENIAL_PROBE_URL": sink.url, **environment},
|
||||
config=_config(directory, callbacks, general),
|
||||
) as proxy,
|
||||
proxy.scenario() as scenario,
|
||||
):
|
||||
|
||||
def agent(suffix: str, version: str, **extra: JsonValue) -> _Agent:
|
||||
name: Final = f"{marker}-{suffix}"
|
||||
return _register(
|
||||
proxy, scenario, {"agent_name": name, "agent_card_params": _card(upstream.url, name, version), **extra}
|
||||
)
|
||||
|
||||
yield _Rig(
|
||||
proxy,
|
||||
_Events(sink),
|
||||
upstream,
|
||||
allowed=agent("allowed", "0.3"),
|
||||
denied=agent("denied", "1.0"),
|
||||
denied_legacy=agent("legacy", "0.3"),
|
||||
bridge=agent(
|
||||
"bridge",
|
||||
"0.3",
|
||||
litellm_params={
|
||||
"custom_llm_provider": "openai",
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"api_base": upstream.url + "/v1",
|
||||
"api_key": "synthetic-bridge-key",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
with _owned_rig(tmp_path_factory.mktemp("a2a-denials")) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=["raising_failure_logger,returning_teapot_logger", "raising_teapot_logger"])
|
||||
def hostile(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
callbacks: Final = tuple(str(request.param).split(","))
|
||||
with _owned_rig(tmp_path_factory.mktemp("a2a-denials-hostile"), callbacks) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=["disable_pre_routing_error_logs", "disable_error_logs"])
|
||||
def flagged(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
with _owned_rig(tmp_path_factory.mktemp("a2a-denials-flagged"), general={str(request.param): True}) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def alerting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[_Rig, Wire]]:
|
||||
with wire_server(lambda request: Reply(body=b"ok", content_type="text/plain")) as alerts:
|
||||
with _owned_rig(
|
||||
tmp_path_factory.mktemp("a2a-denials-alerting"),
|
||||
general={"alerting": ["slack"], "alert_types": ["llm_exceptions"]},
|
||||
environment={"SLACK_WEBHOOK_URL": alerts.url, "DEFAULT_FLUSH_INTERVAL_SECONDS": "1"},
|
||||
) as owned:
|
||||
yield owned, alerts
|
||||
|
||||
|
||||
def _caller(rig: _Rig, scenario: Scenario, grant: str = "team") -> _Caller:
|
||||
alias: Final = "a2adeny-" + uuid.uuid4().hex[:12]
|
||||
user: Final = scenario.user()
|
||||
allowed: Final[list[JsonValue]] = [rig.allowed.agent_id]
|
||||
if grant == "personal":
|
||||
return _Caller(
|
||||
scenario.key(user_id=user, key_alias=alias, object_permission={"agents": allowed}), alias, user, None
|
||||
)
|
||||
if grant == "access group":
|
||||
created: Final = rig.proxy.request(
|
||||
"POST", "/v1/access_group", {"access_group_name": alias, "access_agent_ids": allowed}
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
group: Final = string_value(JSON_OBJECT.validate_json(created.content)["access_group_id"])
|
||||
scenario.cleanups.callback(rig.proxy.request, "DELETE", f"/v1/access_group/{group}")
|
||||
return _Caller(scenario.key(user_id=user, key_alias=alias, access_group_ids=[group]), alias, user, None)
|
||||
team: Final = scenario.team(object_permission={"agents": allowed})
|
||||
return _Caller(scenario.key(user_id=user, team_id=team, key_alias=alias), alias, user, team)
|
||||
|
||||
|
||||
def _params(method: str, message_id: str, version: str | None) -> dict[str, JsonValue]:
|
||||
if method.startswith("tasks/pushNotificationConfig"):
|
||||
return {"taskId": "t1", "pushNotificationConfig": {"url": "https://callback.example.com/hook"}}
|
||||
if method in {"tasks/get", "tasks/cancel", "GetTask", "CancelTask"}:
|
||||
return {"id": "t1"}
|
||||
if version == "1.0":
|
||||
return {"message": {"messageId": message_id, "role": "ROLE_USER", "parts": [{"text": "x"}]}}
|
||||
return {
|
||||
"message": {
|
||||
"kind": "message",
|
||||
"messageId": message_id,
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "x"}],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _send(
|
||||
rig: _Rig,
|
||||
key: str,
|
||||
path: str,
|
||||
*,
|
||||
method: str = "SendMessage",
|
||||
version: str | None = "1.0",
|
||||
request_id: JsonValue = None,
|
||||
message_id: str | None = None,
|
||||
extra: Mapping[str, JsonValue] = _NO_SETTINGS,
|
||||
) -> httpx.Response:
|
||||
headers: Final = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
**({"A2A-Version": version} if version is not None else {}),
|
||||
**({"Accept": "text/event-stream"} if method in _STREAM_METHODS else {}),
|
||||
}
|
||||
body: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": "r-" + uuid.uuid4().hex if request_id is None else request_id,
|
||||
"method": method,
|
||||
"params": _params(method, message_id or uuid.uuid4().hex, version),
|
||||
**extra,
|
||||
}
|
||||
return rig.proxy.client.post(path, headers=headers, json=body)
|
||||
|
||||
|
||||
def _assert_never_dispatched(rig: _Rig, scenario: Scenario, denied_message: str) -> None:
|
||||
control_key: Final = scenario.key(object_permission={"agents": [rig.allowed.agent_id]})
|
||||
control_message: Final = "control-" + uuid.uuid4().hex
|
||||
granted: Final = _send(
|
||||
rig,
|
||||
control_key,
|
||||
f"/a2a/{rig.allowed.agent_id}",
|
||||
method="message/send",
|
||||
version=None,
|
||||
message_id=control_message,
|
||||
)
|
||||
assert granted.status_code == 200, granted.text
|
||||
received: Final = tuple(request.body.decode() for request in rig.upstream.drain())
|
||||
assert any(control_message in body for body in received), received
|
||||
assert not any(denied_message in body for body in received), received
|
||||
|
||||
|
||||
def _metadata(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
value: Final = row["metadata"]
|
||||
return JSON_OBJECT.validate_json(value) if isinstance(value, str) else object_value(value)
|
||||
|
||||
|
||||
def _rows(rig: _Rig, caller: _Caller) -> list[dict[str, JsonValue]]:
|
||||
return rig.proxy.rows(_ROWS, (caller.hashed,))
|
||||
|
||||
|
||||
def _hooks(events: tuple[_Event, ...]) -> tuple[str, ...]:
|
||||
return tuple(sorted(event.hook for event in events))
|
||||
|
||||
|
||||
def _spend(rig: _Rig, caller: _Caller) -> tuple[JsonValue, ...]:
|
||||
key: Final = object_value(rig.proxy.get("/key/info", {"key": caller.key})["info"])["spend"]
|
||||
user: Final = object_value(rig.proxy.get("/user/info", {"user_id": caller.user_id})["user_info"])["spend"]
|
||||
if caller.team_id is None:
|
||||
return key, user
|
||||
team: Final = object_value(rig.proxy.get("/team/info", {"team_id": caller.team_id})["team_info"])["spend"]
|
||||
return key, user, team
|
||||
|
||||
|
||||
def _assert_one_failure(rig: _Rig, caller: _Caller, *, status: int, agent_id: str, model: str) -> dict[str, JsonValue]:
|
||||
events: Final = held(
|
||||
lambda: rig.events.of(caller.alias),
|
||||
lambda found: _hooks(found) == _ONE_FAILURE,
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
logged: Final = next(event for event in events if event.hook == "async_failure")
|
||||
assert (logged.call_type, logged.status, logged.model) == ("asend_message", "failure", model), logged
|
||||
assert (logged.key_hash, logged.user_id, logged.team_id) == (caller.hashed, caller.user_id, caller.team_id), logged
|
||||
hooked: Final = next(event for event in events if event.hook == "post_call_failure")
|
||||
assert (hooked.status_code, hooked.agent_id, hooked.key_hash) == (status, agent_id, caller.hashed), hooked
|
||||
assert (hooked.user_id, hooked.team_id) == (caller.user_id, caller.team_id), hooked
|
||||
rows: Final = held(lambda: _rows(rig, caller), lambda found: len(found) == 1, holding=_SETTLE, seconds=_CONVERGE)
|
||||
row: Final = rows[0]
|
||||
observed: Final = (row["status"], row["call_type"], row["agent_id"], row["api_key"], row["user"], row["spend"])
|
||||
assert observed == ("failure", "asend_message", agent_id, caller.hashed, caller.user_id, 0.0), row
|
||||
assert (row["team_id"] or None) == caller.team_id, row
|
||||
assert object_value(_metadata(row)["error_information"])["error_code"] == str(status), row
|
||||
assert all(spend == 0.0 for spend in _spend(rig, caller)), _spend(rig, caller)
|
||||
return row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method", "version", "legacy"),
|
||||
[
|
||||
pytest.param("/a2a/{}", "SendMessage", "1.0", False, id="regression-send-1.0"),
|
||||
pytest.param("/a2a/{}/message/send", "SendMessage", "1.0", False, id="regression-send-alias"),
|
||||
pytest.param("/v1/a2a/{}/message/send", "SendMessage", "1.0", False, id="regression-send-v1-alias"),
|
||||
pytest.param("/a2a/{}", "SendStreamingMessage", "1.0", False, id="regression-stream-1.0"),
|
||||
pytest.param("/a2a/{}", "GetTask", "1.0", False, id="regression-get-task-1.0"),
|
||||
pytest.param("/a2a/{}", "CancelTask", "1.0", False, id="regression-cancel-task-1.0"),
|
||||
pytest.param("/a2a/{}", "message/send", None, True, id="regression-send-0.3"),
|
||||
pytest.param("/a2a/{}", "message/stream", None, True, id="regression-stream-0.3"),
|
||||
pytest.param("/a2a/{}", "tasks/get", None, True, id="regression-tasks-get-0.3"),
|
||||
pytest.param("/a2a/{}", "tasks/cancel", None, True, id="regression-tasks-cancel-0.3"),
|
||||
pytest.param("/a2a/{}", "tasks/pushNotificationConfig/set", None, True, id="regression-push-set-0.3"),
|
||||
],
|
||||
)
|
||||
def test_denied_agent_keeps_its_403_logs_one_zero_spend_failure_for_the_caller_and_never_dispatches(
|
||||
rig: _Rig, route: str, method: str, version: str | None, legacy: bool
|
||||
) -> None:
|
||||
target: Final = rig.denied_legacy if legacy else rig.denied
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
message: Final = "denied-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig, caller.key, route.format(target.agent_id), method=method, version=version, message_id=message
|
||||
)
|
||||
assert (response.status_code, response.content) == (403, _refusal(target.agent_id, response))
|
||||
assert response.headers["content-type"] == "application/json"
|
||||
_assert_never_dispatched(rig, scenario, message)
|
||||
_assert_one_failure(rig, caller, status=403, agent_id=target.agent_id, model=f"a2a_agent/{target.name}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("grant", ["personal", "team", "access group"], ids=lambda grant: "regression-" + grant)
|
||||
def test_denied_agent_failure_is_attributed_to_the_authenticated_key_for_every_grant_kind(
|
||||
rig: _Rig, grant: str
|
||||
) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario, grant)
|
||||
granted: Final = _send(rig, caller.key, f"/a2a/{rig.allowed.agent_id}", method="message/send", version=None)
|
||||
assert granted.status_code == 200, granted.text
|
||||
response: Final = _send(rig, caller.key, f"/a2a/{rig.denied.agent_id}")
|
||||
assert (response.status_code, response.content) == (403, _refusal(rig.denied.agent_id, response))
|
||||
events: Final = held(
|
||||
lambda: rig.events.of(caller.alias),
|
||||
lambda found: _hooks(found) == tuple(sorted((*_ONE_FAILURE, "async_success"))),
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
failure: Final = next(event for event in events if event.hook == "async_failure")
|
||||
assert (failure.key_hash, failure.user_id, failure.team_id) == (caller.hashed, caller.user_id, caller.team_id)
|
||||
rows: Final = held(
|
||||
lambda: _rows(rig, caller), lambda found: len(found) == 2, holding=_SETTLE, seconds=_CONVERGE
|
||||
)
|
||||
failed: Final = tuple(row for row in rows if row["status"] == "failure")
|
||||
assert [(row["agent_id"], row["user"], row["spend"]) for row in failed] == [
|
||||
(rig.denied.agent_id, caller.user_id, 0.0)
|
||||
], rows
|
||||
|
||||
|
||||
def test_denied_completion_bridge_agent_logs_one_failure_and_never_reaches_the_bridged_model(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
message: Final = "bridge-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig, caller.key, f"/a2a/{rig.bridge.agent_id}", method="message/send", version=None, message_id=message
|
||||
)
|
||||
assert (response.status_code, response.content) == (403, _refusal(rig.bridge.agent_id, response))
|
||||
_assert_never_dispatched(rig, scenario, message)
|
||||
_assert_one_failure(rig, caller, status=403, agent_id=rig.bridge.agent_id, model=f"a2a_agent/{rig.bridge.name}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "method", "version"),
|
||||
[
|
||||
pytest.param("/a2a/{}", "SendMessage", "1.0", id="regression-send"),
|
||||
pytest.param("/a2a/{}/message/send", "SendMessage", "1.0", id="regression-send-alias"),
|
||||
pytest.param("/v1/a2a/{}/message/send", "SendMessage", "1.0", id="regression-send-v1-alias"),
|
||||
pytest.param("/a2a/{}", "message/stream", None, id="regression-stream-0.3"),
|
||||
],
|
||||
)
|
||||
def test_unknown_agent_keeps_its_jsonrpc_404_and_logs_one_zero_spend_failure_for_the_caller(
|
||||
rig: _Rig, route: str, method: str, version: str | None
|
||||
) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
request_id: Final = "u-" + uuid.uuid4().hex
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
message: Final = "unknown-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig,
|
||||
caller.key,
|
||||
route.format(unknown),
|
||||
method=method,
|
||||
version=version,
|
||||
request_id=request_id,
|
||||
message_id=message,
|
||||
)
|
||||
assert (response.status_code, response.content) == (404, _not_found(request_id, unknown))
|
||||
assert response.headers["content-type"] == "application/json"
|
||||
_assert_never_dispatched(rig, scenario, message)
|
||||
_assert_one_failure(rig, caller, status=404, agent_id=unknown, model=f"a2a_agent/{unknown}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("request_id", [0, 7, "abc"], ids=lambda value: f"regression-id-{value!r}")
|
||||
def test_unknown_agent_echoes_the_jsonrpc_id_with_its_type_and_logs_one_failure(
|
||||
rig: _Rig, request_id: JsonValue
|
||||
) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
response: Final = _send(rig, caller.key, f"/a2a/{unknown}", request_id=request_id)
|
||||
assert (response.status_code, response.content) == (404, _not_found(request_id, unknown))
|
||||
assert type(JSON_OBJECT.validate_json(response.content)["id"]) is type(request_id)
|
||||
_assert_one_failure(rig, caller, status=404, agent_id=unknown, model=f"a2a_agent/{unknown}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["denied", "unknown"], ids=lambda outcome: "regression-" + outcome)
|
||||
def test_client_supplied_attribution_cannot_rewrite_a_denial_failure_row(rig: _Rig, outcome: str) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
target: Final = rig.denied.agent_id if outcome == "denied" else unknown
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
other: Final = _caller(rig, scenario)
|
||||
success_id: Final = "s-" + uuid.uuid4().hex
|
||||
granted: Final = _send(
|
||||
rig, other.key, f"/a2a/{rig.allowed.agent_id}", method="message/send", version=None, request_id=success_id
|
||||
)
|
||||
assert granted.status_code == 200, granted.text
|
||||
success_rows: Final = eventually(lambda: _rows(rig, other), lambda found: len(found) == 1, seconds=_CONVERGE)
|
||||
success_request: Final = success_rows[0]["request_id"]
|
||||
request_id: Final = "spoof-" + uuid.uuid4().hex
|
||||
spoof: Final[dict[str, JsonValue]] = {
|
||||
"agent_id": "spoof",
|
||||
"metadata": {
|
||||
"agent_id": "spoof",
|
||||
"user_api_key_alias": "spoof",
|
||||
"user_api_key_team_id": "spoof",
|
||||
"tags": ["x"],
|
||||
},
|
||||
"litellm_call_id": success_request,
|
||||
}
|
||||
body: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": "SendMessage",
|
||||
"params": {**_params("SendMessage", uuid.uuid4().hex, "1.0"), "litellm_call_id": "spoof2"},
|
||||
**spoof,
|
||||
}
|
||||
response: Final = rig.proxy.client.post(
|
||||
f"/a2a/{target}", headers={"Authorization": f"Bearer {caller.key}", "A2A-Version": "1.0"}, json=body
|
||||
)
|
||||
expected: Final = _refusal(target, response) if outcome == "denied" else _not_found(request_id, unknown)
|
||||
assert (response.status_code, response.content) == (403 if outcome == "denied" else 404, expected)
|
||||
model: Final = f"a2a_agent/{rig.denied.name}" if outcome == "denied" else f"a2a_agent/{unknown}"
|
||||
row: Final = _assert_one_failure(rig, caller, status=response.status_code, agent_id=target, model=model)
|
||||
assert row["request_id"] not in {success_request, "spoof2", request_id}, row
|
||||
metadata: Final = _metadata(row)
|
||||
assert (metadata["user_api_key_alias"], metadata["user_api_key_team_id"]) == (caller.alias, caller.team_id)
|
||||
assert rig.events.of("spoof") == ()
|
||||
after: Final = _rows(rig, other)
|
||||
assert [(item["request_id"], item["status"]) for item in after] == [(success_request, "success")], after
|
||||
|
||||
|
||||
def test_two_keys_reusing_one_jsonrpc_id_get_one_denial_row_each(rig: _Rig) -> None:
|
||||
shared_id: Final = "same-" + uuid.uuid4().hex
|
||||
with rig.proxy.scenario() as scenario:
|
||||
callers: Final = (_caller(rig, scenario), _caller(rig, scenario))
|
||||
for caller in callers:
|
||||
response = _send(rig, caller.key, f"/a2a/{rig.denied.agent_id}", request_id=shared_id)
|
||||
assert (response.status_code, response.content) == (403, _refusal(rig.denied.agent_id, response))
|
||||
rows: Final = tuple(
|
||||
_assert_one_failure(
|
||||
rig, caller, status=403, agent_id=rig.denied.agent_id, model=f"a2a_agent/{rig.denied.name}"
|
||||
)
|
||||
for caller in callers
|
||||
)
|
||||
assert rows[0]["request_id"] != rows[1]["request_id"], rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["denied", "unknown"])
|
||||
def test_failing_or_transforming_failure_callbacks_never_change_the_denial_response(
|
||||
hostile: _Rig, outcome: str
|
||||
) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
request_id: Final = "h-" + uuid.uuid4().hex
|
||||
target: Final = hostile.denied.agent_id if outcome == "denied" else unknown
|
||||
with hostile.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(hostile, scenario)
|
||||
response: Final = _send(hostile, caller.key, f"/a2a/{target}", request_id=request_id)
|
||||
expected: Final = (403, _refusal(target, response)) if outcome == "denied" else (404, _not_found(request_id, unknown))
|
||||
assert (response.status_code, response.content) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["denied", "unknown"], ids=lambda outcome: "regression-" + outcome)
|
||||
def test_failing_or_transforming_failure_callbacks_do_not_stop_the_denial_from_being_logged(
|
||||
hostile: _Rig, outcome: str
|
||||
) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
target: Final = hostile.denied.agent_id if outcome == "denied" else unknown
|
||||
with hostile.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(hostile, scenario)
|
||||
response: Final = _send(hostile, caller.key, f"/a2a/{target}")
|
||||
assert response.status_code == (403 if outcome == "denied" else 404), response.text
|
||||
model: Final = f"a2a_agent/{hostile.denied.name}" if outcome == "denied" else f"a2a_agent/{unknown}"
|
||||
_assert_one_failure(hostile, caller, status=response.status_code, agent_id=target, model=model)
|
||||
|
||||
|
||||
def test_denials_follow_the_model_denied_chat_logging_under_error_log_flags(flagged: _Rig) -> None:
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
request_id: Final = "f-" + uuid.uuid4().hex
|
||||
with flagged.proxy.scenario() as scenario:
|
||||
chat_caller: Final = _caller(flagged, scenario)
|
||||
denied_caller: Final = _caller(flagged, scenario)
|
||||
unknown_caller: Final = _caller(flagged, scenario)
|
||||
chat_alias: Final = "a2adeny-chat-" + uuid.uuid4().hex[:8]
|
||||
chat_key: Final = scenario.key(
|
||||
user_id=chat_caller.user_id, team_id=chat_caller.team_id, key_alias=chat_alias, models=["openai/gpt-5-mini"]
|
||||
)
|
||||
chat: Final = flagged.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": "openai/gpt-5.4-mini", "messages": [{"role": "user", "content": "x"}]},
|
||||
key=chat_key,
|
||||
)
|
||||
assert chat.status_code in {401, 403}, chat.text
|
||||
denied: Final = _send(flagged, denied_caller.key, f"/a2a/{flagged.denied.agent_id}")
|
||||
assert (denied.status_code, denied.content) == (403, _refusal(flagged.denied.agent_id, denied))
|
||||
missing: Final = _send(flagged, unknown_caller.key, f"/a2a/{unknown}", request_id=request_id)
|
||||
assert (missing.status_code, missing.content) == (404, _not_found(request_id, unknown))
|
||||
_assert_never_dispatched(flagged, scenario, "never-sent")
|
||||
chat_hooks: Final = held(
|
||||
lambda: _hooks(flagged.events.of(chat_alias)),
|
||||
lambda found: len(found) > 0,
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
assert _hooks(flagged.events.of(denied_caller.alias)) == chat_hooks
|
||||
assert _hooks(flagged.events.of(unknown_caller.alias)) == chat_hooks
|
||||
control: Final = _caller(flagged, scenario)
|
||||
granted: Final = _send(
|
||||
flagged, control.key, f"/a2a/{flagged.allowed.agent_id}", method="message/send", version=None
|
||||
)
|
||||
assert granted.status_code == 200, granted.text
|
||||
eventually(lambda: _rows(flagged, control), lambda found: len(found) == 1, seconds=_CONVERGE)
|
||||
assert (_rows(flagged, denied_caller), _rows(flagged, unknown_caller)) == ([], [])
|
||||
|
||||
|
||||
def test_client_error_denials_never_raise_an_llm_exceptions_alert(alerting: tuple[_Rig, Wire]) -> None:
|
||||
rig, alerts = alerting
|
||||
unknown: Final = "no-such-" + uuid.uuid4().hex
|
||||
marker: Final = "boom-" + uuid.uuid4().hex
|
||||
|
||||
def failing(request: Request) -> Reply:
|
||||
return Reply(
|
||||
status=500, body=_compact({"error": {"message": f"upstream exploded {marker}", "type": "server_error"}})
|
||||
)
|
||||
|
||||
with rig.proxy.scenario() as scenario, wire_server(failing) as broken:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
assert _send(rig, caller.key, f"/a2a/{rig.denied.agent_id}").status_code == 403
|
||||
assert _send(rig, caller.key, f"/a2a/{unknown}").status_code == 404
|
||||
model: Final = scenario.model(api_base=broken.url + "/v1", num_retries=0)
|
||||
failed: Final = rig.proxy.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "x"}]}
|
||||
)
|
||||
assert failed.status_code >= 500, failed.text
|
||||
texts: Final = eventually(
|
||||
lambda: tuple(JSON_OBJECT.validate_json(request.body).get("text") for request in alerts.drain()),
|
||||
lambda found: any(marker in str(text) for text in found),
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
assert not any(rig.denied.agent_id in str(text) or unknown in str(text) for text in texts), texts
|
||||
|
||||
|
||||
def test_guard_model_denied_chat_logs_one_acompletion_failure(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
user: Final = scenario.user()
|
||||
alias: Final = "a2adeny-chat-" + uuid.uuid4().hex[:8]
|
||||
key: Final = scenario.key(user_id=user, key_alias=alias, models=["openai/gpt-5-mini"])
|
||||
caller: Final = _Caller(key, alias, user, None)
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": "openai/gpt-5.4-mini", "messages": [{"role": "user", "content": "x"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 403, response.text
|
||||
assert "key_model_access_denied" in response.text
|
||||
events: Final = held(
|
||||
lambda: rig.events.of(alias),
|
||||
lambda found: _hooks(found) == _ONE_FAILURE,
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
assert {event.status_code for event in events if event.hook == "post_call_failure"} == {403}
|
||||
rows: Final = held(
|
||||
lambda: _rows(rig, caller), lambda found: len(found) == 1, holding=_SETTLE, seconds=_CONVERGE
|
||||
)
|
||||
assert (rows[0]["status"], rows[0]["call_type"], rows[0]["spend"]) == ("failure", "acompletion", 0.0), rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "call_type"), [("message/send", "asend_message"), ("message/stream", "asend_message_streaming")]
|
||||
)
|
||||
def test_guard_granted_agent_succeeds_with_one_success_row_and_no_failure(
|
||||
rig: _Rig, method: str, call_type: str
|
||||
) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
request_id: Final = "g-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig, caller.key, f"/a2a/{rig.allowed.agent_id}", method=method, version=None, request_id=request_id
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
if method == "message/stream":
|
||||
assert response.text.startswith("data: ") and request_id in response.text, response.text
|
||||
else:
|
||||
assert object_value(JSON_OBJECT.validate_json(response.content)["result"])["parts"] == [
|
||||
{"kind": "text", "text": "granted"}
|
||||
]
|
||||
held(
|
||||
lambda: rig.events.of(caller.alias),
|
||||
lambda found: _hooks(found) == ("async_success",),
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
rows: Final = held(
|
||||
lambda: _rows(rig, caller), lambda found: len(found) == 1, holding=_SETTLE, seconds=_CONVERGE
|
||||
)
|
||||
assert (rows[0]["status"], rows[0]["call_type"], rows[0]["agent_id"]) == (
|
||||
"success",
|
||||
call_type,
|
||||
rig.allowed.agent_id,
|
||||
), rows
|
||||
|
||||
|
||||
def test_guard_zero_budget_key_is_refused_once_without_double_logging(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
user: Final = scenario.user()
|
||||
alias: Final = "a2adeny-budget-" + uuid.uuid4().hex[:8]
|
||||
key: Final = scenario.key(
|
||||
user_id=user, key_alias=alias, max_budget=0, object_permission={"agents": [rig.allowed.agent_id]}
|
||||
)
|
||||
caller: Final = _Caller(key, alias, user, None)
|
||||
message: Final = "budget-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig, key, f"/a2a/{rig.allowed.agent_id}", method="message/send", version=None, message_id=message
|
||||
)
|
||||
assert response.status_code == 422, response.text
|
||||
assert "Budget has been exceeded" in response.text
|
||||
_assert_never_dispatched(rig, scenario, message)
|
||||
held(
|
||||
lambda: rig.events.of(alias),
|
||||
lambda found: _hooks(found) == _ONE_FAILURE,
|
||||
holding=_SETTLE,
|
||||
seconds=_CONVERGE,
|
||||
)
|
||||
rows: Final = held(
|
||||
lambda: _rows(rig, caller), lambda found: len(found) == 1, holding=_SETTLE, seconds=_CONVERGE
|
||||
)
|
||||
assert (rows[0]["status"], rows[0]["call_type"]) == ("failure", "asend_message"), rows
|
||||
|
||||
|
||||
def test_guard_agent_card_discovery_keeps_its_refusals_and_writes_no_failure_row(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
headers: Final = {"Authorization": f"Bearer {caller.key}"}
|
||||
denied_card: Final = rig.proxy.client.get(
|
||||
f"/a2a/{rig.denied.agent_id}/.well-known/agent-card.json", headers=headers
|
||||
)
|
||||
assert (denied_card.status_code, denied_card.content) == (403, _refusal(rig.denied.agent_id))
|
||||
missing: Final = rig.proxy.client.get("/a2a/no-such-card/.well-known/agent-card.json", headers=headers)
|
||||
assert (missing.status_code, missing.content) == (404, _compact({"detail": "Agent 'no-such-card' not found"}))
|
||||
granted: Final = _send(rig, caller.key, f"/a2a/{rig.allowed.agent_id}", method="message/send", version=None)
|
||||
assert granted.status_code == 200, granted.text
|
||||
rows: Final = held(
|
||||
lambda: _rows(rig, caller), lambda found: len(found) == 1, holding=_SETTLE, seconds=_CONVERGE
|
||||
)
|
||||
assert [row["status"] for row in rows] == ["success"], rows
|
||||
assert _hooks(rig.events.of(caller.alias)) == ("async_success",)
|
||||
|
||||
|
||||
def test_guard_jsonrpc_1_0_envelope_keeps_its_invalid_request_body(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
caller: Final = _caller(rig, scenario)
|
||||
request_id: Final = "v-" + uuid.uuid4().hex
|
||||
response: Final = _send(
|
||||
rig,
|
||||
caller.key,
|
||||
f"/a2a/{rig.allowed.agent_id}",
|
||||
method="message/send",
|
||||
version=None,
|
||||
request_id=request_id,
|
||||
extra={"jsonrpc": "1.0"},
|
||||
)
|
||||
assert (response.status_code, response.content) == (
|
||||
400,
|
||||
_compact(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"error": {"code": -32600, "message": "Invalid Request: jsonrpc must be '2.0'"},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
|
@ -221,8 +221,13 @@ def test_agent_access_default_deny_denies_ungranted_key_but_keeps_explicit_grant
|
|||
denied_sends: Final = tuple(
|
||||
_send(candidate, agent_id, key_a, f"denied-{index}-{uuid.uuid4().hex}") for index in range(8)
|
||||
)
|
||||
_assert_denial(denied_sends[0], agent_id)
|
||||
assert all(response.json() == denied_sends[0].json() for response in denied_sends), [
|
||||
assert denied_sends[0].status_code == 403, denied_sends[0].text
|
||||
assert denied_sends[0].json()["error"] == {
|
||||
"code": -32000,
|
||||
"message": f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
|
||||
"data": denied_sends[0].json()["error"]["data"],
|
||||
}, denied_sends[0].text
|
||||
assert all(response.json()["error"] == denied_sends[0].json()["error"] for response in denied_sends), [
|
||||
response.text for response in denied_sends
|
||||
]
|
||||
denied_completion: Final = _chat_completion(candidate, agent_name, key_a)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue