diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 764dd3ed3e0..69b528a29f8 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -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) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 15413140d9e..0fb482cbbd9 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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 diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 9d7c9864e62..628486f377f 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -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, + ) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 731b39e681f..db43f7cd052 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -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(): diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 166ad916ea2..f20762789dd 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -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 ( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8..0c3fb3d8550 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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() diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 36d48d49c7e..aa10a80f815 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -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, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 490b393aa68..913310952cd 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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" diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index d2db9357349..13542f22f88 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -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. diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index ff31d7a65c4..096ad1c30dd 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -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) diff --git a/tests/integration/_support/interception_vendor.py b/tests/integration/_support/interception_vendor.py new file mode 100644 index 00000000000..3ca40087441 --- /dev/null +++ b/tests/integration/_support/interception_vendor.py @@ -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 + ), + ) diff --git a/tests/integration/sdk/test_interception_converted_stream_options_sdk.py b/tests/integration/sdk/test_interception_converted_stream_options_sdk.py new file mode 100644 index 00000000000..17876fe605b --- /dev/null +++ b/tests/integration/sdk/test_interception_converted_stream_options_sdk.py @@ -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) diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py new file mode 100644 index 00000000000..46a4211b121 --- /dev/null +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -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 diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index 9d308ac1989..c037b6be80d 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -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={ diff --git a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py index bb55ba57680..064b94942a3 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py @@ -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. diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 0d7f735e6e5..e891706df13 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -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 [] + ) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py index 282373d4b6d..03553e65ad4 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py @@ -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" diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 9daa60bbf88..eb2d62a8d71 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -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, diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 25d2d11b5ca..347840b5c5d 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index f3456f6b60c..6f4b1816150 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -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 = { diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 2ba736d8ae0..cf04b9ff91e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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={ diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 034d501483c..467fe206089 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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 diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 16c9b875636..d27dc733c78 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -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 = (