mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
Merge 01ab5bb7a1 into 24b8889965
This commit is contained in:
yucheng-berri
GitHub
commit
378360bf2c
7 changed files with 983 additions and 63 deletions
|
|
@ -4617,15 +4617,15 @@ class ProxyBaseLLMRequestProcessing:
|
|||
"""
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
|
||||
# is needed. When no callback overrides ``async_post_call_streaming_hook``,
|
||||
# no CustomGuardrail is active, and cost injection is disabled, the
|
||||
# per-chunk hook returns the chunk unchanged, ``str_so_far`` is never
|
||||
# consumed, and cost injection is a no-op -- so the per-chunk coroutine
|
||||
# await, response-string materialization, and cost-injection call are
|
||||
# pure overhead on the streaming hot path (the default config).
|
||||
# is needed. When no callback overrides ``async_post_call_streaming_hook``
|
||||
# and cost injection is disabled, the per-chunk hook returns the chunk
|
||||
# unchanged, ``str_so_far`` is never consumed, and cost injection is a
|
||||
# no-op -- so the per-chunk coroutine await, response-string
|
||||
# materialization, and cost-injection call are pure overhead on the
|
||||
# streaming hot path (the default config).
|
||||
caps: Final = ProxyLogging.callback_capabilities()
|
||||
cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False))
|
||||
fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled
|
||||
fast_path = not caps.has_streaming_chunk_override and not cost_injection_enabled
|
||||
debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG)
|
||||
stream_completed = False
|
||||
client_disconnected = False
|
||||
|
|
|
|||
|
|
@ -10545,10 +10545,10 @@ async def async_data_generator(
|
|||
# ``async_post_call_streaming_iterator_hook`` or has
|
||||
# ``apply_guardrail``; the per-chunk hook (which builds ``str_so_far``
|
||||
# and calls ``async_post_call_streaming_hook``) is only needed when
|
||||
# there is an active CustomGuardrail or a class that overrides the
|
||||
# per-chunk hook. Coalescing them into a single flag forced wasted
|
||||
# ``get_response_string`` work per chunk on every deployment that
|
||||
# happened to ship a streaming-iterator override (the default).
|
||||
# a class overrides the per-chunk hook. Coalescing them into a single
|
||||
# flag forced wasted ``get_response_string`` work per chunk on every
|
||||
# deployment that happened to ship a streaming-iterator override (the
|
||||
# default).
|
||||
needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook()
|
||||
is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream"))
|
||||
strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage"))
|
||||
|
|
|
|||
|
|
@ -1659,6 +1659,7 @@ class _CallbackCapabilities:
|
|||
# every string entry in ``litellm.callbacks``.
|
||||
resolved_callbacks: tuple[object, ...] = field(default_factory=tuple)
|
||||
listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple)
|
||||
streaming_chunk_overrides: tuple[CustomLogger, ...] = field(default_factory=tuple)
|
||||
|
||||
|
||||
def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool:
|
||||
|
|
@ -1666,6 +1667,18 @@ def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool:
|
|||
return any(hook_name in klass.__dict__ for klass in leaf_to_base)
|
||||
|
||||
|
||||
def _overrides_streaming_chunk_hook(callback: CustomLogger) -> bool:
|
||||
base_streaming_hook: Final = CustomLogger.async_post_call_streaming_hook
|
||||
cls_streaming_hook: Final = getattr(type(callback), "async_post_call_streaming_hook", base_streaming_hook)
|
||||
return getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr(
|
||||
base_streaming_hook, "__func__", base_streaming_hook
|
||||
)
|
||||
|
||||
|
||||
def _assigns_streaming_chunk_hook(callback: CustomLogger) -> bool:
|
||||
return "async_post_call_streaming_hook" in vars(callback)
|
||||
|
||||
|
||||
def _overrides_moderation_hook(callback: CustomLogger) -> bool:
|
||||
return _overrides_hook(callback, "async_moderation_hook")
|
||||
|
||||
|
|
@ -3369,7 +3382,6 @@ class ProxyLogging:
|
|||
|
||||
has_post_call_response_headers = False
|
||||
has_iterator_override = False
|
||||
has_streaming_chunk_override = False
|
||||
has_guardrail = False
|
||||
has_pre_call_override = False
|
||||
has_content_enforcer = False
|
||||
|
|
@ -3408,33 +3420,28 @@ class ProxyLogging:
|
|||
iterator_overrides.append((resolved, "override"))
|
||||
elif "apply_guardrail" in cls_attrs and not getattr(resolved, "use_native_lifecycle_hooks", False):
|
||||
iterator_overrides.append((resolved, "apply_guardrail"))
|
||||
# Walk the MRO for ``async_post_call_streaming_hook`` rather than
|
||||
# using the leaf-class ``__dict__`` check used by the other flags:
|
||||
# before this PR the hook was unconditionally invoked, so a
|
||||
# callback that inherits an override from an intermediate parent
|
||||
# (e.g. a vendor base class providing the override, with the
|
||||
# registered class adding nothing else) MUST still be detected.
|
||||
# A leaf-class miss here would silently drop the inherited hook.
|
||||
base_streaming_hook = CustomLogger.async_post_call_streaming_hook
|
||||
cls_streaming_hook = getattr(
|
||||
cls,
|
||||
"async_post_call_streaming_hook",
|
||||
base_streaming_hook,
|
||||
)
|
||||
if getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr(
|
||||
base_streaming_hook, "__func__", base_streaming_hook
|
||||
):
|
||||
has_streaming_chunk_override = True
|
||||
if "async_pre_call_hook" in cls_attrs:
|
||||
has_pre_call_override = True
|
||||
if resolved.enforces_request_content is True:
|
||||
has_content_enforcer = True
|
||||
|
||||
# A hook set on the instance runs per chunk only when a guardrail or a
|
||||
# class-level override is present, the cases where the per-chunk loop
|
||||
# used to reach it.
|
||||
reaches_instance_hooks: Final = has_guardrail or any(
|
||||
_overrides_streaming_chunk_hook(callback) for callback in resolved_callbacks
|
||||
)
|
||||
streaming_chunk_overrides: Final = tuple(
|
||||
callback
|
||||
for callback in resolved_callbacks
|
||||
if _overrides_streaming_chunk_hook(callback)
|
||||
or (reaches_instance_hooks and _assigns_streaming_chunk_hook(callback))
|
||||
)
|
||||
caps: Final = _CallbackCapabilities(
|
||||
has_post_call_response_headers=has_post_call_response_headers,
|
||||
has_iterator_override=has_iterator_override
|
||||
or any(kind == "apply_guardrail" for _, kind in iterator_overrides),
|
||||
has_streaming_chunk_override=has_streaming_chunk_override,
|
||||
has_streaming_chunk_override=bool(streaming_chunk_overrides),
|
||||
has_guardrail=has_guardrail,
|
||||
has_pre_call_override=has_pre_call_override,
|
||||
has_content_enforcer=has_content_enforcer,
|
||||
|
|
@ -3444,6 +3451,7 @@ class ProxyLogging:
|
|||
listed_models_filters=tuple(
|
||||
callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models")
|
||||
),
|
||||
streaming_chunk_overrides=streaming_chunk_overrides,
|
||||
)
|
||||
# Limit cache to handle test churn without leaking; production
|
||||
# callback lists are stable so this rarely grows past 1 entry.
|
||||
|
|
@ -3477,8 +3485,7 @@ class ProxyLogging:
|
|||
def has_streaming_chunk_hook_overrides() -> bool:
|
||||
"""True iff any callback overrides ``async_post_call_streaming_hook``
|
||||
(the per-chunk hook, distinct from the iterator wrapper)."""
|
||||
caps: Final = ProxyLogging.callback_capabilities()
|
||||
return caps.has_streaming_chunk_override or caps.has_guardrail
|
||||
return ProxyLogging.callback_capabilities().has_streaming_chunk_override
|
||||
|
||||
def needs_iterator_wrap(self) -> bool:
|
||||
"""Whether ``async_data_generator`` needs to wrap the upstream stream
|
||||
|
|
@ -3492,8 +3499,7 @@ class ProxyLogging:
|
|||
``_apply_streaming_chunk_hooks`` for every emitted chunk. Instance
|
||||
method for the same reason as :py:meth:`needs_iterator_wrap`.
|
||||
"""
|
||||
caps: Final = ProxyLogging.callback_capabilities()
|
||||
return caps.has_streaming_chunk_override or caps.has_guardrail
|
||||
return ProxyLogging.callback_capabilities().has_streaming_chunk_override
|
||||
|
||||
@staticmethod
|
||||
def has_during_call_guardrails() -> bool:
|
||||
|
|
@ -4455,12 +4461,11 @@ class ProxyLogging:
|
|||
"""
|
||||
# Per-chunk fast path: skip the response-string materialization and
|
||||
# callback scan when no configured callback overrides
|
||||
# ``async_post_call_streaming_hook`` AND no CustomGuardrail is
|
||||
# active. ``get_response_string`` walks every choice/delta on the
|
||||
# chunk so paying it per chunk for no-op callbacks dominated stream
|
||||
# CPU time even after the iterator-chain fix.
|
||||
# ``async_post_call_streaming_hook``. ``get_response_string`` walks
|
||||
# every choice/delta on the chunk so paying it per chunk for no-op
|
||||
# callbacks dominated stream CPU time even after the iterator-chain fix.
|
||||
caps: Final = ProxyLogging.callback_capabilities()
|
||||
if not caps.has_streaming_chunk_override and not caps.has_guardrail:
|
||||
if not caps.has_streaming_chunk_override:
|
||||
return response
|
||||
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
|
@ -4481,9 +4486,8 @@ class ProxyLogging:
|
|||
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
|
||||
)
|
||||
|
||||
for callback in litellm.callbacks:
|
||||
for callback in caps.streaming_chunk_overrides:
|
||||
try:
|
||||
_callback: CustomLogger | None = None
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
if callback.guardrail_name in pipeline_gated:
|
||||
continue
|
||||
|
|
@ -4505,26 +4509,17 @@ class ProxyLogging:
|
|||
is not True
|
||||
):
|
||||
continue
|
||||
if isinstance(callback, str):
|
||||
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
cast(_custom_logger_compatible_callbacks_literal, callback)
|
||||
)
|
||||
if str_so_far is not None:
|
||||
complete_response = str_so_far + response_str
|
||||
else:
|
||||
_callback = callback
|
||||
if _callback is not None and isinstance(_callback, CustomLogger):
|
||||
if str_so_far is not None:
|
||||
complete_response = str_so_far + response_str
|
||||
else:
|
||||
complete_response = response_str
|
||||
callback_response: (
|
||||
ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
|
||||
)
|
||||
callback_response = await _callback.async_post_call_streaming_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=complete_response,
|
||||
)
|
||||
if callback_response is not None:
|
||||
response = callback_response
|
||||
complete_response = response_str
|
||||
callback_response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
|
||||
callback_response = await callback.async_post_call_streaming_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=complete_response,
|
||||
)
|
||||
if callback_response is not None:
|
||||
response = callback_response
|
||||
except Exception as e:
|
||||
raise e
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -0,0 +1,736 @@
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = [
|
||||
pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
|
||||
pytest.mark.xdist_group("request-only-guardrail-stream-fast-path"),
|
||||
]
|
||||
|
||||
Endpoint: TypeAlias = Literal["chat", "messages", "responses"]
|
||||
Client: TypeAlias = Literal["sync", "async"]
|
||||
|
||||
_MODEL: Final = "request-only-guardrail-stream"
|
||||
_MESSAGES_MODEL: Final = "request-only-guardrail-messages"
|
||||
_GUARDED_MODEL: Final = "request-only-guardrail-guarded"
|
||||
_GUARDRAIL: Final = "request-only-check"
|
||||
_CHUNK_GUARDRAIL: Final = "per-chunk-check"
|
||||
_OPT_IN: Final = "opt-in-chunk-check"
|
||||
_INHERITED: Final = "inherited-chunk-check"
|
||||
_RAISING: Final = "raising-chunk-check"
|
||||
_INSTANCE: Final = "instance-chunk-check"
|
||||
_CLASS_LOGGER: Final = "class-chunk-logger"
|
||||
_INSTANCE_LOGGER: Final = "instance-chunk-logger"
|
||||
_CHUNK_ERROR: Final = "Chunk check failed"
|
||||
_MODULE: Final = "request_only_guardrail"
|
||||
_LOG_ENV: Final = "REQUEST_ONLY_GUARDRAIL_LOG"
|
||||
_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses")
|
||||
_PATHS: Final[Mapping[Endpoint, str]] = {
|
||||
"chat": "/v1/chat/completions",
|
||||
"messages": "/v1/messages",
|
||||
"responses": "/v1/responses",
|
||||
}
|
||||
_PIECES: Final = tuple(f"word{index} " for index in range(24))
|
||||
_MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})")
|
||||
_DROP: Final = re.compile(rb"drop-after-([0-9]+)")
|
||||
_MODULE_SOURCE: Final = f'''import json
|
||||
import os
|
||||
import re
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
MARKER = re.compile("marker-([0-9a-f]{{32}})")
|
||||
|
||||
|
||||
def _marker_in(value):
|
||||
found = MARKER.search(str(value))
|
||||
return None if found is None else found.group(1)
|
||||
|
||||
|
||||
def _log(marker, event):
|
||||
with open(os.environ["{_LOG_ENV}"], "a") as log:
|
||||
log.write(json.dumps({{"marker": marker, "event": event}}) + "\\n")
|
||||
|
||||
|
||||
class RequestOnlyGuardrail(CustomGuardrail):
|
||||
def should_run_guardrail(self, data, event_type):
|
||||
marker = _marker_in(data.get("messages") or data.get("input"))
|
||||
if marker is not None:
|
||||
_log(marker, str(getattr(event_type, "value", event_type)))
|
||||
return super().should_run_guardrail(data=data, event_type=event_type)
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
return data
|
||||
|
||||
|
||||
class PerChunkGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
|
||||
_log(_marker_in(response), "chunk:" + self.guardrail_name)
|
||||
return None
|
||||
|
||||
|
||||
class InheritedPerChunkGuardrail(PerChunkGuardrail):
|
||||
pass
|
||||
|
||||
|
||||
class InstancePerChunkGuardrail(CustomGuardrail):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.async_post_call_streaming_hook = self._see_chunk
|
||||
|
||||
async def _see_chunk(self, user_api_key_dict, response):
|
||||
_log(_marker_in(response), "chunk:" + self.guardrail_name)
|
||||
return None
|
||||
|
||||
|
||||
class RaisingPerChunkGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
|
||||
if "word3 " in response:
|
||||
raise HTTPException(status_code=400, detail={{"error": "{_CHUNK_ERROR}"}})
|
||||
return None
|
||||
|
||||
|
||||
class ClassChunkLogger(CustomLogger):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
|
||||
_log(_marker_in(response), "chunk:{_CLASS_LOGGER}")
|
||||
return None
|
||||
|
||||
|
||||
class InstanceChunkLogger(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.async_post_call_streaming_hook = self._see_chunk
|
||||
|
||||
async def _see_chunk(self, user_api_key_dict, response):
|
||||
_log(_marker_in(response), "chunk:{_INSTANCE_LOGGER}")
|
||||
return None
|
||||
|
||||
|
||||
class_chunk_logger = ClassChunkLogger()
|
||||
instance_chunk_logger = InstanceChunkLogger()
|
||||
'''
|
||||
|
||||
|
||||
def _text(marker: str) -> str:
|
||||
return f"marker-{marker} " + "".join(_PIECES)
|
||||
|
||||
|
||||
def _pieces(marker: str) -> tuple[str, ...]:
|
||||
return (f"marker-{marker} ", *_PIECES)
|
||||
|
||||
|
||||
def _sse(name: str | None, payload: Mapping[str, JsonValue]) -> bytes:
|
||||
prefix: Final = f"event: {name}\n" if name is not None else ""
|
||||
return f"{prefix}data: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _chat_chunk(identity: str, delta: Mapping[str, JsonValue], finish: str | None = None) -> bytes:
|
||||
return _sse(
|
||||
None,
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _chat_stream(marker: str) -> tuple[bytes, ...]:
|
||||
identity: Final = f"chatcmpl-{marker}"
|
||||
first, *rest = _pieces(marker)
|
||||
return (
|
||||
_chat_chunk(identity, {"role": "assistant", "content": first}),
|
||||
*(_chat_chunk(identity, {"content": piece}) for piece in rest),
|
||||
_chat_chunk(identity, {}, finish="stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
|
||||
|
||||
def _message(marker: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"msg_{marker}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"content": [{"type": "text", "text": _text(marker)}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 2, "output_tokens": 25},
|
||||
}
|
||||
|
||||
|
||||
def _messages_stream(marker: str) -> tuple[bytes, ...]:
|
||||
started: Final = {**_message(marker), "content": [], "stop_reason": None}
|
||||
events: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{"type": "message_start", "message": {**started, "usage": {"input_tokens": 2, "output_tokens": 0}}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
*(
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}}
|
||||
for piece in _pieces(marker)
|
||||
),
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 25}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
return tuple(_sse(string_value(event["type"]), event) for event in events)
|
||||
|
||||
|
||||
def _response_output(marker: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"type": "message",
|
||||
"id": f"msg_{marker}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": _text(marker), "annotations": []}],
|
||||
}
|
||||
|
||||
|
||||
def _response(marker: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"resp_{marker}",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [_response_output(marker)],
|
||||
"usage": {"input_tokens": 2, "output_tokens": 25, "total_tokens": 27},
|
||||
}
|
||||
|
||||
|
||||
def _responses_stream(marker: str) -> tuple[bytes, ...]:
|
||||
response: Final = _response(marker)
|
||||
output: Final = _response_output(marker)
|
||||
part: Final[dict[str, JsonValue]] = {"type": "output_text", "text": _text(marker), "annotations": []}
|
||||
position: Final[dict[str, JsonValue]] = {"item_id": f"msg_{marker}", "output_index": 0, "content_index": 0}
|
||||
started: Final = {**response, "status": "in_progress", "output": []}
|
||||
events: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{"type": "response.created", "response": started},
|
||||
{"type": "response.in_progress", "response": started},
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"item": {**output, "status": "in_progress", "content": []},
|
||||
"output_index": 0,
|
||||
},
|
||||
{"type": "response.content_part.added", **position, "part": {**part, "text": ""}},
|
||||
*({"type": "response.output_text.delta", **position, "delta": piece} for piece in _pieces(marker)),
|
||||
{"type": "response.output_text.done", **position, "text": _text(marker)},
|
||||
{"type": "response.content_part.done", **position, "part": part},
|
||||
{"type": "response.output_item.done", "item": output, "output_index": 0},
|
||||
{"type": "response.completed", "response": response},
|
||||
)
|
||||
return tuple(
|
||||
_sse(string_value(event["type"]), {**event, "sequence_number": index}) for index, event in enumerate(events)
|
||||
)
|
||||
|
||||
|
||||
def _chat_completion(marker: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"chatcmpl-{marker}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": _text(marker)}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 2, "completion_tokens": 25, "total_tokens": 27},
|
||||
}
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
found: Final = _MARKER.search(request.body)
|
||||
assert found is not None, request.body
|
||||
marker: Final = found.group(1).decode()
|
||||
streamed: Final = JSON_OBJECT.validate_json(request.body).get("stream") is True
|
||||
if request.target.endswith("/messages"):
|
||||
if not streamed:
|
||||
return Reply(body=json.dumps(_message(marker)).encode())
|
||||
return Reply(content_type="text/event-stream", chunks=_messages_stream(marker))
|
||||
if request.target.endswith("/responses"):
|
||||
if not streamed:
|
||||
return Reply(body=json.dumps(_response(marker)).encode())
|
||||
return Reply(content_type="text/event-stream", chunks=_responses_stream(marker))
|
||||
assert request.target.endswith("/chat/completions"), request.target
|
||||
if not streamed:
|
||||
return Reply(body=json.dumps(_chat_completion(marker)).encode())
|
||||
drop: Final = _DROP.search(request.body)
|
||||
if drop is not None:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=_chat_stream(marker),
|
||||
abort_after=int(drop.group(1)),
|
||||
pause_between_chunks=0.02,
|
||||
)
|
||||
return Reply(content_type="text/event-stream", chunks=_chat_stream(marker))
|
||||
|
||||
|
||||
def _guardrail(name: str, guardrail_class: str, mode: str, *, default_on: bool = True) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {"guardrail": f"{_MODULE}.{guardrail_class}", "mode": mode, "default_on": default_on},
|
||||
}
|
||||
|
||||
|
||||
def _deployment(name: str, model: str, api_base: str, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-stream-key", **extra},
|
||||
}
|
||||
|
||||
|
||||
def _config(upstream_url: str, guardrails: list[JsonValue], settings: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
return {
|
||||
**base,
|
||||
"litellm_settings": {**object_value(base["litellm_settings"]), **settings},
|
||||
"guardrails": guardrails,
|
||||
"model_list": [
|
||||
_deployment(_MODEL, "openai/gpt-4o-mini", f"{upstream_url}/v1"),
|
||||
_deployment(_MESSAGES_MODEL, "anthropic/claude-sonnet-4-5-20250929", upstream_url),
|
||||
_deployment(_GUARDED_MODEL, "openai/gpt-4o-mini", f"{upstream_url}/v1", guardrails=[_OPT_IN]),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: OwnedProxy
|
||||
log: Path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rig(
|
||||
gateway: Gateway,
|
||||
directory: Path,
|
||||
guardrails: list[JsonValue],
|
||||
settings: Mapping[str, JsonValue] | None = None,
|
||||
) -> Iterator[Rig]:
|
||||
with wire_server(_respond) as upstream:
|
||||
(directory / f"{_MODULE}.py").write_text(_MODULE_SOURCE)
|
||||
config_path: Final = directory / "request-only-guardrail-stream-fast-path.yaml"
|
||||
config_path.write_text(yaml.safe_dump(_config(upstream.url, guardrails, settings or {})))
|
||||
log: Final = directory / "guardrail-asks.jsonl"
|
||||
log.touch()
|
||||
with owned_proxy_process(gateway, directory, {_LOG_ENV: str(log)}, config=config_path, workers=2) as proxy:
|
||||
yield Rig(proxy, log)
|
||||
|
||||
|
||||
def _request_only() -> list[JsonValue]:
|
||||
return [_guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call")]
|
||||
|
||||
|
||||
def _mixed() -> list[JsonValue]:
|
||||
return [*_request_only(), _guardrail(_CHUNK_GUARDRAIL, "PerChunkGuardrail", "post_call")]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
session_gateway, tmp_path_factory.mktemp("request-only-guardrail-stream-fast-path"), _request_only()
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def mixed_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
session_gateway, tmp_path_factory.mktemp("request-only-guardrail-beside-per-chunk-guardrail"), _mixed()
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def opt_in_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
session_gateway,
|
||||
tmp_path_factory.mktemp("request-only-guardrail-beside-opt-in-per-chunk-guardrails"),
|
||||
[
|
||||
*_request_only(),
|
||||
_guardrail(_OPT_IN, "PerChunkGuardrail", "post_call", default_on=False),
|
||||
_guardrail(_INHERITED, "InheritedPerChunkGuardrail", "post_call", default_on=False),
|
||||
_guardrail(_RAISING, "RaisingPerChunkGuardrail", "post_call", default_on=False),
|
||||
],
|
||||
{"include_cost_in_streaming_usage": True},
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def instance_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
session_gateway,
|
||||
tmp_path_factory.mktemp("request-only-guardrail-beside-instance-per-chunk-hook"),
|
||||
[*_request_only(), _guardrail(_INSTANCE, "InstancePerChunkGuardrail", "post_call")],
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def logger_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
session_gateway,
|
||||
tmp_path_factory.mktemp("per-chunk-loggers-without-guardrails"),
|
||||
[],
|
||||
{"callbacks": [f"{_MODULE}.class_chunk_logger", f"{_MODULE}.instance_chunk_logger"]},
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def chaos_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(session_gateway, tmp_path_factory.mktemp("request-only-guardrail-stream-chaos"), _mixed()) as built:
|
||||
yield built
|
||||
|
||||
|
||||
def _events(rig: Rig, marker: str) -> tuple[str, ...]:
|
||||
entries: Final = (JSON_OBJECT.validate_json(line) for line in rig.log.read_text().splitlines() if line)
|
||||
return tuple(string_value(entry["event"]) for entry in entries if entry["marker"] == marker)
|
||||
|
||||
|
||||
def _prompt(marker: str, directive: str) -> str:
|
||||
return f"marker-{marker} {directive}".strip()
|
||||
|
||||
|
||||
def _body(
|
||||
endpoint: Endpoint,
|
||||
marker: str,
|
||||
*,
|
||||
stream: bool = True,
|
||||
model: str = _MODEL,
|
||||
guardrails: Sequence[str] = (),
|
||||
directive: str = "",
|
||||
) -> dict[str, JsonValue]:
|
||||
requested: Final[dict[str, JsonValue]] = {"guardrails": list(guardrails)} if guardrails else {}
|
||||
if endpoint == "responses":
|
||||
return {"model": model, "stream": stream, "input": _prompt(marker, directive), **requested}
|
||||
messages: Final[list[JsonValue]] = [{"role": "user", "content": _prompt(marker, directive)}]
|
||||
if endpoint == "messages":
|
||||
return {"model": _MESSAGES_MODEL, "stream": stream, "max_tokens": 64, "messages": messages, **requested}
|
||||
return {"model": model, "stream": stream, "messages": messages, "num_retries": 0, **requested}
|
||||
|
||||
|
||||
def _payloads(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
|
||||
|
||||
def _event_text(event: Mapping[str, JsonValue], endpoint: Endpoint) -> str:
|
||||
if endpoint == "responses":
|
||||
return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else ""
|
||||
if endpoint == "messages":
|
||||
if event.get("type") != "content_block_delta":
|
||||
return ""
|
||||
return string_value(object_value(event["delta"]).get("text", ""))
|
||||
return "".join(
|
||||
string_value(object_value(object_value(choice).get("delta", {})).get("content") or "")
|
||||
for choice in list_value(event.get("choices") or [])
|
||||
)
|
||||
|
||||
|
||||
def _streamed_text(text: str, endpoint: Endpoint) -> str:
|
||||
return "".join(_event_text(event, endpoint) for event in _payloads(text))
|
||||
|
||||
|
||||
def _send(
|
||||
rig: Rig,
|
||||
endpoint: Endpoint,
|
||||
marker: str,
|
||||
*,
|
||||
stream: bool = True,
|
||||
model: str = _MODEL,
|
||||
guardrails: Sequence[str] = (),
|
||||
directive: str = "",
|
||||
) -> httpx.Response:
|
||||
body: Final = _body(endpoint, marker, stream=stream, model=model, guardrails=guardrails, directive=directive)
|
||||
return rig.proxy.gateway.request("POST", _PATHS[endpoint], body)
|
||||
|
||||
|
||||
def _stream(rig: Rig, endpoint: Endpoint, marker: str, *, model: str = _MODEL, guardrails: Sequence[str] = ()) -> str:
|
||||
response: Final = _send(rig, endpoint, marker, model=model, guardrails=guardrails)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
assert _streamed_text(response.text, endpoint) == _text(marker), response.text
|
||||
return response.text
|
||||
|
||||
|
||||
def _delivered_chunks(text: str) -> int:
|
||||
return sum(1 for payload in _payloads(text) if "choices" in payload)
|
||||
|
||||
|
||||
def _assert_every_chunk_seen(rig: Rig, marker: str, guardrail: str, text: str) -> None:
|
||||
seen: Final = _events(rig, marker).count(f"chunk:{guardrail}")
|
||||
assert seen >= _delivered_chunks(text) > 0, (seen, _delivered_chunks(text), _events(rig, marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
|
||||
def test_a_request_only_guardrail_is_not_asked_about_each_streamed_chunk(rig: Rig, endpoint: Endpoint) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_stream(rig, endpoint, marker)
|
||||
asks: Final = _events(rig, marker)
|
||||
assert "pre_call" in asks, asks
|
||||
assert asks.count("post_call") == 0, f"asked about the response {asks.count('post_call')} times: {asks}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
|
||||
def test_a_request_only_guardrail_answers_non_streaming_calls_on_every_endpoint(rig: Rig, endpoint: Endpoint) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = _send(rig, endpoint, marker, stream=False)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _text(marker) in response.text, response.text
|
||||
asks: Final = _events(rig, marker)
|
||||
assert "pre_call" in asks, asks
|
||||
assert asks.count("post_call") <= 1, asks
|
||||
|
||||
|
||||
def _openai_text(client: openai.OpenAI, endpoint: Endpoint, marker: str) -> str:
|
||||
if endpoint == "responses":
|
||||
events: Final = client.responses.create(model=_MODEL, input=_prompt(marker, ""), stream=True)
|
||||
return "".join(event.delta for event in events if event.type == "response.output_text.delta")
|
||||
chunks: Final = client.chat.completions.create(
|
||||
model=_MODEL, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True
|
||||
)
|
||||
return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
|
||||
|
||||
|
||||
async def _async_openai_text(client: openai.AsyncOpenAI, endpoint: Endpoint, marker: str) -> str:
|
||||
pieces: Final[list[str]] = []
|
||||
if endpoint == "responses":
|
||||
async for event in await client.responses.create(model=_MODEL, input=_prompt(marker, ""), stream=True):
|
||||
if event.type == "response.output_text.delta":
|
||||
pieces.append(event.delta)
|
||||
return "".join(pieces)
|
||||
async for chunk in await client.chat.completions.create(
|
||||
model=_MODEL, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True
|
||||
):
|
||||
pieces.extend(choice.delta.content or "" for choice in chunk.choices)
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def _anthropic_text(client: anthropic.Anthropic, marker: str) -> str:
|
||||
events: Final = client.messages.create(
|
||||
model=_MESSAGES_MODEL, max_tokens=64, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True
|
||||
)
|
||||
return "".join(
|
||||
event.delta.text for event in events if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
)
|
||||
|
||||
|
||||
async def _async_anthropic_text(client: anthropic.AsyncAnthropic, marker: str) -> str:
|
||||
pieces: Final[list[str]] = []
|
||||
async for event in await client.messages.create(
|
||||
model=_MESSAGES_MODEL, max_tokens=64, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True
|
||||
):
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta":
|
||||
pieces.append(event.delta.text)
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def _sync_sdk_text(gateway: Gateway, endpoint: Endpoint, marker: str) -> str:
|
||||
origin: Final = str(gateway.client.base_url).rstrip("/")
|
||||
with httpx.Client(timeout=30, trust_env=False) as transport:
|
||||
if endpoint == "messages":
|
||||
with anthropic.Anthropic(
|
||||
api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport
|
||||
) as client:
|
||||
return _anthropic_text(client, marker)
|
||||
with openai.OpenAI(
|
||||
api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport
|
||||
) as client:
|
||||
return _openai_text(client, endpoint, marker)
|
||||
|
||||
|
||||
def _async_sdk_text(gateway: Gateway, endpoint: Endpoint, marker: str) -> str:
|
||||
origin: Final = str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
async def collect() -> str:
|
||||
async with httpx.AsyncClient(timeout=30, trust_env=False) as transport:
|
||||
if endpoint == "messages":
|
||||
async with anthropic.AsyncAnthropic(
|
||||
api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport
|
||||
) as client:
|
||||
return await _async_anthropic_text(client, marker)
|
||||
async with openai.AsyncOpenAI(
|
||||
api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport
|
||||
) as client:
|
||||
return await _async_openai_text(client, endpoint, marker)
|
||||
|
||||
return asyncio.run(collect())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ["sync", "async"])
|
||||
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
|
||||
def test_sdk_streams_through_a_request_only_guardrail_without_per_chunk_asks(
|
||||
rig: Rig, endpoint: Endpoint, client: Client
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
read: Final = _sync_sdk_text if client == "sync" else _async_sdk_text
|
||||
assert read(rig.proxy.gateway, endpoint, marker) == _text(marker)
|
||||
asks: Final = _events(rig, marker)
|
||||
assert "pre_call" in asks, asks
|
||||
assert asks.count("post_call") == 0, asks
|
||||
|
||||
|
||||
def test_a_request_only_guardrail_stream_writes_one_success_spend_row_naming_it(rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = _send(rig, "chat", marker)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _streamed_text(response.text, "chat") == _text(marker), response.text
|
||||
call_id: Final = response.headers["x-litellm-call-id"]
|
||||
rows: Final = eventually(
|
||||
lambda: rig.proxy.gateway.rows(
|
||||
'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', (call_id,)
|
||||
),
|
||||
lambda found: len(found) >= 1,
|
||||
seconds=30,
|
||||
)
|
||||
assert [row["status"] for row in rows] == ["success"], rows
|
||||
assert _GUARDRAIL in list_value(object_value(rows[0]["metadata"]).get("applied_guardrails") or []), rows
|
||||
|
||||
|
||||
def test_a_per_chunk_guardrail_still_sees_every_chunk_without_asking_the_request_only_one(mixed_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(mixed_rig, "chat", marker)
|
||||
asks: Final = _events(mixed_rig, marker)
|
||||
assert "pre_call" in asks, asks
|
||||
assert asks.count("post_call") <= 1, f"asked about the response {asks.count('post_call')} times: {asks}"
|
||||
_assert_every_chunk_seen(mixed_rig, marker, _CHUNK_GUARDRAIL, text)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ["messages", "responses"])
|
||||
def test_a_per_chunk_guardrail_beside_a_request_only_one_keeps_other_endpoints_whole(
|
||||
mixed_rig: Rig, endpoint: Endpoint
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_stream(mixed_rig, endpoint, marker)
|
||||
asks: Final = _events(mixed_rig, marker)
|
||||
assert "pre_call" in asks, asks
|
||||
assert asks.count("post_call") <= 1, asks
|
||||
|
||||
|
||||
def test_an_opt_in_per_chunk_guardrail_sees_every_chunk_only_when_requested(opt_in_rig: Rig) -> None:
|
||||
requested: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(opt_in_rig, "chat", requested, guardrails=[_OPT_IN])
|
||||
_assert_every_chunk_seen(opt_in_rig, requested, _OPT_IN, text)
|
||||
assert _events(opt_in_rig, requested).count("post_call") <= 1, _events(opt_in_rig, requested)
|
||||
|
||||
skipped: Final = uuid.uuid4().hex
|
||||
_stream(opt_in_rig, "chat", skipped)
|
||||
assert not any(event.startswith("chunk:") for event in _events(opt_in_rig, skipped)), _events(opt_in_rig, skipped)
|
||||
assert _events(opt_in_rig, skipped).count("post_call") <= 1, _events(opt_in_rig, skipped)
|
||||
|
||||
|
||||
def test_an_inherited_per_chunk_hook_still_sees_every_chunk(opt_in_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(opt_in_rig, "chat", marker, guardrails=[_INHERITED])
|
||||
_assert_every_chunk_seen(opt_in_rig, marker, _INHERITED, text)
|
||||
|
||||
|
||||
def test_a_per_chunk_hook_set_on_the_instance_still_sees_every_chunk(instance_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(instance_rig, "chat", marker)
|
||||
_assert_every_chunk_seen(instance_rig, marker, _INSTANCE, text)
|
||||
assert _events(instance_rig, marker).count("post_call") <= 1, _events(instance_rig, marker)
|
||||
|
||||
|
||||
def test_a_per_chunk_hook_set_on_the_instance_runs_beside_a_class_hook_without_guardrails(logger_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(logger_rig, "chat", marker)
|
||||
_assert_every_chunk_seen(logger_rig, marker, _CLASS_LOGGER, text)
|
||||
_assert_every_chunk_seen(logger_rig, marker, _INSTANCE_LOGGER, text)
|
||||
|
||||
|
||||
def test_a_model_level_per_chunk_guardrail_sees_every_chunk(opt_in_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
text: Final = _stream(opt_in_rig, "chat", marker, model=_GUARDED_MODEL)
|
||||
_assert_every_chunk_seen(opt_in_rig, marker, _OPT_IN, text)
|
||||
|
||||
|
||||
def test_a_raising_per_chunk_guardrail_ends_the_stream_with_its_error(opt_in_rig: Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = _send(opt_in_rig, "chat", marker, guardrails=[_RAISING])
|
||||
assert response.status_code == 200, response.text
|
||||
lines: Final = [line for line in response.text.splitlines() if line.startswith("data: ")]
|
||||
assert lines[-1] == "data: [DONE]" and lines.count("data: [DONE]") == 1, response.text
|
||||
payloads: Final = _payloads(response.text)
|
||||
assert "error" in payloads[-1] and all("error" not in payload for payload in payloads[:-1]), response.text
|
||||
assert object_value(payloads[-1]["error"])["message"] == _CHUNK_ERROR, payloads[-1]
|
||||
delivered: Final = "".join(_event_text(payload, "chat") for payload in payloads[:-1])
|
||||
assert delivered and _text(marker).startswith(delivered) and "word3 " not in delivered, response.text
|
||||
assert _events(opt_in_rig, marker).count("post_call") <= 1, _events(opt_in_rig, marker)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
|
||||
def test_cost_injection_keeps_streams_whole_beside_a_request_only_guardrail(
|
||||
opt_in_rig: Rig, endpoint: Endpoint
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_stream(opt_in_rig, endpoint, marker)
|
||||
assert _events(opt_in_rig, marker).count("post_call") <= 1, _events(opt_in_rig, marker)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Sent:
|
||||
endpoint: Endpoint
|
||||
marker: str
|
||||
dropped: bool
|
||||
response: httpx.Response
|
||||
|
||||
|
||||
def _burst_member(rig: Rig, endpoint: Endpoint, dropped: bool) -> Sent:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
directive: Final = "drop-after-8" if dropped else ""
|
||||
return Sent(endpoint, marker, dropped, _send(rig, endpoint, marker, directive=directive))
|
||||
|
||||
|
||||
def _assert_burst_member(rig: Rig, sent: Sent) -> None:
|
||||
assert sent.response.status_code == 200, sent.response.text
|
||||
text: Final = _streamed_text(sent.response.text, sent.endpoint)
|
||||
if sent.dropped:
|
||||
assert text and _text(sent.marker).startswith(text) and text != _text(sent.marker), sent.response.text
|
||||
else:
|
||||
assert text == _text(sent.marker), sent.response.text
|
||||
asks: Final = _events(rig, sent.marker)
|
||||
assert "pre_call" in asks and asks.count("post_call") <= 1, asks
|
||||
if sent.endpoint == "chat":
|
||||
_assert_every_chunk_seen(rig, sent.marker, _CHUNK_GUARDRAIL, sent.response.text)
|
||||
|
||||
|
||||
def test_a_mixed_burst_with_dropped_upstreams_and_a_killed_worker_keeps_every_guardrail_contract(
|
||||
chaos_rig: Rig,
|
||||
) -> None:
|
||||
plan: Final[tuple[tuple[Endpoint, bool], ...]] = (
|
||||
*(("chat", index % 2 == 0) for index in range(8)),
|
||||
*(("messages", False) for _ in range(8)),
|
||||
*(("responses", False) for _ in range(8)),
|
||||
)
|
||||
with ThreadPoolExecutor(max_workers=len(plan)) as pool:
|
||||
burst: Final = tuple(pool.map(lambda member: _burst_member(chaos_rig, *member), plan))
|
||||
for sent in burst:
|
||||
_assert_burst_member(chaos_rig, sent)
|
||||
|
||||
victim: Final = min(chaos_rig.proxy.worker_pids())
|
||||
replacement: Final = chaos_rig.proxy.kill_worker(victim)
|
||||
assert replacement in chaos_rig.proxy.worker_pids(), (replacement, chaos_rig.proxy.worker_pids())
|
||||
for _ in range(8):
|
||||
_assert_burst_member(chaos_rig, _burst_member(chaos_rig, "chat", False))
|
||||
|
|
@ -39,6 +39,7 @@ from litellm.constants import (
|
|||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.opentelemetry import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
|
@ -5183,6 +5184,35 @@ class TestAsyncStreamingDataGeneratorFastPath:
|
|||
assert out == chunks # bytes pass through return_sse_chunk untouched
|
||||
hook_spy.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
hook_spy: Final = AsyncMock(side_effect=lambda **kw: kw["response"])
|
||||
monkeypatch.setattr(proxy_logging_obj, "async_post_call_streaming_hook", hook_spy)
|
||||
|
||||
chunks: Final = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"]
|
||||
out: Final = [
|
||||
c
|
||||
async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=self._aiter(chunks),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
request_data={"model": "claude-x"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
|
||||
serialize_error=lambda e: "data: error\n\n",
|
||||
)
|
||||
]
|
||||
|
||||
assert out == chunks
|
||||
hook_spy.assert_not_awaited()
|
||||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_path_runs_per_chunk_hook(self, monkeypatch):
|
||||
"""A callback that overrides async_post_call_streaming_hook forces the
|
||||
|
|
|
|||
|
|
@ -9,12 +9,14 @@ Covers ``_callback_capabilities`` (the cached deriver),
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging, _CallbackCapabilities
|
||||
|
||||
|
||||
|
|
@ -37,11 +39,26 @@ class _OverridesPerChunk(CustomLogger):
|
|||
return None
|
||||
|
||||
|
||||
class _PerChunkGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _OverridesPreCall(CustomLogger):
|
||||
async def async_pre_call_hook(self, *args, **kwargs): # type: ignore[override]
|
||||
return None
|
||||
|
||||
|
||||
def _instance_per_chunk_logger() -> CustomLogger:
|
||||
logger: Final = _PlainLogger()
|
||||
|
||||
async def per_chunk(user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
|
||||
return None
|
||||
|
||||
vars(logger)["async_post_call_streaming_hook"] = per_chunk
|
||||
return logger
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_caps_cache():
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
|
@ -195,10 +212,18 @@ def test_has_streaming_chunk_hook_overrides_truth_table(monkeypatch, mock_callba
|
|||
monkeypatch.setattr(litellm, "callbacks", [_OverridesIterator()])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["only_iterator_false"] = ProxyLogging.has_streaming_chunk_hook_overrides()
|
||||
monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["request_only_guardrail_false"] = ProxyLogging.has_streaming_chunk_hook_overrides()
|
||||
monkeypatch.setattr(litellm, "callbacks", [_PerChunkGuardrail(guardrail_name="g", event_hook="post_call")])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["per_chunk_guardrail_true"] = ProxyLogging.has_streaming_chunk_hook_overrides()
|
||||
assert snapshot == {
|
||||
"empty_false": False,
|
||||
"per_chunk_override_true": True,
|
||||
"only_iterator_false": False,
|
||||
"request_only_guardrail_false": False,
|
||||
"per_chunk_guardrail_true": True,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -251,10 +276,30 @@ def test_needs_per_chunk_streaming_hook_truth_table(proxy_logging, monkeypatch,
|
|||
monkeypatch.setattr(litellm, "callbacks", [_OverridesIterator()])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["only_iter_override_false"] = proxy_logging.needs_per_chunk_streaming_hook()
|
||||
monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["request_only_guardrail_false"] = proxy_logging.needs_per_chunk_streaming_hook()
|
||||
monkeypatch.setattr(litellm, "callbacks", [_PerChunkGuardrail(guardrail_name="g", event_hook="post_call")])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["per_chunk_guardrail_true"] = proxy_logging.needs_per_chunk_streaming_hook()
|
||||
monkeypatch.setattr(litellm, "callbacks", [_instance_per_chunk_logger()])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["instance_hook_alone_false"] = proxy_logging.needs_per_chunk_streaming_hook()
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[CustomGuardrail(guardrail_name="g", event_hook="pre_call"), _instance_per_chunk_logger()],
|
||||
)
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
snapshot["instance_hook_beside_guardrail_true"] = proxy_logging.needs_per_chunk_streaming_hook()
|
||||
assert snapshot == {
|
||||
"empty_false": False,
|
||||
"per_chunk_override_true": True,
|
||||
"only_iter_override_false": False,
|
||||
"request_only_guardrail_false": False,
|
||||
"per_chunk_guardrail_true": True,
|
||||
"instance_hook_alone_false": False,
|
||||
"instance_hook_beside_guardrail_true": True,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Final, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.native.streaming_iterator import (
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import Usage
|
||||
from litellm.types.utils import ModelResponseStream, Usage
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -384,6 +384,120 @@ async def test_async_post_call_streaming_hook_invokes_per_chunk_callback(proxy_l
|
|||
assert out.startswith("modified-")
|
||||
|
||||
|
||||
class _RequestOnlyGuardrail(CustomGuardrail):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(guardrail_name="request-only", event_hook=GuardrailEventHooks.pre_call, default_on=True)
|
||||
self.asked: List[GuardrailEventHooks] = []
|
||||
|
||||
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
self.asked.append(event_type)
|
||||
return super().should_run_guardrail(data=data, event_type=event_type)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardrail(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrail: Final = _RequestOnlyGuardrail()
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="rid",
|
||||
choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}],
|
||||
created=0,
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
out: Final = await proxy_logging.async_post_call_streaming_hook(
|
||||
data={"messages": [{"role": "user", "content": "hi"}]},
|
||||
response=chunk,
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
)
|
||||
assert out is chunk
|
||||
assert guardrail.asked == []
|
||||
|
||||
|
||||
class _PerChunkRewriter(CustomLogger):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> str:
|
||||
return f"modified-{response}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_beside_a_per_chunk_callback(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrail: Final = _RequestOnlyGuardrail()
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail, _PerChunkRewriter()])
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="rid",
|
||||
choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}],
|
||||
created=0,
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
out: Final = await proxy_logging.async_post_call_streaming_hook(
|
||||
data={"messages": [{"role": "user", "content": "hi"}]},
|
||||
response=chunk,
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
str_so_far="say ",
|
||||
)
|
||||
assert out == "modified-say hi"
|
||||
assert guardrail.asked == []
|
||||
|
||||
|
||||
def _instance_rewriter() -> CustomLogger:
|
||||
rewriter: Final = CustomLogger()
|
||||
|
||||
async def per_chunk(user_api_key_dict: UserAPIKeyAuth, response: str) -> str:
|
||||
return f"modified-{response}"
|
||||
|
||||
vars(rewriter)["async_post_call_streaming_hook"] = per_chunk
|
||||
return rewriter
|
||||
|
||||
|
||||
class _ClassChunkLogger(CustomLogger):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
async def _say_hi_through_chunk_hooks(proxy_logging: ProxyLogging, user_api_key_dict: UserAPIKeyAuth) -> object:
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="rid",
|
||||
choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}],
|
||||
created=0,
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
return await proxy_logging.async_post_call_streaming_hook(
|
||||
data={"messages": [{"role": "user", "content": "hi"}]},
|
||||
response=chunk,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
str_so_far="say ",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_guardrail(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrail: Final = _RequestOnlyGuardrail()
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail, _instance_rewriter()])
|
||||
assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"
|
||||
assert guardrail.asked == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_class_hook(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [_ClassChunkLogger(), _instance_rewriter()])
|
||||
assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_streaming_hook_callback_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
class _Per(CustomLogger):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue