diff --git a/litellm/constants.py b/litellm/constants.py index 42ad624026c..d02114f56b1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -158,7 +158,7 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_ DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250) HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( - "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400 + "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 60, 1, 86400 ) # Maximum wall-clock seconds a streaming response is allowed to run. @@ -627,6 +627,7 @@ STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS: Final = 60 BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = get_env_int("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024) # Anthropic's Messages API rejects thinking.budget_tokens < 1024. ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: Final = 1024 +ANTHROPIC_BILLING_METADATA_PREFIX: Final = "x-anthropic-billing-header:" REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_float("REPLICATE_POLLING_DELAY_SECONDS", 0.5) DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = get_env_int("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 128000) DEFAULT_OCI_CHAT_MAX_TOKENS: Final = 4096 diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 92e7c84031c..865247afc63 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( + ANTHROPIC_BILLING_METADATA_PREFIX, ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES, DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, @@ -2066,7 +2067,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if not system_message_block["content"]: continue if self.should_strip_billing_metadata() and system_message_block["content"].startswith( - "x-anthropic-billing-header:" + ANTHROPIC_BILLING_METADATA_PREFIX ): continue anthropic_system_message_content = AnthropicSystemMessageContent( @@ -2086,7 +2087,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self.should_strip_billing_metadata() and _content.get("type") == "text" and text_value - and text_value.startswith("x-anthropic-billing-header:") + and text_value.startswith(ANTHROPIC_BILLING_METADATA_PREFIX) ): continue anthropic_system_message_content = AnthropicSystemMessageContent( diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index f1b97f3ebeb..c086618e829 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -17,6 +17,7 @@ from litellm.anthropic_beta_headers_manager import ( has_beta_headers_config, is_beta_header_supported_by_bundled_config, ) +from litellm.constants import ANTHROPIC_BILLING_METADATA_PREFIX from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -222,7 +223,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): """ if isinstance(system_param, str): # If it's a string and starts with billing header, filter it out - if system_param.startswith("x-anthropic-billing-header:"): + if system_param.startswith(ANTHROPIC_BILLING_METADATA_PREFIX): return None return system_param elif isinstance(system_param, list): @@ -233,7 +234,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): text = content_block.get("text", "") content_type = content_block.get("type", "") # Skip text blocks that start with billing header - if content_type == "text" and text.startswith("x-anthropic-billing-header:"): + if content_type == "text" and text.startswith(ANTHROPIC_BILLING_METADATA_PREFIX): continue filtered_list.append(content_block) else: diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 26ae30635d4..990c8869de3 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, ) +from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_reasoning_transformation import GPT_REASONING_SERIES_MARKERS from litellm.types.llms.azure import ( @@ -98,6 +99,22 @@ class AzureOpenAIConfig(BaseConfig): if key != "self" and value is not None: setattr(self.__class__, key, value) + def translate_developer_role_to_system_role( + self, + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, + ) -> Sequence[AllMessageValues]: + """ + Azure OpenAI accepts `system` messages at any position, so a developer message + only changes role. Declared here so `AzureOpenAIGPT5Config` resolves to it + ahead of the OpenAI-compatible hoist it also inherits. + """ + return map_developer_role_to_system_role(messages=messages) + @classmethod def get_config(cls): return super().get_config() diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index a9ec1f97d20..ef8d080027f 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -10,14 +10,23 @@ from abc import ABC, abstractmethod from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, TypeAlias from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # narrows a message to its system member by the role literal from litellm._logging import verbose_logger +from litellm.constants import ANTHROPIC_BILLING_METADATA_PREFIX from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionCachedContent, + ChatCompletionSystemMessage, + ChatCompletionTextObject, + ChatCompletionToolCallChunk, + ChatCompletionUserMessage, +) from litellm.types.utils import ( CountContextManagement, Message, @@ -458,7 +467,7 @@ def type_to_response_format_param( def map_developer_role_to_system_role( - messages: list[AllMessageValues], + messages: Sequence[AllMessageValues], ) -> list[AllMessageValues]: """ Translate `developer` role to `system` role for non-OpenAI providers. @@ -473,3 +482,226 @@ def map_developer_role_to_system_role( else: new_messages.append(m) return new_messages + + +SystemMessageContent: TypeAlias = str | list[object] + +_system_content_adapter: Final = TypeAdapter[SystemMessageContent | None](SystemMessageContent | None) + + +def _system_content(message: ChatCompletionSystemMessage) -> SystemMessageContent | None: + if "content" not in message: + return "" + try: + content: Final = _system_content_adapter.validate_python(message["content"]) + except ValidationError: + return None + return "" if content is None else content + + +def _message_cache_control(message: ChatCompletionSystemMessage) -> ChatCompletionCachedContent | None: + if "cache_control" not in message: + return None + return message["cache_control"] + + +_SystemPart: TypeAlias = tuple[SystemMessageContent, ChatCompletionCachedContent | None] + + +def _with_message_cache_control( + blocks: Sequence[object], cache_control: ChatCompletionCachedContent | None +) -> tuple[object, ...]: + last: Final = as_str_mapping(blocks[-1]) if blocks else None + if last is None or cache_control is None or "cache_control" in last: + return tuple(blocks) + return (*blocks[:-1], {**last, "cache_control": cache_control}) + + +def _as_system_message(message: AllMessageValues) -> AllMessageValues: + if message["role"] != "developer": + return message + content_only: Final[ChatCompletionSystemMessage] = {"role": "system", "content": message["content"]} + named: Final[ChatCompletionSystemMessage] = ( + content_only if "name" not in message else {**content_only, "name": message["name"]} + ) + if "cache_control" not in message: + return named + cached: Final[ChatCompletionSystemMessage] = {**named, "cache_control": message["cache_control"]} + return cached + + +def _text_blocks( + content: SystemMessageContent, cache_control: ChatCompletionCachedContent | None +) -> tuple[object, ...]: + if not isinstance(content, str): + return _with_message_cache_control(content, cache_control) + if not content: + return () + if cache_control: + cached_block: Final[ChatCompletionTextObject] = { + "type": "text", + "text": content, + "cache_control": cache_control, + } + return (cached_block,) + text_block: Final[ChatCompletionTextObject] = {"type": "text", "text": content} + return (text_block,) + + +def _lone_text_block_text(blocks: Sequence[object]) -> str | None: + match blocks: + case [{"type": "text", "text": str() as text} as block] if block == {"type": "text", "text": text}: + return text + case _: + return None + + +def _plain_text(content: SystemMessageContent, cache_control: ChatCompletionCachedContent | None) -> str | None: + if cache_control: + return None + text: Final = content if isinstance(content, str) else _lone_text_block_text(content) + if text is None or text.startswith(ANTHROPIC_BILLING_METADATA_PREFIX): + return None + return text + + +def _is_system_message( + message: AllMessageValues, +) -> TypeIs[ChatCompletionSystemMessage]: # guard-ok: the role literal picks the TypedDict member + return message["role"] == "system" + + +def _with_leading_separator(blocks: tuple[object, ...]) -> tuple[object, ...]: + first: Final = as_str_mapping(blocks[0]) + text: Final = None if first is None else first.get("text") + if first is None or not isinstance(text, str) or text.startswith(ANTHROPIC_BILLING_METADATA_PREFIX): + return blocks + return ({**first, "text": f"\n\n{text}"}, *blocks[1:]) + + +def _merged_system_content(parts: Sequence[_SystemPart]) -> SystemMessageContent: + texts: Final = tuple(itertools.starmap(_plain_text, parts)) + if all(text is not None for text in texts): + return "\n\n".join(text for text in texts if text) + block_runs: Final = tuple(blocks for blocks in itertools.starmap(_text_blocks, parts) if blocks) + separated: Final = (*block_runs[:1], *map(_with_leading_separator, block_runs[1:])) + return list(itertools.chain.from_iterable(separated)) + + +def _merged_system_message( + run: Sequence[ChatCompletionSystemMessage], parts: Sequence[_SystemPart] +) -> ChatCompletionSystemMessage: + content: Final = _merged_system_content(parts) + named: Final = next((message for message in reversed(run) if "name" in message), None) + if named is None: + merged: Final[ChatCompletionSystemMessage] = {"role": "system", "content": content} + return merged + with_name: Final[ChatCompletionSystemMessage] = {"role": "system", "content": content, "name": named["name"]} + return with_name + + +def _merged_run(run: tuple[AllMessageValues, ...]) -> tuple[AllMessageValues, ...]: + system_run: Final = tuple(message for message in run if _is_system_message(message)) + if len(system_run) < 2: + return run + contents: Final = tuple((_system_content(message), _message_cache_control(message)) for message in system_run) + parts: Final = tuple((content, cache_control) for content, cache_control in contents if content is not None) + if len(parts) != len(contents): + return run + return (_merged_system_message(system_run, parts),) + + +def _merged_system_runs(messages: Sequence[AllMessageValues]) -> Iterator[AllMessageValues]: + runs: Final = (tuple(run) for _, run in itertools.groupby(messages, key=_is_system_message)) + return itertools.chain.from_iterable(map(_merged_run, runs)) + + +_INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) + + +def _leading_instruction_block_length(roles: Sequence[str | None]) -> int: + return next((index for index, role in enumerate(roles) if role not in _INSTRUCTION_ROLES), len(roles)) + + +def _closing_instruction_block_start(roles: Sequence[str | None], leading_length: int) -> int: + last_conversation_index: Final = next( + (index for index in range(len(roles) - 1, leading_length - 1, -1) if roles[index] not in _INSTRUCTION_ROLES), + None, + ) + if last_conversation_index is None or roles[last_conversation_index] != "assistant": + return len(roles) + return last_conversation_index + 1 + + +def hoisted_instruction_indices(roles: Sequence[str | None]) -> tuple[int, ...]: + """ + Indices of the leading instruction block plus every developer item after it, except + a developer block that closes the conversation right after an assistant turn, which + stays in place so the request does not end on the assistant's turn. Items without a + role (tool calls and their outputs) count as conversation. + """ + leading_length: Final = _leading_instruction_block_length(roles) + closing_start: Final = _closing_instruction_block_start(roles, leading_length) + moved: Final = tuple(index for index in range(leading_length, closing_start) if roles[index] == "developer") + return (*range(leading_length), *moved) + + +def closing_instruction_indices(roles: Sequence[str | None]) -> frozenset[int]: + closing_start: Final = _closing_instruction_block_start(roles, _leading_instruction_block_length(roles)) + return frozenset(range(closing_start, len(roles))) + + +def _role_name(role: object) -> str | None: + return role if isinstance(role, str) else None + + +def _message_roles(messages: Sequence[AllMessageValues]) -> tuple[str | None, ...]: + return tuple(_role_name(message["role"]) for message in messages) + + +def _hoisted_message_order(messages: Sequence[AllMessageValues]) -> tuple[int, ...]: + hoisted: Final = hoisted_instruction_indices(_message_roles(messages)) + hoisted_set: Final = frozenset(hoisted) + return (*hoisted, *(index for index in range(len(messages)) if index not in hoisted_set)) + + +def _developer_message_as_user(message: AllMessageValues) -> AllMessageValues: + if message["role"] != "developer": + return message + copied: Final = copy.deepcopy(message) + if "cache_control" not in copied: + as_user: Final[ChatCompletionUserMessage] = {"role": "user", "content": copied["content"]} + return as_user + cached_as_user: Final[ChatCompletionUserMessage] = { + "role": "user", + "content": copied["content"], + "cache_control": copied["cache_control"], + } + return cached_as_user + + +def merge_system_message_runs(messages: Sequence[AllMessageValues]) -> Sequence[AllMessageValues]: + """ + Translate `developer` role to `system` role in place and fold each run of + consecutive system messages into one, for chat templates that render a single + system message per run. + """ + return tuple(_merged_system_runs(tuple(map(_as_system_message, messages)))) + + +def hoist_developer_messages_into_leading_system_message( + messages: Sequence[AllMessageValues], +) -> Sequence[AllMessageValues]: + """ + Place instructions for an OpenAI-compatible model whose chat template allows a single + system message and only at the start: developer messages after the first user turn move + into the leading system block, a developer message that closes the conversation right + after an assistant turn becomes a user turn so the request does not end on the + assistant's turn, and each run of consecutive system messages folds into one. + """ + order: Final = _hoisted_message_order(messages) + closing: Final = closing_instruction_indices(_message_roles(messages)) + placed: Final = tuple( + _developer_message_as_user(messages[index]) if index in closing else messages[index] for index in order + ) + return merge_system_message_runs(placed) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 207a4705c92..b87a783e244 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -4,7 +4,7 @@ Common base config for all LLM providers import types from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Union import httpx @@ -199,12 +199,19 @@ class BaseConfig(ABC): def translate_developer_role_to_system_role( self, - messages: list[AllMessageValues], - ) -> list[AllMessageValues]: + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, + ) -> Sequence[AllMessageValues]: """ Translate `developer` role to `system` role for non-OpenAI providers. - Overriden by OpenAI/Azure + Overridden by OpenAI, which also places the instructions in one leading system + message for an OpenAI-compatible model declared to reject a later system + message, and by the o-series and Azure configs. """ return map_developer_role_to_system_role(messages=messages) diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index db234c09930..86d9df9618a 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -1,10 +1,9 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final from urllib.parse import unquote import httpx -from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.fireworks_ai.common_utils import ( @@ -15,12 +14,18 @@ from litellm.llms.fireworks_ai.common_utils import ( with_fireworks_session_affinity, without_caller_user, ) +from litellm.llms.openai.responses.instruction_items import developer_items_as_system, instruction_items_folded from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ResponseInputParam from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.utils import ( + declared_mid_conversation_system_support, + deployment_declared_mid_conversation_system_support, + deployment_model_info, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -45,90 +50,6 @@ def _forwarded_user_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, ) -_INSTRUCTION_ROLES: Final = frozenset({"system", "developer"}) - - -def _role(item: ResponseInputItemParam) -> str | None: - match item: - case {"role": str(role)}: - return role - case _: - return None - - -def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam: - if "role" not in item or item["role"] != "developer": - return item - return EasyInputMessageParam(role="system", content=item["content"], type="message") - - -def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam: - if isinstance(input, str): - return input - return [_developer_item_as_system(item) for item in input] - - -def _text_part(part: ResponseInputContentParam) -> str | None: - match part: - case {"type": "input_text", "text": str(text)}: - return text - case _: - return None - - -def _text_only_content(item: ResponseInputItemParam) -> str | None: - match item: - case {"role": "system" | "developer", "content": str(text)}: - return text - case {"role": "system" | "developer", "content": [*parts]}: - texts: Final = tuple(map(_text_part, parts)) - return None if any(text is None for text in texts) else "\n\n".join(text for text in texts if text) - case _: - return None - - -def _leading_instruction_block_length(roles: Sequence[str | None]) -> int: - return next((index for index, role in enumerate(roles) if role not in _INSTRUCTION_ROLES), len(roles)) - - -def _closing_instruction_block_start(roles: Sequence[str | None], leading_length: int) -> int: - last_conversation_index: Final = next( - (index for index in range(len(roles) - 1, leading_length - 1, -1) if roles[index] not in _INSTRUCTION_ROLES), - None, - ) - if last_conversation_index is None or roles[last_conversation_index] != "assistant": - return len(roles) - return last_conversation_index + 1 - - -def _hoisted_indices(roles: Sequence[str | None]) -> tuple[int, ...]: - leading_length: Final = _leading_instruction_block_length(roles) - closing_start: Final = _closing_instruction_block_start(roles, leading_length) - return tuple( - index for index, role in enumerate(roles[:closing_start]) if index < leading_length or role == "developer" - ) - - -def _with_instruction_items_folded( - input: str | ResponseInputParam, instructions: str | None -) -> tuple[str | None, str | ResponseInputParam]: - if isinstance(input, str): - return instructions, input - items: Final = tuple(input) - folded: Final = MappingProxyType( - { - index: text - for index in _hoisted_indices(tuple(map(_role, items))) - if (text := _text_only_content(items[index])) is not None - } - ) - joined: Final = "\n\n".join(chunk for chunk in (instructions, *folded.values()) if chunk) - return ( - instructions if not folded else joined or None, - [_developer_item_as_system(item) for index, item in enumerate(items) if index not in folded], - ) - - class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -158,29 +79,28 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): self, model: str, input: str | ResponseInputParam, - response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature + response_api_optional_request_params: dict[str, object], # mutable-ok: overrides the base class signature litellm_params: GenericLiteLLMParams, - headers: dict, # mutable-ok: overrides the base class signature - ) -> dict: # mutable-ok: overrides the base class signature - instructions_param: Final[object] = response_api_optional_request_params.get("instructions") + headers: dict[str, str], # mutable-ok: overrides the base class signature + ) -> dict[str, object]: # mutable-ok: overrides the base class signature validated_input: Final = self._validate_input_param(input) - instructions, folded_input = ( - _with_instruction_items_folded(validated_input, instructions_param) - if isinstance(instructions_param, str | None) - else (instructions_param, _developer_items_as_system(validated_input)) - ) - instruction_entries: Final = () if instructions is None else (("instructions", instructions),) - folded_params: Final = { - key: value - for key, value in ( - *((key, value) for key, value in response_api_optional_request_params.items() if key != "instructions"), - *instruction_entries, + model_info: Final = deployment_model_info(litellm_params) + declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info) + folded: Final = ( + None + if deployment_declared_mid_conversation_system_support(model_info) is True + else instruction_items_folded( + validated_input, response_api_optional_request_params, "user" if declared is False else "system" ) - } + ) + placed_input, placed_params = folded or ( + developer_items_as_system(validated_input), + response_api_optional_request_params, + ) request: Final = super().transform_responses_api_request( model=resolve_fireworks_resource_name(model), - input=folded_input, - response_api_optional_request_params=folded_params, + input=placed_input, + response_api_optional_request_params=dict(placed_params), litellm_params=litellm_params, headers=headers, ) diff --git a/litellm/llms/hosted_vllm/responses/transformation.py b/litellm/llms/hosted_vllm/responses/transformation.py index 8489bc0da84..0831ed3b37a 100644 --- a/litellm/llms/hosted_vllm/responses/transformation.py +++ b/litellm/llms/hosted_vllm/responses/transformation.py @@ -9,10 +9,13 @@ the chat completions → responses conversion pipeline. from typing import Final from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.openai.responses.instruction_items import instruction_items_folded from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponseInputParam from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.utils import declared_mid_conversation_system_support, deployment_model_info class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig): @@ -69,6 +72,31 @@ class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig): return join_api_base(api_base, "v1/responses") + def transform_responses_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: dict[str, object], # mutable-ok: overrides the base class signature + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: overrides the base class signature + ) -> dict[str, object]: # mutable-ok: overrides the base class signature + declared: Final = declared_mid_conversation_system_support( + model, self.custom_llm_provider.value, deployment_model_info(litellm_params) + ) + folded: Final = ( + instruction_items_folded(self._validate_input_param(input), response_api_optional_request_params, "user") + if declared is False + else None + ) + placed_input, placed_params = folded or (input, response_api_optional_request_params) + return super().transform_responses_api_request( + model=model, + input=placed_input, + response_api_optional_request_params=dict(placed_params), + litellm_params=litellm_params, + headers=headers, + ) + def supports_native_websocket(self) -> bool: """Hosted vLLM does not support native WebSocket for Responses API""" return False diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index b7396e8fb03..22af51aaabe 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -36,7 +36,11 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( ) from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import ( + BaseLLMModelInfo, + hoist_developer_messages_into_leading_system_message, + map_developer_role_to_system_role, +) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -58,7 +62,7 @@ from litellm.types.utils import ( ModelResponseStream, chat_completion_tool_call_from_dict, ) -from litellm.utils import convert_to_model_response_object +from litellm.utils import convert_to_model_response_object, declared_mid_conversation_system_support from ..common_utils import OpenAIError from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config @@ -75,6 +79,45 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def targets_openai_hosted_endpoint(custom_llm_provider: str | None, api_base: str | None) -> bool: + if custom_llm_provider != "openai": + return False + resolved_api_base: Final = ( + api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") + ) + if not resolved_api_base: + return True + hostname: Final = urlparse(resolved_api_base).hostname + if hostname is None: + return True + return hostname == "openai.com" or hostname.endswith(".openai.com") + + +def _roles(messages: Sequence[AllMessageValues]) -> tuple[str, ...]: + return tuple(message["role"] for message in messages) + + +def translate_developer_role_for_openai_compatible_endpoint( + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, +) -> Sequence[AllMessageValues]: + in_place: Final = map_developer_role_to_system_role(messages=messages) + if targets_openai_hosted_endpoint(custom_llm_provider, api_base): + return in_place + placed: Final = hoist_developer_messages_into_leading_system_message(messages=messages) + if _roles(placed) == _roles(in_place): + return in_place + if declared_mid_conversation_system_support(model, custom_llm_provider, model_info) is not False: + return in_place + return placed + + _OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) _CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _DECODED_JSON: Final = TypeAdapter(object) @@ -437,20 +480,30 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) return messages, tools + def translate_developer_role_to_system_role( + self, + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, + ) -> Sequence[AllMessageValues]: + """ + A developer message only changes role, except on a model whose deployment or cost + map declares `supports_mid_conversation_system` false, where the instructions are + placed in one leading system message for its chat template. + """ + return translate_developer_role_for_openai_compatible_endpoint( + messages, model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, model_info=model_info + ) + def _targets_openai_hosted_endpoint( self, custom_llm_provider: str | None, api_base: str | None, ) -> bool: - if custom_llm_provider != "openai": - return False - resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") - if not resolved_api_base: - return True - hostname: Final = urlparse(resolved_api_base).hostname - if hostname is None: - return True - return hostname == "openai.com" or hostname.endswith(".openai.com") + return targets_openai_hosted_endpoint(custom_llm_provider, api_base) def _should_preserve_cache_control_for_endpoint( self, diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index c7552d47985..64285aafdb0 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -11,7 +11,7 @@ Translations handled by LiteLLM: - Logprobs => drop param (if user opts in to dropping param) """ -from collections.abc import Coroutine +from collections.abc import Coroutine, Sequence from typing import Final, Literal, cast, overload import litellm @@ -37,7 +37,15 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): def get_config(cls): return super().get_config() - def translate_developer_role_to_system_role(self, messages: list[AllMessageValues]) -> list[AllMessageValues]: + def translate_developer_role_to_system_role( + self, + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, + ) -> Sequence[AllMessageValues]: """ O-series models support `developer` role. """ diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 5a3ab8aebea..af2e85c5f20 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,6 +1,6 @@ import time import types -from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Final, Literal, Optional, cast import httpx @@ -50,7 +50,11 @@ from litellm.utils import ( from ...types.llms.openai import * from ..base import BaseLLM from .chat.gpt_reasoning_transformation import OpenAIGPT5Config -from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig +from .chat.gpt_transformation import ( + OpenAIGPTConfig, + OpenAIUnknownModelConfig, + translate_developer_role_for_openai_compatible_endpoint, +) from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, @@ -206,6 +210,19 @@ class OpenAIConfig(BaseConfig): else: return self._gpt_config_for_model(model).get_supported_openai_params(model=model) + def translate_developer_role_to_system_role( + self, + messages: Sequence[AllMessageValues], + *, + model: str, + custom_llm_provider: str | None, + api_base: str | None, + model_info: object, + ) -> Sequence[AllMessageValues]: + return translate_developer_role_for_openai_compatible_endpoint( + messages, model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, model_info=model_info + ) + def _gpt_config_for_model(self, model: str) -> OpenAIGPTConfig: if type(self) is OpenAIConfig and not OpenAIGPTConfig.is_openai_catalog_model(model): return OpenAIUnknownModelConfig() diff --git a/litellm/llms/openai/responses/instruction_items.py b/litellm/llms/openai/responses/instruction_items.py new file mode 100644 index 00000000000..82bf4329b34 --- /dev/null +++ b/litellm/llms/openai/responses/instruction_items.py @@ -0,0 +1,98 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam + +from litellm.llms.base_llm.base_utils import hoisted_instruction_indices +from litellm.types.llms.openai import ResponseInputParam + +UnfoldedDeveloperRole: TypeAlias = Literal["system", "user"] + +_INSTRUCTION_ROLES: Final = frozenset({"system", "developer"}) + + +def _role(item: ResponseInputItemParam) -> str | None: + match item: + case {"role": str(role)}: + return role + case _: + return None + + +def _developer_item_as(item: ResponseInputItemParam, role: UnfoldedDeveloperRole) -> ResponseInputItemParam: + if "role" not in item or item["role"] != "developer" or "content" not in item: + return item + return EasyInputMessageParam(role=role, content=item["content"], type="message") + + +def developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam: + if isinstance(input, str): + return input + return [_developer_item_as(item, "system") for item in input] + + +def _text_part(part: ResponseInputContentParam) -> str | None: + match part: + case {"type": "input_text", "text": str(text)}: + return text + case _: + return None + + +def _text_only_content(item: ResponseInputItemParam) -> str | None: + match item: + case {"role": "system" | "developer", "content": str(text)}: + return text + case {"role": "system" | "developer", "content": [*parts]}: + texts: Final = tuple(map(_text_part, parts)) + return None if any(text is None for text in texts) else "\n\n".join(text for text in texts if text) + case _: + return None + + +def _folded_instructions_and_items( + items: tuple[ResponseInputItemParam, ...], instructions: str | None, unfolded_developer_role: UnfoldedDeveloperRole +) -> tuple[str | None, ResponseInputParam]: + folded: Final = MappingProxyType( + { + index: text + for index in hoisted_instruction_indices(tuple(map(_role, items))) + if (text := _text_only_content(items[index])) is not None + } + ) + joined: Final = "\n\n".join(chunk for chunk in (instructions, *folded.values()) if chunk) + return ( + instructions if not folded else joined or None, + [_developer_item_as(item, unfolded_developer_role) for index, item in enumerate(items) if index not in folded], + ) + + +def instruction_items_folded( + input: str | ResponseInputParam, + params: Mapping[str, object], + unfolded_developer_role: UnfoldedDeveloperRole, +) -> tuple[ResponseInputParam, Mapping[str, object]] | None: + """ + For a model whose chat template takes a single leading system message, the + leading system and developer items and every later text developer item folded into + top-level `instructions`, the developer items left over sent with + `unfolded_developer_role`, or None when the request needs no folding. + """ + instructions: Final = params.get("instructions") + if isinstance(input, str) or not isinstance(instructions, str | None): + return None + items: Final = tuple(input) + if not any(_role(item) in _INSTRUCTION_ROLES for item in items): + return None + folded_instructions, folded_items = _folded_instructions_and_items(items, instructions, unfolded_developer_role) + instruction_entries: Final = () if folded_instructions is None else (("instructions", folded_instructions),) + return folded_items, MappingProxyType( + { + key: value + for key, value in ( + *((key, value) for key, value in params.items() if key != "instructions"), + *instruction_entries, + ) + } + ) diff --git a/litellm/main.py b/litellm/main.py index 996e93589fe..a3a63d5a83a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5703,15 +5703,16 @@ def completion( ), ) + deployment_model_info: Final = str_keyed_mapping( + cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg + ) ## RESPONSES API BRIDGE LOGIC ## - check early and normalize model name responses_api_model_info, model = responses_api_bridge_check( model=model, custom_llm_provider=custom_llm_provider, web_search_options=web_search_options, api_base=api_base, - deployment_model_info=str_keyed_mapping( - cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg - ), + deployment_model_info=deployment_model_info, ) routing_kwargs: Final = cast("Mapping[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict drop_unsupported_reasoning: Final = ( @@ -5808,7 +5809,15 @@ def completion( ) if provider_config is not None: - translated_messages: Final = provider_config.translate_developer_role_to_system_role(messages=messages) + translated_messages: Final = list( + provider_config.translate_developer_role_to_system_role( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + model_info=deployment_model_info, + ) + ) history_for_route: Final = _history_without_unanswered_tool_search_calls( translated_messages, custom_llm_provider, provider_config ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 50b3560f05b..aa00ea73f8c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -63655,6 +63655,7 @@ "output_cost_per_token_priority": 9e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -63801,6 +63802,7 @@ "output_cost_per_token_priority": 9e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 01aa69ba5d2..dd964a74735 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -7816,7 +7816,7 @@ "paths": { "/claude-code/marketplace.json": { "get": { - "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` or an access group on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", + "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout a key the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` or an access group on the key or its team.\n\nWith `general_settings.claude_code_marketplace_auth_required: true` the key must come\nin the `Authorization` or `x-litellm-api-key` header (Claude Code sends them from the\nmarketplace source's `headers`), `?key=` is ignored, and a caller with no valid key\ngets a 401 with `WWW-Authenticate: Bearer`.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", "operationId": "get_marketplace_claude_code_marketplace_json_get", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index df139a9b84d..dd873ae49dd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2833,6 +2833,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, stores request messages and responses in spend logs. Default is False.", ) + claude_code_marketplace_auth_required: bool = Field( + False, + description="If True, GET /claude-code/marketplace.json requires a LiteLLM key in the Authorization or x-litellm-api-key header (?key= is not accepted) and answers 401 with WWW-Authenticate: Bearer without one, and GET /public/skill_hub returns an empty list. Default is False, which serves the enabled plugins to anyone.", + ) disable_requester_ip_address_logging: bool | None = Field( None, description="If True, new spend log rows store requester_ip_address as null in the column, the metadata and the stored request. The stored request also drops client-IP forwarding headers such as x-forwarded-for, and IP addresses in its other header values and in stored error messages are masked. Guardrails, metrics and callbacks that read the standard logging payload still receive the IP. Default is False.", diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index b0088d93112..fc82686947d 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -7,7 +7,8 @@ Actual plugin files are hosted on GitHub/GitLab/Bitbucket or as a zip archive on any HTTPS host (S3, Artifactory, a static file server). Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated; `?key=` adds the key's granted skills) +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (anonymous unless + general_settings.claude_code_marketplace_auth_required; a key adds its granted skills) /claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) /claude-code/plugins - GET - List plugins visible to the key (enabled, plus granted disabled ones) /claude-code/plugins/{name} - GET - Get plugin details (403 on a disabled plugin the key is not granted) @@ -31,14 +32,17 @@ from pydantic import TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( SkillVisibility, resolve_skill_visibility, ) +from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.claude_code_marketplace_auth import marketplace_auth_required from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository +from litellm.types.proxy.auth.special_headers import SpecialHeaders from litellm.types.proxy.claude_code_endpoints import ( InstallationPreference, ListPluginsResponse, @@ -109,6 +113,41 @@ async def get_prisma_client() -> object: _get_prisma_client: Final = get_prisma_client +async def _header_key_caller(request: Request) -> UserAPIKeyAuth: + try: + caller: Final = await user_api_key_auth( + request=request, + api_key=request.headers.get(SpecialHeaders.openai_authorization.value) or "", + custom_litellm_key_header=request.headers.get(SpecialHeaders.custom_litellm_api_key.value), + ) + except ProxyException as e: + if e.code != "401": + raise + raise ProxyException( + message=e.message, + type=e.type, + param=e.param, + code=401, + headers={**e.headers, "WWW-Authenticate": "Bearer"}, + ) from e + if caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID: + raise ProxyException( + message="The LiteLLM key could not be verified because the database is unavailable", + type=ProxyErrorTypes.no_db_connection, + param=None, + code=503, + ) + return caller + + +async def _marketplace_caller(request: Request, key: str | None) -> UserAPIKeyAuth | None: + if marketplace_auth_required(): + return await _header_key_caller(request) + if key: + return await user_api_key_auth(request=request, api_key=f"Bearer {key}") + return None + + @router.get( "/claude-code/marketplace.json", tags=["Claude Code Marketplace"], @@ -121,10 +160,15 @@ async def get_marketplace(request: Request, key: str | None = None): - claude plugin marketplace add - claude plugin install @ - Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` + Without a key the catalog holds the enabled (public) plugins. With `?key=sk-...` the key is authenticated and the catalog also holds the disabled plugins granted to it through `object_permission.skills` or an access group on the key or its team. + With `general_settings.claude_code_marketplace_auth_required: true` the key must come + in the `Authorization` or `x-litellm-api-key` header (Claude Code sends them from the + marketplace source's `headers`), `?key=` is ignored, and a caller with no valid key + gets a 401 with `WWW-Authenticate: Bearer`. + Returns: Marketplace catalog with list of available plugins and their git sources. @@ -137,10 +181,7 @@ async def get_marketplace(request: Request, key: str | None = None): """ try: prisma_client: Final = await get_prisma_client() - - caller: Final[UserAPIKeyAuth | None] = ( - await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None - ) + caller: Final = await _marketplace_caller(request, key) visibility: Final[SkillVisibility] = await resolve_skill_visibility(caller) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( where=visibility.where() @@ -201,7 +242,7 @@ async def get_marketplace(request: Request, key: str | None = None): verbose_proxy_logger.exception("Error generating marketplace: %s", e) raise HTTPException( status_code=500, - detail={"error": f"Failed to generate marketplace: {e}"}, + detail={"error": "Failed to generate marketplace"}, ) diff --git a/litellm/proxy/common_utils/claude_code_marketplace_auth.py b/litellm/proxy/common_utils/claude_code_marketplace_auth.py new file mode 100644 index 00000000000..0cfb3061744 --- /dev/null +++ b/litellm/proxy/common_utils/claude_code_marketplace_auth.py @@ -0,0 +1,17 @@ +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +_AUTH_REQUIRED_SETTING: Final[TypeAdapter[bool | None]] = TypeAdapter(bool | None) + + +def marketplace_auth_required() -> bool: + from litellm.proxy.proxy_server import general_settings_view + + try: + return ( + _AUTH_REQUIRED_SETTING.validate_python(general_settings_view().get("claude_code_marketplace_auth_required")) + is True + ) + except ValidationError: + return True diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 792ef8a3f8c..36f424507d7 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -4,6 +4,7 @@ from typing import Final from fastapi import APIRouter +from litellm.proxy.common_utils.claude_code_marketplace_auth import marketplace_auth_required from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint from litellm.proxy.mcp_server.stdio_gate import is_mcp_stdio_enabled from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( @@ -51,4 +52,5 @@ async def get_ui_config(): workers=proxy_config.worker_registry if is_control_plane else [], mcp_stdio_enabled=is_mcp_stdio_enabled(), use_si_terminology=use_si_terminology, + claude_code_marketplace_auth_required=marketplace_auth_required(), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d034a72934..091e05610ba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6978,6 +6978,7 @@ class ProxyConfig: for name in ( "vector_store_deny_by_default", "search_tool_deny_by_default", + "claude_code_marketplace_auth_required", "vertex_ai_pass_through_drop_anthropic_betas", "request_tag_mappings", ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 9f33e0d2adc..b8cdea0dea4 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -367,19 +367,23 @@ async def get_mcp_servers(): tags=["public", "Claude Code Marketplace"], ) async def public_skill_hub(): - """Return enabled (public) Claude Code skills — no auth required.""" + """Return enabled (public) Claude Code skills, or none when the marketplace requires a key.""" from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( MANIFEST_FIELDS, get_installation_preference, - get_prisma_client, ) + from litellm.proxy.common_utils.claude_code_marketplace_auth import marketplace_auth_required + from litellm.proxy.proxy_server import prisma_client from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, PluginListItem, ) + if marketplace_auth_required(): + return ListPluginsResponse(plugins=[], count=0) + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - prisma_client: Final = await get_prisma_client() plugins: Final = await _plugin_table(prisma_client).find_many(where={"enabled": True}) items: Final = [] for plugin in plugins: @@ -406,7 +410,8 @@ async def public_skill_hub(): ) return ListPluginsResponse(plugins=items, count=len(items)) except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + verbose_logger.exception("Error listing public skills: %s", e) + raise HTTPException(status_code=500, detail="Failed to list skills") @router.get( diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index d396c79b844..ca6ada0d1ec 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -15,4 +15,5 @@ class UiDiscoveryEndpoints(LiteLLMBaseModel): is_control_plane: bool = False mcp_stdio_enabled: bool = False use_si_terminology: bool = False + claude_code_marketplace_auth_required: bool = False workers: list[WorkerRegistryEntry] = Field(default=[]) diff --git a/litellm/utils.py b/litellm/utils.py index 26ce38d63b3..adf195668ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -29,8 +29,8 @@ import threading import time import traceback from contextlib import nullcontext -from dataclasses import dataclass, field -from functools import lru_cache, wraps +from dataclasses import dataclass, field, replace +from functools import lru_cache, partial, wraps from importlib import resources from inspect import iscoroutine from io import StringIO @@ -220,6 +220,7 @@ def _get_cached_audio_utils(): from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper as _TextCompletionStreamWrapper +from litellm.types.llms.base_model import LiteLLMBaseModel from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -2683,9 +2684,99 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: return None +@dataclass(frozen=True, slots=True) +class _HuggingFaceTokenizerLoaded: + tokenizer: HuggingFace + + @dataclass(frozen=True, slots=True) class _HuggingFaceTokenizerLoadFailure: reason: str + retry_at: float + retry_pid: int | None = None + + +_HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] +_HuggingFaceTokenizerLoadOutcome: TypeAlias = _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure + + +def _start_in_thread_joined_at_exit(work: Callable[[], None]) -> None: + threading.Thread(target=work, name="litellm-huggingface-tokenizer-retry", daemon=False).start() + + +@dataclass(slots=True) +class _HuggingFaceTokenizerLoadOutcomes: + clock: Callable[[], float] = time.monotonic + start_retry: Callable[[Callable[[], None]], None] = _start_in_thread_joined_at_exit + pid: Callable[[], int] = os.getpid + lock: threading.Lock = field(default_factory=threading.Lock) + by_tokenizer: Mapping[_HuggingFaceTokenizerKey, _HuggingFaceTokenizerLoadOutcome] = field( + default_factory=lambda: MappingProxyType({}) + ) + + def known( + self, key: _HuggingFaceTokenizerKey, retry_when_due: Callable[[], None] + ) -> _HuggingFaceTokenizerLoadOutcome | None: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if ( + not isinstance(known, _HuggingFaceTokenizerLoadFailure) + or known.retry_pid == self.pid() + or self.clock() < known.retry_at + ): + return known + retrying: Final = replace(known, retry_pid=self.pid()) + self._remember(key, retrying) + try: + self.start_retry(retry_when_due) + except RuntimeError as error: + return self.remember_failed_retry(key, str(error)) + return retrying + + def remember_loaded(self, key: _HuggingFaceTokenizerKey, tokenizer: HuggingFace) -> _HuggingFaceTokenizerLoaded: + loaded: Final = _HuggingFaceTokenizerLoaded(tokenizer) + with self.lock: + self._remember(key, loaded) + return loaded + + def remember_first_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if known is not None: + return known + failure: Final = self._failure(reason) + self._remember(key, failure) + return failure + + def remember_failed_retry(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if isinstance(known, _HuggingFaceTokenizerLoaded): + return known + failure: Final = self._failure(reason) + self._remember(key, failure) + return failure + + def end_retry(self, key: _HuggingFaceTokenizerKey) -> None: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if isinstance(known, _HuggingFaceTokenizerLoadFailure) and known.retry_pid is not None: + self._remember(key, self._failure(known.reason)) + + def flush_cache(self) -> None: + with self.lock: + self.by_tokenizer = MappingProxyType({}) + + def _failure(self, reason: str) -> _HuggingFaceTokenizerLoadFailure: + return _HuggingFaceTokenizerLoadFailure( + reason=reason, retry_at=self.clock() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + ) + + def _remember(self, key: _HuggingFaceTokenizerKey, outcome: _HuggingFaceTokenizerLoadOutcome) -> None: + self.by_tokenizer = MappingProxyType({**self.by_tokenizer, key: outcome}) + + +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES: Final = _HuggingFaceTokenizerLoadOutcomes() def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: @@ -2697,30 +2788,28 @@ def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: ) -def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: +def _return_huggingface_tokenizer( + model: str, outcomes: _HuggingFaceTokenizerLoadOutcomes = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES +) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: kind: Final = huggingface_tokenizer_kind(model) if kind is None or (kind not in _BUNDLED_TOKENIZER_LOAD_LOCKS and is_offline_mode()): return None - backend: Final = _huggingface_tokenizer_backend() - key: Final = _huggingface_tokenizer_failure_key(kind, backend) + key: Final = (kind, _huggingface_tokenizer_backend()) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - cached_failure: Final = _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.get_cache(key) - if isinstance(cached_failure, _HuggingFaceTokenizerLoadFailure): - return cached_failure - tokenizer: Final = _load_huggingface_tokenizer_recording_failure(kind, backend, key) - if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): - _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) - return tokenizer - return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + known: Final = outcomes.known(key, partial(_retry_huggingface_tokenizer_load, model, key, outcomes)) + outcome: Final = ( + known + if known is not None + else _load_huggingface_tokenizer_remembering_outcome(model, key, outcomes, outcomes.remember_first_failure) + ) + match outcome: + case _HuggingFaceTokenizerLoaded(tokenizer=tokenizer): + return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + case _HuggingFaceTokenizerLoadFailure(): + return outcome -def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: - return f"{kind}:{backend.name}" - - -@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) -def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decision) -> HuggingFace: - """One tokenizer per kind and backend; `backend` is the cache key, the dispatch re-derives it.""" +def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> HuggingFace: match kind: case "cohere": return tokenizer_dispatch.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") @@ -2740,19 +2829,34 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio assert_never(kind) -def _load_huggingface_tokenizer_recording_failure( - kind: HuggingFaceTokenizerKind, backend: Decision, key: str -) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: +def _load_huggingface_tokenizer_remembering_outcome( + model: str, + key: _HuggingFaceTokenizerKey, + outcomes: _HuggingFaceTokenizerLoadOutcomes, + remember_failure: Callable[[_HuggingFaceTokenizerKey, str], _HuggingFaceTokenizerLoadOutcome], +) -> _HuggingFaceTokenizerLoadOutcome: + kind, _ = key try: - return _load_huggingface_tokenizer(kind, backend) + tokenizer: Final = _load_huggingface_tokenizer(kind) except Exception as error: from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - failure: Final = _HuggingFaceTokenizerLoadFailure(reason=str(error)) - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.set_cache(key, failure, ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) - return failure + outcome: Final = remember_failure(key, str(error)) + if isinstance(outcome, _HuggingFaceTokenizerLoadFailure): + _warn_huggingface_tokenizer_fallback(model, outcome.reason) + return outcome + return outcomes.remember_loaded(key, tokenizer) + + +def _retry_huggingface_tokenizer_load( + model: str, key: _HuggingFaceTokenizerKey, outcomes: _HuggingFaceTokenizerLoadOutcomes +) -> None: + try: + _load_huggingface_tokenizer_remembering_outcome(model, key, outcomes, outcomes.remember_failed_retry) + finally: + outcomes.end_retry(key) def encode(model="", text="", custom_tokenizer: dict | None = None): @@ -3398,6 +3502,42 @@ def supports_mid_conversation_system(model: str, custom_llm_provider: str | None ) +class _DeploymentModelInfo(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + model_info: Mapping[str, object] | None = None + + +def deployment_model_info(litellm_params: object) -> Mapping[str, object] | None: + """The deployment's own ``model_info`` read off its litellm params.""" + return _DeploymentModelInfo.model_validate(litellm_params).model_info + + +def deployment_declared_mid_conversation_system_support(model_info: object) -> bool | None: + match model_info: + case {"supports_mid_conversation_system": bool(declared)}: + return declared + case _: + return None + + +def declared_mid_conversation_system_support( + model: str, custom_llm_provider: str | None, model_info: object +) -> bool | None: + """ + The `supports_mid_conversation_system` the deployment's own model_info declares, else the one the + cost map declares for the model, else None when neither sets it. + """ + deployment_declared: Final = deployment_declared_mid_conversation_system_support(model_info) + if deployment_declared is not None: + return deployment_declared + if is_explicitly_disabled_factory(model, custom_llm_provider, "supports_mid_conversation_system"): + return False + if supports_mid_conversation_system(model, custom_llm_provider): + return True + return None + + def supports_native_structured_output(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports native structured outputs and return a boolean value. @@ -8346,26 +8486,6 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache -class _HuggingFaceTokenizerLoadFailureCache(Protocol): - def get_cache(self, key: str) -> object | None: ... - - def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ... - - def flush_cache(self) -> None: ... - - -_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = ( - cast( # cast-ok: cache methods are untyped - _HuggingFaceTokenizerLoadFailureCache, - InMemoryCache( - max_size_in_memory=16, - default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, - clock=lambda: time.monotonic(), - ), - ) -) - - class AvailableModelsCache(InMemoryCache): def __init__(self, ttl_seconds: int = 300, max_size: int = 1000): super().__init__(ttl_seconds, max_size) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 50b3560f05b..aa00ea73f8c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -63655,6 +63655,7 @@ "output_cost_per_token_priority": 9e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -63801,6 +63802,7 @@ "output_cost_per_token_priority": 9e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3070ebb0235..ed7dd5ec146 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -378,6 +378,20 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def declare_mid_conversation_system(monkeypatch: pytest.MonkeyPatch) -> Callable[[str, bool], None]: + def declare(model_cost_key: str, supported: bool) -> None: + entry: Final = { + "litellm_provider": model_cost_key.split("/", 1)[0], + "mode": "chat", + "supports_mid_conversation_system": supported, + } + monkeypatch.setitem(litellm.model_cost, model_cost_key, entry) + litellm.get_model_info.cache_clear() + + return declare + + NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock"] | None diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 000e3993bc2..d97970753f8 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -14,6 +14,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain from pathlib import Path +from queue import SimpleQueue from typing import Final, cast from unittest.mock import MagicMock, sentinel @@ -45,6 +46,7 @@ from litellm.litellm_core_utils.token_counter import ( from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.utils import SelectTokenizerResponse from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -978,16 +980,17 @@ def test_token_counter(): import unittest from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, + _HuggingFaceTokenizerLoadFailure, + _HuggingFaceTokenizerLoadOutcomes, + _return_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding, ) # Clear the cache at module load to ensure clean state -_load_huggingface_tokenizer.cache_clear() -_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() class TestTokenizerSelection(unittest.TestCase): @@ -998,8 +1001,7 @@ class TestTokenizerSelection(unittest.TestCase): @lru_cache, which can cause cache hits from previous tests when running with --dist=loadscope (tests from same file run on same worker). """ - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): @@ -1087,64 +1089,23 @@ class TestTokenizerSelection(unittest.TestCase): def test_failed_huggingface_tokenizer_download_is_not_retried_on_every_count( self, mock_from_pretrained: MagicMock ): - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: - with patch("litellm.utils.time.monotonic", return_value=0.0): - mock_from_pretrained.side_effect = RuntimeError("HF unavailable") - model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" - text: Final = "Count the tokens for this message." - expected: Final = len(encoding.encode(text, disallowed_special=())) - - first_count: Final = litellm.token_counter(model=model, text=text) - first_attempts: Final = mock_from_pretrained.call_count - second_count: Final = litellm.token_counter(model=model, text=text) - - assert first_attempts > 0 - assert mock_from_pretrained.call_count == first_attempts - assert first_count == expected - assert second_count == expected - finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() - - @patch("litellm.utils.time.monotonic", return_value=299.0) - @patch("litellm.utils.tokenizer_dispatch.from_pretrained") - def test_failed_huggingface_tokenizer_download_is_retried_after_retry_window( - self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock - ): - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() - try: - def from_pretrained_side_effect(_: str) -> object: - if mock_from_pretrained.call_count == 1: - mock_monotonic.return_value = 322.0 - raise RuntimeError("HF unavailable") - return sentinel.tokenizer - - mock_from_pretrained.side_effect = from_pretrained_side_effect + mock_from_pretrained.side_effect = RuntimeError("HF unavailable") model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + text: Final = "Count the tokens for this message." + expected: Final = len(encoding.encode(text, disallowed_special=())) - first_result: Final = _select_tokenizer_helper(model) - assert first_result["type"] == "openai_tokenizer" - mock_monotonic.return_value = 323.0 - second_result: Final = _select_tokenizer_helper(model) - assert second_result["type"] == "openai_tokenizer", "tokenizer was retried at t=323" - assert mock_from_pretrained.call_count == 1 + first_count: Final = litellm.token_counter(model=model, text=text) + first_attempts: Final = mock_from_pretrained.call_count + second_count: Final = litellm.token_counter(model=model, text=text) - mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 - third_result: Final = _select_tokenizer_helper(model) - assert third_result["type"] == "openai_tokenizer" - assert mock_from_pretrained.call_count == 1 - - mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1 - retried_result: Final = _select_tokenizer_helper(model) - assert retried_result["type"] == "huggingface_tokenizer" - assert retried_result["tokenizer"] is sentinel.tokenizer - assert mock_from_pretrained.call_count == 2 + assert first_attempts > 0 + assert mock_from_pretrained.call_count == first_attempts + assert first_count == expected + assert second_count == expected finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() _CODE_SAMPLE: Final = " for item in items:\n total += item.price * item.quantity\n" * 20 @@ -1212,8 +1173,7 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: @pytest.mark.parametrize("model", ["llama-2-7b", "meta-llama/llama-3-8b-instruct"]) def test_offline_mode_counts_downloaded_tokenizer_models_with_tiktoken(monkeypatch: pytest.MonkeyPatch, model: str): - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") with patch("litellm.utils.tokenizer_dispatch.from_pretrained") as from_pretrained: @@ -1232,8 +1192,7 @@ def test_offline_mode_keeps_deepseek_v4_on_its_bundled_tokenizer(monkeypatch: py def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() start: Final = threading.Barrier(8) def first_tokenizer(_: int) -> object: @@ -1247,8 +1206,7 @@ def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplog: pytest.LogCaptureFixture): - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" try: with ( @@ -1260,15 +1218,346 @@ def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplo assert types == ("openai_tokenizer",) * 3 assert len(fallback_warnings) == 1 finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_download_that_fails_after_another_count_loaded_the_tokenizer_counts_with_it(): + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + slow_download_started: Final = threading.Event() + loading_count_finished: Final = threading.Event() + + def download(_: str) -> object: + if threading.current_thread().name.startswith("slow-download"): + slow_download_started.set() + loading_count_finished.wait(timeout=5) + raise RuntimeError("HF rate limited") + return sentinel.tokenizer + + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download): + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="slow-download") as pool: + slow_count: Final = pool.submit(_select_tokenizer_helper, model) + assert slow_download_started.wait(timeout=5) + loading_count: Final = _select_tokenizer_helper(model) + loading_count_finished.set() + slow_count_tokenizer: Final = slow_count.result(timeout=5)["tokenizer"] + later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) + assert slow_count_tokenizer is sentinel.tokenizer + assert loading_count["tokenizer"] is sentinel.tokenizer + assert tuple(count["tokenizer"] for count in later_counts) == (sentinel.tokenizer,) * 3 + finally: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_loaded_tokenizer_serves_later_counts_without_downloading_it_again(): + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", return_value=sentinel.tokenizer): + loaded: Final = _select_tokenizer_helper(model) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF rate limited") + ) as from_pretrained: + later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) + assert loaded["tokenizer"] is sentinel.tokenizer + assert tuple(count["tokenizer"] for count in later_counts) == (sentinel.tokenizer,) * 3 + assert from_pretrained.call_count == 0 + finally: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def _run_now(retry: Callable[[], None]) -> None: + retry() + + +def _start_retry_resolving(ended: Future[None]) -> Callable[[Callable[[], None]], None]: + def start(retry: Callable[[], None]) -> None: + def run_then_resolve() -> None: + retry() + ended.set_result(None) + + threading.Thread(target=run_then_resolve, daemon=True).start() + + return start + + +def _eventually(probe: Callable[[], bool]) -> bool: + deadline: Final = time.monotonic() + 5 + while time.monotonic() < deadline: + if probe(): + return True + time.sleep(0.01) + return False + + +def test_failed_download_is_retried_once_per_window_until_it_loads(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=299.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_run_now) + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + clock.return_value = 322.0 + raise RuntimeError("HF unavailable") + if from_pretrained.call_count == 2: + raise RuntimeError("HF rate limited") + return sentinel.tokenizer + + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained: + first: Final = _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 + inside_the_window: Final = _return_huggingface_tokenizer(model, outcomes) + downloads_inside_the_window: Final = from_pretrained.call_count + clock.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + claimant: Final = _return_huggingface_tokenizer(model, outcomes) + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + downloads_inside_the_next_window: Final = from_pretrained.call_count + clock.return_value = 322.0 + 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + after_the_second_retry: Final = _return_huggingface_tokenizer(model, outcomes) + assert isinstance(first, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_window, _HuggingFaceTokenizerLoadFailure) + assert downloads_inside_the_window == 1 + assert isinstance(claimant, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert downloads_inside_the_next_window == 2 + assert from_pretrained.call_count == 3 + assert after_the_second_retry == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_retry_after_the_window_runs_in_the_background_while_every_count_answers_at_once(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock) + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + return sentinel.tokenizer + + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 3 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(4)) + assert retry_started.wait(timeout=5) + clock.return_value = 9 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + later_counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(2)) + downloads_while_retry_runs: Final = from_pretrained.call_count + retry_released.set() + retried: Final = _eventually( + lambda: _return_huggingface_tokenizer(model, outcomes) + == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + ) + assert all( + isinstance(count, _HuggingFaceTokenizerLoadFailure) + for count in counts_while_retry_runs + later_counts_while_retry_runs + ) + assert downloads_while_retry_runs == 2 + assert retried + assert from_pretrained.call_count == 2 + + +def test_download_from_before_the_first_failure_ending_during_a_retry_leaves_it_the_only_retry(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock) + early_download_started: Final = threading.Event() + early_download_released: Final = threading.Event() + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + early_download_started.set() + early_download_released.wait(timeout=5) + raise RuntimeError("HF unavailable") + if from_pretrained.call_count == 2: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + return sentinel.tokenizer + + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, + ThreadPoolExecutor(max_workers=1) as early_pool, + ): + early_count: Final = early_pool.submit(_return_huggingface_tokenizer, model, outcomes) + assert early_download_started.wait(timeout=5) + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 + _return_huggingface_tokenizer(model, outcomes) + assert retry_started.wait(timeout=5) + early_download_released.set() + early: Final = early_count.result(timeout=5) + clock.return_value = 5 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(3)) + downloads_while_retry_runs: Final = from_pretrained.call_count + retry_released.set() + retried: Final = _eventually( + lambda: _return_huggingface_tokenizer(model, outcomes) + == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + ) + assert isinstance(early, _HuggingFaceTokenizerLoadFailure) + assert all(isinstance(count, _HuggingFaceTokenizerLoadFailure) for count in counts_while_retry_runs) + assert downloads_while_retry_runs == 3 + assert retried + + +def test_retry_failing_after_an_earlier_download_loaded_the_tokenizer_keeps_it(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + retry_ended: Final = Future[None]() + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_start_retry_resolving(retry_ended)) + early_download_started: Final = threading.Event() + early_download_released: Final = threading.Event() + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + early_download_started.set() + early_download_released.wait(timeout=5) + return sentinel.tokenizer + if from_pretrained.call_count == 2: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + raise RuntimeError("HF rate limited") + + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, + ThreadPoolExecutor(max_workers=1) as early_pool, + ): + early_count: Final = early_pool.submit(_return_huggingface_tokenizer, model, outcomes) + assert early_download_started.wait(timeout=5) + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + assert retry_started.wait(timeout=5) + early_download_released.set() + early: Final = early_count.result(timeout=5) + retry_released.set() + retry_ended.result(timeout=5) + later_counts: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(3)) + loaded: Final = {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + assert early == loaded + assert from_pretrained.call_count == 3 + assert later_counts == (loaded,) * 3 + + +class _NativePanic(BaseException): + pass + + +def _fork_guard_error() -> BaseException: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + return ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + +def _native_panic() -> BaseException: + return _NativePanic("tokenizer parser panicked") + + +@pytest.mark.parametrize("escaping_error", (_fork_guard_error, _native_panic), ids=("fork_guard", "native_panic")) +def test_retry_that_raises_past_failure_handling_still_retries_after_the_next_window( + escaping_error: Callable[[], BaseException], +): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + error: Final = escaping_error() + retries: Final = SimpleQueue[Future[None]]() + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="tokenizer-retry") as pool: + + def start_retry(retry: Callable[[], None]) -> None: + retries.put(pool.submit(retry)) + + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_retry) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + from_pretrained.side_effect = error + _return_huggingface_tokenizer(model, outcomes) + raised: Final = retries.get(timeout=5).exception(timeout=5) + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + no_retry_inside_the_next_window: Final = retries.empty() + clock.return_value = 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + retries.get(timeout=5).result(timeout=5) + after_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + assert raised is error + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert no_retry_inside_the_next_window + assert from_pretrained.call_count == 3 + assert after_the_next_window == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_process_forked_during_a_retry_starts_its_own_retry_instead_of_waiting_on_the_parents(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + pid: Final = MagicMock(return_value=100) + start_retry: Final = MagicMock() + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_retry, pid=pid) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + _return_huggingface_tokenizer(model, outcomes) + retries_started_by_the_parent: Final = start_retry.call_count + pid.return_value = 101 + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + _return_huggingface_tokenizer(model, outcomes) + start_retry.call_args.args[0]() + in_the_forked_process: Final = _return_huggingface_tokenizer(model, outcomes) + assert retries_started_by_the_parent == 1 + assert start_retry.call_count == 2 + assert in_the_forked_process == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_retry_thread_that_cannot_start_answers_the_count_and_retries_after_the_next_window(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + thread_starts: Final = MagicMock(side_effect=(RuntimeError("can't start new thread"), None)) + + def start_thread_then_retry(retry: Callable[[], None]) -> None: + thread_starts() + retry() + + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_thread_then_retry) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + when_the_thread_cannot_start: Final = _return_huggingface_tokenizer(model, outcomes) + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + after_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + assert isinstance(when_the_thread_cannot_start, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert from_pretrained.call_count == 2 + assert after_the_next_window == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered(): from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" forked: Final = ForkedAfterNativeRuntimeStarted("native runtime started before fork") try: @@ -1280,8 +1569,7 @@ def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered assert loaded["type"] == "huggingface_tokenizer" assert loaded["tokenizer"] is sentinel.tokenizer finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() def test_token_counter_with_anthropic_tool_use(): @@ -2325,6 +2613,27 @@ tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) """ +HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +import threading +import time +from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer +retry_reached_hub = threading.Event() +def answer_once_the_process_is_exiting(request): + retry_reached_hub.set() + while threading.main_thread().is_alive(): + time.sleep(0.01) + time.sleep(0.5) + return handle(request) +now = [0.0] +outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0]) +model = "meta-llama/llama-3-70b-instruct" +_return_huggingface_tokenizer(model, outcomes) +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting))) +now[0] = 3600.0 +_return_huggingface_tokenizer(model, outcomes) +print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)})) +""" + def _word_level_tokenizer_json( pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None @@ -2387,6 +2696,10 @@ def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_t } +def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None: + assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True} + + def test_empty_custom_tokenizer_uses_model_tokenizer() -> None: text_value: Final = "A tokenizer fallback should preserve the model encoding." custom_count: Final = _get_exact_count_function("gpt-3.5-turbo", {})(text_value) diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 54185c08c10..343b67699d3 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -402,14 +402,12 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): result = _select_tokenizer_helper("llama-2") @@ -418,8 +416,7 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) @@ -477,14 +474,12 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029" secret = "sk-" + "x" * 48 failure = OSError("download failed\r\nforged-error\x1b[31m api_key=" + secret) @@ -511,8 +506,7 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): assert "REDACTED" in message assert "install tokenizers and huggingface-hub" in message finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index 53df6610913..f6dae2c304f 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -783,3 +783,26 @@ def test_responses_api_version_leaves_the_chat_route_and_body_untouched(respx_mo assert request.url.params["api-version"] == "2024-12-01-preview" assert json.loads(request.content) == {"model": "gpt-5.4-nano", "messages": messages} assert response.choices[0].message.content == "Hello" + + +def _developer_message_after_a_user_turn() -> list[dict[str, str]]: + return [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.mark.parametrize("config", [AzureOpenAIConfig(), AzureOpenAIGPT5Config()]) +def test_azure_keeps_a_later_developer_message_in_place_as_system(config: AzureOpenAIConfig) -> None: + messages: Final = config.translate_developer_role_to_system_role( + messages=_developer_message_after_a_user_turn(), + model="my-deployment", + custom_llm_provider="azure", + api_base="https://example-resource.openai.azure.com", + model_info=None, + ) + assert [message["role"] for message in messages] == ["system", "user", "assistant", "system", "user"] + assert messages[3]["content"] == "Answer with exactly one word." 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 ccac6ffc309..42dc222e420 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,5 +1,6 @@ import json import traceback +from collections.abc import Callable from typing import Final from unittest.mock import MagicMock, patch @@ -1098,3 +1099,25 @@ def test_map_azure_model_group(model_group_header, expected_model): config = AzureAICohereConfig() assert config._map_azure_model_group(model_group_header) == expected_model + + +def test_translate_developer_role_hoists_a_later_developer_message_for_a_model_that_rejects_it( + declare_mid_conversation_system: Callable[[str, bool], None], +) -> None: + declare_mid_conversation_system("azure_ai/system-first-model", False) + messages: Final = AzureAIStudioConfig().translate_developer_role_to_system_role( + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + model="system-first-model", + custom_llm_provider="azure_ai", + api_base="https://example.services.ai.azure.com/models", + model_info=None, + ) + + assert [message["role"] for message in messages] == ["system", "user", "assistant", "user"] + assert messages[0]["content"] == "You are terse.\n\nAnswer with exactly one word." diff --git a/tests/unit/llms/base_llm/test_base_utils.py b/tests/unit/llms/base_llm/test_base_utils.py new file mode 100644 index 00000000000..7052fd04242 --- /dev/null +++ b/tests/unit/llms/base_llm/test_base_utils.py @@ -0,0 +1,494 @@ +from typing import Final + +from litellm.llms.base_llm.base_utils import ( + hoist_developer_messages_into_leading_system_message, + merge_system_message_runs, +) + + +class TestMergeSystemMessageRuns: + def test_later_developer_message_stays_in_place_while_the_leading_run_folds(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are Codex"}, + {"role": "developer", "content": "Sandbox: read-only."}, + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(merge_system_message_runs(messages)) == [ + {"role": "system", "content": "You are Codex\n\nSandbox: read-only."}, + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_run_with_content_that_is_neither_text_nor_blocks_is_forwarded_unmerged(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are Codex"}, + {"role": "developer", "content": {"text": "not a content list"}}, + {"role": "user", "content": "Hi there"}, + ] + assert list(merge_system_message_runs(messages)) == [ + {"role": "system", "content": "You are Codex"}, + {"role": "system", "content": {"text": "not a content list"}}, + {"role": "user", "content": "Hi there"}, + ] + + def test_hoisted_developer_message_with_content_that_is_neither_text_nor_blocks_is_forwarded_unmerged(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are Codex"}, + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": 42}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "You are Codex"}, + {"role": "system", "content": 42}, + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +class TestHoistDeveloperMessagesIntoLeadingSystemMessage: + def test_single_developer_message(self) -> None: + messages: Final = [ + {"role": "developer", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + + def test_consecutive_developer_messages_merge(self) -> None: + messages: Final = [ + {"role": "developer", "content": "Rule 1: Be concise."}, + {"role": "developer", "content": "Rule 2: Respond in Markdown."}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Rule 1: Be concise.\n\nRule 2: Respond in Markdown."}, + {"role": "user", "content": "Hello"}, + ] + + def test_system_followed_by_developer_merges(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base instructions."}, + {"role": "developer", "content": "Developer override."}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Base instructions.\n\nDeveloper override."}, + {"role": "user", "content": "Hello"}, + ] + + def test_developer_followed_by_system_merges(self) -> None: + messages: Final = [ + {"role": "developer", "content": "Developer instruction."}, + {"role": "system", "content": "System instruction."}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Developer instruction.\n\nSystem instruction."}, + {"role": "user", "content": "Hello"}, + ] + + def test_consecutive_native_system_messages_merge(self) -> None: + messages: Final = [ + {"role": "system", "content": "System part 1."}, + {"role": "system", "content": "System part 2."}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "System part 1.\n\nSystem part 2."}, + {"role": "user", "content": "Hello"}, + ] + + def test_empty_string_side_does_not_pad(self) -> None: + assert list( + hoist_developer_messages_into_leading_system_message( + [{"role": "system", "content": ""}, {"role": "developer", "content": "Rules"}] + ) + ) == [{"role": "system", "content": "Rules"}] + assert list( + hoist_developer_messages_into_leading_system_message( + [{"role": "system", "content": "Rules"}, {"role": "developer", "content": ""}] + ) + ) == [{"role": "system", "content": "Rules"}] + + def test_empty_string_message_adds_no_block_when_the_run_merges_as_blocks(self) -> None: + messages: Final = [ + {"role": "system", "content": ""}, + {"role": "developer", "content": [{"type": "text", "text": "Rule 1"}, {"type": "text", "text": "Rule 2"}]}, + {"role": "user", "content": "Hi"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": [{"type": "text", "text": "Rule 1"}, {"type": "text", "text": "Rule 2"}]}, + {"role": "user", "content": "Hi"}, + ] + + def test_two_cached_string_messages_keep_one_breakpoint_each(self) -> None: + messages: Final = [ + {"role": "system", "content": "A", "cache_control": {"type": "ephemeral", "ttl": "5m"}}, + {"role": "system", "content": "B", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "A", "cache_control": {"type": "ephemeral", "ttl": "5m"}}, + {"type": "text", "text": "\n\nB", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + ], + }, + ] + + def test_cache_control_supplied_only_by_the_later_message_marks_only_its_block(self) -> None: + messages: Final = [ + {"role": "system", "content": "A"}, + {"role": "developer", "content": "B", "cache_control": {"type": "ephemeral"}}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "A"}, + {"type": "text", "text": "\n\nB", "cache_control": {"type": "ephemeral"}}, + ], + }, + ] + + def test_cache_control_of_the_first_message_stays_on_its_own_block(self) -> None: + messages: Final = [ + {"role": "system", "content": "Cached prompt", "cache_control": {"type": "ephemeral"}}, + {"role": "developer", "content": "Additional instructions"}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Cached prompt", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "\n\nAdditional instructions"}, + ], + }, + {"role": "user", "content": "Hello"}, + ] + + def test_cached_string_message_merged_with_block_content_keeps_its_breakpoint(self) -> None: + messages: Final = [ + {"role": "system", "content": "Cached prompt", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": [{"type": "text", "text": "Answer with exactly one word."}]}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Cached prompt", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "\n\nAnswer with exactly one word."}, + ], + }, + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_plain_messages_ahead_of_a_cached_one_keep_their_own_blocks(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base"}, + {"role": "developer", "content": "Rule"}, + {"role": "developer", "content": "Cached", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Base"}, + {"type": "text", "text": "\n\nRule"}, + {"type": "text", "text": "\n\nCached", "cache_control": {"type": "ephemeral"}}, + ], + }, + {"role": "user", "content": "Hello"}, + ] + + def test_name_of_the_last_named_message_survives_a_long_run(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base", "name": "first"}, + *({"role": "developer", "content": f"Rule {index}"} for index in range(2000)), + {"role": "developer", "content": "Last", "name": "last"}, + {"role": "user", "content": "Hello"}, + ] + merged: Final = list(hoist_developer_messages_into_leading_system_message(messages)) + assert len(merged) == 2 + assert merged[0]["name"] == "last" + assert merged[0]["content"] == "\n\n".join(["Base", *(f"Rule {index}" for index in range(2000)), "Last"]) + + def test_billing_metadata_keeps_its_own_block_when_merged(self) -> None: + billing: Final = {"role": "system", "content": "x-anthropic-billing-header: cc_version=1.0;"} + advisory: Final = {"role": "system", "content": "Guardrail advisory: request was flagged"} + assert list(hoist_developer_messages_into_leading_system_message([billing, advisory])) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1.0;"}, + {"type": "text", "text": "\n\nGuardrail advisory: request was flagged"}, + ], + } + ] + assert list(hoist_developer_messages_into_leading_system_message([advisory, billing])) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Guardrail advisory: request was flagged"}, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1.0;"}, + ], + } + ] + + def test_developer_message_after_a_user_turn_is_hoisted_into_the_leading_system_message(self) -> None: + messages: Final = [ + {"role": "system", "content": "System turn 1"}, + {"role": "user", "content": "User turn 1"}, + {"role": "developer", "content": "Developer turn 2"}, + {"role": "user", "content": "User turn 2"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "System turn 1\n\nDeveloper turn 2"}, + {"role": "user", "content": "User turn 1"}, + {"role": "user", "content": "User turn 2"}, + ] + + def test_hoisted_developer_message_keeps_only_the_fields_a_system_message_has(self) -> None: + messages: Final = [ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word.", "name": "codex", "extra_key": "x"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Answer with exactly one word.", "name": "codex"}, + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_developer_message_in_second_position_without_a_leading_system_message_is_hoisted(self) -> None: + messages: Final = [ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Answer with exactly one word."}, + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_every_later_developer_message_is_hoisted_in_order_across_assistant_and_tool_turns(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base"}, + {"role": "user", "content": "Turn 1"}, + {"role": "assistant", "content": "Reply 1"}, + {"role": "developer", "content": "Update A"}, + {"role": "user", "content": "Turn 2"}, + {"role": "developer", "content": "Update B"}, + {"role": "user", "content": "Turn 3"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Base\n\nUpdate A\n\nUpdate B"}, + {"role": "user", "content": "Turn 1"}, + {"role": "assistant", "content": "Reply 1"}, + {"role": "user", "content": "Turn 2"}, + {"role": "user", "content": "Turn 3"}, + ] + + def test_developer_message_that_closes_the_conversation_after_an_assistant_turn_becomes_a_user_turn(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello! How can I help?"}, + {"role": "developer", "content": "Reply with the single word PONG and nothing else."}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello! How can I help?"}, + {"role": "user", "content": "Reply with the single word PONG and nothing else."}, + ] + + def test_closing_developer_message_keeps_its_cache_mark_as_a_user_turn(self) -> None: + messages: Final = [ + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello! How can I help?"}, + {"role": "developer", "content": "Reply with PONG.", "cache_control": {"type": "ephemeral"}}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages))[-1] == { + "role": "user", + "content": "Reply with PONG.", + "cache_control": {"type": "ephemeral"}, + } + + def test_only_the_closing_developer_run_after_an_assistant_turn_becomes_user_turns(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base"}, + {"role": "user", "content": "Turn 1"}, + {"role": "developer", "content": "Update A"}, + {"role": "user", "content": "Turn 2"}, + {"role": "assistant", "content": "Reply 2"}, + {"role": "developer", "content": "Closing B"}, + {"role": "developer", "content": "Closing C"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Base\n\nUpdate A"}, + {"role": "user", "content": "Turn 1"}, + {"role": "user", "content": "Turn 2"}, + {"role": "assistant", "content": "Reply 2"}, + {"role": "user", "content": "Closing B"}, + {"role": "user", "content": "Closing C"}, + ] + + def test_developer_message_that_closes_the_conversation_after_a_user_turn_is_hoisted(self) -> None: + messages: Final = [ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Reply with exactly one word: the capital of France"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Reply with exactly one word: the capital of France"}, + {"role": "user", "content": "Hi there"}, + ] + + def test_developer_message_that_closes_the_conversation_after_a_tool_result_is_hoisted(self) -> None: + messages: Final = [ + {"role": "system", "content": "Base"}, + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Paris"}, + {"role": "developer", "content": "Answer with exactly one word."}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Base\n\nAnswer with exactly one word."}, + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Paris"}, + ] + + def test_hoisted_single_text_block_developer_message_joins_string_instructions_as_text(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are Codex"}, + {"role": "user", "content": [{"type": "text", "text": "# AGENTS.md"}]}, + {"role": "developer", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "You are Codex\n\n"}, + {"role": "user", "content": [{"type": "text", "text": "# AGENTS.md"}]}, + {"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}, + ] + + def test_multi_block_developer_message_keeps_its_blocks_after_string_instructions(self) -> None: + messages: Final = [ + {"role": "system", "content": "You are Codex"}, + {"role": "user", "content": "Hi"}, + {"role": "developer", "content": [{"type": "text", "text": "Rule 1"}, {"type": "text", "text": "Rule 2"}]}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are Codex"}, + {"type": "text", "text": "\n\nRule 1"}, + {"type": "text", "text": "Rule 2"}, + ], + }, + {"role": "user", "content": "Hi"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_system_message_without_content_folds_into_the_run_instead_of_raising(self) -> None: + messages: Final = [ + {"role": "system", "content": "Rules"}, + {"role": "system"}, + {"role": "system", "content": None}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Rules"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_message_level_cache_control_of_a_block_content_member_lands_on_its_last_block(self) -> None: + messages: Final = [ + { + "role": "system", + "content": [{"type": "text", "text": "Rule 1"}, {"type": "text", "text": "Rule 2"}], + "cache_control": {"type": "ephemeral"}, + }, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Rule 1"}, + {"type": "text", "text": "Rule 2", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "\n\nAnswer with exactly one word."}, + ], + }, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_message_level_cache_control_does_not_override_a_block_that_already_carries_one(self) -> None: + messages: Final = [ + { + "role": "system", + "content": [{"type": "text", "text": "Rule 1", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + "cache_control": {"type": "ephemeral"}, + }, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Rule 1", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + {"type": "text", "text": "\n\nAnswer with exactly one word."}, + ], + }, + {"role": "user", "content": "What is the capital of France?"}, + ] + + def test_non_consecutive_native_system_messages_stay_where_the_client_put_them(self) -> None: + messages: Final = [ + {"role": "system", "content": "System turn 1"}, + {"role": "user", "content": "User turn 1"}, + {"role": "system", "content": "System turn 2"}, + {"role": "user", "content": "User turn 2"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == messages + + def test_list_content_merge_preserves_block_level_cache_control(self) -> None: + messages: Final = [ + {"role": "system", "content": [{"type": "text", "text": "Cached", "cache_control": {"type": "ephemeral"}}]}, + {"role": "developer", "content": [{"type": "text", "text": "Extra"}]}, + {"role": "user", "content": "Hello"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Cached", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "\n\nExtra"}, + ], + }, + {"role": "user", "content": "Hello"}, + ] diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 61e5545b8bf..28f4fdead9e 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1,4 +1,5 @@ import json +from collections.abc import Callable from typing import Final from unittest.mock import MagicMock, patch @@ -2358,3 +2359,86 @@ def test_global_disable_flag_with_transform_messages_helper( ) assert json.loads(route.calls.last.request.content)["messages"] == [{"role": "user", "content": content}] + + +_CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN: Final = ( + {"role": "system", "content": "You are Codex"}, + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, +) + + +def test_translate_developer_role_keeps_a_later_developer_message_in_place_for_an_undeclared_fireworks_model() -> None: + messages: Final = FireworksAIConfig().translate_developer_role_to_system_role( + messages=_CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN, + model="accounts/fireworks/models/unlisted-model", + custom_llm_provider="fireworks_ai", + api_base=None, + model_info=None, + ) + + assert list(messages) == [ + {"role": "system", "content": "You are Codex"}, + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +def test_translate_developer_role_leaves_a_leading_system_and_developer_pair_unmerged_for_an_undeclared_model() -> None: + messages: Final = FireworksAIConfig().translate_developer_role_to_system_role( + messages=[ + {"role": "system", "content": "You are Codex"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + model="accounts/fireworks/models/unlisted-model", + custom_llm_provider="fireworks_ai", + api_base=None, + model_info=None, + ) + + assert list(messages) == [ + {"role": "system", "content": "You are Codex"}, + {"role": "system", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.mark.parametrize("model", ["qwen3p8-max", "accounts/fireworks/models/qwen3p8-max"]) +def test_translate_developer_role_hoists_a_later_developer_message_for_the_system_first_qwen3p8_model( + local_model_cost_map: None, model: str +) -> None: + messages: Final = FireworksAIConfig().translate_developer_role_to_system_role( + messages=_CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN, + model=model, + custom_llm_provider="fireworks_ai", + api_base=None, + model_info=None, + ) + + assert list(messages) == [ + {"role": "system", "content": "You are Codex\n\nAnswer with exactly one word."}, + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.mark.parametrize( + ("deployment_declared", "expected_roles"), + [(False, ["system", "user", "user"]), (True, ["system", "user", "system", "user"])], +) +def test_translate_developer_role_follows_the_deployment_model_info_over_the_cost_map( + declare_mid_conversation_system: Callable[[str, bool], None], deployment_declared: bool, expected_roles: list[str] +) -> None: + declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/shared-backend", not deployment_declared) + messages: Final = FireworksAIConfig().translate_developer_role_to_system_role( + messages=_CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN, + model="accounts/fireworks/models/shared-backend", + custom_llm_provider="fireworks_ai", + api_base=None, + model_info={"id": "deployment-a", "supports_mid_conversation_system": deployment_declared}, + ) + + assert [message["role"] for message in messages] == expected_roles diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index 878e0285cd1..7fc6ac2c436 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -1,5 +1,5 @@ import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping from types import MappingProxyType from typing import Final, TypedDict, cast from unittest.mock import MagicMock, patch @@ -7,6 +7,7 @@ from urllib.parse import quote import httpx import pytest +import respx from openai.types.responses import ( EasyInputMessage, ResponseFunctionToolCall, @@ -81,6 +82,14 @@ def _mock_http_client(response_body: Mapping[str, object]) -> MagicMock: return client +def _fireworks_responses_route(model: str) -> respx.Route: + return respx.post(FIREWORKS_RESPONSES_URL).mock(return_value=httpx.Response(200, json=_fireworks_response(model))) + + +def _routed_body(route: respx.Route) -> Mapping[str, object]: + return json.loads(route.calls.last.request.content) + + def _sent_request(client: MagicMock) -> tuple[str, Mapping[str, str], Mapping[str, object]]: kwargs: Final = client.post.call_args.kwargs body: Final = kwargs["json"] if "json" in kwargs else json.loads(kwargs["data"]) @@ -182,6 +191,60 @@ def test_responses_call_folds_developer_items_into_instructions() -> None: ) +@respx.mock +def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_accepts_later_system_items() -> None: + route: Final = _fireworks_responses_route("accounts/fireworks/models/accepts-later-system") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/accepts-later-system", + instructions="You are Codex", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + model_info={"id": "deployment-1", "supports_mid_conversation_system": True}, + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "You are Codex" + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +@respx.mock +@pytest.mark.parametrize( + ("model_info", "closing_role"), + [(None, "system"), ({"id": "deployment-1", "supports_mid_conversation_system": False}, "user")], +) +def test_responses_call_sends_the_closing_developer_item_as_a_user_turn_only_for_a_declared_system_first_deployment( + model_info: Mapping[str, object] | None, closing_role: str +) -> None: + route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "developer", "content": "Reply with PONG."}, + ], + model_info=model_info, + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": closing_role, "content": "Reply with PONG.", "type": "message"}, + ) + + def test_responses_call_folds_instructions_and_developer_item_into_instructions_with_reasoning_replayed() -> None: client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) with patch(HTTPX_CLIENT_FACTORY, return_value=client): @@ -667,3 +730,40 @@ def test_responses_call_forwards_litellm_user_id_as_user( _, _, body = _sent_request(client) assert body.get("user") == expected_user assert "fireworks_forward_user_id" not in body + + +@respx.mock +@pytest.mark.parametrize( + ("model_info", "expected_instructions", "expected_input"), + [ + (None, "Answer in one word.", ({"role": "user", "content": "What is the capital of France?"},)), + ( + {"id": "deployment-1", "supports_mid_conversation_system": True}, + None, + ( + {"role": "user", "content": "What is the capital of France?"}, + {"role": "system", "content": "Answer in one word.", "type": "message"}, + ), + ), + ], +) +def test_responses_call_forwards_developer_items_in_place_only_when_the_deployment_itself_declares_true( + declare_mid_conversation_system: Callable[[str, bool], None], + model_info: Mapping[str, object] | None, + expected_instructions: str | None, + expected_input: tuple[Mapping[str, object], ...], +) -> None: + declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/shared-backend", True) + route: Final = _fireworks_responses_route("accounts/fireworks/models/shared-backend") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/shared-backend", + input=[ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "developer", "content": "Answer in one word."}, + ], + model_info=model_info, + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body.get("instructions") == expected_instructions + assert tuple(body["input"]) == expected_input diff --git a/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py index e32f876a6b9..2f63ef070e6 100644 --- a/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py +++ b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py @@ -8,9 +8,13 @@ hosted_vllm (and any OpenAI-compatible provider using add_provider_specific_para """ import json +from collections.abc import Callable, Mapping +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +import respx import litellm @@ -183,3 +187,92 @@ def test_hosted_vllm_validate_environment_custom_api_key(): ) assert headers.get("Authorization") == "Bearer my-custom-key" + + +_LATER_DEVELOPER_INPUT: Final = ( + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, +) + + +_VLLM_RESPONSES_URL: Final = "https://test-vllm.example.com/v1/responses" + + +def _vllm_responses_route(output_text: str) -> respx.Route: + return respx.post(_VLLM_RESPONSES_URL).mock( + return_value=httpx.Response(200, json=_make_mock_responses_api_response(output_text)) + ) + + +def _sent_body(route: respx.Route) -> Mapping[str, object]: + return json.loads(route.calls.last.request.content) + + +@respx.mock +def test_hosted_vllm_responses_folds_a_later_developer_item_into_instructions_for_a_system_first_model( + declare_mid_conversation_system: Callable[[str, bool], None], +) -> None: + declare_mid_conversation_system("hosted_vllm/Qwen/Qwen3.6-27B", False) + route: Final = _vllm_responses_route("Paris") + litellm.responses( + model="hosted_vllm/Qwen/Qwen3.6-27B", + instructions="You are a terse assistant.", + input=list(_LATER_DEVELOPER_INPUT), + api_base="https://test-vllm.example.com/v1", + ) + body: Final = _sent_body(route) + assert body["instructions"] == "You are a terse assistant.\n\nAnswer with exactly one word." + assert tuple(body["input"]) == (_LATER_DEVELOPER_INPUT[0], _LATER_DEVELOPER_INPUT[2]) + + +@respx.mock +def test_hosted_vllm_responses_folds_for_a_deployment_declared_system_first_and_sends_a_closing_developer_as_user() -> None: + assistant_turn: Final = {"type": "message", "role": "assistant", "content": "Paris"} + route: Final = _vllm_responses_route("PONG") + litellm.responses( + model="hosted_vllm/Qwen/Qwen3-8B", + instructions="You are a terse assistant.", + input=[*_LATER_DEVELOPER_INPUT, assistant_turn, {"role": "developer", "content": "Reply with PONG."}], + model_info={"id": "deployment-1", "supports_mid_conversation_system": False}, + api_base="https://test-vllm.example.com/v1", + ) + body: Final = _sent_body(route) + assert body["instructions"] == "You are a terse assistant.\n\nAnswer with exactly one word." + assert tuple(body["input"]) == ( + _LATER_DEVELOPER_INPUT[0], + _LATER_DEVELOPER_INPUT[2], + assistant_turn, + {"role": "user", "content": "Reply with PONG.", "type": "message"}, + ) + + +@respx.mock +def test_hosted_vllm_responses_forwards_a_later_developer_item_in_place_for_an_undeclared_model() -> None: + route: Final = _vllm_responses_route("Paris") + litellm.responses( + model="hosted_vllm/Qwen/Qwen3-8B", + instructions="You are a terse assistant.", + input=list(_LATER_DEVELOPER_INPUT), + api_base="https://test-vllm.example.com/v1", + ) + body: Final = _sent_body(route) + assert body["instructions"] == "You are a terse assistant." + assert tuple(body["input"]) == _LATER_DEVELOPER_INPUT + + +@respx.mock +def test_hosted_vllm_responses_forwards_a_developer_item_without_content_on_a_deployment_declared_system_first() -> None: + route: Final = _vllm_responses_route("Paris") + input_items: Final = ( + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Paris."}]}, + {"type": "message", "role": "developer"}, + ) + litellm.responses( + model="hosted_vllm/Qwen/Qwen3-8B", + input=list(input_items), + model_info={"id": "deployment-1", "supports_mid_conversation_system": False}, + api_base="https://test-vllm.example.com/v1", + ) + assert tuple(_sent_body(route)["input"]) == input_items diff --git a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py index 04c2ce153c6..35e66c9742b 100644 --- a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py +++ b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py @@ -10,7 +10,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.huggingface.common_utils import _fetch_inference_provider_mapping -from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, _load_huggingface_tokenizer +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct" HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo" @@ -84,13 +84,11 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator monkeypatch.delenv("HF_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() _fetch_inference_provider_mapping.cache_clear() yield _fetch_inference_provider_mapping.cache_clear() - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() PROVIDER_MAPPING_RESPONSE = { diff --git a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py index 9b8788d3769..8a690a70a97 100644 --- a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py +++ b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py @@ -639,3 +639,25 @@ def test_litellm_gateway_from_sdk_with_thinking_param(respx_mock): ) e = exc_info.value assert "Connection error." in str(e) + + +def test_translate_developer_role_leaves_a_later_developer_message_in_place_for_the_downstream_proxy() -> None: + translated: Final = LiteLLMProxyChatConfig().translate_developer_role_to_system_role( + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + model="fireworks-qwen", + custom_llm_provider="litellm_proxy", + api_base="http://inner-proxy:4000", + model_info=None, + ) + + assert list(translated) == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 70436662ef0..5b85a94b24d 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -2,7 +2,9 @@ Tests for OpenAI GPT transformation (litellm/llms/openai/chat/gpt_transformation.py) """ import base64 +import copy import json +from collections.abc import Callable from typing import Final import pytest @@ -1820,3 +1822,97 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: ) assert request["tools"][0] is tool + + +def _developer_message_after_a_user_turn() -> list[dict[str, str]]: + return [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.mark.parametrize( + ("config", "custom_llm_provider", "api_base", "declared", "expected_roles"), + [ + (OpenAIGPTConfig(), "openai", "https://api.openai.com/v1", False, ["system", "user", "assistant", "system", "user"]), + (OpenAIGPT5Config(), "openai", "https://api.openai.com/v1", False, ["system", "user", "assistant", "system", "user"]), + (OpenAIGPTConfig(), "openai", "http://vllm.internal:8000/v1", None, ["system", "user", "assistant", "system", "user"]), + (OpenAIGPTConfig(), "openai", "http://vllm.internal:8000/v1", False, ["system", "user", "assistant", "user"]), + (OpenAIGPTConfig(), "groq", None, None, ["system", "user", "assistant", "system", "user"]), + (OpenAIGPTConfig(), "groq", None, False, ["system", "user", "assistant", "user"]), + ], +) +def test_translate_developer_role_hoists_a_later_developer_message_only_for_a_model_off_openai_that_rejects_it( + declare_mid_conversation_system: Callable[[str, bool], None], + config: OpenAIGPTConfig, + custom_llm_provider: str, + api_base: str | None, + declared: bool | None, + expected_roles: list[str], +) -> None: + if declared is not None: + declare_mid_conversation_system(f"{custom_llm_provider}/system-first-model", declared) + messages: Final = config.translate_developer_role_to_system_role( + messages=_developer_message_after_a_user_turn(), + model="system-first-model", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + model_info=None, + ) + assert [message["role"] for message in messages] == expected_roles + assert "Answer with exactly one word." in [message["content"] for message in messages if message["role"] == "system"][-1] + + +@pytest.mark.parametrize("role", [["user"], {"name": "user"}]) +@pytest.mark.parametrize( + ("declared", "expected_order"), + [(None, ("caller", "instruction")), (False, ("instruction", "caller"))], +) +def test_translate_developer_role_places_a_message_with_a_non_string_role_like_any_conversation_turn( + declare_mid_conversation_system: Callable[[str, bool], None], + role: object, + declared: bool | None, + expected_order: tuple[str, str], +) -> None: + if declared is not None: + declare_mid_conversation_system("groq/system-first-model", declared) + caller_turn: Final = {"role": role, "content": "Hi there"} + by_name: Final = {"caller": caller_turn, "instruction": {"role": "system", "content": "Answer in one word."}} + messages: Final = OpenAIGPTConfig().translate_developer_role_to_system_role( + messages=[caller_turn, {"role": "developer", "content": "Answer in one word."}], + model="system-first-model", + custom_llm_provider="groq", + api_base=None, + model_info=None, + ) + assert list(messages) == [by_name[name] for name in expected_order] + + +@pytest.mark.respx(assert_all_called=True) +def test_completion_leaves_the_callers_closing_developer_message_unchanged_on_a_system_first_deployment( + respx_mock: respx.MockRouter, +) -> None: + chat: Final = respx_mock.post("http://vllm.internal:8000/v1/chat/completions").mock( + return_value=Response(200, json=_OPENAI_CHAT_REPLY) + ) + messages: Final = [ + {"role": "user", "content": "Describe the image."}, + {"role": "assistant", "content": "Send it."}, + { + "role": "developer", + "content": [{"type": "text", "text": "Look at this."}, {"type": "image_url", "image_url": "data:image/png;base64,aW1hZ2U="}], + }, + ] + sent: Final = copy.deepcopy(messages) + litellm.completion( + model="openai/system-first-model", + api_base="http://vllm.internal:8000/v1", + api_key="sk-test", + messages=messages, + model_info={"id": "deployment-1", "supports_mid_conversation_system": False}, + ) + assert json.loads(chat.calls.last.request.content)["messages"][-1]["role"] == "user" + assert messages == sent diff --git a/tests/unit/llms/openai/test_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py index 121fab66c15..e78e1a6f12b 100644 --- a/tests/unit/llms/openai/test_o_series_transformation.py +++ b/tests/unit/llms/openai/test_o_series_transformation.py @@ -181,3 +181,20 @@ def test_litellm_responses(): print("response: ", response) assert isinstance(response.usage.completion_tokens_details, CompletionTokensDetails) + + +def test_translate_developer_role_keeps_developer_messages_untouched() -> None: + messages: Final = [ + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + assert ( + OpenAIOSeriesConfig().translate_developer_role_to_system_role( + messages=messages, + model="my-reasoning-deployment", + custom_llm_provider="openai", + api_base=None, + model_info=None, + ) + == messages + ) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 11d7c7d6d3c..5b27a446fb3 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -2,7 +2,7 @@ import asyncio import importlib import json import os -from collections.abc import Mapping +from collections.abc import Callable, Mapping from types import ModuleType from typing import Final, Literal, Protocol from unittest.mock import MagicMock, Mock, patch @@ -24,6 +24,7 @@ from litellm.llms.openai.openai import ( AsyncCursorPage, MessageData, OpenAIChatCompletion, + OpenAIConfig, OpenAIMessage as Message, Run, SyncCursorPage, @@ -2010,3 +2011,40 @@ async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_head assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True} assert sent.headers["x-caller"] == "unit" assert sent.headers["authorization"] == "Bearer sk-test" + + +def _developer_message_after_a_user_turn() -> list[dict[str, str]]: + return [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hi there"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + + +@pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "declared", "expected_roles"), + [ + ("openai", "https://api.openai.com/v1", False, ["system", "user", "assistant", "system", "user"]), + ("azure_ai", "https://example.services.ai.azure.com/models", None, ["system", "user", "assistant", "system", "user"]), + ("azure_ai", "https://example.services.ai.azure.com/models", False, ["system", "user", "assistant", "user"]), + ], +) +def test_openai_config_hoists_a_later_developer_message_only_for_a_model_off_openai_that_rejects_it( + declare_mid_conversation_system: Callable[[str, bool], None], + custom_llm_provider: str, + api_base: str, + declared: bool | None, + expected_roles: list[str], +) -> None: + if declared is not None: + declare_mid_conversation_system(f"{custom_llm_provider}/system-first-model", declared) + messages: Final = OpenAIConfig().translate_developer_role_to_system_role( + messages=_developer_message_after_a_user_turn(), + model="system-first-model", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + model_info=None, + ) + assert [message["role"] for message in messages] == expected_roles diff --git a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 6f426a4a3a7..9c693ff1c3e 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -5,15 +5,19 @@ Covers the git-subdir and archive source types added alongside the existing gith """ import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from fastapi import HTTPException from pydantic import ValidationError import litellm -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth +from litellm.caching.caching import DualCache +from litellm.proxy._types import CommonProxyErrors, LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints import claude_code_marketplace from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( delete_plugin, @@ -25,7 +29,7 @@ from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketp register_plugin, update_plugin, ) -from litellm.proxy.proxy_server import LitellmUserRoles +from litellm.proxy.proxy_server import LitellmUserRoles, hash_token from litellm.types.proxy.claude_code_endpoints import ( MarketplaceResponse, RegisterPluginRequest, @@ -964,3 +968,206 @@ def test_marketplace_response_rejects_unknown_installation_preference(): with pytest.raises(ValidationError): MarketplaceResponse.model_validate(catalog) + + +def _require_marketplace_auth(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "general_settings", {"claude_code_marketplace_auth_required": True} + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache()) + + +async def _get_marketplace_over_http( + headers: Mapping[str, str] = MappingProxyType({}), params: Mapping[str, str] = MappingProxyType({}) +) -> httpx.Response: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=litellm.proxy.proxy_server.app), base_url="http://test" + ) as client: + return await client.get("/claude-code/marketplace.json", headers=dict(headers), params=dict(params)) + + +@pytest.mark.asyncio +async def test_marketplace_auth_required_refuses_an_anonymous_caller(monkeypatch: pytest.MonkeyPatch) -> None: + _require_marketplace_auth(monkeypatch) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http() + + assert response.status_code == 401 + assert response.headers["www-authenticate"] == "Bearer" + assert "public-skill" not in response.text + + +@pytest.mark.asyncio +async def test_marketplace_auth_required_refuses_a_key_in_the_query_string(monkeypatch: pytest.MonkeyPatch) -> None: + _require_marketplace_auth(monkeypatch) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http(params={"key": MASTER_KEY}) + + assert response.status_code == 401 + assert response.headers["www-authenticate"] == "Bearer" + assert "public-skill" not in response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers", + [ + {"Authorization": f"Bearer {MASTER_KEY}"}, + {"x-litellm-api-key": f"Bearer {MASTER_KEY}"}, + ], +) +async def test_marketplace_auth_required_serves_a_caller_with_a_valid_key( + monkeypatch: pytest.MonkeyPatch, headers: Mapping[str, str] +) -> None: + _require_marketplace_auth(monkeypatch) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http(headers) + + assert response.status_code == 200 + assert sorted(plugin["name"] for plugin in response.json()["plugins"]) == ["private-skill", "public-skill"] + + +@pytest.mark.asyncio +async def test_marketplace_auth_required_refuses_an_unknown_bearer_key(monkeypatch: pytest.MonkeyPatch) -> None: + _require_marketplace_auth(monkeypatch) + litellm.proxy.proxy_server.prisma_client.get_data = AsyncMock(return_value=None) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http({"Authorization": "Bearer sk-not-a-litellm-key"}) + + assert response.status_code == 401 + assert response.headers["www-authenticate"] == "Bearer" + assert "public-skill" not in response.text + + +_LIMITED_KEY: Final = "sk-limited" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limited_key_auth, status_code", + [ + ( + UserAPIKeyAuth( + token=hash_token(_LIMITED_KEY), + user_role=LitellmUserRoles.INTERNAL_USER, + allowed_routes=["/chat/completions"], + ), + 403, + ), + ( + UserAPIKeyAuth( + token=hash_token(_LIMITED_KEY), user_role=LitellmUserRoles.INTERNAL_USER, spend=2.0, max_budget=1.0 + ), + 422, + ), + ], + ids=["route-not-allowed", "over-budget"], +) +async def test_marketplace_auth_required_keeps_a_valid_key_refusal_status( + monkeypatch: pytest.MonkeyPatch, limited_key_auth: UserAPIKeyAuth, status_code: int +) -> None: + _require_marketplace_auth(monkeypatch) + litellm.proxy.proxy_server.user_api_key_cache.set_cache(key=hash_token(_LIMITED_KEY), value=limited_key_auth) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http({"Authorization": f"Bearer {_LIMITED_KEY}"}) + + assert response.status_code == status_code + assert "www-authenticate" not in response.headers + assert "public-skill" not in response.text + + +@pytest.mark.asyncio +async def test_marketplace_without_auth_required_ignores_an_unrelated_authorization_header(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {}) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http({"Authorization": "Bearer an-identity-aware-proxy-token"}) + + assert response.status_code == 200 + assert [plugin["name"] for plugin in response.json()["plugins"]] == ["public-skill"] + + +@pytest.mark.asyncio +async def test_marketplace_500_body_carries_no_exception_text(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {}) + table: Final = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + table.find_many = AsyncMock(side_effect=RuntimeError("could not reach postgres://admin:hunter2@db.internal")) + + response: Final = await _get_marketplace_over_http() + + assert response.status_code == 500 + assert "hunter2" not in response.text + assert "db.internal" not in response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["true", "True", 1, "a typo"]) +async def test_marketplace_auth_required_treats_a_quoted_or_unreadable_setting_as_on( + monkeypatch: pytest.MonkeyPatch, setting: object +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "general_settings", {"claude_code_marketplace_auth_required": setting} + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache()) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http() + + assert response.status_code == 401 + assert "public-skill" not in response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["false", False, 0, None]) +async def test_marketplace_auth_required_off_values_keep_the_public_catalog( + monkeypatch: pytest.MonkeyPatch, setting: object +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "general_settings", {"claude_code_marketplace_auth_required": setting} + ) + await _register_public_and_private_plugins() + + response: Final = await _get_marketplace_over_http() + + assert response.status_code == 200 + assert [plugin["name"] for plugin in response.json()["plugins"]] == ["public-skill"] + + +@pytest.mark.asyncio +async def test_marketplace_auth_required_refuses_the_database_outage_fallback_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "general_settings", + {"claude_code_marketplace_auth_required": True, "allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache()) + await _register_public_and_private_plugins() + litellm.proxy.proxy_server.prisma_client.get_data = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + litellm.proxy.proxy_server.prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + response: Final = await _get_marketplace_over_http({"Authorization": "Bearer sk-not-a-litellm-key"}) + + assert response.status_code == 503 + assert "www-authenticate" not in response.headers + assert "public-skill" not in response.text + + +@pytest.mark.asyncio +async def test_marketplace_without_a_database_answers_db_not_connected_for_a_query_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {}) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache()) + + response: Final = await _get_marketplace_over_http(params={"key": "sk-not-the-master-key"}) + + assert response.status_code == 500 + assert CommonProxyErrors.db_not_connected_error.value in response.text diff --git a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index d338f38d862..6cdfd8139dc 100644 --- a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -523,3 +523,33 @@ def test_ui_config_tells_the_dashboard_whether_to_use_si_terminology( assert response.status_code == 200 assert response.json()["use_si_terminology"] is expected + + +@pytest.mark.parametrize( + ("settings", "expected"), + [ + ({}, False), + ({"claude_code_marketplace_auth_required": False}, False), + ({"claude_code_marketplace_auth_required": True}, True), + ({"claude_code_marketplace_auth_required": "true"}, True), + ({"claude_code_marketplace_auth_required": "maybe"}, True), + ], +) +def test_ui_config_reports_whether_the_marketplace_requires_a_key_as_the_marketplace_reads_it( + monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], expected: bool +) -> None: + import litellm.proxy.proxy_server + + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", settings) + app: Final[FastAPI] = FastAPI() + app.include_router(router) + + with ( + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), + ): + response: Final = TestClient(app).get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["claude_code_marketplace_auth_required"] is expected diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 20a4ca87d42..23dbe86ee5f 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -2660,7 +2660,9 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke @pytest.mark.asyncio -@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"]) +@pytest.mark.parametrize( + "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "claude_code_marketplace_auth_required"] +) @pytest.mark.parametrize(("yaml_value", "expected"), [("true", True), ("false", False)]) async def test_load_config_yaml_deny_by_default_is_boolean( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str, expected: bool @@ -2678,7 +2680,9 @@ async def test_load_config_yaml_deny_by_default_is_boolean( @pytest.mark.asyncio -@pytest.mark.parametrize("setting", ["vector_store_deny_by_default", "search_tool_deny_by_default"]) +@pytest.mark.parametrize( + "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "claude_code_marketplace_auth_required"] +) @pytest.mark.parametrize("yaml_value", ["", "enabled"], ids=["null", "string"]) async def test_load_config_rejects_non_boolean_deny_by_default( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, setting: str, yaml_value: str diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 072ef838e78..e694bca3c8d 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1,6 +1,7 @@ import asyncio import json import re +from collections.abc import Mapping from datetime import datetime, timezone from importlib.resources import files from itertools import chain @@ -8,6 +9,7 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient @@ -18,6 +20,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.public_endpoints import router from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets +from litellm.types.proxy.common_proxy_errors import CommonProxyErrors from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) @@ -1782,3 +1785,72 @@ def test_public_skill_hub_echoes_registered_installation_preference(preference): assert [(p["name"], p["installation_preference"]) for p in response.json()["plugins"]] == [ ("security-review", preference) ] + + +_INTERNAL_PLUGIN: Final = SimpleNamespace( + id="plugin-id", + name="deploy-helper", + enabled=True, + created_at=None, + updated_at=None, + manifest_json=json.dumps( + { + "source": { + "source": "git-subdir", + "url": "https://gitlab.internal.example.com/platform/claude-skills.git", + "path": "plugins/deploy-helper", + } + } + ), +) + + +def _skill_hub_response(find_many: AsyncMock, general_settings: Mapping[str, object]) -> httpx.Response: + table: Final = MagicMock() + table.find_many = find_many + prisma: Final = MagicMock() + prisma.db.litellm_claudecodeplugintable = table + app: Final = FastAPI() + app.include_router(router) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", dict(general_settings)), + ): + return TestClient(app).get("/public/skill_hub") + + +def test_public_skill_hub_is_empty_when_marketplace_auth_is_required() -> None: + find_many: Final = AsyncMock(return_value=[_INTERNAL_PLUGIN]) + + public: Final = _skill_hub_response(find_many, {}) + gated: Final = _skill_hub_response(find_many, {"claude_code_marketplace_auth_required": True}) + + assert [plugin["name"] for plugin in public.json()["plugins"]] == ["deploy-helper"] + assert gated.status_code == 200 + assert gated.json()["plugins"] == [] + assert "gitlab.internal.example.com" not in gated.text + + +def test_public_skill_hub_500_body_carries_no_exception_text() -> None: + response: Final = _skill_hub_response( + AsyncMock(side_effect=RuntimeError("could not reach postgres://admin:hunter2@db.internal")), {} + ) + + assert response.status_code == 500 + assert "hunter2" not in response.text + assert "db.internal" not in response.text + + +def test_public_skill_hub_without_a_database_answers_db_not_connected() -> None: + app: Final = FastAPI() + app.include_router(router) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + response: Final = TestClient(app).get("/public/skill_hub") + + assert response.status_code == 500 + assert response.json() == {"detail": CommonProxyErrors.db_not_connected_error.value} diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 665a540e462..2870016a5e9 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10869,6 +10869,44 @@ def test_get_optional_params_keeps_a_tool_config_cache_point_out_of_the_openai_e assert "cache_control_injection_points" not in optional_params +@pytest.mark.parametrize( + ("cost_map_declared", "deployment_declared", "expected"), + [ + (None, None, None), + (False, None, False), + (True, None, True), + (True, False, False), + (False, True, True), + (None, False, False), + ], +) +def test_declared_mid_conversation_system_support_prefers_the_deployment_over_the_cost_map( + declare_mid_conversation_system: Callable[[str, bool], None], + cost_map_declared: bool | None, + deployment_declared: bool | None, + expected: bool | None, +) -> None: + if cost_map_declared is not None: + declare_mid_conversation_system("hosted_vllm/mid-conversation-probe", cost_map_declared) + deployment_model_info: Final = ( + {"id": "deployment-1"} + if deployment_declared is None + else {"id": "deployment-1", "supports_mid_conversation_system": deployment_declared} + ) + assert ( + litellm.utils.declared_mid_conversation_system_support( + "mid-conversation-probe", "hosted_vllm", deployment_model_info + ) + is expected + ) + + +def test_declared_mid_conversation_system_support_is_none_when_the_model_lookup_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(litellm.model_cost, "hosted_vllm/unreadable-probe", "not a model info entry") + litellm.get_model_info.cache_clear() + assert litellm.utils.declared_mid_conversation_system_support("unreadable-probe", "hosted_vllm", None) is None + + def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_litellm_proxy(): """A litellm_proxy deployment calls the downstream proxy through the OpenAI SDK, and that proxy applies cache_control_injection_points itself, so the point must ride along in extra_body.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.test.tsx new file mode 100644 index 00000000000..0cffcb82644 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.test.tsx @@ -0,0 +1,54 @@ +import React from "react"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { useClaudeCodeMarketplaceRequiresKey } from "./useClaudeCodeMarketplaceRequiresKey"; + +const renderRequiresKey = (respond: () => Promise) => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(respond); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const { result } = renderHook(() => useClaudeCodeMarketplaceRequiresKey(), { + wrapper: ({ children }: { children: React.ReactNode }) => ( + {children} + ), + }); + const settled = () => + waitFor(() => + expect( + client + .getQueryCache() + .getAll() + .map((query) => query.state.status), + ).toEqual(["success"]), + ); + return { result, fetchSpy, settled }; +}; + +describe("useClaudeCodeMarketplaceRequiresKey", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it.each([ + [{ claude_code_marketplace_auth_required: false }, false], + [{ claude_code_marketplace_auth_required: true }, true], + [{}, true], + ])("reads %o from the UI config as requires key=%s", async (config, requiresKey) => { + const { result, fetchSpy, settled } = renderRequiresKey( + async () => new Response(JSON.stringify(config), { status: 200 }), + ); + + await settled(); + expect(result.current).toBe(requiresKey); + expect(fetchSpy.mock.calls.map(([request]) => new URL((request as Request).url).pathname)).toEqual([ + "/.well-known/litellm-ui-config", + ]); + }); + + it("assumes a key is needed while the proxy has not answered yet", async () => { + const { result, fetchSpy } = renderRequiresKey(() => new Promise(() => {})); + + await waitFor(() => expect(fetchSpy).toHaveBeenCalled()); + expect(result.current).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.ts new file mode 100644 index 00000000000..dec7098bf35 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey.ts @@ -0,0 +1,4 @@ +import { $api } from "@/lib/http/api"; + +export const useClaudeCodeMarketplaceRequiresKey = (): boolean => + $api.useQuery("get", "/.well-known/litellm-ui-config").data?.claude_code_marketplace_auth_required !== false; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index 5a4b8840c3e..ca1ee99cc27 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -16,8 +16,22 @@ import { } from "./helpers"; describe("buildMarketplaceSettingsSnippet", () => { + it("sends the key placeholder as a bearer header when the proxy requires a key", () => { + expect(JSON.parse(buildMarketplaceSettingsSnippet("https://proxy.example.com", true))).toEqual({ + extraKnownMarketplaces: { + litellm: { + source: { + source: "url", + url: "https://proxy.example.com/claude-code/marketplace.json", + headers: { Authorization: "Bearer " }, + }, + }, + }, + }); + }); + it("nests the url under a source object so Claude Code accepts the marketplace", () => { - expect(JSON.parse(buildMarketplaceSettingsSnippet("https://proxy.example.com"))).toEqual({ + expect(JSON.parse(buildMarketplaceSettingsSnippet("https://proxy.example.com", false))).toEqual({ extraKnownMarketplaces: { litellm: { source: { diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 52f860dc0ec..519d49ca704 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -240,7 +240,9 @@ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourceP * bare `"url"` string, so the url/source pair is nested one level deeper. The key must be * "litellm" to match the name the proxy returns in marketplace.json. */ -export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => +export const MARKETPLACE_KEY_PLACEHOLDER = ""; + +export const buildMarketplaceSettingsSnippet = (proxyOrigin: string, requiresKey: boolean): string => JSON.stringify( { extraKnownMarketplaces: { @@ -248,6 +250,7 @@ export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => source: { source: "url", url: `${proxyOrigin}/claude-code/marketplace.json`, + ...(requiresKey ? { headers: { Authorization: `Bearer ${MARKETPLACE_KEY_PLACEHOLDER}` } } : {}), }, }, }, diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx index 92e18514f66..951f30fc9cc 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx @@ -1,10 +1,20 @@ -import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { Plugin } from "./types"; import SkillDetail from "./skill_detail"; +const { mockUseMarketplaceRequiresKey } = vi.hoisted(() => ({ mockUseMarketplaceRequiresKey: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey", () => ({ + useClaudeCodeMarketplaceRequiresKey: () => mockUseMarketplaceRequiresKey(), +})); + +beforeEach(() => { + mockUseMarketplaceRequiresKey.mockReturnValue(false); +}); + const buildSkill = (source: Plugin["source"]): Plugin => ({ id: "plugin-id", name: "my-skill", @@ -55,3 +65,27 @@ describe("SkillDetail source", () => { expect(screen.queryByRole("link")).not.toBeInTheDocument(); }); }); + +describe("SkillDetail marketplace setup", () => { + const renderSetupTab = () => { + render(); + fireEvent.click(screen.getByText("How to Use")); + fireEvent.click(screen.getByText(/See one-time setup/)); + }; + + it("offers the add command and a keyless snippet when the marketplace is open", () => { + renderSetupTab(); + expect(screen.getByText(/^\/plugin marketplace add .*\/claude-code\/marketplace\.json$/)).toBeInTheDocument(); + expect(screen.queryByText(/Bearer/)).not.toBeInTheDocument(); + expect(screen.queryByText(/\/plugin marketplace remove/)).not.toBeInTheDocument(); + }); + + it("replaces the add command with a snippet carrying the key header when the marketplace requires a key", () => { + mockUseMarketplaceRequiresKey.mockReturnValue(true); + renderSetupTab(); + expect(screen.queryByText(/\/plugin marketplace add/)).not.toBeInTheDocument(); + expect(screen.getByText(/"Authorization": "Bearer "/)).toBeInTheDocument(); + expect(screen.getByText("/plugin marketplace remove litellm")).toBeInTheDocument(); + expect(screen.getByText(/uninstalls the plugins you installed from it/)).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 9b03c25aade..658a0472699 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -1,7 +1,14 @@ import React, { useState } from "react"; import { ArrowLeft, Check, Copy, Link2 } from "lucide-react"; +import { useClaudeCodeMarketplaceRequiresKey } from "@/app/(dashboard)/hooks/claudeCodeMarketplace/useClaudeCodeMarketplaceRequiresKey"; import { cn } from "@/lib/cva.config"; -import { buildMarketplaceSettingsSnippet, formatInstallCommand, getSourceDisplayText, getSourceLink } from "./helpers"; +import { + buildMarketplaceSettingsSnippet, + formatInstallCommand, + getSourceDisplayText, + getSourceLink, + MARKETPLACE_KEY_PLACEHOLDER, +} from "./helpers"; import { Plugin } from "./types"; const SkillSource: React.FC<{ source: Plugin["source"] }> = ({ source }) => { @@ -45,6 +52,7 @@ interface SkillDetailProps { const SkillDetail: React.FC = ({ skill, onBack }) => { const [activeTab, setActiveTab] = useState("overview"); const [copiedKey, setCopiedKey] = useState(null); + const marketplaceRequiresKey = useClaudeCodeMarketplaceRequiresKey(); const copyToClipboard = (text: string, key: string) => { navigator.clipboard.writeText(text); @@ -56,6 +64,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { const settingsSnippet = buildMarketplaceSettingsSnippet( typeof window !== "undefined" ? window.location.origin : "", + marketplaceRequiresKey, ); const detailRows = [ @@ -226,38 +235,65 @@ const SkillDetail: React.FC = ({ skill, onBack }) => {

One-time marketplace setup

- {/* Option 1: single command — fastest path for most users */} -

- Run this command in Claude Code to register the marketplace: -

-
-
- Run in Claude Code - -
- 
-              {`/plugin marketplace add ${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`}
-
-
+ {marketplaceRequiresKey ? ( + <> +
+

+ If you added this marketplace before without a key, remove it first in Claude Code. Removing it also + uninstalls the plugins you installed from it, so install them again from /plugin after you add the + snippet: +

+ 
+                  /plugin marketplace remove litellm
+
+
+

+ This proxy requires a LiteLLM key to read the marketplace. Add this to{" "} + ~/.claude/settings.json and replace{" "} + {MARKETPLACE_KEY_PLACEHOLDER} with + your key: +

+ + ) : ( + <> + {/* Option 1: single command — fastest path for most users */} +

+ Run this command in Claude Code to register the marketplace: +

+
+
+ Run in Claude Code + +
+ 
+                  {`/plugin marketplace add ${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`}
+
+
- {/* Option 2: settings.json — for persistent config or managed deployments. - extraKnownMarketplaces requires source to be a nested object, not a flat string. */} -

- Or add this to ~/.claude/settings.json{" "} - for a persistent configuration: -

+ {/* Option 2: settings.json — for persistent config or managed deployments. + extraKnownMarketplaces requires source to be a nested object, not a flat string. */} +

+ Or add this to{" "} + ~/.claude/settings.json for a + persistent configuration: +

+ + )}
~/.claude/settings.json diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2147afc4931..f114c1c73d2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2408,10 +2408,15 @@ export interface paths { * - claude plugin marketplace add * - claude plugin install @ * - * Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` + * Without a key the catalog holds the enabled (public) plugins. With `?key=sk-...` * the key is authenticated and the catalog also holds the disabled plugins granted * to it through `object_permission.skills` or an access group on the key or its team. * + * With `general_settings.claude_code_marketplace_auth_required: true` the key must come + * in the `Authorization` or `x-litellm-api-key` header (Claude Code sends them from the + * marketplace source's `headers`), `?key=` is ignored, and a caller with no valid key + * gets a 401 with `WWW-Authenticate: Bearer`. + * * Returns: * Marketplace catalog with list of available plugins and their git sources. * @@ -13786,7 +13791,7 @@ export interface paths { }; /** * Public Skill Hub - * @description Return enabled (public) Claude Code skills — no auth required. + * @description Return enabled (public) Claude Code skills, or none when the marketplace requires a key. */ get: operations["public_skill_hub_public_skill_hub_get"]; put?: never; @@ -29890,6 +29895,12 @@ export interface components { claude_code_gateway_managed_settings?: { [key: string]: unknown; } | null; + /** + * Claude Code Marketplace Auth Required + * @description If True, GET /claude-code/marketplace.json requires a LiteLLM key in the Authorization or x-litellm-api-key header (?key= is not accepted) and answers 401 with WWW-Authenticate: Bearer without one, and GET /public/skill_hub returns an empty list. Default is False, which serves the enabled plugins to anyone. + * @default false + */ + claude_code_marketplace_auth_required: boolean; /** * Completion Model * @description proxy level default model for all chat completion calls @@ -48183,6 +48194,11 @@ export interface components { admin_ui_disabled: boolean; /** Auto Redirect To Sso */ auto_redirect_to_sso: boolean; + /** + * Claude Code Marketplace Auth Required + * @default false + */ + claude_code_marketplace_auth_required: boolean; /** Gateway Base Url */ gateway_base_url?: string | null; /**