diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index fc1f1830129..45c82589d2d 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -23,6 +23,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( cache_control_as_prompt_cache_breakpoint, with_prompt_cache_breakpoint, @@ -175,7 +176,10 @@ def _hosted_entry_flag(entry: Mapping[str, object], resolve_provider: Callable[[ def _hosted_openai_dialect_flag( - model: str, custom_llm_provider: str | None, resolve_provider: Callable[[str], str | None] + model: str, + custom_llm_provider: str | None, + resolve_provider: Callable[[str], str | None], + model_cost: Mapping[str, Mapping[str, object]] | None = None, ) -> bool | None: """ Explicit opt-in for an OpenAI-shaped deployment served by another provider. @@ -194,14 +198,19 @@ def _hosted_openai_dialect_flag( """ import litellm - exact_entry: Final = litellm.model_cost.get(model) - exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, dict) else None + cost_map: Final = ( + cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: the live map holds catalog entries + if model_cost is None + else model_cost + ) + exact_entry: Final = cost_map.get(model) + exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, Mapping) else None if custom_llm_provider is None and (exact_provider == "openai" or ("/" not in model and exact_provider is None)): return None provider: Final = custom_llm_provider or resolve_provider(model) if provider is None or provider == "openai": return None - if isinstance(exact_entry, dict) and exact_provider == provider: + if isinstance(exact_entry, Mapping) and exact_provider == provider: return _hosted_entry_flag(exact_entry, lambda: provider) from litellm.utils import get_potential_model_names @@ -214,8 +223,8 @@ def _hosted_openai_dialect_flag( names["stripped_model_name"], names["provider_prefixed_model_name"], ) - entries: Final = (litellm.model_cost.get(candidate) for candidate in candidates) - flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, dict)) + entries: Final = (cost_map.get(candidate) for candidate in candidates) + flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping)) return next((flag for flag in flags if flag is not None), None) @@ -537,6 +546,18 @@ class AnthropicCacheControlHook(CustomPromptManagement): return False return prompt_cache_options is not None or targets_openai_api(api_base) + @staticmethod + def catalog_targets_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str) -> bool: + return ( + _hosted_openai_dialect_flag( + model, + custom_llm_provider, + AnthropicCacheControlHook._resolve_provider, + GetModelCostMap.loaded_model_cost_map(), + ) + is True + ) + @staticmethod def _resolve_provider(model: str) -> str | None: return _served_provider_for_model(model) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index d382d653a53..5381447f28f 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -6,15 +6,19 @@ from typing import TYPE_CHECKING, Final, cast import httpx from httpx import Response +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.prompt_templates.common_utils import ( content_parts, convert_content_list_to_str, drop_tool_reference_parts_from_tool_messages, filter_value_from_dict, is_pdf_file_part, + prompt_cache_breakpoint_for_wire, + with_prompt_cache_breakpoint, ) from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import ( @@ -27,7 +31,7 @@ from litellm.llms.openai.common_utils import drop_params_from_unprocessable_enti from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, PromptCacheBreakpoint from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField from litellm.utils import add_path_to_api_base, supports_tool_choice @@ -62,6 +66,72 @@ def _has_list_only_content_part(message: Mapping[str, object]) -> bool: return any(_is_list_only_content_part(part) for part in content_parts(message)) +_CONTENT_PARTS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _text_parts(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + content: Final = message.get("content") + if not isinstance(content, list): + return () + try: + parts: Final = _CONTENT_PARTS_ADAPTER.validate_python(content) + except ValidationError: + return () + return tuple(part for part in parts if part.get("type") == "text") + + +def _with_only_text_parts(message: AllMessageValues) -> AllMessageValues: + if _has_list_only_content_part(message) or not (text_parts := _text_parts(message)): + return message + text_only: Final = {**message, "content": list(text_parts)} + return cast(AllMessageValues, text_only) # cast-ok: the same message, text parts only + + +_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def _foundry_breakpoint(marker: object) -> PromptCacheBreakpoint | None: + if prompt_cache_breakpoint_for_wire(marker, drop_params=True) is None: + return None + return PromptCacheBreakpoint(mode="explicit") + + +def _part_with_foundry_breakpoint(part: object) -> object: + if not isinstance(part, Mapping): + return part + fields: Final = _FIELDS.validate_python(part) + if "prompt_cache_breakpoint" not in fields: + return fields + unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"} + return with_prompt_cache_breakpoint(unmarked, _foundry_breakpoint(fields["prompt_cache_breakpoint"])) + + +def _with_breakpoints_for_foundry(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + renamed: Final = AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(_with_only_text_parts(message)) + parts: Final = content_parts(renamed) if drop_params else () + if not parts: + return renamed + foundry_message: Final = {**renamed, "content": [_part_with_foundry_breakpoint(part) for part in parts]} + return cast(AllMessageValues, foundry_message) # cast-ok: the same message, its breakpoints in Foundry's shape + + +def _deployment_reads_breakpoints(model: str, litellm_params: Mapping[str, object]) -> bool: + base_model: Final = litellm_params.get("base_model") + named_models: Final = (base_model, model) if isinstance(base_model, str) and base_model else (model,) + model_info: Final = litellm_params.get("model_info") + if not isinstance(model_info, Mapping): + return any( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai") for name in named_models + ) + deployment_flag: Final = _FIELDS.validate_python(model_info).get("supports_prompt_cache_breakpoint") + if isinstance(deployment_flag, bool): + return deployment_flag + return any( + AnthropicCacheControlHook.catalog_targets_openai_prompt_cache_breakpoint(name, "azure_ai") + for name in named_models + ) + + class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract usage: Final = chunk.get("usage") @@ -251,6 +321,8 @@ class AzureAIStudioConfig(OpenAIConfig): self, messages: list[AllMessageValues], model: str, + reads_breakpoints: bool | None = None, + drop_params: bool = False, ) -> list: """ - Azure AI Studio doesn't support content as a list. This handles: @@ -260,13 +332,31 @@ class AzureAIStudioConfig(OpenAIConfig): Azure AI Foundry backends set additionalProperties=false and reject these with "Extra inputs are not permitted", which breaks multi-turn Anthropic-format clients that echo thinking blocks back as history. - 2. Transforms list content to a string. + 2. Transforms list content to a string, except on deployments that read OpenAI explicit + prompt cache breakpoints (the deployment's own model_info opt-in, else the cost-map entry + of its base_model or name, as loaded and without a sibling deployment's opt-in): only a + content part can carry the breakpoint (a message-level one is silently ignored), so those + keep their text parts, with each Anthropic `cache_control` mark renamed to a breakpoint the + way the openai provider renames it. Every other part is dropped, exactly as the string + conversion drops it, before the marks are renamed, so a message-level mark lands on the + last text part. Under drop_params each breakpoint goes out as `{"mode": "explicit"}`, the only + shape Foundry takes, and a malformed one is dropped. 3. If message contains an image, audio, or PDF file part, send as is (user-intended) Operates on a deep copy so the caller's messages keep their thinking blocks and provider metadata, which a fallback to another provider still needs. """ - stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) + deployment_reads_breakpoints: Final = ( + AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(model, "azure_ai") + if reads_breakpoints is None + else reads_breakpoints + ) + copied_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) + stripped_messages: Final = ( + [_with_breakpoints_for_foundry(message, drop_params) for message in copied_messages] + if deployment_reads_breakpoints + else copied_messages + ) for message in stripped_messages: message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: @@ -275,6 +365,10 @@ class AzureAIStudioConfig(OpenAIConfig): if _has_list_only_content_part(message): continue + if deployment_reads_breakpoints and (text_parts := _text_parts(message)): + message_dict["content"] = list(text_parts) + continue + texts = convert_content_list_to_str(message=message) if texts: message["content"] = texts @@ -329,7 +423,16 @@ class AzureAIStudioConfig(OpenAIConfig): if extra_body and isinstance(extra_body, dict): optional_params.update(extra_body) optional_params.pop("max_retries", None) - return super().transform_request(model, messages, optional_params, litellm_params, headers) + return { + "model": model, + "messages": self._transform_messages( + messages, + model, + _deployment_reads_breakpoints(model, _FIELDS.validate_python(litellm_params)), + bool(litellm_params.get("drop_params") or litellm.drop_params), + ), + **optional_params, + } def transform_response( self, diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index 3df73d3b767..d470605b1e5 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -153,15 +153,16 @@ def _foundry_deployment(rig: _Rig, scenario: Scenario, *, bare: bool) -> str: return name -def _assert_foundry_chat_wire(received: Request, prompt: str, *, options: JsonValue | None) -> None: +def _assert_foundry_chat_wire(received: Request, user_content: JsonValue, *, options: JsonValue | None) -> None: body: Final = body_of(received) + marked_system: Final = [{"type": "text", "text": SYSTEM, "prompt_cache_breakpoint": EXPLICIT}] assert urlsplit(received.target).netloc == _FOUNDRY_HOST, received.target assert urlsplit(received.target).path == _FOUNDRY_CHAT_PATH, received.target - assert "prompt_cache_breakpoint" not in received.body.decode(), received.body assert body["model"] == _FOUNDRY_MODEL, received.body - assert body["messages"] == [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}], ( - received.body - ) + assert body["messages"] == [ + {"role": "system", "content": marked_system}, + {"role": "user", "content": user_content}, + ], received.body assert body.get("prompt_cache_options") == options, received.body @@ -185,7 +186,7 @@ def test_b9_b12_foundry_gpt6_responses_reach_the_openai_dialect_on_the_foundry_h @pytest.mark.timeout(_CELL_TIMEOUT) -def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig) -> None: +def test_b10_foundry_gpt6_messages_bridge_to_chat_keeps_the_injected_breakpoint(rig: _Rig) -> None: marker: Final = fresh_marker() with rig.gateway.scenario() as scenario: name: Final = _foundry_deployment(rig, scenario, bare=False) @@ -193,12 +194,12 @@ def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig rig.gateway, rig.wire, "messages", request_body("messages", name, prompt_text(marker)) ) assert_answered(outcome, marker) - _assert_foundry_chat_wire(received, prompt_text(marker), options=IMPLICIT) + _assert_foundry_chat_wire(received, [{"type": "text", "text": prompt_text(marker)}], options=IMPLICIT) success_row(name, marker) @pytest.mark.timeout(_CELL_TIMEOUT) -def test_b11_foundry_gpt6_chat_stays_on_the_plain_chat_wire(rig: _Rig) -> None: +def test_b11_foundry_gpt6_chat_stays_on_the_chat_wire_with_the_injected_breakpoint(rig: _Rig) -> None: marker: Final = fresh_marker() with rig.gateway.scenario() as scenario: name: Final = _foundry_deployment(rig, scenario, bare=False) diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 6269fd367fe..145d45317c3 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1,6 +1,7 @@ import json import traceback from collections.abc import Callable +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import respx from httpx import Response import litellm +from litellm import Router from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.azure_ai.azure_model_router.transformation import ( @@ -880,6 +882,360 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages(): } +def test_azure_ai_keeps_prompt_cache_breakpoint_for_a_model_that_supports_it(_local_model_cost_map): + """ + Foundry GPT-6 deployments take OpenAI explicit prompt caching, and the breakpoint only + counts on a content part, so a system message carrying one must reach the request body as + the content-part list it arrived in while the Anthropic-only fields are still stripped. + """ + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "stable instructions", + "prompt_cache_breakpoint": {"mode": "explicit"}, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][1]["content"] == [{"type": "text", "text": "hi"}] + assert request["prompt_cache_options"] == {"mode": "explicit"} + + +def test_azure_ai_still_flattens_list_content_for_a_model_without_the_breakpoint_flag(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-4o", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "stable ", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "instructions"}, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == "stable instructions" + assert request["messages"][1]["content"] == "hi" + assert not _find_key_anywhere(request["messages"], "prompt_cache_breakpoint") + + +def test_azure_ai_drops_thinking_parts_from_history_while_keeping_the_breakpoint(_local_model_cost_map): + """ + A multi-turn client can echo an assistant turn back with thinking parts inside its content + list. Foundry rejects those parts, and the string conversion always dropped them, so a model + that keeps content parts must still drop them while the text parts and the breakpoint stay. + """ + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ], + }, + {"role": "user", "content": "read the file"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "The user wants me to read a file.", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"}, + {"type": "text", "text": "Reading it now."}, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "thanks"}]}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][2]["content"] == [{"type": "text", "text": "Reading it now."}] + assert request["messages"][3]["content"] == [{"type": "text", "text": "thanks"}] + assert not _find_key_anywhere(request["messages"], "thinking") + assert not _find_key_anywhere(request["messages"], "data") + + + +def test_azure_ai_renames_a_cache_control_mark_to_a_breakpoint_for_a_model_that_reads_them(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "system", + "content": [{"type": "text", "text": "stable instructions", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "hi", "cache_control": {"type": "ephemeral"}}, + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request["messages"][1]["content"] == [ + {"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert not _find_key_anywhere(request["messages"], "cache_control") + + +def test_azure_ai_drops_a_file_part_next_to_the_breakpoint_like_the_string_conversion(_local_model_cost_map): + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "file", "file": {"file_id": "file-held-by-another-provider"}}, + ], + } + ] + + kept: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + flattened: Final = AzureAIStudioConfig().transform_request( + model="gpt-4o", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + + assert kept["messages"][0]["content"] == [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert flattened["messages"][0]["content"] == "summarize" + + +def test_azure_ai_sends_a_pdf_part_next_to_the_breakpoint_as_is(_local_model_cost_map): + content: Final = [ + {"type": "file", "file": {"filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}}, + {"type": "text", "text": "What is the codeword?", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ] + + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[{"role": "user", "content": content}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [{"role": "user", "content": content}] + + +def test_azure_ai_puts_a_message_level_mark_on_the_last_text_part_when_a_file_part_follows(_local_model_cost_map): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + { + "role": "user", + "cache_control": {"type": "ephemeral"}, + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "file", "file": {"file_id": "file-held-by-another-provider"}}, + ], + } + ], + optional_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][0]["content"] == [ + {"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + + +@pytest.mark.parametrize( + ("registered", "keeps_parts"), + [ + ({"azure_ai/qa-foundry-alias": {"litellm_provider": "azure_ai", "supports_prompt_cache_breakpoint": True}}, True), + ({"qa-foundry-alias": {"litellm_provider": "openai", "supports_prompt_cache_breakpoint": True}}, False), + ], + ids=["azure_ai-entry-opts-in", "openai-entry-does-not-speak-for-azure_ai"], +) +def test_azure_ai_keeps_parts_only_when_the_azure_ai_entry_reads_breakpoints( + _local_model_cost_map, registered: dict[str, dict[str, object]], keeps_parts: bool +): + litellm.register_model(model_cost=registered, persist_across_reloads=False, add_to_provider_model_sets=False) + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="qa-foundry-alias", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params={"custom_llm_provider": "azure_ai", "drop_params": True}, + headers={}, + ) + + assert request["messages"][0]["content"] == ([breakpoint_part] if keeps_parts else "stable") + + +@pytest.mark.parametrize( + "registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")] +) +def test_azure_ai_deployment_reads_only_its_own_breakpoint_flag_not_a_siblings_sharing_its_model_string( + _local_model_cost_map, respx_mock: respx.MockRouter, registration_order: tuple[str, ...] +) -> None: + breakpoint_flags: Final = MappingProxyType( + { + "opted-in": {"supports_prompt_cache_breakpoint": True}, + "opted-out": {"supports_prompt_cache_breakpoint": False}, + } + ) + def deployment(host: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": host, + "litellm_params": { + "model": "azure_ai/qa-shared-deployment", + "api_base": f"https://{host}.services.ai.azure.com/models", + "api_key": "azure-ai-test-key", + }, + "model_info": model_info, + } + + routes: Final = { + host: respx_mock.post(host=f"{host}.services.ai.azure.com", path="/models/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-azure-ai", + "object": "chat.completion", + "created": 1700000000, + "model": "qa-shared-deployment", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + }, + ) + ) + for host in registration_order + } + router: Final = Router( + model_list=[ + deployment(host, {"id": f"{host}-id", **breakpoint_flags.get(host, {})}) for host in registration_order + ] + ) + marked_system: Final = { + "role": "system", + "content": [{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}], + } + + for host in registration_order: + router.completion(model=host, messages=[marked_system, {"role": "user", "content": "hi"}]) + + def sent_system_content(host: str) -> object: + return json.loads(routes[host].calls.last.request.content)["messages"][0]["content"] + + assert sent_system_content("opted-in") == [ + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert sent_system_content("opted-out") == "stable" + assert sent_system_content("unflagged") == "stable" + + +@pytest.mark.parametrize( + "litellm_params", + [ + {"base_model": "azure_ai/gpt-6-astra"}, + {"base_model": "azure_ai/gpt-6-astra", "model_info": {"id": "my-astra-deployment-id"}}, + ], +) +def test_azure_ai_keeps_parts_for_a_custom_deployment_name_whose_base_model_reads_breakpoints( + _local_model_cost_map, litellm_params: dict[str, object] +): + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="my-astra-deployment", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [breakpoint_part] + + +@pytest.mark.parametrize( + "litellm_params", + [ + {"base_model": "azure/cost-tracking-row"}, + {"base_model": "azure/cost-tracking-row", "model_info": {"id": "astra-deployment-id"}}, + ], +) +def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_read_breakpoints( + _local_model_cost_map, litellm_params: dict[str, object] +): + breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}} + + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[{"role": "system", "content": [breakpoint_part]}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [breakpoint_part] + + +@pytest.mark.parametrize( + "sent_marker, litellm_params, wire_part", + [ + ( + {"mode": "explicit", "ttl": "30m"}, + {"drop_params": True}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ( + {"mode": "explicit", "ttl": "30m"}, + {}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}, + ), + ], +) +def test_azure_ai_reshapes_breakpoints_for_foundry_only_under_drop_params( + _local_model_cost_map, + sent_marker: dict[str, object], + litellm_params: dict[str, object], + wire_part: dict[str, object], +): + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-6-astra", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "stable", "prompt_cache_breakpoint": sent_marker}]} + ], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + assert request["messages"][0]["content"] == [wire_part] + + def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text(): messages: Final = [ {"role": "user", "content": "load the WebFetch tool"},