mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge pull request #50326 from BerriAI/litellm_backport_1108_2_fixes_lts_1_108_x
fix(backport): bring #50132, #50218, #50274 and #50276 onto lts/1.108.x
This commit is contained in:
Kerry Lu
GitHub
commit
612587a91d
25 changed files with 870 additions and 28 deletions
|
|
@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
def redacts_messages_itself(self) -> bool:
|
||||
return False
|
||||
|
||||
def exempt_from_message_redaction(self) -> bool:
|
||||
"""Router pre-call checks that key routing off the request and never log it return True, so
|
||||
`turn_off_message_logging` leaves the success event they receive unredacted."""
|
||||
return False
|
||||
|
||||
def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict:
|
||||
"""
|
||||
Redacts or excludes fields from StandardLoggingPayload before callbacks receive it.
|
||||
|
|
|
|||
|
|
@ -3751,7 +3751,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if isinstance(callback, CustomLogger): # custom logger class
|
||||
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
|
||||
|
||||
model_call_details: dict = self.model_call_details
|
||||
model_call_details: dict = (
|
||||
self._model_call_details_with_request_messages()
|
||||
if callback.exempt_from_message_redaction()
|
||||
else self.model_call_details
|
||||
)
|
||||
##################################
|
||||
# call redaction hook for custom logger
|
||||
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
|
||||
|
|
@ -4289,6 +4293,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
) -> None:
|
||||
executor.submit(self.success_handler, result, start_time, end_time, cache_hit)
|
||||
|
||||
def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict
|
||||
"""The success-event payload with the request messages put back where `turn_off_message_logging`
|
||||
redacted them in place, for a callback that routes on the request and never logs it."""
|
||||
request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is
|
||||
standard_logging_object: Final = self.model_call_details.get("standard_logging_object")
|
||||
return {
|
||||
**self.model_call_details,
|
||||
"messages": request_messages,
|
||||
**(
|
||||
{"standard_logging_object": {**standard_logging_object, "messages": request_messages}}
|
||||
if isinstance(standard_logging_object, dict)
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
def _should_run_sync_callbacks_for_async_calls(self) -> bool:
|
||||
"""
|
||||
Returns:
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import inspect
|
||||
import time
|
||||
import types
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
|
||||
|
|
@ -351,6 +352,11 @@ _EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None)
|
|||
_EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None)
|
||||
_NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({})
|
||||
_SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body"))
|
||||
_CLIENT_API_KEY_ONLY_AUTH: Final[Mapping[str, Mapping[str, bool]]] = (
|
||||
types.MappingProxyType({"security": types.MappingProxyType({"bearer_auth": True})})
|
||||
if "security" in inspect.signature(make_request_options).parameters
|
||||
else types.MappingProxyType({})
|
||||
)
|
||||
|
||||
|
||||
def _raw_chat_response(
|
||||
|
|
@ -370,6 +376,7 @@ def _request_without_sdk_transform(
|
|||
extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")),
|
||||
extra_body=data.get("extra_body"),
|
||||
timeout=timeout,
|
||||
**_CLIENT_API_KEY_ONLY_AUTH, # pyright: ignore[reportArgumentType] # openai<3 make_request_options has no security kwarg
|
||||
)
|
||||
return body, options
|
||||
|
||||
|
|
|
|||
|
|
@ -55,3 +55,10 @@ class ToolDiscoveryQueue:
|
|||
"""Return and clear all pending rows."""
|
||||
items, self._pending = list(self._pending.values()), {}
|
||||
return items
|
||||
|
||||
def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None:
|
||||
"""Put rows from a failed write back, ahead of and merged with anything queued since the flush."""
|
||||
since_flush: Final = self._pending
|
||||
self._pending = {}
|
||||
for item in (*items, *since_flush.values()):
|
||||
self.add_update(item)
|
||||
|
|
|
|||
|
|
@ -31,6 +31,24 @@ _CONNECTION_CAPACITY_PHRASES: Final = (
|
|||
"too many connections for database",
|
||||
"remaining connection slots are reserved",
|
||||
)
|
||||
_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03"
|
||||
_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = (
|
||||
"the database system is starting up",
|
||||
"the database system is shutting down",
|
||||
"the database system is not yet accepting connections",
|
||||
"the database system is not accepting connections",
|
||||
"the database system is in recovery mode",
|
||||
)
|
||||
_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"})
|
||||
_SERVER_TERMINATED_PHRASES: Final = (
|
||||
"terminating connection due to administrator command",
|
||||
"terminating connection due to immediate shutdown command",
|
||||
"terminating connection because of unexpected sigquit signal",
|
||||
"terminating connection because of crash of another server process",
|
||||
"terminating connection due to idle-session timeout",
|
||||
)
|
||||
_UNREACHABLE_PRISMA_CODE: Final = "P1001"
|
||||
_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server")
|
||||
_PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
|
||||
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
|
||||
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
|
||||
|
|
@ -74,6 +92,22 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]:
|
|||
return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException))
|
||||
|
||||
|
||||
def _prisma_error_code(e: Exception) -> str | None:
|
||||
import prisma.engine.errors
|
||||
|
||||
if isinstance(e, _exception_types(prisma.errors.DataError)):
|
||||
code: Final[object] = getattr(e, "code", None)
|
||||
return code if isinstance(code, str) else None
|
||||
if not isinstance(e, _exception_types(prisma.engine.errors.EngineRequestError)):
|
||||
return None
|
||||
try:
|
||||
body: Final = _DATABASE_ERROR_META.validate_json(str(e).partition(": ")[2])
|
||||
except ValidationError:
|
||||
return None
|
||||
engine_code: Final = body.get("error_code")
|
||||
return engine_code if isinstance(engine_code, str) else None
|
||||
|
||||
|
||||
class PrismaDBExceptionHandler:
|
||||
"""
|
||||
Class to handle DB Exceptions or Connection Errors
|
||||
|
|
@ -123,6 +157,8 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -271,6 +307,39 @@ class PrismaDBExceptionHandler:
|
|||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_database_not_accepting_connections_error(e: Exception) -> bool:
|
||||
"""True iff Postgres refused the connection because it is starting up, shutting
|
||||
down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement
|
||||
runs, so nothing the write carried was applied."""
|
||||
import prisma
|
||||
|
||||
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||
return False
|
||||
if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE:
|
||||
return True
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_database_lifecycle_error(e: Exception) -> bool:
|
||||
"""True iff Postgres failed the statement because of a server lifecycle event:
|
||||
a refused connection during startup, shutdown or recovery (57P03), or a session
|
||||
the server terminated (57P01, 57P02, 57P05). The server comes back on its own,
|
||||
so the rows are not at fault. A terminated session may have committed before it
|
||||
was cut, so only ``is_database_not_accepting_connections_error`` proves a
|
||||
non-idempotent write never applied."""
|
||||
import prisma
|
||||
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||
return False
|
||||
if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES:
|
||||
return True
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_lock_timeout_error(e: Exception) -> bool:
|
||||
"""True iff Postgres cancelled a statement for exceeding ``lock_timeout``
|
||||
|
|
@ -388,6 +457,8 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_lifecycle_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e):
|
||||
return True
|
||||
if "cached plan must not change result type" in str(e).lower():
|
||||
|
|
@ -662,3 +733,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool:
|
|||
return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
|
||||
sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
|
||||
return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES
|
||||
|
||||
|
||||
def db_write_never_sent(e: Exception) -> bool:
|
||||
"""Whether the database could not take the write at all, because it was unreachable (P1001),
|
||||
refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again."""
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE:
|
||||
return True
|
||||
if not PrismaDBExceptionHandler.is_prisma_error(e):
|
||||
return False
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES)
|
||||
|
|
|
|||
|
|
@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
|
|||
ToolDiscoveryQueue,
|
||||
aggregate_tool_discoveries,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
call_with_db_reconnect_retry,
|
||||
db_write_failure_is_requeue_safe,
|
||||
db_write_never_sent,
|
||||
)
|
||||
from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
|
@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove
|
|||
return
|
||||
try:
|
||||
await batch_upsert_tools(prisma_client=prisma_client, items=items)
|
||||
except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail
|
||||
verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e)
|
||||
except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail
|
||||
if db_write_never_sent(e):
|
||||
tool_discovery_queue.requeue(items)
|
||||
verbose_proxy_logger.warning(
|
||||
"Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e)
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.error(
|
||||
"Tool discovery - dropped %d tool rows: the failed statement may have applied or the database "
|
||||
"refused the data, so re-sending it is not safe. Error: %s",
|
||||
len(items),
|
||||
str(e),
|
||||
)
|
||||
|
||||
|
||||
async def commit_tool_discoveries(
|
||||
|
|
|
|||
|
|
@ -198,6 +198,7 @@ from litellm.proxy.db.db_url_settings import (
|
|||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
db_write_never_sent,
|
||||
)
|
||||
from litellm.proxy.db.health_check_latest import (
|
||||
LatestHealthCheckRow,
|
||||
|
|
@ -8500,9 +8501,20 @@ async def update_daily_tag_spend(
|
|||
|
||||
def _rollup_batch_never_applied(e: Exception) -> bool:
|
||||
"""True when a drained rollup batch provably never reached its rows: Postgres had no
|
||||
connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row
|
||||
lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued."""
|
||||
return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e)
|
||||
connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or
|
||||
cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case
|
||||
re-sending the batch cannot double-count, so it is requeued."""
|
||||
return (
|
||||
PrismaDBExceptionHandler.is_database_capacity_error(e)
|
||||
or db_write_never_sent(e)
|
||||
or PrismaDBExceptionHandler.is_lock_timeout_error(e)
|
||||
)
|
||||
|
||||
|
||||
def _rollup_failure_stops_job(e: Exception) -> bool:
|
||||
"""True when the database cannot take any write, so the job stops instead of draining
|
||||
the queues after this one into the same failure."""
|
||||
return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e)
|
||||
|
||||
|
||||
async def update_spend_logs_job(
|
||||
|
|
@ -8593,12 +8605,12 @@ async def _run_spend_logs_job(
|
|||
*prisma_client.tool_usage_transactions,
|
||||
]
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - database out of connections or rollup row lock timed out during tool usage "
|
||||
"Spend tracking - database unavailable or rollup row lock timed out during tool usage "
|
||||
"flush, requeued %s tool usage transactions for the next flush: %s",
|
||||
len(tool_usage_to_process),
|
||||
tool_tracking_err,
|
||||
)
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err):
|
||||
if _rollup_failure_stops_job(tool_tracking_err):
|
||||
raise
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -8625,12 +8637,12 @@ async def _run_spend_logs_job(
|
|||
*prisma_client.model_usage_transactions,
|
||||
]
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - database out of connections or rollup row lock timed out during model usage "
|
||||
"Spend tracking - database unavailable or rollup row lock timed out during model usage "
|
||||
"flush, requeued %s model usage transactions for the next flush: %s",
|
||||
len(model_usage_to_process),
|
||||
model_usage_err,
|
||||
)
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err):
|
||||
if _rollup_failure_stops_job(model_usage_err):
|
||||
raise
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -8807,6 +8819,7 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
|||
return (
|
||||
PrismaDBExceptionHandler.is_database_transport_error(e)
|
||||
or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e)
|
||||
or PrismaDBExceptionHandler.is_database_lifecycle_error(e)
|
||||
or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -138,6 +138,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
|
||||
return healthy_deployments
|
||||
|
||||
def exempt_from_message_redaction(self) -> bool:
|
||||
return True
|
||||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: _PromptCachingSuccessEventKwargs,
|
||||
|
|
|
|||
|
|
@ -653,6 +653,14 @@ class CallTypes(str, Enum):
|
|||
retrieve_fine_tuning_job = "retrieve_fine_tuning_job"
|
||||
responses = "responses"
|
||||
aresponses = "aresponses"
|
||||
get_responses = "get_responses"
|
||||
aget_responses = "aget_responses"
|
||||
delete_responses = "delete_responses"
|
||||
adelete_responses = "adelete_responses"
|
||||
cancel_responses = "cancel_responses"
|
||||
acancel_responses = "acancel_responses"
|
||||
compact_responses = "compact_responses"
|
||||
acompact_responses = "acompact_responses"
|
||||
alist_input_items = "alist_input_items"
|
||||
llm_passthrough_route = "llm_passthrough_route"
|
||||
allm_passthrough_route = "allm_passthrough_route"
|
||||
|
|
@ -754,6 +762,8 @@ CallTypesLiteral = Literal[
|
|||
"send_message",
|
||||
"aresponses",
|
||||
"responses",
|
||||
"compact_responses",
|
||||
"acompact_responses",
|
||||
"acreate_skill",
|
||||
"acreate_realtime_client_secret",
|
||||
"arealtime_calls",
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ def _payload(kwargs: Mapping[str, object], response_obj: object) -> str:
|
|||
"request": request,
|
||||
"response": response_obj.model_dump() if isinstance(response_obj, BaseModel) else response_obj,
|
||||
"standard_response": standard.get("response") if isinstance(standard, Mapping) else None,
|
||||
"standard": standard,
|
||||
},
|
||||
default=str,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,114 @@
|
|||
"""GET, compact, DELETE and cancel on /v1/responses answer 200 when a during_call guardrail is on by default
|
||||
(LIT-9770): the hook used to raise `'acompact_responses' is not a valid CallTypes` and the proxy returned 500"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, string_value
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.provider import SharedProvider
|
||||
from integration._support.wire import Reply
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_MODEL: Final = "openai/gpt-5.4-mini"
|
||||
|
||||
|
||||
def _input() -> list[JsonValue]:
|
||||
return [
|
||||
{"role": "user", "content": f"Remember the number {uuid.uuid4().hex}."},
|
||||
{"role": "assistant", "content": "Got it."},
|
||||
]
|
||||
|
||||
|
||||
def _response(response_id: str) -> Reply:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": _MODEL.removeprefix("openai/"),
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{uuid.uuid4().hex}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Got it.", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": False,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17},
|
||||
}
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _deleted(response_id: str) -> Reply:
|
||||
return Reply(body=json.dumps({"id": response_id, "object": "response.deleted", "deleted": True}).encode())
|
||||
|
||||
|
||||
def _guarded_config(path: Path) -> Path:
|
||||
base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
guardrail: Final[dict[str, JsonValue]] = {
|
||||
"guardrail_name": f"content-filter-{uuid.uuid4().hex}",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "during_call",
|
||||
"default_on": True,
|
||||
"blocked_words": [{"keyword": "forbiddenword", "action": "BLOCK"}],
|
||||
},
|
||||
}
|
||||
path.write_text(yaml.safe_dump({**base, "guardrails": [guardrail]}))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("during-call-guardrail")
|
||||
with owned_proxy(session_gateway, directory, {}, config=_guarded_config(directory / "config.yaml")) as candidate:
|
||||
yield candidate
|
||||
|
||||
|
||||
def _create(guarded: Gateway, provider: SharedProvider, upstream_id: str) -> str:
|
||||
provider.expect(_response(upstream_id))
|
||||
created: Final = guarded.request("POST", "/v1/responses", {"model": _MODEL, "input": _input(), "store": True})
|
||||
assert created.status_code == 200, created.text
|
||||
(upstream,) = provider.received()
|
||||
assert upstream.target.endswith("/responses"), upstream.target
|
||||
return string_value(_JSON_OBJECT.validate_json(created.content)["id"])
|
||||
|
||||
|
||||
def _lifecycle(guarded: Gateway, provider: SharedProvider, operation: str) -> tuple[httpx.Response, str]:
|
||||
if operation == "compact":
|
||||
provider.expect(_response(f"resp_{uuid.uuid4().hex}"))
|
||||
return guarded.request(
|
||||
"POST", "/v1/responses/compact", {"model": _MODEL, "input": _input()}
|
||||
), "/responses/compact"
|
||||
upstream_id: Final = f"resp_{uuid.uuid4().hex}"
|
||||
managed_id: Final = _create(guarded, provider, upstream_id)
|
||||
provider.expect(_deleted(upstream_id) if operation == "delete" else _response(upstream_id))
|
||||
if operation == "get":
|
||||
return guarded.request("GET", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}"
|
||||
if operation == "delete":
|
||||
return guarded.request("DELETE", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}"
|
||||
return guarded.request("POST", f"/v1/responses/{managed_id}/cancel"), f"/responses/{upstream_id}/cancel"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["get", "compact", "delete", "cancel"])
|
||||
def test_responses_lifecycle_answers_200_with_a_during_call_guardrail_on(
|
||||
guarded: Gateway, provider: SharedProvider, operation: str
|
||||
) -> None:
|
||||
response, upstream_target = _lifecycle(guarded, provider, operation)
|
||||
assert response.status_code == 200, response.text
|
||||
(upstream,) = provider.received()
|
||||
assert upstream.target.endswith(upstream_target), upstream.target
|
||||
|
|
@ -68,6 +68,7 @@ def pytest_addoption(parser: pytest.Parser) -> None:
|
|||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line("markers", "integration: owned real-service integration contracts")
|
||||
config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced")
|
||||
config.addinivalue_line("markers", "shares_worker_state: owned proxies here share the worker DB and Redis")
|
||||
config.stash[REPORTS] = []
|
||||
config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files")))
|
||||
config.pluginmanager.register(IntegrationReportPlugin(config))
|
||||
|
|
|
|||
|
|
@ -0,0 +1,95 @@
|
|||
import collections
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import eventually, gateway_from_environment, object_value
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.redis_process import owned_redis
|
||||
from integration._support.wire import Reply, Wire, wire_server
|
||||
from integration.routing._prompt_cache_pin_support import (
|
||||
DEFAULT_MARK,
|
||||
anthropic_deployment,
|
||||
ask_turn,
|
||||
chat_body,
|
||||
claude_peer,
|
||||
new_marker,
|
||||
one_new_pin,
|
||||
pin_config,
|
||||
post_chat,
|
||||
read_pins,
|
||||
system_turn,
|
||||
)
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from redis import Redis
|
||||
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300)
|
||||
|
||||
_GROUP: Final = "pin-logging-off"
|
||||
_LEFT: Final = "pin-logging-off-left"
|
||||
_RIGHT: Final = "pin-logging-off-right"
|
||||
_FOLLOW_UPS: Final = 20
|
||||
_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
owned: OwnedProxy
|
||||
redis: "Redis[bytes]"
|
||||
logger_sink: Wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_redis(directory) as cache,
|
||||
wire_server(claude_peer) as left,
|
||||
wire_server(claude_peer) as right,
|
||||
wire_server(lambda request: Reply()) as logger_sink,
|
||||
):
|
||||
config: Final = pin_config(
|
||||
directory,
|
||||
[anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)],
|
||||
litellm_settings={
|
||||
"turn_off_message_logging": True,
|
||||
"callbacks": ["integration._support.redaction_probe.probe"],
|
||||
"cache": False,
|
||||
"enable_redis_auth_cache": False,
|
||||
},
|
||||
)
|
||||
environment: Final = {
|
||||
"REDIS_HOST": cache.host,
|
||||
"REDIS_PORT": str(cache.port),
|
||||
"REDACTION_PROBE_URL": logger_sink.url,
|
||||
}
|
||||
with (
|
||||
owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned,
|
||||
Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis,
|
||||
):
|
||||
yield _Rig(owned, redis, logger_sink)
|
||||
|
||||
|
||||
def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None:
|
||||
before: Final = frozenset(read_pins(rig.redis))
|
||||
body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())])
|
||||
|
||||
_, first_served = post_chat(rig.owned.gateway, body)
|
||||
_, pin = one_new_pin(rig.redis, before)
|
||||
assert pin.model_id == first_served, (pin, first_served)
|
||||
|
||||
served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS))
|
||||
assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served)
|
||||
|
||||
eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30)
|
||||
logged: Final = tuple(
|
||||
object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"]
|
||||
for request in rig.logger_sink.drain()
|
||||
)
|
||||
assert len(logged) == 1 + _FOLLOW_UPS, len(logged)
|
||||
assert all(messages == _REDACTED_MESSAGES for messages in logged), logged
|
||||
|
|
@ -456,6 +456,10 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client():
|
|||
"add_message",
|
||||
"arun_thread_stream",
|
||||
"aresponses",
|
||||
"aget_responses",
|
||||
"adelete_responses",
|
||||
"acancel_responses",
|
||||
"acompact_responses",
|
||||
"aresponses_websocket",
|
||||
"alist_input_items",
|
||||
"acreate_fine_tuning_job",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -1937,3 +1939,49 @@ async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_head
|
|||
assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True}
|
||||
assert sent.headers["x-caller"] == "unit"
|
||||
assert sent.headers["authorization"] == "Bearer sk-test"
|
||||
|
||||
|
||||
_SDK_READS_OPENAI_ADMIN_KEY: Final = "admin_api_key" in inspect.signature(AsyncOpenAI.__init__).parameters
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _SDK_READS_OPENAI_ADMIN_KEY, reason="openai<3 has no OPENAI_ADMIN_KEY fallback to leak")
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_never_sends_the_openai_admin_key_for_a_deployment_without_a_key(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.setenv("OPENAI_ADMIN_KEY", "sk-admin-only")
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
outbound: Final = asyncio.Queue()
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
outbound.put_nowait(request.headers.get("authorization"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-no-admin-key",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-5.4-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
|
||||
monkeypatch.setattr(litellm, "aclient_session", http_client)
|
||||
with contextlib.suppress(litellm.AuthenticationError, litellm.InternalServerError):
|
||||
await asyncio.wait_for(
|
||||
litellm.acompletion(
|
||||
model="openai/gpt-5.4-mini",
|
||||
api_key="",
|
||||
messages=[{"role": "user", "content": "who signs this request"}],
|
||||
num_retries=0,
|
||||
max_retries=0,
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
sent: Final = tuple(outbound.get_nowait() for _ in range(outbound.qsize()))
|
||||
assert all(authorization is None for authorization in sent), sent
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
Unit tests for ToolDiscoveryQueue.
|
||||
"""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
|
||||
|
|
@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes():
|
|||
{"tool_name": "get_weather", "origin": "user_defined", "call_count": 5},
|
||||
{"tool_name": "search", "call_count": 1},
|
||||
)
|
||||
|
||||
|
||||
def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None:
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"})
|
||||
queue.add_update({"tool_name": "search", "origin": "user_defined"})
|
||||
failed_write = queue.flush()
|
||||
queue.add_update({"tool_name": "translate", "origin": "user_defined"})
|
||||
queue.add_update({"tool_name": "search", "origin": "user_defined"})
|
||||
|
||||
queue.requeue(failed_write)
|
||||
|
||||
assert queue.flush() == [
|
||||
{"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1},
|
||||
{"tool_name": "search", "origin": "user_defined", "call_count": 2},
|
||||
{"tool_name": "translate", "origin": "user_defined", "call_count": 1},
|
||||
]
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from prisma.errors import (
|
|||
|
||||
import litellm
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent
|
||||
|
||||
|
||||
# Test is_database_connection_error method
|
||||
|
|
@ -966,3 +966,137 @@ def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None
|
|||
)
|
||||
def test_is_statement_timeout_error(error: Exception, expected: bool) -> None:
|
||||
assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected
|
||||
|
||||
|
||||
def _connector_error(message: str) -> DataError:
|
||||
return DataError(
|
||||
data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}}
|
||||
)
|
||||
|
||||
|
||||
def _batched_postgres_error(sqlstate: str, message: str) -> DataError:
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: "
|
||||
f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})'
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _localized_error(sqlstate: str) -> DataError:
|
||||
return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}})
|
||||
|
||||
|
||||
_NOT_ACCEPTING_CONNECTIONS: Final = (
|
||||
_connector_error("the database system is starting up"),
|
||||
_connector_error("the database system is shutting down"),
|
||||
_connector_error("the database system is not yet accepting connections"),
|
||||
_connector_error("the database system is not accepting connections"),
|
||||
_connector_error("the database system is in recovery mode"),
|
||||
_localized_error("57P03"),
|
||||
_batched_postgres_error("57P03", "the database system is starting up"),
|
||||
EngineRequestError(
|
||||
MagicMock(status=500),
|
||||
'{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is '
|
||||
'starting up","backtrace":null}',
|
||||
),
|
||||
RawQueryError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`",
|
||||
"meta": {"code": "57P03", "message": "FATAL: the database system is starting up"},
|
||||
"error_code": "P2010",
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
_SERVER_TERMINATED: Final = (
|
||||
_connector_error("terminating connection due to administrator command"),
|
||||
_connector_error("terminating connection due to immediate shutdown command"),
|
||||
_connector_error("terminating connection because of unexpected SIGQUIT signal"),
|
||||
_connector_error("terminating connection because of crash of another server process"),
|
||||
_connector_error("terminating connection due to idle-session timeout"),
|
||||
_localized_error("57P01"),
|
||||
_localized_error("57P02"),
|
||||
_localized_error("57P05"),
|
||||
_batched_postgres_error("57P01", "terminating connection due to administrator command"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "connection_error"),
|
||||
[
|
||||
*((error, True) for error in _NOT_ACCEPTING_CONNECTIONS),
|
||||
*((error, False) for error in _SERVER_TERMINATED),
|
||||
],
|
||||
)
|
||||
def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None:
|
||||
"""A server that is restarting, shutting down or cutting sessions is a transient outage:
|
||||
auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison.
|
||||
Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback,
|
||||
one terminated session does not prove the database is down."""
|
||||
assert {
|
||||
"lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error),
|
||||
"service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error),
|
||||
"connection_error": PrismaDBExceptionHandler.is_database_connection_error(error),
|
||||
"permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error),
|
||||
} == {
|
||||
"lifecycle": True,
|
||||
"service_unavailable": True,
|
||||
"connection_error": connection_error,
|
||||
"permanent": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "never_sent"),
|
||||
[
|
||||
*((error, True) for error in _NOT_ACCEPTING_CONNECTIONS),
|
||||
*((error, False) for error in _SERVER_TERMINATED),
|
||||
(DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True),
|
||||
(DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True),
|
||||
(_capacity_error("sorry, too many clients already"), True),
|
||||
(_localized_error("53300"), True),
|
||||
(
|
||||
DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"error_code": "P2024",
|
||||
"message": "Timed out fetching a new connection from the connection pool",
|
||||
}
|
||||
}
|
||||
),
|
||||
True,
|
||||
),
|
||||
(httpx.ReadTimeout("read timed out"), False),
|
||||
(_localized_error("57014"), False),
|
||||
],
|
||||
)
|
||||
def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None:
|
||||
"""A refused, unreachable or full server provably applied nothing, while a terminated session
|
||||
or a timeout may have committed before the failure, so only the former lets a non-idempotent
|
||||
write be replayed."""
|
||||
assert db_write_never_sent(error) is never_sent
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
_localized_error("57P04"),
|
||||
_localized_error("57014"),
|
||||
_localized_error("23505"),
|
||||
DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}),
|
||||
RuntimeError("the database system is starting up"),
|
||||
],
|
||||
)
|
||||
def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None:
|
||||
assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False
|
||||
assert db_write_never_sent(error) is False
|
||||
|
||||
|
||||
def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None:
|
||||
wrapped: Final = RuntimeError("spend flush failed")
|
||||
wrapped.__cause__ = _connector_error("the database system is starting up")
|
||||
assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True
|
||||
|
|
|
|||
|
|
@ -69,9 +69,7 @@ def _make_prisma(
|
|||
prisma.db.litellm_tooltable = MagicMock()
|
||||
prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return)
|
||||
prisma.db.execute_raw = AsyncMock(return_value=1)
|
||||
prisma.db.litellm_tooltable.find_many = AsyncMock(
|
||||
return_value=find_many_rows if find_many_rows is not None else []
|
||||
)
|
||||
prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else [])
|
||||
prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row)
|
||||
prisma.db.litellm_verificationtoken = MagicMock()
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows))
|
||||
|
|
@ -227,9 +225,7 @@ async def test_get_tool_found():
|
|||
result = await get_tool(prisma, "my_tool")
|
||||
assert result is not None
|
||||
assert result.tool_name == "my_tool"
|
||||
prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(
|
||||
where={"tool_name": "my_tool"}
|
||||
)
|
||||
prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool():
|
|||
updated_by="admin",
|
||||
)
|
||||
prisma = _make_prisma(find_unique_row=row)
|
||||
result = await update_tool_policy(
|
||||
prisma, "my_tool", updated_by="admin", input_policy="blocked"
|
||||
)
|
||||
result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked")
|
||||
assert result is not None
|
||||
assert result.input_policy == "blocked"
|
||||
prisma.db.litellm_tooltable.upsert.assert_awaited_once()
|
||||
|
|
@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably
|
|||
left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])]
|
||||
expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else []
|
||||
assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("failure", "requeued"),
|
||||
[
|
||||
(_postgres_rejection("57P03"), True),
|
||||
(_postgres_rejection("53300"), True),
|
||||
(_postgres_rejection("57P01"), False),
|
||||
(_postgres_rejection("54000"), False),
|
||||
(_postgres_rejection("23502"), False),
|
||||
],
|
||||
ids=[
|
||||
"restarting_server_keeps_the_rows",
|
||||
"full_server_keeps_the_rows",
|
||||
"terminated_session_may_have_counted_so_the_rows_are_dropped",
|
||||
"oversized_value_is_dropped",
|
||||
"rejected_row_is_dropped",
|
||||
],
|
||||
)
|
||||
async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply(
|
||||
failure: Exception, requeued: bool
|
||||
):
|
||||
queue = ToolDiscoveryQueue()
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.execute_raw = AsyncMock(side_effect=failure)
|
||||
|
||||
await commit_tool_discoveries(
|
||||
prisma_client=prisma_client,
|
||||
tool_discovery_queue=queue,
|
||||
redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()),
|
||||
pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()),
|
||||
use_redis_buffer=False,
|
||||
)
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
|
||||
|
||||
expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}]
|
||||
assert queue.flush() == expected
|
||||
|
|
|
|||
|
|
@ -3053,3 +3053,21 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none():
|
|||
)
|
||||
|
||||
assert len(guardrail.apply_calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_moderation_hook_lets_responses_compact_through_without_a_translation() -> None:
|
||||
handler: Final = UnifiedLLMGuardrails()
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
data: Final = {
|
||||
"guardrail_to_apply": guardrail,
|
||||
"model": "gpt-5.4-mini",
|
||||
"input": [{"role": "user", "content": "Remember the number 42."}],
|
||||
}
|
||||
|
||||
result: Final = await handler.async_moderation_hook(
|
||||
data=data, user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), call_type="acompact_responses"
|
||||
)
|
||||
|
||||
assert result is data
|
||||
assert guardrail.apply_calls == []
|
||||
|
|
|
|||
|
|
@ -23,9 +23,9 @@ import pytest
|
|||
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
|
||||
from litellm.proxy.utils import (
|
||||
MAX_SPEND_LOG_DRAIN_ITERATIONS,
|
||||
drain_spend_logs_queue,
|
||||
monitor_spend_logs_queue,
|
||||
raise_failed_update_spend_exception,
|
||||
drain_spend_logs_queue,
|
||||
recover_parked_spend_logs,
|
||||
update_daily_tag_spend,
|
||||
update_spend,
|
||||
|
|
@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o
|
|||
} == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []}
|
||||
|
||||
|
||||
def _postgres_restarting() -> Exception:
|
||||
"""prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after
|
||||
a restart: a base ``DataError`` carrying only the connector message."""
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Error in connector: Error querying the database: FATAL: the database system is starting up",
|
||||
"backtrace": None,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _postgres_unreachable() -> Exception:
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Can't reach database server at `db.internal`:`5432`",
|
||||
"error_code": "P1001",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _postgres_terminated_session() -> Exception:
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Error in connector: Error querying the database: FATAL: terminating connection due to "
|
||||
"administrator command",
|
||||
"meta": {"code": "57P01", "message": "terminating connection due to administrator command"},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"])
|
||||
async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts(
|
||||
error: Exception,
|
||||
mock_prisma_client: MagicMock,
|
||||
make_spend_log_row: Callable[..., dict[str, object]],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A restarting server refuses or cuts every statement, so the batch is not bisected
|
||||
into per-row drops: it is retried whole and then goes back to the queue head."""
|
||||
monkeypatch.setattr(asyncio, "sleep", AsyncMock())
|
||||
proxy_logging: Final = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [
|
||||
make_spend_log_row(request_id="r1"),
|
||||
make_spend_log_row(request_id="r2"),
|
||||
make_spend_log_row(request_id="r3"),
|
||||
make_spend_log_row(request_id="r4"),
|
||||
]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error)
|
||||
|
||||
with suppress(Exception):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
)
|
||||
|
||||
assert {
|
||||
"create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count,
|
||||
"queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions],
|
||||
} == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]}
|
||||
|
||||
|
||||
def _tool_usage_transaction(request_id: str) -> object:
|
||||
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
|
||||
|
||||
|
|
@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out
|
|||
} == {"queue_after": [first, second], "autorouter_queue_after": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"])
|
||||
async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes(
|
||||
error: Exception,
|
||||
mock_prisma_client: MagicMock,
|
||||
) -> None:
|
||||
"""A server that is restarting or unreachable never saw the rollup batch, so the tool
|
||||
usage increments are kept for the next flush, and the job stops before draining the
|
||||
auto-router queue into the same outage."""
|
||||
mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error)
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2")
|
||||
mock_prisma_client.tool_usage_transactions = [first, second]
|
||||
mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")]
|
||||
|
||||
with suppress(Exception):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert {
|
||||
"queue_after": mock_prisma_client.tool_usage_transactions,
|
||||
"autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions),
|
||||
} == {"queue_after": [first, second], "autorouter_queue_after": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rollups_drop_a_batch_whose_session_the_server_terminated(
|
||||
mock_prisma_client: MagicMock,
|
||||
) -> None:
|
||||
"""A terminated session (57P01) may have committed before it was cut, and the rollup
|
||||
increments are not idempotent, so the batch is dropped rather than counted twice."""
|
||||
mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session())
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")]
|
||||
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert mock_prisma_client.tool_usage_transactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied(
|
||||
mock_prisma_client: MagicMock,
|
||||
|
|
|
|||
|
|
@ -1236,3 +1236,58 @@ async def test_router_with_prompt_caching(anthropic_messages):
|
|||
print("response=", response)
|
||||
|
||||
assert response._hidden_params["model_id"] == initial_model_id
|
||||
|
||||
|
||||
class _LoggedMessagesProbe(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
self.logged_messages: list[object] = []
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"]))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages(
|
||||
monkeypatch, local_model_cost_map
|
||||
):
|
||||
"""
|
||||
`turn_off_message_logging` redacted the shared success-event payload before the check read it, so the
|
||||
pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the
|
||||
group. The pin must land whatever the flag says, while a real logger alongside still sees only the
|
||||
redacted messages.
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False)
|
||||
probe = _LoggedMessagesProbe()
|
||||
monkeypatch.setattr(litellm, "callbacks", [probe])
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": MODEL_GROUP_ALIAS,
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
for model_id in ("dep-1", "dep-2")
|
||||
],
|
||||
optional_pre_call_checks=["prompt_caching"],
|
||||
)
|
||||
messages = _messages(word_count=5000)
|
||||
|
||||
first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok")
|
||||
served_by = first._hidden_params["model_id"]
|
||||
|
||||
affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None)
|
||||
pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key))
|
||||
assert pin == {"model_id": served_by}, pin
|
||||
|
||||
subsequent = [
|
||||
(await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[
|
||||
"model_id"
|
||||
]
|
||||
for _ in range(4)
|
||||
]
|
||||
assert subsequent == [served_by] * 4
|
||||
|
||||
await _eventually(lambda: len(probe.logged_messages) == 5)
|
||||
redacted = [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ from litellm.utils import (
|
|||
_snapshot_exception_for_hook,
|
||||
async_post_call_failure_deployment_hook,
|
||||
async_post_call_success_deployment_hook,
|
||||
async_pre_call_deployment_hook,
|
||||
calculate_max_parallel_requests,
|
||||
client,
|
||||
get_non_default_completion_params,
|
||||
|
|
@ -10052,3 +10053,25 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_
|
|||
)
|
||||
|
||||
assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]}
|
||||
|
||||
|
||||
class _RecordingPreCallDeploymentHook(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seen_call_types: tuple[CallTypes | None, ...] = ()
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, object], call_type: CallTypes | None
|
||||
) -> dict[str, object] | None:
|
||||
self.seen_call_types = (*self.seen_call_types, call_type)
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_hands_responses_compact_its_call_type(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
hook: Final = _RecordingPreCallDeploymentHook()
|
||||
monkeypatch.setattr(litellm, "callbacks", [hook])
|
||||
|
||||
await async_pre_call_deployment_hook({"model": "gpt-5.4-mini"}, "acompact_responses")
|
||||
|
||||
assert hook.seen_call_types == (CallTypes.acompact_responses,)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,12 @@
|
|||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value
|
||||
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content
|
||||
from litellm.types.utils import (
|
||||
CAPTURE_MESSAGE_CONTENT_VALUES,
|
||||
CallTypes,
|
||||
CaptureMessageContent,
|
||||
captures_span_content,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"])
|
||||
|
|
@ -45,3 +50,8 @@ def test_capture_message_content_constructor_accepts_any_string() -> None:
|
|||
)
|
||||
def test_captures_span_content(mode: str | None, captures: bool) -> None:
|
||||
assert captures_span_content(mode) is captures
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["compact_responses", "acompact_responses"])
|
||||
def test_responses_compact_call_types_round_trip_their_public_names(call_type: str) -> None:
|
||||
assert CallTypes(call_type).value == call_type
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -28370,7 +28370,7 @@ export interface components {
|
|||
* CallTypes
|
||||
* @enum {string}
|
||||
*/
|
||||
CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill";
|
||||
CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "get_responses" | "aget_responses" | "delete_responses" | "adelete_responses" | "cancel_responses" | "acancel_responses" | "compact_responses" | "acompact_responses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill";
|
||||
/** CallbackDelete */
|
||||
CallbackDelete: {
|
||||
/** Callback Name */
|
||||
|
|
|
|||
6
uv.lock
generated
6
uv.lock
generated
|
|
@ -6194,11 +6194,11 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "oauthlib"
|
||||
version = "4.0.0"
|
||||
version = "4.0.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7a/d8/a1bcc8ba112a627f8ffbdc212a78ce18d3ac07e91a5ca65d27918eee25a1/oauthlib-4.0.0.tar.gz", hash = "sha256:efb274799819440f95b4ab3b818869f1ce9ae26c5beacba0201d1a1b76b54f86", size = 187232, upload-time = "2026-09-28T06:01:18.77Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/33/17/2bb072de0b990303c985f26869fd47bb179024a051f47e7f9d208eb92f8e/oauthlib-4.0.1.tar.gz", hash = "sha256:3d73febdb6faa00f277c81fb04be7b7586fb96fc3b7f0b9c772dada69a25ede7", size = 191219, upload-time = "2026-10-31T09:52:58.008Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d9/f4/78229a1066068ca14fc60fb26cf7381cabe4382261392b90e5f9552722d4/oauthlib-4.0.0-py3-none-any.whl", hash = "sha256:624c28c13a0a59cabf9747dfa52af63be3e512a7f2714df16e91b5b3a145e6cd", size = 159715, upload-time = "2026-09-28T06:01:17.008Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/df/ef/3b8adde1b32a6f540758bff734d3442c621734c83708e5f02b0ba2f43b84/oauthlib-4.0.1-py3-none-any.whl", hash = "sha256:498a298876e996301950a5b053d12c31628bb7b895db03666430c69ce182eb61", size = 161934, upload-time = "2026-10-31T09:52:56.206Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue