mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(websearch_interception): drop stream_options when forcing stream=False and replay it on the final stream (#45433)
* fix(websearch_interception): drop stream_options when forcing stream=False and replay it on the final stream * test(websearch_interception): type the conversion helpers and drive the replay test through mock_response * test(websearch_interception): pin the replay test's values as Final and record follow-ups immutably * fix(proxy): strip every agentic-loop control key from client bodies, the web search keys included * fix(interception): move stream_options off forced non-stream calls for code interpreter and headroom, and keep adapters from replaying nested calls * test: annotate new interception test locals with Final * fix(responses): drop params a deployment hook removed on sync Responses calls * test: drive the sync Responses stream_options regression through a mock transport and real Azure config * test: type the agentic loop recorder and capture the sent request on an object * test(integration): cover interception stream_options on chat, responses, messages and the SDK * test(integration): assert the follow-up error text and release a held burst on failure * test(integration): pin the follow-up failure contract and wait for the replacement worker before teardown --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
4f08cf250b
commit
13928bdabe
23 changed files with 1662 additions and 72 deletions
|
|
@ -26,9 +26,11 @@ from litellm.types.integrations.code_interpreter_interception import (
|
|||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
CODE_INTERPRETER_INTERCEPTION_PREFIX,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
as_converted_stream,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -290,8 +292,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
|
|||
else:
|
||||
kwargs[_SANDBOX_KEY] = uuid.uuid4().hex
|
||||
if kwargs.get("stream"):
|
||||
kwargs["stream"] = False
|
||||
kwargs[_CONVERTED_STREAM_KEY] = True
|
||||
kwargs.update(as_converted_stream(kwargs, CODE_INTERPRETER_INTERCEPTION_PREFIX))
|
||||
kwargs.pop("stream_options", None)
|
||||
self._write_interception_metadata(kwargs)
|
||||
|
||||
function_tool: Final = self._get_function_tool(call_type=call_type)
|
||||
|
|
|
|||
|
|
@ -40,8 +40,12 @@ from litellm.llms.base_llm.search.transformation import SearchResponse
|
|||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
RESPONSES_AGENTIC_SURFACE,
|
||||
WEBSEARCH_CONVERTED_STREAM_KEY,
|
||||
WEBSEARCH_INTERCEPTION_PREFIX,
|
||||
WEBSEARCH_STREAM_OPTIONS_KEY,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
as_converted_stream,
|
||||
)
|
||||
from litellm.types.integrations.websearch_interception import (
|
||||
AnthropicSearchQuery,
|
||||
|
|
@ -489,8 +493,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
if kwargs.get("stream"):
|
||||
verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
|
||||
kwargs["stream"] = False
|
||||
kwargs["_websearch_interception_converted_stream"] = True
|
||||
return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
|
||||
|
||||
return kwargs
|
||||
|
||||
|
|
@ -511,8 +514,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
if kwargs.get("stream"):
|
||||
verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
|
||||
converted_kwargs["stream"] = False
|
||||
converted_kwargs["_websearch_interception_converted_stream"] = True
|
||||
return as_converted_stream(converted_kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
|
||||
|
||||
return converted_kwargs
|
||||
|
||||
|
|
@ -670,7 +672,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if kwargs.get("stream"):
|
||||
verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False")
|
||||
kwargs["stream"] = False
|
||||
kwargs["_websearch_interception_converted_stream"] = True
|
||||
kwargs[WEBSEARCH_CONVERTED_STREAM_KEY] = True
|
||||
if "stream_options" in kwargs:
|
||||
kwargs[WEBSEARCH_STREAM_OPTIONS_KEY] = kwargs["stream_options"]
|
||||
del kwargs["stream_options"]
|
||||
|
||||
return kwargs
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,9 @@ import json
|
|||
from collections.abc import Mapping
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -17,11 +19,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
HEADROOM_CONVERTED_STREAM_KEY,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
converted_stream_requested,
|
||||
is_interception_internal_key,
|
||||
stashed_stream_options,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
|
@ -57,10 +60,14 @@ def _post_hook_overridden(callback: CustomLogger) -> bool:
|
|||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool:
|
||||
return bool(
|
||||
kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY)
|
||||
)
|
||||
_STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
||||
|
||||
def _stashed_stream_options(kwargs: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _STREAM_OPTIONS_ADAPTER.validate_python(stashed_stream_options(kwargs))
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _coerce_int(value: object, default: int) -> int:
|
||||
|
|
@ -100,23 +107,28 @@ def _check_agentic_loop_safety(
|
|||
return fingerprint
|
||||
|
||||
|
||||
def _wrap_response_as_fake_stream(
|
||||
def _converted_stream_replay(
|
||||
response: object,
|
||||
*,
|
||||
kwargs: Mapping[str, object],
|
||||
depth: int,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
logging_obj: object,
|
||||
) -> object:
|
||||
) -> CustomStreamWrapper | None:
|
||||
if depth or not converted_stream_requested(kwargs):
|
||||
return None
|
||||
if isinstance(response, CustomStreamWrapper):
|
||||
return response
|
||||
if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject):
|
||||
return response
|
||||
return None
|
||||
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=MockResponseIterator(model_response=response),
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
stream_options=_stashed_stream_options(kwargs),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -213,14 +225,15 @@ async def _execute_chat_completion_agentic_plan(
|
|||
model,
|
||||
str(e),
|
||||
)
|
||||
if _converted_stream_requested(kwargs) and not depth:
|
||||
return _wrap_response_as_fake_stream(
|
||||
response_followup,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return response_followup
|
||||
replay: Final = _converted_stream_replay(
|
||||
response_followup,
|
||||
kwargs=kwargs,
|
||||
depth=depth,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return response_followup if replay is None else replay
|
||||
finally:
|
||||
try:
|
||||
await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs)
|
||||
|
|
@ -240,7 +253,7 @@ async def maybe_run_chat_completion_agentic_loop(
|
|||
model: str,
|
||||
messages: list,
|
||||
optional_params: dict,
|
||||
kwargs: dict,
|
||||
kwargs: dict[str, object],
|
||||
logging_obj: object,
|
||||
custom_llm_provider: str,
|
||||
stream: bool,
|
||||
|
|
@ -343,14 +356,11 @@ async def maybe_run_chat_completion_agentic_loop(
|
|||
str(e),
|
||||
)
|
||||
|
||||
if _converted_stream_requested(kwargs) and not depth:
|
||||
return cast(
|
||||
"ModelResponse | CustomStreamWrapper",
|
||||
_wrap_response_as_fake_stream(
|
||||
response,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
),
|
||||
)
|
||||
return None
|
||||
return _converted_stream_replay(
|
||||
response,
|
||||
kwargs=kwargs,
|
||||
depth=depth,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.llms.anthropic.pass_through.utils import (
|
|||
litellm_logging_obj_from_kwargs,
|
||||
local_model_name,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -564,7 +565,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
# Maintainability: when adding a new Anthropic-only request param to
|
||||
# ``AnthropicMessagesRequestOptionalParams``, also extend
|
||||
# ``ANTHROPIC_ONLY_REQUEST_KEYS`` here so it doesn't silently leak.
|
||||
excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | {"anthropic_messages"}
|
||||
excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | CONVERTED_STREAM_KEYS | {"anthropic_messages"}
|
||||
# NOTE: extra_kwargs was already coerced from None to {} at the top of
|
||||
# this method (line ~220). It is guaranteed to be a dict here.
|
||||
for key, value in extra_kwargs.items():
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Coroutine, Mapping
|
|||
from typing import Final, TypeAlias
|
||||
|
||||
import litellm
|
||||
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
|
||||
from litellm.types.llms.anthropic import (
|
||||
AllAnthropicMessageValues,
|
||||
AllAnthropicToolsValues,
|
||||
|
|
@ -126,7 +127,7 @@ def _build_responses_kwargs(
|
|||
responses_kwargs["stream"] = True
|
||||
|
||||
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
|
||||
excluded: Final = frozenset(("anthropic_messages",))
|
||||
excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages"}
|
||||
for key, value in forwarded_kwargs.items():
|
||||
if key == "litellm_logging_obj" and value is not None:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
|
|||
|
|
@ -2495,12 +2495,12 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
optional_param_names: Final = _responses_api_optional_request_param_names()
|
||||
updated_response_params: Final = {
|
||||
**response_api_optional_request_params,
|
||||
**{key: value for key, value in response_api_optional_request_params.items() if key in modified_kwargs},
|
||||
**{key: value for key, value in modified_kwargs.items() if key in optional_param_names},
|
||||
}
|
||||
updated_litellm_params: Final = GenericLiteLLMParams(
|
||||
**{
|
||||
**dict(litellm_params),
|
||||
**{key: value for key, value in dict(litellm_params).items() if key in modified_kwargs},
|
||||
**{
|
||||
key: value
|
||||
for key, value in modified_kwargs.items()
|
||||
|
|
|
|||
|
|
@ -42,9 +42,10 @@ from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
HEADROOM_CONVERTED_STREAM_KEY,
|
||||
HEADROOM_INTERCEPTION_PREFIX,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
as_converted_stream,
|
||||
)
|
||||
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -883,18 +884,14 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
call_type: CallTypes | None,
|
||||
) -> dict[str, object] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
|
||||
base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
|
||||
effective: Final = base_result if base_result is not None else kwargs
|
||||
if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
|
||||
return base_result
|
||||
effective: Final = kwargs if base_result is None else _REQUEST_DATA_ADAPTER.validate_python(base_result)
|
||||
if not effective.get("stream") or effective.get("background"):
|
||||
return base_result
|
||||
if not has_headroom_retrieve_tool(effective.get("tools")):
|
||||
return base_result
|
||||
return {
|
||||
**effective,
|
||||
"stream": False,
|
||||
HEADROOM_CONVERTED_STREAM_KEY: True,
|
||||
}
|
||||
return as_converted_stream(effective, HEADROOM_INTERCEPTION_PREFIX)
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -223,6 +223,7 @@ def sanitize_for_log(value: object) -> str:
|
|||
|
||||
from litellm.router import Router
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
|
||||
from litellm.types.services import ServiceTypes
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -342,14 +343,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
|
|||
# re-wrapped as a synthetic stream the caller never asked for, or raise the
|
||||
# loop ceiling to drive many upstream model calls and sandbox executions
|
||||
# from a single request.
|
||||
"_agentic_loop_depth",
|
||||
"_agentic_loop_fingerprints",
|
||||
"_code_interpreter_interception_active",
|
||||
"_code_interpreter_interception_converted_stream",
|
||||
"_code_interpreter_interception_sandbox_key",
|
||||
"_code_interpreter_interception_session_scoped",
|
||||
"_headroom_interception_converted_stream",
|
||||
"max_agentic_loops",
|
||||
*AGENTIC_LOOP_KWARG_NAMES,
|
||||
# Recomputed below from the actual caller-controlled timeout sources (headers and
|
||||
# body fields); a client-forged value here would let a request either dodge cooldown
|
||||
# protection on a real deployment failure or force a false "not caller-controlled"
|
||||
|
|
|
|||
|
|
@ -10,10 +10,20 @@ CHAT_COMPLETION_AGENTIC_SURFACE: Final[AgenticSurface] = "chat_completions"
|
|||
RESPONSES_AGENTIC_SURFACE: Final[AgenticSurface] = "responses"
|
||||
CODE_INTERPRETER_INTERCEPTION_PREFIX: Final = "_code_interpreter_interception"
|
||||
HEADROOM_INTERCEPTION_PREFIX: Final = "_headroom_interception"
|
||||
WEBSEARCH_INTERCEPTION_PREFIX: Final = "_websearch_interception"
|
||||
HEADROOM_CONVERTED_STREAM_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream"
|
||||
WEBSEARCH_CONVERTED_STREAM_KEY: Final = f"{WEBSEARCH_INTERCEPTION_PREFIX}_converted_stream"
|
||||
WEBSEARCH_STREAM_OPTIONS_KEY: Final = f"{WEBSEARCH_INTERCEPTION_PREFIX}_stream_options"
|
||||
CODE_INTERPRETER_STREAM_OPTIONS_KEY: Final = f"{CODE_INTERPRETER_INTERCEPTION_PREFIX}_stream_options"
|
||||
HEADROOM_STREAM_OPTIONS_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_stream_options"
|
||||
STREAM_OPTIONS_STASH_KEYS: Final = (
|
||||
CODE_INTERPRETER_STREAM_OPTIONS_KEY,
|
||||
HEADROOM_STREAM_OPTIONS_KEY,
|
||||
WEBSEARCH_STREAM_OPTIONS_KEY,
|
||||
)
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset(
|
||||
(
|
||||
"_websearch_interception",
|
||||
WEBSEARCH_INTERCEPTION_PREFIX,
|
||||
"_compression_interception",
|
||||
HEADROOM_INTERCEPTION_PREFIX,
|
||||
)
|
||||
|
|
@ -40,6 +50,23 @@ def converted_stream_requested(params: Mapping[str, object]) -> bool:
|
|||
return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS)
|
||||
|
||||
|
||||
def as_converted_stream(
|
||||
params: Mapping[str, object], prefix: str
|
||||
) -> dict[str, object]: # mutable-ok: deployment hooks hand back the kwargs dict their caller keeps editing
|
||||
stream_options: Final = params.get("stream_options")
|
||||
stash: Final = {} if stream_options is None else {f"{prefix}_stream_options": stream_options}
|
||||
return {
|
||||
**{key: value for key, value in params.items() if key != "stream_options"},
|
||||
"stream": False,
|
||||
f"{prefix}_converted_stream": True,
|
||||
**stash,
|
||||
}
|
||||
|
||||
|
||||
def stashed_stream_options(params: Mapping[str, object]) -> object:
|
||||
return next((params[key] for key in STREAM_OPTIONS_STASH_KEYS if params.get(key) is not None), None)
|
||||
|
||||
|
||||
class AgenticLoopSafetyError(ValueError):
|
||||
"""
|
||||
Raised when an agentic-loop safety rail refuses a rerun.
|
||||
|
|
|
|||
|
|
@ -355,15 +355,24 @@ class AgenticLoopState:
|
|||
code_interpreter_converted_stream: bool | None = field(
|
||||
default=None, metadata=wire("_code_interpreter_interception_converted_stream")
|
||||
)
|
||||
code_interpreter_stream_options: Mapping[str, object] | None = field(
|
||||
default=None, metadata=wire("_code_interpreter_interception_stream_options")
|
||||
)
|
||||
websearch_emit_native_blocks: bool | None = field(
|
||||
default=None, metadata=wire("_websearch_interception_emit_native_blocks")
|
||||
)
|
||||
websearch_converted_stream: bool | None = field(
|
||||
default=None, metadata=wire("_websearch_interception_converted_stream")
|
||||
)
|
||||
websearch_stream_options: Mapping[str, object] | None = field(
|
||||
default=None, metadata=wire("_websearch_interception_stream_options")
|
||||
)
|
||||
headroom_converted_stream: bool | None = field(
|
||||
default=None, metadata=wire("_headroom_interception_converted_stream")
|
||||
)
|
||||
headroom_stream_options: Mapping[str, object] | None = field(
|
||||
default=None, metadata=wire("_headroom_interception_stream_options")
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
|
|
|
|||
245
tests/integration/_support/interception_vendor.py
Normal file
245
tests/integration/_support/interception_vendor.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal, TypeAlias
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from integration._support.openai_wire import chat_stream, responses_stream
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
AZURE_KEY: Final = "synthetic-azure-key"
|
||||
TAVILY_KEY: Final = "synthetic-tavily-key"
|
||||
API_VERSION: Final = "2025-04-01-preview"
|
||||
SEARCH_TOOL: Final = "litellm_web_search"
|
||||
SEARCH_TARGET: Final = "/tavily/search"
|
||||
STREAM_OPTIONS_REFUSAL: Final = "The 'stream_options' parameter is only allowed when 'stream' is enabled."
|
||||
FOLLOWUP_FAILURE: Final = "scripted follow-up failure"
|
||||
FIRST_USAGE: Final = {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}
|
||||
FOLLOWUP_USAGE: Final = {"prompt_tokens": 23, "completion_tokens": 5, "total_tokens": 28}
|
||||
RESPONSES_FIRST_USAGE: Final = {"input_tokens": 13, "output_tokens": 6, "total_tokens": 19}
|
||||
RESPONSES_FOLLOWUP_USAGE: Final = {"input_tokens": 29, "output_tokens": 4, "total_tokens": 33}
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
_CHAT_TARGET: Final = re.compile(r"^/openai/deployments/([^/]+)/chat/completions$")
|
||||
|
||||
Mode: TypeAlias = Literal["search", "plain", "search-down", "followup-fails", "missing"]
|
||||
|
||||
|
||||
def deployment(mode: Mode) -> str:
|
||||
return f"dep-{uuid.uuid4().hex}-{mode}"
|
||||
|
||||
|
||||
def answer(name: str) -> str:
|
||||
return f"answer-{name}"
|
||||
|
||||
|
||||
def plain(name: str) -> str:
|
||||
return f"plain-{name}"
|
||||
|
||||
|
||||
def query(name: str) -> str:
|
||||
return f"q-{name}"
|
||||
|
||||
|
||||
def _minted(prefix: str) -> str:
|
||||
return f"{prefix}{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def _azure_error(status: int, message: str, *, param: str | None = None) -> Reply:
|
||||
return Reply(
|
||||
status=status,
|
||||
body=json.dumps(
|
||||
{"error": {"message": message, "type": "invalid_request_error", "param": param, "code": None}}
|
||||
).encode(),
|
||||
)
|
||||
|
||||
|
||||
def _refusal(body: Mapping[str, JsonValue]) -> Reply | None:
|
||||
if body.get("stream_options") is not None and body.get("stream") is not True:
|
||||
return _azure_error(400, STREAM_OPTIONS_REFUSAL, param="stream_options")
|
||||
unknown: Final = tuple(key for key in body if key.startswith("_"))
|
||||
if unknown:
|
||||
return _azure_error(400, f"Unrecognized request argument supplied: {unknown[0]}")
|
||||
return None
|
||||
|
||||
|
||||
def _offers_search(body: Mapping[str, JsonValue]) -> bool:
|
||||
tools: Final = _ITEMS.validate_python(body.get("tools") or [])
|
||||
return any(
|
||||
tool.get("name") == SEARCH_TOOL
|
||||
or JSON_OBJECT.validate_python(tool.get("function") or {}).get("name") == SEARCH_TOOL
|
||||
for tool in tools
|
||||
)
|
||||
|
||||
|
||||
def _chat_completion(identity: str, message: Mapping[str, JsonValue], finish: str, usage: Mapping[str, int]) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o",
|
||||
"choices": [{"index": 0, "message": dict(message), "finish_reason": finish}],
|
||||
"usage": dict(usage),
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _chat(name: str, body: Mapping[str, JsonValue]) -> Reply:
|
||||
if body.get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=chat_stream(_minted(f"chatcmpl-{name}-s-"), "gpt-4o", plain(name))
|
||||
)
|
||||
messages: Final = _ITEMS.validate_python(body["messages"])
|
||||
if any(message.get("role") == "tool" for message in messages):
|
||||
if name.endswith("-followup-fails"):
|
||||
return _azure_error(500, FOLLOWUP_FAILURE)
|
||||
return _chat_completion(
|
||||
_minted(f"chatcmpl-{name}-2-"), {"role": "assistant", "content": answer(name)}, "stop", FOLLOWUP_USAGE
|
||||
)
|
||||
if name.endswith("-plain") or not _offers_search(body):
|
||||
return _chat_completion(
|
||||
_minted(f"chatcmpl-{name}-1-"), {"role": "assistant", "content": plain(name)}, "stop", FIRST_USAGE
|
||||
)
|
||||
call: Final = {
|
||||
"id": f"call_{name[4:36]}",
|
||||
"type": "function",
|
||||
"function": {"name": SEARCH_TOOL, "arguments": json.dumps({"query": query(name)})},
|
||||
}
|
||||
message: Final = JSON_OBJECT.validate_python({"role": "assistant", "content": None, "tool_calls": [call]})
|
||||
return _chat_completion(_minted(f"chatcmpl-{name}-1-"), message, "tool_calls", FIRST_USAGE)
|
||||
|
||||
|
||||
def _response(identity: str, name: str, output: Sequence[Mapping[str, JsonValue]], usage: Mapping[str, int]) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": name,
|
||||
"output": [dict(item) for item in output],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": dict(usage),
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _text_item(identity: str, text: str) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
||||
}
|
||||
|
||||
|
||||
def _responses(body: Mapping[str, JsonValue]) -> Reply:
|
||||
name: Final = str(body["model"])
|
||||
if body.get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=responses_stream(_minted(f"resp_{name}_s_"), name, plain(name))
|
||||
)
|
||||
items: Final = _ITEMS.validate_python(body["input"]) if isinstance(body.get("input"), list) else []
|
||||
if any(item.get("type") == "function_call_output" for item in items):
|
||||
if name.endswith("-followup-fails"):
|
||||
return _azure_error(500, FOLLOWUP_FAILURE)
|
||||
return _response(
|
||||
_minted(f"resp_{name}_2_"), name, (_text_item(f"{name}_2", answer(name)),), RESPONSES_FOLLOWUP_USAGE
|
||||
)
|
||||
if name.endswith("-plain") or not _offers_search(body):
|
||||
return _response(
|
||||
_minted(f"resp_{name}_1_"), name, (_text_item(f"{name}_1", plain(name)),), RESPONSES_FIRST_USAGE
|
||||
)
|
||||
call: Final = {
|
||||
"type": "function_call",
|
||||
"id": f"fc_{name[4:36]}",
|
||||
"call_id": f"call_{name[4:36]}",
|
||||
"name": SEARCH_TOOL,
|
||||
"arguments": json.dumps({"query": query(name)}),
|
||||
"status": "completed",
|
||||
}
|
||||
return _response(_minted(f"resp_{name}_1_"), name, (call,), RESPONSES_FIRST_USAGE)
|
||||
|
||||
|
||||
def _search(request: Request, body: Mapping[str, JsonValue]) -> Reply:
|
||||
if request.headers.get("authorization") != f"Bearer {TAVILY_KEY}":
|
||||
return Reply(status=401, body=b'{"detail":{"error":"Unauthorized"}}')
|
||||
text: Final = str(body.get("query", ""))
|
||||
if text.endswith("-search-down"):
|
||||
return Reply(status=500, body=b'{"detail":{"error":"scripted search outage"}}')
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"query": text,
|
||||
"results": [
|
||||
{"title": f"title-{text}", "url": "https://example.test/result", "content": f"snippet-{text}"}
|
||||
],
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
target: Final = urlsplit(request.target).path
|
||||
if request.method == "GET":
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
body: Final = JSON_OBJECT.validate_json(request.body or b"{}")
|
||||
if target == SEARCH_TARGET:
|
||||
return _search(request, body)
|
||||
if request.headers.get("api-key") != AZURE_KEY:
|
||||
return Reply(
|
||||
status=401, body=b'{"error":{"code":"401","message":"Access denied due to invalid subscription key."}}'
|
||||
)
|
||||
chat: Final = _CHAT_TARGET.match(target)
|
||||
if chat is not None and chat.group(1).endswith("-missing"):
|
||||
return Reply(
|
||||
status=404,
|
||||
body=b'{"error":{"code":"DeploymentNotFound","message":"The API deployment for this resource does not exist."}}',
|
||||
)
|
||||
refused: Final = _refusal(body)
|
||||
if refused is not None:
|
||||
return refused
|
||||
if chat is not None:
|
||||
return _chat(chat.group(1), body)
|
||||
if target.endswith("/responses"):
|
||||
return _responses(body)
|
||||
return Reply(status=404, body=b'{"error":{"code":"404","message":"Resource not found"}}')
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Received:
|
||||
model_calls: tuple[Mapping[str, JsonValue], ...]
|
||||
searches: tuple[Mapping[str, JsonValue], ...]
|
||||
|
||||
|
||||
def _named(request: Request, name: str) -> bool:
|
||||
return name in request.target or name in request.body.decode(errors="replace")
|
||||
|
||||
|
||||
def received(drained: Sequence[Request], name: str) -> Received:
|
||||
requests: Final = tuple(request for request in drained if request.method == "POST" and _named(request, name))
|
||||
return Received(
|
||||
model_calls=tuple(
|
||||
JSON_OBJECT.validate_json(request.body)
|
||||
for request in requests
|
||||
if urlsplit(request.target).path != SEARCH_TARGET
|
||||
),
|
||||
searches=tuple(
|
||||
JSON_OBJECT.validate_json(request.body)
|
||||
for request in requests
|
||||
if urlsplit(request.target).path == SEARCH_TARGET
|
||||
),
|
||||
)
|
||||
|
|
@ -0,0 +1,73 @@
|
|||
from collections.abc import Iterator, Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support import interception_vendor as iv
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from openai.types.responses import WebSearchToolParam
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
|
||||
from litellm.responses.streaming_iterator import (
|
||||
MockResponsesAPIStreamingIterator,
|
||||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
)
|
||||
from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamingResponse
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wire(monkeypatch: pytest.MonkeyPatch) -> Iterator[Wire]:
|
||||
monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"])])
|
||||
with wire_server(iv.respond) as served:
|
||||
yield served
|
||||
|
||||
|
||||
def _text(events: Sequence[ResponsesAPIStreamingResponse]) -> str:
|
||||
return "".join(event.delta for event in events if isinstance(event, OutputTextDeltaEvent))
|
||||
|
||||
|
||||
def _assert_one_clean_call(wire: Wire, name: str) -> None:
|
||||
received: Final = iv.received(wire.drain(), name)
|
||||
(call,) = received.model_calls
|
||||
assert call.get("stream") is not True and "stream_options" not in call, call
|
||||
assert [key for key in call if key.startswith("_")] == [], call
|
||||
assert received.searches == (), received
|
||||
|
||||
|
||||
def test_sync_responses_stream_with_web_search_keeps_stream_options_off_azure(wire: Wire) -> None:
|
||||
name: Final = iv.deployment("plain")
|
||||
stream: Final = litellm.responses(
|
||||
model=f"azure/{name}",
|
||||
input=f"Search for {name}",
|
||||
tools=[WebSearchToolParam(type="web_search")],
|
||||
stream=True,
|
||||
stream_options={"include_obfuscation": True},
|
||||
api_base=wire.url,
|
||||
api_key=iv.AZURE_KEY,
|
||||
api_version=iv.API_VERSION,
|
||||
num_retries=0,
|
||||
)
|
||||
assert isinstance(stream, (SyncResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator)), stream
|
||||
events: Final = tuple(stream)
|
||||
assert _text(events) == iv.plain(name), events
|
||||
_assert_one_clean_call(wire, name)
|
||||
|
||||
|
||||
async def test_async_responses_stream_with_web_search_keeps_stream_options_off_azure(wire: Wire) -> None:
|
||||
name: Final = iv.deployment("plain")
|
||||
stream: Final = await litellm.aresponses(
|
||||
model=f"azure/{name}",
|
||||
input=f"Search for {name}",
|
||||
tools=[WebSearchToolParam(type="web_search")],
|
||||
stream=True,
|
||||
stream_options={"include_obfuscation": True},
|
||||
api_base=wire.url,
|
||||
api_key=iv.AZURE_KEY,
|
||||
api_version=iv.API_VERSION,
|
||||
num_retries=0,
|
||||
)
|
||||
assert isinstance(stream, (ResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator)), stream
|
||||
events: Final = tuple([event async for event in stream])
|
||||
assert _text(events) == iv.plain(name), events
|
||||
_assert_one_clean_call(wire, name)
|
||||
|
|
@ -0,0 +1,905 @@
|
|||
import asyncio
|
||||
import json
|
||||
import signal
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from anthropic.types import RawContentBlockDeltaEvent, TextDelta, WebSearchTool20250305Param
|
||||
from integration._support import interception_vendor as iv
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import (
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai.types.chat import ChatCompletionChunk, ChatCompletionFunctionToolParam, ChatCompletionUserMessageParam
|
||||
from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent, WebSearchToolParam
|
||||
from pydantic import JsonValue
|
||||
|
||||
_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
|
||||
pytestmark: Final = pytest.mark.timeout(_CELL_SECONDS)
|
||||
|
||||
_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}})
|
||||
_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType({"include_usage": True})
|
||||
_OBFUSCATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"include_obfuscation": True})
|
||||
_PARAMETERS: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}
|
||||
)
|
||||
_SEARCH_FUNCTION: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": iv.SEARCH_TOOL, "description": "Search the web", "parameters": dict(_PARAMETERS)},
|
||||
}
|
||||
)
|
||||
_CODE_INTERPRETER: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"type": "code_interpreter", "container": {"type": "auto"}}
|
||||
)
|
||||
_HEADROOM_RETRIEVE: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "headroom_retrieve", "description": "Retrieve content", "parameters": dict(_PARAMETERS)},
|
||||
}
|
||||
)
|
||||
_RESPONSES_SEARCH: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "web_search"})
|
||||
_BURST_MODEL: Final = f"interception-burst-{uuid.uuid4().hex}"
|
||||
_BURST_DEPLOYMENT: Final = iv.deployment("search")
|
||||
_STARTED_WORKER: Final = "Started server process ["
|
||||
_STARTUP_COMPLETE: Final = "Application startup complete."
|
||||
_MALFORMED: Final = (
|
||||
pytest.param(5, id="int"),
|
||||
pytest.param([{"include_usage": True}], id="list"),
|
||||
pytest.param("", id="empty-string"),
|
||||
pytest.param("x" * 5120, id="5kb-string"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
wire: Wire
|
||||
proxy: OwnedProxy
|
||||
|
||||
@property
|
||||
def gateway(self) -> Gateway:
|
||||
return self.proxy.gateway
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return str(self.proxy.gateway.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Deployment:
|
||||
model: str
|
||||
name: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Stream:
|
||||
status: int
|
||||
text: str
|
||||
|
||||
@property
|
||||
def frames(self) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
iv.JSON_OBJECT.validate_json(line[6:]) for line in self.text.splitlines() if line.startswith("data: {")
|
||||
)
|
||||
|
||||
@property
|
||||
def content(self) -> str:
|
||||
return "".join(_delta_text(frame) for frame in self.frames)
|
||||
|
||||
@property
|
||||
def ids(self) -> frozenset[str]:
|
||||
return frozenset(string_value(frame["id"]) for frame in self.frames)
|
||||
|
||||
@property
|
||||
def usages(self) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(object_value(frame["usage"]) for frame in self.frames if frame.get("usage") is not None)
|
||||
|
||||
@property
|
||||
def deltas(self) -> str:
|
||||
return "".join(
|
||||
string_value(frame["delta"]) for frame in self.frames if frame.get("type") == "response.output_text.delta"
|
||||
)
|
||||
|
||||
@property
|
||||
def completed(self) -> Mapping[str, JsonValue]:
|
||||
(frame,) = [frame for frame in self.frames if frame.get("type") == "response.completed"]
|
||||
return object_value(frame["response"])
|
||||
|
||||
|
||||
def _delta_text(frame: Mapping[str, JsonValue]) -> str:
|
||||
choices: Final = rv.ITEMS.validate_python(frame.get("choices") or [])
|
||||
return "".join(_text_of(object_value(choice.get("delta") or {})) for choice in choices)
|
||||
|
||||
|
||||
def _text_of(delta: Mapping[str, JsonValue]) -> str:
|
||||
content: Final = delta.get("content")
|
||||
return content if isinstance(content, str) else ""
|
||||
|
||||
|
||||
def _config(
|
||||
directory: Path,
|
||||
wire: Wire,
|
||||
general_settings: Mapping[str, JsonValue],
|
||||
models: Sequence[Mapping[str, JsonValue]] = (),
|
||||
) -> Path:
|
||||
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
|
||||
config: Final = {
|
||||
**base,
|
||||
"model_list": [*rv.ITEMS.validate_python(base["model_list"]), *(dict(model) for model in models)],
|
||||
"general_settings": {**object_value(base["general_settings"]), **general_settings},
|
||||
"litellm_settings": {
|
||||
**object_value(base["litellm_settings"]),
|
||||
"callbacks": ["websearch_interception", "code_interpreter_interception"],
|
||||
"websearch_interception_params": {"enabled_providers": ["azure"], "search_tool_name": "integration-search"},
|
||||
"code_interpreter_interception_params": {"enabled_providers": ["azure"]},
|
||||
},
|
||||
"search_tools": [
|
||||
{
|
||||
"search_tool_name": "integration-search",
|
||||
"litellm_params": {
|
||||
"search_provider": "tavily",
|
||||
"api_key": iv.TAVILY_KEY,
|
||||
"api_base": f"{wire.url}/tavily",
|
||||
},
|
||||
}
|
||||
],
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": "integration-headroom",
|
||||
"litellm_params": {
|
||||
"guardrail": "headroom",
|
||||
"mode": "post_call",
|
||||
"api_base": "http://127.0.0.1:9",
|
||||
"unreachable_fallback": "fail_open",
|
||||
"default_on": False,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
path: Final = directory / f"interception-{uuid.uuid4().hex}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _interception_proxy(
|
||||
directory: Path,
|
||||
respond: Callable[[Request], Reply],
|
||||
general_settings: Mapping[str, JsonValue],
|
||||
models: Callable[[Wire], Sequence[Mapping[str, JsonValue]]] = lambda _: (),
|
||||
) -> Generator[_Rig, None, None]:
|
||||
with gateway_from_environment() as gateway, wire_server(respond) as wire:
|
||||
config: Final = _config(directory, wire, general_settings, models(wire))
|
||||
with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
|
||||
yield _Rig(wire, owned)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
with _interception_proxy(tmp_path_factory.mktemp("interception-rig"), iv.respond, {}) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scenario(rig: _Rig) -> Iterator[Scenario]:
|
||||
rig.wire.drain()
|
||||
with rig.gateway.scenario() as opened:
|
||||
yield opened
|
||||
|
||||
|
||||
def _deploy(rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY) -> _Deployment:
|
||||
name: Final = iv.deployment(mode)
|
||||
model: Final = scenario.model(
|
||||
model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION
|
||||
)
|
||||
return _Deployment(model, name)
|
||||
|
||||
|
||||
def _chat_body(
|
||||
deployment: _Deployment,
|
||||
*,
|
||||
stream: bool = True,
|
||||
tool: Mapping[str, JsonValue] = _SEARCH_FUNCTION,
|
||||
extra: Mapping[str, JsonValue] = _NO_CACHE,
|
||||
) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"model": deployment.model,
|
||||
"messages": [{"role": "user", "content": f"Search for {deployment.name}"}],
|
||||
"tools": [dict(tool)],
|
||||
"stream": stream,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def _responses_body(deployment: _Deployment, extra: Mapping[str, JsonValue] = _NO_CACHE) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"model": deployment.model,
|
||||
"input": f"Search for {deployment.name}",
|
||||
"tools": [dict(_RESPONSES_SEARCH)],
|
||||
"stream": True,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def _post(rig: _Rig, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> _Stream:
|
||||
response: Final = rig.gateway.request("POST", path, body, key=key)
|
||||
return _Stream(response.status_code, response.text)
|
||||
|
||||
|
||||
def _chat(rig: _Rig, body: Mapping[str, JsonValue]) -> _Stream:
|
||||
return _post(rig, "/v1/chat/completions", body)
|
||||
|
||||
|
||||
def _responses(rig: _Rig, body: Mapping[str, JsonValue]) -> _Stream:
|
||||
return _post(rig, "/v1/responses", body)
|
||||
|
||||
|
||||
def _upstream(rig: _Rig, deployment: _Deployment) -> iv.Received:
|
||||
return iv.received(rig.wire.drain(), deployment.name)
|
||||
|
||||
|
||||
def _leaks(call: Mapping[str, JsonValue]) -> bool:
|
||||
return "stream_options" in call or call.get("stream") is True or any(key.startswith("_") for key in call)
|
||||
|
||||
|
||||
def _assert_upstream(received: iv.Received, *, model_calls: int, searches: int) -> None:
|
||||
assert len(received.model_calls) == model_calls, received
|
||||
assert len(received.searches) == searches, received
|
||||
assert [call for call in received.model_calls if _leaks(call)] == [], received.model_calls
|
||||
|
||||
|
||||
def _totals(usage: Mapping[str, JsonValue]) -> int:
|
||||
total: Final = usage["total_tokens"]
|
||||
assert isinstance(total, int), usage
|
||||
return total
|
||||
|
||||
|
||||
def _assert_restreamed(stream: _Stream, text: str, *, usage: Mapping[str, int] | None) -> str:
|
||||
assert stream.status == 200, stream.text
|
||||
assert stream.content == text, stream.text
|
||||
assert [_totals(item) for item in stream.usages] == ([] if usage is None else [usage["total_tokens"]]), stream.text
|
||||
(identity,) = stream.ids
|
||||
return identity
|
||||
|
||||
|
||||
def _assert_responses_restreamed(stream: _Stream, text: str) -> str:
|
||||
assert stream.status == 200, stream.text
|
||||
assert stream.deltas == text, stream.text
|
||||
completed: Final = stream.completed
|
||||
assert _totals(object_value(completed["usage"])) == iv.RESPONSES_FOLLOWUP_USAGE["total_tokens"], completed
|
||||
return string_value(completed["id"])
|
||||
|
||||
|
||||
def _spend_rows(model: str, expected: int) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)
|
||||
),
|
||||
lambda rows: sum(1 for row in rows if row["status"] == "success") >= expected,
|
||||
seconds=70,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _assert_landed_once(model: str, identities: frozenset[str], *, rows: int) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
landed: Final = _spend_rows(model, rows)
|
||||
assert len(landed) == rows, landed
|
||||
assert frozenset(string_value(row["request_id"]) for row in landed if row["status"] == "success") == identities, (
|
||||
landed
|
||||
)
|
||||
return landed
|
||||
|
||||
|
||||
def _assert_response_landed_once(model: str, identity: str) -> None:
|
||||
(row,) = _spend_rows(model, 1)
|
||||
assert row["status"] == "success" and rv.same_response(string_value(row["request_id"]), identity), row
|
||||
|
||||
|
||||
def _openai(rig: _Rig) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _async_openai(rig: _Rig) -> openai.AsyncOpenAI:
|
||||
return openai.AsyncOpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _sdk_tool() -> ChatCompletionFunctionToolParam:
|
||||
return ChatCompletionFunctionToolParam(
|
||||
type="function",
|
||||
function={"name": iv.SEARCH_TOOL, "description": "Search the web", "parameters": dict(_PARAMETERS)},
|
||||
)
|
||||
|
||||
|
||||
def _sdk_message(deployment: _Deployment) -> ChatCompletionUserMessageParam:
|
||||
return ChatCompletionUserMessageParam(role="user", content=f"Search for {deployment.name}")
|
||||
|
||||
|
||||
def _chunk_text(chunk: ChatCompletionChunk) -> str:
|
||||
return "".join(choice.delta.content or "" for choice in chunk.choices)
|
||||
|
||||
|
||||
def _assert_sdk_chunks(chunks: Sequence[ChatCompletionChunk], deployment: _Deployment) -> str:
|
||||
assert "".join(map(_chunk_text, chunks)) == iv.answer(deployment.name), chunks
|
||||
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
|
||||
assert [usage.total_tokens for usage in usages] == [iv.FOLLOWUP_USAGE["total_tokens"]], chunks
|
||||
(identity,) = {chunk.id for chunk in chunks}
|
||||
return identity
|
||||
|
||||
|
||||
def _assert_search_turn(rig: _Rig, deployment: _Deployment, identity: str) -> None:
|
||||
received: Final = _upstream(rig, deployment)
|
||||
_assert_upstream(received, model_calls=2, searches=1)
|
||||
assert [search["query"] for search in received.searches] == [iv.query(deployment.name)], received.searches
|
||||
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
|
||||
|
||||
|
||||
def test_openai_sdk_chat_stream_through_web_search_is_restreamed_with_its_usage_chunk(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
chunks: Final = tuple(
|
||||
_openai(rig).chat.completions.create(
|
||||
model=deployment.model,
|
||||
messages=[_sdk_message(deployment)],
|
||||
tools=[_sdk_tool()],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
)
|
||||
_assert_search_turn(rig, deployment, _assert_sdk_chunks(chunks, deployment))
|
||||
|
||||
|
||||
async def test_async_openai_sdk_chat_stream_through_web_search_is_restreamed_with_its_usage_chunk(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = await _async_openai(rig).chat.completions.create(
|
||||
model=deployment.model,
|
||||
messages=[_sdk_message(deployment)],
|
||||
tools=[_sdk_tool()],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
chunks: Final = tuple([chunk async for chunk in stream])
|
||||
_assert_search_turn(rig, deployment, _assert_sdk_chunks(chunks, deployment))
|
||||
|
||||
|
||||
def test_raw_chat_stream_without_stream_options_answers_without_a_usage_chunk(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
identity: Final = _assert_restreamed(_chat(rig, _chat_body(deployment)), iv.answer(deployment.name), usage=None)
|
||||
_assert_search_turn(rig, deployment, identity)
|
||||
|
||||
|
||||
def test_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
events: Final = tuple(
|
||||
_openai(rig).responses.create(
|
||||
model=deployment.model,
|
||||
input=f"Search for {deployment.name}",
|
||||
tools=[WebSearchToolParam(type="web_search")],
|
||||
stream=True,
|
||||
stream_options={"include_obfuscation": True},
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
)
|
||||
text: Final = "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent))
|
||||
(completed,) = [event for event in events if isinstance(event, ResponseCompletedEvent)]
|
||||
assert text == iv.answer(deployment.name), events
|
||||
usage: Final = completed.response.usage
|
||||
assert usage is not None and usage.total_tokens == iv.RESPONSES_FOLLOWUP_USAGE["total_tokens"], completed
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_response_landed_once(deployment.model, completed.response.id)
|
||||
|
||||
|
||||
async def test_async_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = await _async_openai(rig).responses.create(
|
||||
model=deployment.model,
|
||||
input=f"Search for {deployment.name}",
|
||||
tools=[WebSearchToolParam(type="web_search")],
|
||||
stream=True,
|
||||
stream_options={"include_obfuscation": True},
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
events: Final = tuple([event async for event in stream])
|
||||
text: Final = "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent))
|
||||
(completed,) = [event for event in events if isinstance(event, ResponseCompletedEvent)]
|
||||
assert text == iv.answer(deployment.name), events
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_response_landed_once(deployment.model, completed.response.id)
|
||||
|
||||
|
||||
def test_raw_responses_stream_without_stream_options_answers(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
identity: Final = _assert_responses_restreamed(
|
||||
_responses(rig, _responses_body(deployment)), iv.answer(deployment.name)
|
||||
)
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_response_landed_once(deployment.model, identity)
|
||||
|
||||
|
||||
def test_openai_sdk_chat_without_streaming_through_web_search_answers_in_json(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
completion: Final = _openai(rig).chat.completions.create(
|
||||
model=deployment.model,
|
||||
messages=[_sdk_message(deployment)],
|
||||
tools=[_sdk_tool()],
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
assert completion.choices[0].message.content == iv.answer(deployment.name), completion
|
||||
assert completion.usage is not None, completion
|
||||
_assert_search_turn(rig, deployment, completion.id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool", "mode"),
|
||||
(
|
||||
pytest.param(_CODE_INTERPRETER, "plain", id="code-interpreter"),
|
||||
pytest.param(_HEADROOM_RETRIEVE, "plain", id="headroom-retrieve"),
|
||||
pytest.param(_SEARCH_FUNCTION, "plain", id="web-search-unused"),
|
||||
),
|
||||
)
|
||||
def test_converted_chat_stream_with_usage_replays_the_single_call_usage(
|
||||
rig: _Rig, scenario: Scenario, tool: Mapping[str, JsonValue], mode: iv.Mode
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, mode)
|
||||
stream: Final = _chat(rig, _chat_body(deployment, tool=tool, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
|
||||
identity: Final = _assert_restreamed(stream, iv.plain(deployment.name), usage=iv.FIRST_USAGE)
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=1, searches=0)
|
||||
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
|
||||
|
||||
|
||||
def test_anthropic_sdk_messages_stream_with_only_web_search_on_azure_answers_from_the_search(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
prompt: Final = f"Search for {deployment.name}"
|
||||
client: Final = anthropic.Anthropic(base_url=rig.base_url, api_key=rig.gateway.key, max_retries=0)
|
||||
events: Final = tuple(
|
||||
client.messages.create(
|
||||
model=deployment.model,
|
||||
max_tokens=64,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
tools=[WebSearchTool20250305Param(type="web_search_20250305", name="web_search")],
|
||||
stream=True,
|
||||
extra_body=dict(_NO_CACHE),
|
||||
)
|
||||
)
|
||||
deltas: Final = tuple(event.delta for event in events if isinstance(event, RawContentBlockDeltaEvent))
|
||||
text: Final = "".join(delta.text for delta in deltas if isinstance(delta, TextDelta))
|
||||
assert text == f"Title: title-{prompt}\nURL: https://example.test/result\nSnippet: snippet-{prompt}", events
|
||||
received: Final = _upstream(rig, deployment)
|
||||
_assert_upstream(received, model_calls=0, searches=1)
|
||||
assert [search.get("query") for search in received.searches] == [prompt], received
|
||||
|
||||
|
||||
def test_per_request_fallback_from_a_missing_deployment_restreams_with_usage(rig: _Rig, scenario: Scenario) -> None:
|
||||
missing: Final = _deploy(rig, scenario, "missing")
|
||||
good: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _chat_body(missing, extra={**_NO_CACHE, "stream_options": dict(_USAGE), "fallbacks": [good.model]})
|
||||
stream: Final = _chat(rig, {**body, "messages": [{"role": "user", "content": "Search the web"}]})
|
||||
identity: Final = _assert_restreamed(stream, iv.answer(good.name), usage=iv.FOLLOWUP_USAGE)
|
||||
drained: Final = rig.wire.drain()
|
||||
_assert_upstream(iv.received(drained, missing.name), model_calls=1, searches=0)
|
||||
_assert_upstream(iv.received(drained, good.name), model_calls=2, searches=1)
|
||||
_assert_landed_once(good.model, frozenset({identity}), rows=1)
|
||||
|
||||
|
||||
def test_chat_stream_cache_hit_replays_the_answer_without_calling_azure(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _chat_body(deployment, extra={"stream_options": dict(_USAGE)})
|
||||
first: Final = _assert_restreamed(_chat(rig, body), iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE)
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
second: Final = _chat(rig, body)
|
||||
assert second.status == 200 and second.content == iv.answer(deployment.name), second.text
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=0, searches=0)
|
||||
rows: Final = _spend_rows(deployment.model, 2)
|
||||
served: Final = tuple(string_value(row["request_id"]) for row in rows if str(row["cache_hit"]) != "True")
|
||||
cached: Final = tuple(string_value(row["request_id"]) for row in rows if str(row["cache_hit"]) == "True")
|
||||
assert served == (first,) and len(cached) == 1 and cached[0].startswith(first), rows
|
||||
|
||||
|
||||
def test_responses_stream_cache_hit_replays_the_answer_without_calling_azure(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _responses_body(deployment, {"stream_options": dict(_OBFUSCATION)})
|
||||
_assert_responses_restreamed(_responses(rig, body), iv.answer(deployment.name))
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
second: Final = _responses(rig, body)
|
||||
assert second.status == 200 and second.deltas == iv.answer(deployment.name), second.text
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=0, searches=0)
|
||||
|
||||
|
||||
def _with_duplicated_stream_options(body: Mapping[str, JsonValue]) -> bytes:
|
||||
return (json.dumps({**body, "stream_options": dict(_USAGE)})[:-1] + ', "stream_options": 5}').encode()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_options", _MALFORMED)
|
||||
def test_malformed_chat_stream_options_never_reach_azure(
|
||||
rig: _Rig, scenario: Scenario, stream_options: JsonValue
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": stream_options}))
|
||||
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
|
||||
_assert_search_turn(rig, deployment, identity)
|
||||
|
||||
|
||||
def test_duplicated_chat_stream_options_key_never_reaches_azure(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
response: Final = rig.gateway.client.post(
|
||||
"/v1/chat/completions",
|
||||
content=_with_duplicated_stream_options(_chat_body(deployment)),
|
||||
headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"},
|
||||
)
|
||||
identity: Final = _assert_restreamed(
|
||||
_Stream(response.status_code, response.text), iv.answer(deployment.name), usage=None
|
||||
)
|
||||
_assert_search_turn(rig, deployment, identity)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_options", _MALFORMED)
|
||||
def test_malformed_responses_stream_options_never_reach_azure(
|
||||
rig: _Rig, scenario: Scenario, stream_options: JsonValue
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = _responses(rig, _responses_body(deployment, {**_NO_CACHE, "stream_options": stream_options}))
|
||||
identity: Final = _assert_responses_restreamed(stream, iv.answer(deployment.name))
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_response_landed_once(deployment.model, identity)
|
||||
|
||||
|
||||
def test_forged_interception_keys_on_a_json_request_never_reach_azure(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
forged: Final = {
|
||||
**_NO_CACHE,
|
||||
"_websearch_interception_stream_options": dict(_USAGE),
|
||||
"_websearch_interception_converted_stream": True,
|
||||
"_headroom_interception_stream_options": dict(_USAGE),
|
||||
"_code_interpreter_interception_stream_options": dict(_USAGE),
|
||||
"max_agentic_loops": 50,
|
||||
}
|
||||
response: Final = rig.gateway.request(
|
||||
"POST", "/v1/chat/completions", _chat_body(deployment, stream=False, extra=forged)
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("application/json"), response.headers
|
||||
body: Final = iv.JSON_OBJECT.validate_json(response.content)
|
||||
(choice,) = rv.ITEMS.validate_python(body["choices"])
|
||||
assert object_value(choice["message"])["content"] == iv.answer(deployment.name), body
|
||||
_assert_search_turn(rig, deployment, string_value(body["id"]))
|
||||
|
||||
|
||||
def test_forged_garbage_stash_on_a_stream_is_replaced_by_the_real_stream_options(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "plain")
|
||||
extra: Final = {**_NO_CACHE, "stream_options": dict(_USAGE), "_code_interpreter_interception_stream_options": "x"}
|
||||
stream: Final = _chat(rig, _chat_body(deployment, tool=_CODE_INTERPRETER, extra=extra))
|
||||
identity: Final = _assert_restreamed(stream, iv.plain(deployment.name), usage=iv.FIRST_USAGE)
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=1, searches=0)
|
||||
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
|
||||
|
||||
|
||||
def test_wrong_azure_key_surfaces_azures_401(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search", api_key="wrong-azure-key")
|
||||
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
|
||||
assert stream.status == 401, stream.text
|
||||
assert "Access denied due to invalid subscription key" in stream.text, stream.text
|
||||
assert len(_upstream(rig, deployment).model_calls) == 1
|
||||
|
||||
|
||||
def _tool_call(call: Mapping[str, JsonValue]) -> tuple[str, Mapping[str, JsonValue]]:
|
||||
function: Final = object_value(call["function"])
|
||||
return string_value(function["name"]), iv.JSON_OBJECT.validate_json(string_value(function["arguments"]))
|
||||
|
||||
|
||||
def _stream_choices(stream: _Stream) -> Iterator[Mapping[str, JsonValue]]:
|
||||
for frame in stream.frames:
|
||||
yield from rv.ITEMS.validate_python(frame.get("choices") or [])
|
||||
|
||||
|
||||
def _streamed_tool_calls(stream: _Stream) -> Iterator[tuple[str, Mapping[str, JsonValue]]]:
|
||||
for choice in _stream_choices(stream):
|
||||
yield from map(
|
||||
_tool_call, rv.ITEMS.validate_python(object_value(choice.get("delta") or {}).get("tool_calls") or [])
|
||||
)
|
||||
|
||||
|
||||
def _handed_back(deployment: _Deployment) -> list[tuple[str, Mapping[str, JsonValue]]]:
|
||||
return [(iv.SEARCH_TOOL, {"query": iv.query(deployment.name)})]
|
||||
|
||||
|
||||
def _assert_succeeded_once(model: str, identity: str) -> None:
|
||||
rows: Final = _spend_rows(model, 1)
|
||||
assert [string_value(row["request_id"]) for row in rows if row["status"] == "success"] == [identity], rows
|
||||
|
||||
|
||||
def test_failed_follow_up_call_hands_the_search_tool_call_back_with_its_usage(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "followup-fails")
|
||||
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
|
||||
assert stream.status == 200 and stream.content == "", stream.text
|
||||
assert list(_streamed_tool_calls(stream)) == _handed_back(deployment), stream.text
|
||||
finished: Final = [choice["finish_reason"] for choice in _stream_choices(stream) if choice.get("finish_reason")]
|
||||
assert finished == ["tool_calls"], stream.text
|
||||
assert [_totals(usage) for usage in stream.usages] == [iv.FIRST_USAGE["total_tokens"]], stream.text
|
||||
(identity,) = stream.ids
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_succeeded_once(deployment.model, identity)
|
||||
assert rig.gateway.chat(_deploy(rig, scenario, "plain").model)["object"] == "chat.completion"
|
||||
|
||||
|
||||
def test_failed_follow_up_call_on_a_json_request_hands_the_search_tool_call_back(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "followup-fails")
|
||||
response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(deployment, stream=False))
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = iv.JSON_OBJECT.validate_json(response.content)
|
||||
(choice,) = rv.ITEMS.validate_python(body["choices"])
|
||||
calls: Final = rv.ITEMS.validate_python(object_value(choice["message"])["tool_calls"])
|
||||
assert [_tool_call(call) for call in calls] == _handed_back(deployment), body
|
||||
assert choice["finish_reason"] == "tool_calls", body
|
||||
assert _totals(object_value(body["usage"])) == iv.FIRST_USAGE["total_tokens"], body
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_succeeded_once(deployment.model, string_value(body["id"]))
|
||||
|
||||
|
||||
def test_search_outage_still_restreams_the_follow_up_answer(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search-down")
|
||||
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
|
||||
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE)
|
||||
_assert_search_turn(rig, deployment, identity)
|
||||
|
||||
|
||||
def test_unknown_virtual_key_is_refused_before_azure(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
|
||||
stream: Final = _post(rig, "/v1/chat/completions", body, key=f"sk-{uuid.uuid4().hex}")
|
||||
assert stream.status == 401, stream.text
|
||||
assert _upstream(rig, deployment).model_calls == ()
|
||||
|
||||
|
||||
def test_json_request_with_its_own_stream_options_still_gets_azures_refusal(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _chat_body(deployment, stream=False, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
|
||||
response: Final = rig.gateway.request("POST", "/v1/chat/completions", body)
|
||||
assert response.status_code == 400, response.text
|
||||
assert iv.STREAM_OPTIONS_REFUSAL in response.text, response.text
|
||||
assert [call.get("stream_options") for call in _upstream(rig, deployment).model_calls] == [dict(_USAGE)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stream_options",
|
||||
(
|
||||
pytest.param({"include_usage": False}, id="include-usage-false"),
|
||||
pytest.param(None, id="null"),
|
||||
pytest.param({}, id="empty"),
|
||||
),
|
||||
)
|
||||
def test_chat_stream_without_a_usage_request_answers_without_a_usage_chunk(
|
||||
rig: _Rig, scenario: Scenario, stream_options: JsonValue
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": stream_options}))
|
||||
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
|
||||
_assert_search_turn(rig, deployment, identity)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_options", (pytest.param(None, id="null"), pytest.param({}, id="empty")))
|
||||
def test_responses_stream_with_empty_stream_options_answers(
|
||||
rig: _Rig, scenario: Scenario, stream_options: JsonValue
|
||||
) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
stream: Final = _responses(rig, _responses_body(deployment, {**_NO_CACHE, "stream_options": stream_options}))
|
||||
identity: Final = _assert_responses_restreamed(stream, iv.answer(deployment.name))
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
|
||||
_assert_response_landed_once(deployment.model, identity)
|
||||
|
||||
|
||||
def test_repeated_identical_streams_each_land_their_own_row(rig: _Rig, scenario: Scenario) -> None:
|
||||
deployment: Final = _deploy(rig, scenario, "search")
|
||||
body: Final = _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
|
||||
identities: Final = frozenset(
|
||||
_assert_restreamed(_chat(rig, body), iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE) for _ in range(3)
|
||||
)
|
||||
assert len(identities) == 3, identities
|
||||
_assert_upstream(_upstream(rig, deployment), model_calls=6, searches=3)
|
||||
_assert_landed_once(deployment.model, identities, rows=3)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
path: str
|
||||
stream: _Stream
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, path: str, body: Mapping[str, JsonValue]) -> _Served:
|
||||
response: Final = await client.post(path, json=body, headers={"Authorization": f"Bearer {key}"})
|
||||
return _Served(path, _Stream(response.status_code, response.text))
|
||||
|
||||
|
||||
async def _burst(
|
||||
base_url: str, key: str, requests: Sequence[tuple[str, Mapping[str, JsonValue]]], *, tolerate: bool = False
|
||||
) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_send(client, key, path, body) for path, body in requests), return_exceptions=tolerate
|
||||
)
|
||||
assert [result for result in results if isinstance(result, BaseException)] == [] or tolerate, results
|
||||
assert all(isinstance(result, _Served | httpx.TransportError) for result in results), results
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
def _chat_request(deployment: _Deployment) -> tuple[str, Mapping[str, JsonValue]]:
|
||||
return "/v1/chat/completions", _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
|
||||
|
||||
|
||||
def _responses_request(deployment: _Deployment) -> tuple[str, Mapping[str, JsonValue]]:
|
||||
return "/v1/responses", _responses_body(deployment, {**_NO_CACHE, "stream_options": dict(_OBFUSCATION)})
|
||||
|
||||
|
||||
async def test_concurrent_burst_across_chat_and_responses_answers_each_and_lands_each_once(
|
||||
rig: _Rig, scenario: Scenario
|
||||
) -> None:
|
||||
chat: Final = _deploy(rig, scenario, "search")
|
||||
responses: Final = _deploy(rig, scenario, "search")
|
||||
outage: Final = _deploy(rig, scenario, "search-down")
|
||||
requests: Final = (
|
||||
*(_chat_request(chat) for _ in range(8)),
|
||||
*(_responses_request(responses) for _ in range(8)),
|
||||
*(_chat_request(outage) for _ in range(4)),
|
||||
)
|
||||
served: Final = await _burst(rig.base_url, rig.gateway.key, requests)
|
||||
chat_ids: Final = frozenset(
|
||||
_assert_restreamed(item.stream, iv.answer(chat.name), usage=iv.FOLLOWUP_USAGE) for item in served[:8]
|
||||
)
|
||||
response_ids: Final = tuple(
|
||||
_assert_responses_restreamed(item.stream, iv.answer(responses.name)) for item in served[8:16]
|
||||
)
|
||||
outage_ids: Final = frozenset(
|
||||
_assert_restreamed(item.stream, iv.answer(outage.name), usage=iv.FOLLOWUP_USAGE) for item in served[16:]
|
||||
)
|
||||
drained: Final = rig.wire.drain()
|
||||
_assert_upstream(iv.received(drained, chat.name), model_calls=16, searches=8)
|
||||
_assert_upstream(iv.received(drained, responses.name), model_calls=16, searches=8)
|
||||
_assert_upstream(iv.received(drained, outage.name), model_calls=8, searches=4)
|
||||
_assert_landed_once(chat.model, chat_ids, rows=8)
|
||||
_assert_landed_once(outage.model, outage_ids, rows=4)
|
||||
rows: Final = _spend_rows(responses.model, 8)
|
||||
assert len(rows) == 8 and len(set(response_ids)) == 8, rows
|
||||
assert all(
|
||||
any(rv.same_response(string_value(row["request_id"]), identity) for row in rows) for identity in response_ids
|
||||
)
|
||||
|
||||
|
||||
def test_always_include_stream_usage_restreams_usage_on_chat_and_strips_it_from_azure_on_responses(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
with _interception_proxy(tmp_path, iv.respond, {"always_include_stream_usage": True}) as owned:
|
||||
with owned.gateway.scenario() as opened:
|
||||
chat: Final = _deploy(owned, opened, "search")
|
||||
responses: Final = _deploy(owned, opened, "search")
|
||||
stream: Final = _chat(owned, _chat_body(chat))
|
||||
identity: Final = _assert_restreamed(stream, iv.answer(chat.name), usage=iv.FOLLOWUP_USAGE)
|
||||
replayed: Final = _assert_responses_restreamed(
|
||||
_responses(owned, _responses_body(responses)), iv.answer(responses.name)
|
||||
)
|
||||
drained: Final = owned.wire.drain()
|
||||
_assert_upstream(iv.received(drained, chat.name), model_calls=2, searches=1)
|
||||
_assert_upstream(iv.received(drained, responses.name), model_calls=2, searches=1)
|
||||
_assert_landed_once(chat.model, frozenset({identity}), rows=1)
|
||||
_assert_response_landed_once(responses.model, replayed)
|
||||
|
||||
|
||||
def _holding(release: threading.Event, held: SimpleQueue[str]) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "POST" and urlsplit(request.target).path.endswith("/chat/completions"):
|
||||
held.put(request.target)
|
||||
assert release.wait(timeout=90), "The burst was never released"
|
||||
return iv.respond(request)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _releasing(release: threading.Event) -> Generator[None, None, None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
def _open_upstream_connections(pid: int, upstream: str) -> int:
|
||||
port: Final = urlsplit(upstream).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
||||
|
||||
def _worker_pids(log: str) -> tuple[int, ...]:
|
||||
return tuple(
|
||||
int(line.split(_STARTED_WORKER, 1)[1].split("]", 1)[0]) for line in log.splitlines() if _STARTED_WORKER in line
|
||||
)
|
||||
|
||||
|
||||
def _burst_models(wire: Wire) -> Sequence[Mapping[str, JsonValue]]:
|
||||
return (
|
||||
{
|
||||
"model_name": _BURST_MODEL,
|
||||
"litellm_params": {
|
||||
"model": f"azure/{_BURST_DEPLOYMENT}",
|
||||
"api_base": wire.url,
|
||||
"api_key": iv.AZURE_KEY,
|
||||
"api_version": iv.API_VERSION,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def test_worker_sigkill_mid_burst_leaves_the_sibling_restreaming_with_usage(tmp_path: Path) -> None:
|
||||
release: Final = threading.Event()
|
||||
held: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
deployment: Final = _Deployment(_BURST_MODEL, _BURST_DEPLOYMENT)
|
||||
with _interception_proxy(tmp_path, _holding(release, held), {}, _burst_models) as owned, _releasing(release):
|
||||
workers: Final = eventually(
|
||||
lambda: _worker_pids(owned.proxy.log.read_text()),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(
|
||||
owned.base_url, owned.gateway.key, tuple(_chat_request(deployment) for _ in range(20)), tolerate=True
|
||||
)
|
||||
)
|
||||
await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, owned.wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == 20, held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
victim: Final = psutil.Process(victim_pid)
|
||||
victim.suspend()
|
||||
victim.send_signal(signal.SIGKILL)
|
||||
release.set()
|
||||
served: Final = await burst
|
||||
assert held_by[survivor_pid] >= 10, held_by
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
identities: Final = frozenset(
|
||||
_assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) for item in served
|
||||
)
|
||||
(follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),))
|
||||
recovered: Final = _assert_restreamed(follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE)
|
||||
rows: Final = _spend_rows(_BURST_MODEL, len(identities) + 1)
|
||||
landed: Final = tuple(string_value(row["request_id"]) for row in rows if row["status"] == "success")
|
||||
assert sorted(landed) == sorted((*identities, recovered)), rows
|
||||
replaced: Final = eventually(
|
||||
lambda: owned.proxy.log.read_text(),
|
||||
lambda log: len(_worker_pids(log)) == 3 and log.count(_STARTUP_COMPLETE) == 3,
|
||||
seconds=2 * graceful_stop_seconds(),
|
||||
)
|
||||
assert victim_pid not in _worker_pids(replaced)[2:], replaced
|
||||
|
|
@ -6,22 +6,26 @@ config and records how it is called.
|
|||
"""
|
||||
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.code_interpreter_interception.handler import (
|
||||
CodeInterpreterInterceptionLogger,
|
||||
LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
_INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY,
|
||||
)
|
||||
from litellm.integrations.code_interpreter_interception.handler import (
|
||||
_SANDBOX_KEY,
|
||||
_SESSION_SCOPED_KEY,
|
||||
LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
CodeInterpreterInterceptionLogger,
|
||||
)
|
||||
from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
CODE_INTERPRETER_STREAM_OPTIONS_KEY,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
|
|
@ -516,6 +520,26 @@ async def test_pre_call_forces_non_stream_for_loop():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", (CallTypes.acompletion, CallTypes.aresponses))
|
||||
async def test_pre_call_moves_stream_options_off_the_forced_non_stream_call(call_type: CallTypes):
|
||||
logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox())
|
||||
stream_options: Final = {"include_usage": True}
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"tools": [{"type": "code_interpreter", "container": {"type": "auto"}}],
|
||||
"custom_llm_provider": "azure",
|
||||
"stream": True,
|
||||
"stream_options": stream_options,
|
||||
}
|
||||
|
||||
out: Final = await logger.async_pre_call_deployment_hook(kwargs, call_type)
|
||||
|
||||
assert out is not None
|
||||
assert out["stream"] is False
|
||||
assert "stream_options" not in out
|
||||
assert out[CODE_INTERPRETER_STREAM_OPTIONS_KEY] == stream_options
|
||||
|
||||
|
||||
async def _build_plan(logger, sandbox, call_id="k1", provider="openai"):
|
||||
return await logger.async_build_agentic_loop_plan(
|
||||
tools={
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ Unit tests for WebSearch Interception Handler
|
|||
Tests the WebSearchInterceptionLogger class and helper functions.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, Final
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock
|
||||
|
||||
import pytest
|
||||
|
|
@ -20,7 +22,7 @@ from litellm.proxy._types import (
|
|||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
|
||||
|
||||
def test_initialize_from_proxy_config():
|
||||
|
|
@ -810,6 +812,71 @@ async def test_deployment_hook_converts_stream_and_logging_obj_syncs():
|
|||
assert logging_obj.stream is False
|
||||
|
||||
|
||||
_STREAM_USAGE_OPTIONS = {"include_usage": True}
|
||||
|
||||
|
||||
_Convert = Callable[[WebSearchInterceptionLogger, dict[str, Any]], Awaitable[dict[str, Any] | None]]
|
||||
|
||||
|
||||
async def _convert_via_deployment_hook_chat(
|
||||
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
|
||||
) -> dict[str, Any] | None:
|
||||
return await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
|
||||
|
||||
|
||||
async def _convert_via_deployment_hook_responses(
|
||||
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
|
||||
) -> dict[str, Any] | None:
|
||||
return await logger.async_pre_call_deployment_hook(
|
||||
kwargs={**kwargs, "tools": [{"type": "web_search"}]}, call_type=CallTypes.aresponses
|
||||
)
|
||||
|
||||
|
||||
async def _convert_via_pre_request_hook(
|
||||
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
|
||||
) -> dict[str, Any] | None:
|
||||
return await logger.async_pre_request_hook(model=kwargs["model"], messages=kwargs["messages"], kwargs=kwargs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"convert",
|
||||
(_convert_via_deployment_hook_chat, _convert_via_deployment_hook_responses, _convert_via_pre_request_hook),
|
||||
ids=("deployment_hook_chat", "deployment_hook_responses", "pre_request_hook"),
|
||||
)
|
||||
async def test_converted_stream_moves_stream_options_out_of_the_provider_request(convert: _Convert):
|
||||
"""Every site that flips a streamed web search request to stream=False must take
|
||||
stream_options with it: a provider that validates its body (Azure chat completions)
|
||||
rejects stream_options on a non-streaming request, so the converted kwargs carry no
|
||||
stream_options and stash the client's value for the replayed stream instead."""
|
||||
logger: Final = WebSearchInterceptionLogger(enabled_providers=["azure"])
|
||||
kwargs: Final = {
|
||||
"model": "azure/gpt-4o",
|
||||
"messages": [{"role": "user", "content": "Search for LiteLLM"}],
|
||||
"custom_llm_provider": "azure",
|
||||
"litellm_params": {"custom_llm_provider": "azure"},
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": LITELLM_WEB_SEARCH_TOOL_NAME, "parameters": {"type": "object"}},
|
||||
}
|
||||
],
|
||||
"stream": True,
|
||||
"stream_options": dict(_STREAM_USAGE_OPTIONS),
|
||||
}
|
||||
|
||||
result: Final = await convert(logger, kwargs)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result["stream"],
|
||||
"stream_options" in result,
|
||||
result["_websearch_interception_converted_stream"],
|
||||
result["_websearch_interception_stream_options"],
|
||||
) == (False, False, True, _STREAM_USAGE_OPTIONS)
|
||||
|
||||
|
||||
|
||||
def test_sync_forced_tool_choice_repoints_converted_web_search():
|
||||
"""Regression (tool_choice 400): a forced tool_choice naming the original
|
||||
web_search tool must be repointed to litellm_web_search after conversion.
|
||||
|
|
|
|||
|
|
@ -20,7 +20,8 @@ removed, so `test_internal_control_fields_never_leak_into_provider_body` proves
|
|||
they stay out of the body even without it.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
import datetime
|
||||
from typing import Any, Dict, Final, List, Optional, Tuple
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -34,7 +35,12 @@ from litellm.integrations.code_interpreter_interception.handler import (
|
|||
from litellm.litellm_core_utils.chat_completion_agentic_loop import (
|
||||
maybe_run_chat_completion_agentic_loop,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CODE_INTERPRETER_INTERCEPTION_PREFIX,
|
||||
HEADROOM_INTERCEPTION_PREFIX,
|
||||
WEBSEARCH_INTERCEPTION_PREFIX,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
|
|
@ -253,6 +259,21 @@ class _GateOnlyLogger(CustomLogger):
|
|||
self.cleanup_calls += 1
|
||||
|
||||
|
||||
class _RecordingGateLogger(_GateOnlyLogger):
|
||||
"""Also keeps every follow-up response the loop hands to the post hook."""
|
||||
|
||||
def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None:
|
||||
super().__init__(plan=plan, tool_calls=tool_calls)
|
||||
self.followups: tuple[ModelResponse, ...] = ()
|
||||
|
||||
async def async_post_agentic_loop_response_hook(
|
||||
self, response: object, plan: AgenticLoopPlan, kwargs: Dict[str, Any]
|
||||
) -> object:
|
||||
assert isinstance(response, ModelResponse)
|
||||
self.followups = (*self.followups, response)
|
||||
return response
|
||||
|
||||
|
||||
def _patched_messages() -> List[Dict[str, Any]]:
|
||||
return [
|
||||
{"role": "user", "content": "what is 6*7?"},
|
||||
|
|
@ -443,3 +464,58 @@ async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callb
|
|||
)
|
||||
|
||||
acompletion_mock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("include_usage", (True, False), ids=("usage_requested", "usage_not_requested"))
|
||||
@pytest.mark.parametrize(
|
||||
"prefix", (WEBSEARCH_INTERCEPTION_PREFIX, CODE_INTERPRETER_INTERCEPTION_PREFIX, HEADROOM_INTERCEPTION_PREFIX)
|
||||
)
|
||||
async def test_converted_stream_replays_the_final_answer_with_the_requested_usage_chunk(
|
||||
include_usage: bool,
|
||||
prefix: str,
|
||||
):
|
||||
"""A streamed web search request runs non-streaming and is replayed as a fake
|
||||
stream at depth 0. The stream_options the interception moved out of the provider
|
||||
request decide whether the replay ends with the usage-only chunk the client asked for."""
|
||||
plan: Final = AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=AgenticLoopRequestPatch(messages=_patched_messages()),
|
||||
)
|
||||
gate: Final = _RecordingGateLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]})
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "what is 6*7?"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
litellm_call_id="call-websearch",
|
||||
function_id="fn-websearch",
|
||||
dynamic_success_callbacks=[gate],
|
||||
)
|
||||
stash: Final = {f"{prefix}_stream_options": {"include_usage": True}} if include_usage else {}
|
||||
|
||||
result: Final = await maybe_run_chat_completion_agentic_loop(
|
||||
response=_tool_call_model_response(),
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "what is 6*7?"}],
|
||||
optional_params={},
|
||||
kwargs={f"{prefix}_converted_stream": True, "mock_response": "done", **stash},
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="azure",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
chunks: Final = [chunk async for chunk in result]
|
||||
content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
|
||||
usages: Final = [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None]
|
||||
(followup,) = gate.followups
|
||||
|
||||
assert content == "done"
|
||||
assert followup.usage.total_tokens > 0
|
||||
assert [(usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) for usage in usages] == (
|
||||
[(followup.usage.prompt_tokens, followup.usage.completion_tokens, followup.usage.total_tokens)]
|
||||
if include_usage
|
||||
else []
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ Tests cover (consolidating PRs #23706 and #22727):
|
|||
|
||||
import os
|
||||
import sys
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -40,6 +41,7 @@ from litellm.llms.anthropic.pass_through.adapters.handler import (
|
|||
ANTHROPIC_ONLY_REQUEST_KEYS,
|
||||
LiteLLMMessagesToCompletionTransformationHandler,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
|
||||
|
||||
MESSAGES = [{"role": "user", "content": "hello"}]
|
||||
|
||||
|
|
@ -246,3 +248,18 @@ class TestPromptCacheOptionsForwarded:
|
|||
result = _call_prepare(extra_kwargs={"prompt_cache_options": {"mode": "explicit"}}, model="gpt-5.6")
|
||||
completion_kwargs = result[0] if isinstance(result, tuple) else result
|
||||
assert completion_kwargs["prompt_cache_options"] == {"mode": "explicit"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("converted_stream_key", sorted(CONVERTED_STREAM_KEYS))
|
||||
async def test_messages_layer_converted_stream_flag_still_gets_a_message_back(converted_stream_key: str):
|
||||
response: Final = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler(
|
||||
max_tokens=16,
|
||||
messages=MESSAGES,
|
||||
model="openai/gpt-4o",
|
||||
stream=False,
|
||||
mock_response="hi",
|
||||
**{converted_stream_key: True},
|
||||
)
|
||||
|
||||
assert response["content"][0]["text"] == "hi"
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import datetime
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -14,6 +15,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.handler import (
|
|||
LiteLLMMessagesToResponsesAPIHandler,
|
||||
_build_responses_kwargs,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS, converted_stream_requested
|
||||
|
||||
MESSAGES = [{"role": "user", "content": "hello"}]
|
||||
CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"})
|
||||
|
|
@ -31,6 +33,18 @@ RESPONSES_SSE_BODY = (
|
|||
)
|
||||
|
||||
|
||||
def test_build_responses_kwargs_never_asks_the_nested_call_to_replay_a_stream():
|
||||
responses_kwargs: Final = _build_responses_kwargs(
|
||||
max_tokens=1024,
|
||||
messages=MESSAGES,
|
||||
model="azure/gpt-5.6-luna",
|
||||
extra_kwargs={"custom_llm_provider": "azure", **{key: True for key in CONVERTED_STREAM_KEYS}},
|
||||
)
|
||||
|
||||
assert not converted_stream_requested(responses_kwargs)
|
||||
assert responses_kwargs["custom_llm_provider"] == "azure"
|
||||
|
||||
|
||||
def test_build_responses_kwargs_derives_prompt_cache_key_from_claude_code_session_id():
|
||||
responses_kwargs = _build_responses_kwargs(
|
||||
max_tokens=1024,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ import json
|
|||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
|
|
@ -18,6 +20,9 @@ from litellm.integrations.code_interpreter_interception.handler import (
|
|||
CodeInterpreterInterceptionLogger,
|
||||
LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
|
||||
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
|
||||
from litellm.llms.base_llm.audio_transcription.transformation import (
|
||||
AudioTranscriptionRequestData,
|
||||
BaseAudioTranscriptionConfig,
|
||||
|
|
@ -353,6 +358,89 @@ def test_response_api_handler_runs_responses_pre_call_hook_before_transform():
|
|||
assert hook_litellm_params.get(_SANDBOX_KEY)
|
||||
|
||||
|
||||
class _AgenticLoopKwargsRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seen_kwargs: Mapping[str, object] = MappingProxyType({})
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict[str, object],
|
||||
) -> tuple[bool, dict[str, object]]:
|
||||
self.seen_kwargs = MappingProxyType(dict(kwargs))
|
||||
return False, {}
|
||||
|
||||
|
||||
class _SentRequestRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.body: Mapping[str, object] = MappingProxyType({})
|
||||
|
||||
def respond(self, request: httpx.Request) -> httpx.Response:
|
||||
self.body = MappingProxyType(json.loads(request.content))
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 0,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "hello", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_response_api_handler_sends_no_stream_options_when_web_search_forces_a_non_streaming_call(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
sent: Final = _SentRequestRecorder()
|
||||
recorder: Final = _AgenticLoopKwargsRecorder()
|
||||
monkeypatch.setattr(
|
||||
litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"]), recorder]
|
||||
)
|
||||
logging_obj: Final = Mock()
|
||||
logging_obj.dynamic_success_callbacks = []
|
||||
|
||||
BaseLLMHTTPHandler().response_api_handler(
|
||||
model="gpt-4o",
|
||||
input="what is new in litellm?",
|
||||
responses_api_provider_config=AzureOpenAIResponsesAPIConfig(),
|
||||
response_api_optional_request_params={
|
||||
"stream": True,
|
||||
"stream_options": {"include_obfuscation": True},
|
||||
"tools": [{"type": "web_search"}],
|
||||
},
|
||||
custom_llm_provider="azure",
|
||||
litellm_params=GenericLiteLLMParams(
|
||||
api_key="sk-test",
|
||||
api_base="https://example.openai.azure.com",
|
||||
api_version="2025-04-01-preview",
|
||||
stream_options={"include_obfuscation": True},
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(sent.respond))),
|
||||
)
|
||||
|
||||
assert sent.body["stream"] is False
|
||||
assert "stream_options" not in sent.body
|
||||
assert "stream_options" not in recorder.seen_kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_response_api_handler_streams_when_provider_transform_adds_stream():
|
||||
handler = BaseLLMHTTPHandler()
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ Tests cover:
|
|||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -41,7 +41,7 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
|
|||
from litellm.proxy.spend_tracking.compression_savings import (
|
||||
extract_compression_saved_tokens,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
GenericGuardrailAPIInputs,
|
||||
|
|
@ -2171,6 +2171,28 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple
|
|||
assert kwargs["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", (CallTypes.acompletion, CallTypes.aresponses))
|
||||
async def test_pre_call_deployment_hook_moves_stream_options_off_the_forced_non_stream_call(
|
||||
guardrail: HeadroomGuardrail, call_type: CallTypes
|
||||
):
|
||||
stream_options: Final = {"include_usage": True}
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "gpt-4o",
|
||||
"stream": True,
|
||||
"stream_options": stream_options,
|
||||
"tools": [_retrieve_tool_definition()],
|
||||
}
|
||||
|
||||
result: Final = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type)
|
||||
|
||||
assert result is not None
|
||||
assert result["stream"] is False
|
||||
assert "stream_options" not in result
|
||||
assert result[HEADROOM_STREAM_OPTIONS_KEY] == stream_options
|
||||
assert kwargs["stream_options"] == stream_options
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail):
|
||||
kwargs = {
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_log
|
|||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
|
||||
from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth
|
||||
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
KeyAndTeamLoggingSettings,
|
||||
LiteLLMProxyRequestSetup,
|
||||
|
|
@ -1276,13 +1277,7 @@ async def test_key_metadata_enable_prompt_caching_promoted_to_request_root(key_v
|
|||
"service_callback",
|
||||
"logger_fn",
|
||||
"litellm_disabled_callbacks",
|
||||
"_agentic_loop_depth",
|
||||
"_agentic_loop_fingerprints",
|
||||
"_code_interpreter_interception_active",
|
||||
"_code_interpreter_interception_converted_stream",
|
||||
"_code_interpreter_interception_sandbox_key",
|
||||
"_headroom_interception_converted_stream",
|
||||
"max_agentic_loops",
|
||||
*AGENTIC_LOOP_KWARG_NAMES,
|
||||
],
|
||||
)
|
||||
async def test_add_litellm_data_to_request_strips_callback_control_fields(
|
||||
|
|
@ -1313,13 +1308,18 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields(
|
|||
"logger_fn": "module.func",
|
||||
"_agentic_loop_depth": 5,
|
||||
"_agentic_loop_fingerprints": ["forged"],
|
||||
"_agentic_loop_api_surface": "responses",
|
||||
"_code_interpreter_interception_active": True,
|
||||
"_code_interpreter_interception_converted_stream": True,
|
||||
"_code_interpreter_interception_sandbox_key": "forged-key",
|
||||
"_code_interpreter_interception_session_scoped": True,
|
||||
"_websearch_interception_emit_native_blocks": True,
|
||||
"_websearch_interception_converted_stream": True,
|
||||
"_websearch_interception_stream_options": {"include_usage": True},
|
||||
"_headroom_interception_converted_stream": True,
|
||||
"max_agentic_loops": 9999,
|
||||
}
|
||||
sample_value = sample_values[control_field]
|
||||
sample_value: Final = sample_values.get(control_field, "forged")
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
|||
from litellm.proxy.utils import is_valid_api_key
|
||||
from litellm.responses.main import aresponses, responses
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, STREAM_OPTIONS_STASH_KEYS
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams, LiteLLM_Params
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -5291,6 +5291,7 @@ def test_function_setup_failure_log_line_shows_outer_not_doomed_ids(monkeypatch)
|
|||
WEBSEARCH_INTERNAL_CONTROL_FIELDS = (
|
||||
"_websearch_interception_emit_native_blocks",
|
||||
"_websearch_interception_converted_stream",
|
||||
"_websearch_interception_stream_options",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -5318,6 +5319,15 @@ def test_websearch_interception_control_fields_never_reach_the_provider():
|
|||
assert set(WEBSEARCH_INTERNAL_CONTROL_FIELDS) <= set(all_litellm_params)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stash_key", STREAM_OPTIONS_STASH_KEYS)
|
||||
def test_converted_stream_options_stash_never_reaches_the_provider(stash_key: str):
|
||||
non_default: Final = get_non_default_completion_params(
|
||||
{"a_real_provider_specific_param": 1, stash_key: {"include_usage": True}}
|
||||
)
|
||||
|
||||
assert non_default == {"a_real_provider_specific_param": 1}
|
||||
|
||||
|
||||
def test_get_litellm_params_keys_never_reach_the_provider():
|
||||
"""Bridges (chat <-> Responses, agentic loop follow-ups) forward litellm_params as
|
||||
`completion()` kwargs. Any key the param builder does not recognize is swept into
|
||||
|
|
|
|||
|
|
@ -227,9 +227,12 @@ AGENTIC_LOOP_STATE_NAMES: Final = (
|
|||
"_code_interpreter_interception_sandbox_key",
|
||||
"_code_interpreter_interception_session_scoped",
|
||||
"_code_interpreter_interception_converted_stream",
|
||||
"_code_interpreter_interception_stream_options",
|
||||
"_websearch_interception_emit_native_blocks",
|
||||
"_websearch_interception_converted_stream",
|
||||
"_websearch_interception_stream_options",
|
||||
"_headroom_interception_converted_stream",
|
||||
"_headroom_interception_stream_options",
|
||||
)
|
||||
|
||||
INTERNAL_STATE_NAMES: Final = (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue