mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(azure_ai)!: keep OpenAI prompt cache breakpoints on the Foundry chat completions path (#45299)
* fix(azure_ai): keep OpenAI prompt cache breakpoints on the Foundry chat completions path * fix(azure_ai): keep only OpenAI content part types when preserving prompt cache breakpoints * fix(azure_ai): share the breakpoint check with the cache-mark step and keep only text parts A deployment's model_info opt-in is registered under its provider-prefixed cost-map key without a litellm_provider, so that prefix now names the provider the flag speaks for. * refactor(azure_ai): read text parts through a walrus and mark new test values Final * fix(azure_ai): read a router-registered breakpoint opt-in only for azure_ai and keep a message mark off dropped file parts * fix(azure_ai): read the breakpoint opt-in from each deployment's own model_info and base_model A router deployment without its own supports_prompt_cache_breakpoint flag now reads the cost map as loaded, so a sibling's opt-in on the same model string no longer turns its requests into 400s. The shared hook no longer reads router registered keys for azure_ai. * fix(azure_ai): send breakpoints in Foundry's shape under drop_params and fall back to the model name * test(integration): expect the injected breakpoint on the Foundry GPT-6 chat wire --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
dbcacc7ab7
commit
4a5c5665fa
4 changed files with 499 additions and 18 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue