mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
Merge remote-tracking branch 'origin/main' into litellm_model_access_denied_message
This commit is contained in:
commit
7267c6bed7
93 changed files with 6232 additions and 402 deletions
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.67"
|
||||
version = "0.1.68"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.67"
|
||||
version = "0.1.68"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,18 @@
|
|||
-- DropIndex
|
||||
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx";
|
||||
|
||||
-- DropIndex
|
||||
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key";
|
||||
|
||||
-- AlterTable
|
||||
-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every
|
||||
-- NULL as distinct, so a nullable column would let multiple unscoped mappings
|
||||
-- collide on the same claim without a constraint violation. The constant
|
||||
-- default is a fast, metadata-only backfill for existing rows, not a rewrite.
|
||||
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT '';
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value");
|
||||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.97"
|
||||
version = "0.4.98"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.97"
|
||||
version = "0.4.98"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -344,6 +344,7 @@ _anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_
|
|||
anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = (
|
||||
"1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None
|
||||
)
|
||||
openai_system_messages_first: bool = False
|
||||
disable_vertex_batch_output_transformation: bool = False
|
||||
extra_spend_tag_headers: Optional[List[str]] = None
|
||||
in_memory_llm_clients_cache: "LLMClientCache"
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from litellm.litellm_core_utils.logging_utils import (
|
|||
_assemble_complete_response_from_streaming_chunks,
|
||||
)
|
||||
from litellm.types.caching import CachedEmbedding
|
||||
from litellm.types.integrations.custom_logger import converted_stream_requested
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -107,17 +108,31 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
|
|||
return "choices" in cached_result
|
||||
|
||||
|
||||
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
|
||||
def _stream_replay_requested(kwargs: Mapping[str, object]) -> bool:
|
||||
if kwargs.get("stream", False) is True:
|
||||
return True
|
||||
return converted_stream_requested(kwargs) and not kwargs.get("_agentic_loop_depth")
|
||||
|
||||
|
||||
def _should_defer_streaming_cache_hit_callbacks(*, cached_result: object) -> bool:
|
||||
"""
|
||||
When stream=True, do not run success callbacks at cache-hit time.
|
||||
When the cache hit is replayed as a stream, do not run success callbacks at cache-hit time.
|
||||
|
||||
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
|
||||
replay uses CachedResponsesAPIStreamingIterator; cached Anthropic Messages
|
||||
replay uses CachedAnthropicMessagesStreamIterator. All invoke logging success
|
||||
handlers when the stream finishes; firing them here too would double-count
|
||||
spend and callback records.
|
||||
spend and callback records. A plain (non-stream) replay logs here, since nothing
|
||||
else will.
|
||||
"""
|
||||
return kwargs.get("stream", False) is True
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import (
|
||||
CachedAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
return isinstance(
|
||||
cached_result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator, CachedAnthropicMessagesStreamIterator)
|
||||
)
|
||||
|
||||
|
||||
def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
|
||||
|
|
@ -267,7 +282,7 @@ class LLMCachingHandler:
|
|||
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
||||
args=args,
|
||||
)
|
||||
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
||||
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
|
||||
# LOG SUCCESS
|
||||
self._async_log_cache_hit_on_callbacks(
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -383,7 +398,7 @@ class LLMCachingHandler:
|
|||
is_async=False,
|
||||
)
|
||||
|
||||
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
||||
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=cached_result,
|
||||
start_time=start_time,
|
||||
|
|
@ -823,7 +838,7 @@ class LLMCachingHandler:
|
|||
if (call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value) and isinstance(
|
||||
cached_result, dict
|
||||
):
|
||||
if kwargs.get("stream", False) is True:
|
||||
if _stream_replay_requested(kwargs):
|
||||
cached_result = self._convert_cached_stream_response(
|
||||
cached_result=cached_result,
|
||||
call_type=call_type,
|
||||
|
|
@ -838,7 +853,7 @@ class LLMCachingHandler:
|
|||
if (
|
||||
call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value
|
||||
) and isinstance(cached_result, dict):
|
||||
if kwargs.get("stream", False) is True:
|
||||
if _stream_replay_requested(kwargs):
|
||||
cached_result = self._convert_cached_stream_response(
|
||||
cached_result=cached_result,
|
||||
call_type=call_type,
|
||||
|
|
@ -893,7 +908,7 @@ class LLMCachingHandler:
|
|||
elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict):
|
||||
use_chat_completion_cache: Final = _is_chat_completion_cached_dict(cached_result)
|
||||
if use_chat_completion_cache:
|
||||
if kwargs.get("stream", False) is True:
|
||||
if _stream_replay_requested(kwargs):
|
||||
bridge_call_type: Final = (
|
||||
CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value
|
||||
)
|
||||
|
|
@ -921,7 +936,7 @@ class LLMCachingHandler:
|
|||
):
|
||||
response_obj._hidden_params["cache_hit"] = True
|
||||
|
||||
if kwargs.get("stream", False) is True:
|
||||
if _stream_replay_requested(kwargs):
|
||||
cached_result = CachedResponsesAPIStreamingIterator(
|
||||
response=response_obj,
|
||||
logging_obj=logging_obj,
|
||||
|
|
|
|||
|
|
@ -1777,6 +1777,7 @@ DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_
|
|||
LENGTH_OF_LITELLM_GENERATED_KEY: Final = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
|
||||
MINIMUM_CUSTOM_KEY_LENGTH: Final = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
|
||||
SECRET_MANAGER_REFRESH_INTERVAL: Final = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
|
||||
OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS: Final = frozenset({"openai", "azure"})
|
||||
LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
|
||||
"default_internal_user_params",
|
||||
"default_team_params",
|
||||
|
|
@ -1794,6 +1795,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
|
|||
# test_general_settings_ui_fields_are_db_overridable enforces that pairing.
|
||||
"enable_anthropic_prompt_caching",
|
||||
"anthropic_prompt_caching_ttl",
|
||||
"openai_system_messages_first",
|
||||
"max_ui_session_budget",
|
||||
"budget_rollover",
|
||||
"mcp_tool_search",
|
||||
|
|
|
|||
|
|
@ -39,6 +39,8 @@ def is_serializable(value):
|
|||
|
||||
|
||||
class LangsmithLogger(CustomBatchLogger):
|
||||
preserve_events_added_during_flush = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
langsmith_api_key: str | None = None,
|
||||
|
|
|
|||
|
|
@ -238,6 +238,8 @@ def get_llm_provider(
|
|||
if dynamic_api_key is not None and not isinstance(dynamic_api_key, str):
|
||||
raise Exception(f"dynamic_api_key needs to be a string. Got type={type(dynamic_api_key).__name__}")
|
||||
return model, custom_llm_provider, dynamic_api_key, api_base
|
||||
if "/" in model and is_registered_custom_provider(provider_prefix):
|
||||
return model.split("/", 1)[1], provider_prefix, dynamic_api_key, api_base
|
||||
# check if api base is a known openai compatible endpoint
|
||||
if api_base:
|
||||
for endpoint in litellm.openai_compatible_endpoints:
|
||||
|
|
@ -536,6 +538,10 @@ def get_llm_provider(
|
|||
)
|
||||
|
||||
|
||||
def is_registered_custom_provider(custom_llm_provider: str | None) -> bool:
|
||||
return any(item["provider"] == custom_llm_provider for item in litellm.custom_provider_map)
|
||||
|
||||
|
||||
def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
return litellm.QwenCloudChatConfig()
|
||||
|
|
|
|||
|
|
@ -2256,6 +2256,22 @@ def drop_tool_reference_parts_from_tool_messages(
|
|||
return [_drop_tool_reference_parts(message) for message in messages] # mutable-ok: pipelines mutate message lists
|
||||
|
||||
|
||||
INSTRUCTION_MESSAGE_ROLES: Final = frozenset({"system", "developer"})
|
||||
|
||||
|
||||
def _is_instruction_message(message: AllMessageValues) -> bool:
|
||||
return message.get("role") in INSTRUCTION_MESSAGE_ROLES
|
||||
|
||||
|
||||
def system_messages_first(
|
||||
messages: list[AllMessageValues], # mutable-ok: message pipelines type messages as mutable lists
|
||||
) -> list[AllMessageValues]: # mutable-ok: message pipelines type messages as mutable lists
|
||||
return [ # mutable-ok: pipelines mutate message lists
|
||||
*(message for message in messages if _is_instruction_message(message)),
|
||||
*(message for message in messages if not _is_instruction_message(message)),
|
||||
]
|
||||
|
||||
|
||||
def _attempt_json_repair(s: str) -> object | None:
|
||||
"""
|
||||
Attempt to repair truncated JSON produced by LLM tool calls.
|
||||
|
|
|
|||
|
|
@ -1500,19 +1500,33 @@ def convert_to_gemini_tool_call_result(
|
|||
return _part
|
||||
|
||||
|
||||
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
|
||||
"""
|
||||
Sanitize tool_use_id to match Anthropic's required pattern: ^[a-zA-Z0-9_-]+$
|
||||
_TOOL_USE_ID_FALLBACK: Final = "tool_use_id"
|
||||
_ANTHROPIC_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
_BEDROCK_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_.:-]")
|
||||
_BEDROCK_TOOL_USE_ID_MAX_LEN: Final = 64
|
||||
_BEDROCK_TOOL_USE_ID_HASH_LEN: Final = 8
|
||||
|
||||
Anthropic requires tool_use_id to only contain alphanumeric characters, underscores, and hyphens.
|
||||
This function replaces any invalid characters with underscores.
|
||||
|
||||
def _replace_invalid_tool_use_id_chars(tool_use_id: str, invalid_chars: re.Pattern[str]) -> str:
|
||||
return invalid_chars.sub("_", tool_use_id) or _TOOL_USE_ID_FALLBACK
|
||||
|
||||
|
||||
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
|
||||
"""Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$."""
|
||||
return _replace_invalid_tool_use_id_chars(tool_use_id, _ANTHROPIC_TOOL_USE_ID_INVALID_CHARS)
|
||||
|
||||
|
||||
def _sanitize_bedrock_tool_use_id(tool_use_id: str) -> str:
|
||||
"""
|
||||
# Replace any character that's not alphanumeric, underscore, or hyphen with underscore
|
||||
sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", tool_use_id)
|
||||
# Ensure it's not empty (fallback to a default if needed)
|
||||
if not sanitized:
|
||||
sanitized = "tool_use_id"
|
||||
return sanitized
|
||||
Bedrock Converse requires toolUseId to match [a-zA-Z0-9_.:-]+ and be at most 64 chars.
|
||||
Ids that need rewriting get a short hash of the original appended so two ids that only
|
||||
differ in a replaced char or past the cut still map to distinct values.
|
||||
"""
|
||||
sanitized: Final = _replace_invalid_tool_use_id_chars(tool_use_id, _BEDROCK_TOOL_USE_ID_INVALID_CHARS)
|
||||
if sanitized == tool_use_id and len(sanitized) <= _BEDROCK_TOOL_USE_ID_MAX_LEN:
|
||||
return sanitized
|
||||
digest: Final = hashlib.sha256(tool_use_id.encode()).hexdigest()[:_BEDROCK_TOOL_USE_ID_HASH_LEN]
|
||||
return f"{sanitized[: _BEDROCK_TOOL_USE_ID_MAX_LEN - _BEDROCK_TOOL_USE_ID_HASH_LEN - 1]}_{digest}"
|
||||
|
||||
|
||||
_ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES: Final = {"application/pdf", "text/plain"}
|
||||
|
|
@ -3661,7 +3675,9 @@ def _convert_to_bedrock_tool_call_invoke(
|
|||
if parsed_objects:
|
||||
# First object keeps the original tool id.
|
||||
for obj_idx, obj in enumerate(parsed_objects):
|
||||
block_id = tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
|
||||
block_id = _sanitize_bedrock_tool_use_id(
|
||||
tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
|
||||
)
|
||||
bedrock_tool = BedrockToolUseBlock(input=obj, name=name, toolUseId=block_id)
|
||||
_parts_list.append(BedrockContentBlock(toolUse=bedrock_tool))
|
||||
# cache_control applies to the whole original
|
||||
|
|
@ -3678,7 +3694,9 @@ def _convert_to_bedrock_tool_call_invoke(
|
|||
# Fallback: no objects extracted — use empty dict.
|
||||
arguments_dict = {}
|
||||
|
||||
bedrock_tool = BedrockToolUseBlock(input=arguments_dict, name=name, toolUseId=tool_id)
|
||||
bedrock_tool = BedrockToolUseBlock(
|
||||
input=arguments_dict, name=name, toolUseId=_sanitize_bedrock_tool_use_id(tool_id)
|
||||
)
|
||||
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
|
||||
_parts_list.append(bedrock_content_block)
|
||||
|
||||
|
|
@ -3849,7 +3867,7 @@ def _convert_to_bedrock_tool_call_result(
|
|||
tool_result_content_blocks, used_search_results = _build_bedrock_tool_result_content_blocks(message)
|
||||
|
||||
message.get("name", "")
|
||||
id: Final = str(message.get("tool_call_id", str(uuid.uuid4())))
|
||||
id: Final = _sanitize_bedrock_tool_use_id(str(message.get("tool_call_id", str(uuid.uuid4()))))
|
||||
|
||||
tool_result: Final = BedrockToolResultBlock(content=tool_result_content_blocks, toolUseId=id)
|
||||
if used_search_results:
|
||||
|
|
|
|||
|
|
@ -539,6 +539,13 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
value: Final = litellm.model_cost.get(model, {}).get(key)
|
||||
return value if isinstance(value, bool) else None
|
||||
|
||||
@staticmethod
|
||||
def supports_fast_mode(model: str, custom_llm_provider: str) -> bool:
|
||||
return (
|
||||
custom_llm_provider == "anthropic"
|
||||
and AnthropicModelInfo._get_exact_model_capability(model, "supports_fast_mode") is True
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_provider_resolved_capability(model: str, key: str, custom_llm_provider: str) -> bool | None:
|
||||
"""Resolve boolean capability ``key`` for ``model`` under the caller's provider.
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
drop_tool_reference_parts_from_tool_messages,
|
||||
flatten_combinators_and_drop_non_python_regex_patterns,
|
||||
hoist_images_from_tool_messages,
|
||||
system_messages_first,
|
||||
tool_with_sanitized_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -276,7 +277,8 @@ class AzureOpenAIConfig(BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages)
|
||||
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
|
||||
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
|
||||
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
|
||||
return {
|
||||
"model": model,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from urllib.parse import urlparse
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
_extract_reasoning_content,
|
||||
|
|
@ -24,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
flatten_combinators_and_drop_non_python_regex_patterns,
|
||||
get_tool_call_names,
|
||||
hoist_images_from_tool_messages,
|
||||
system_messages_first,
|
||||
tool_with_sanitized_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
|
|
@ -463,6 +465,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
]
|
||||
return MappingProxyType({"tools": sanitized})
|
||||
|
||||
def _prompt_cache_ordered_messages(
|
||||
self, messages: list[AllMessageValues], litellm_params: Mapping[str, object]
|
||||
) -> list[AllMessageValues]:
|
||||
if not litellm.openai_system_messages_first:
|
||||
return messages
|
||||
if litellm_params.get("custom_llm_provider") not in OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS:
|
||||
return messages
|
||||
return system_messages_first(messages)
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -477,7 +488,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
Returns:
|
||||
dict: The transformed request. Sent as the body of the API call.
|
||||
"""
|
||||
messages = self._transform_messages(messages=messages, model=model)
|
||||
messages = self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model
|
||||
)
|
||||
if not self._should_preserve_cache_control_for_endpoint(
|
||||
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
|
||||
):
|
||||
|
|
@ -506,7 +519,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
transformed_messages = await self._transform_messages(messages=messages, model=model, is_async=True)
|
||||
transformed_messages = await self._transform_messages(
|
||||
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model, is_async=True
|
||||
)
|
||||
if not self._should_preserve_cache_control_for_endpoint(
|
||||
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
|
||||
):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -15226,6 +15226,17 @@
|
|||
"title": "Jwt Claim Value",
|
||||
"type": "string"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"key": {
|
||||
"title": "Key",
|
||||
"type": "string"
|
||||
|
|
@ -15310,6 +15321,17 @@
|
|||
"title": "Jwt Claim Value",
|
||||
"type": "string"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"updated_at": {
|
||||
"format": "date-time",
|
||||
"title": "Updated At",
|
||||
|
|
@ -15366,6 +15388,17 @@
|
|||
],
|
||||
"title": "Is Active"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"key": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -4499,12 +4499,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
|||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
key: str
|
||||
jwt_issuer: str | None = None
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
||||
id: str
|
||||
key: str | None = None
|
||||
jwt_issuer: str | None = None
|
||||
description: str | None = None
|
||||
is_active: bool | None = None
|
||||
|
||||
|
|
@ -4515,6 +4517,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
|
||||
id: str
|
||||
jwt_issuer: str | None = None
|
||||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
description: str | None = None
|
||||
|
|
|
|||
|
|
@ -150,6 +150,7 @@ class _PrismaDictableRow(Protocol):
|
|||
|
||||
class _PrismaJWTKeyMappingRow(Protocol):
|
||||
token: str
|
||||
jwt_issuer: str
|
||||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
|
||||
|
|
@ -3603,9 +3604,18 @@ async def _fetch_key_object_from_db_with_reconnect(
|
|||
raise
|
||||
|
||||
|
||||
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
|
||||
"""Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping."""
|
||||
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
|
||||
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
|
||||
"""Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
|
||||
issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
|
||||
|
||||
Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
|
||||
for one issuer's claim value can never be served to a different issuer whose claim
|
||||
value happens to collide. Unchanged for the global scope, keeping the single-issuer
|
||||
(no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
|
||||
"""
|
||||
if not jwt_issuer:
|
||||
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
|
||||
return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
|
|
@ -3617,7 +3627,7 @@ async def get_jwt_key_mapping_cache_keys_for_token(
|
|||
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
|
||||
where={"token": hashed_token}
|
||||
)
|
||||
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings)
|
||||
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
|
|
@ -3625,9 +3635,14 @@ async def get_jwt_key_mapping_object(
|
|||
jwt_claim_name: str,
|
||||
jwt_claim_value: str,
|
||||
prisma_client: PrismaClient,
|
||||
jwt_issuer: str | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Lookup a JWT-to-virtual-key mapping from the database.
|
||||
Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
|
||||
``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
|
||||
back to the global scope itself -- a caller that wants "issuer-scoped mapping,
|
||||
else the global one" queries both scopes itself, so each result can be cached
|
||||
under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
|
||||
|
||||
Returns the hashed token (str) if a matching active mapping is found, else None.
|
||||
"""
|
||||
|
|
@ -3635,6 +3650,7 @@ async def get_jwt_key_mapping_object(
|
|||
where={
|
||||
"jwt_claim_name": jwt_claim_name,
|
||||
"jwt_claim_value": jwt_claim_value,
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"is_active": True,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -269,6 +269,17 @@ class _TokenTeamModels(Protocol):
|
|||
def team_models(self) -> list[str]: ...
|
||||
|
||||
|
||||
class _RawCacheRead(Protocol):
|
||||
async def async_get_cache(self, *, key: str) -> object: ...
|
||||
|
||||
|
||||
def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
|
||||
"""View an untyped cache object's ``async_get_cache`` as returning ``object``
|
||||
instead of ``Any``, so a caller can ``isinstance``-narrow it without paying
|
||||
the ``reportAny`` cost of the underlying (unannotated) cache implementation."""
|
||||
return cache
|
||||
|
||||
|
||||
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
|
||||
return valid_token.team_models
|
||||
|
||||
|
|
@ -842,6 +853,7 @@ class _PendingAutoRegister(NamedTuple):
|
|||
claim_field: str
|
||||
claim_value: str
|
||||
cache_key: str
|
||||
jwt_issuer: str | None = None
|
||||
|
||||
|
||||
async def _auto_register_jwt_mapping(
|
||||
|
|
@ -853,6 +865,7 @@ async def _auto_register_jwt_mapping(
|
|||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
cache_key: str,
|
||||
jwt_issuer: str | None = None,
|
||||
team_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
org_id: str | None = None,
|
||||
|
|
@ -905,6 +918,7 @@ async def _auto_register_jwt_mapping(
|
|||
try:
|
||||
await prisma_client.db.litellm_jwtkeymapping.create(
|
||||
data={
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"jwt_claim_name": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
"token": token_hash,
|
||||
|
|
@ -939,6 +953,7 @@ async def _auto_register_jwt_mapping(
|
|||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=jwt_issuer,
|
||||
)
|
||||
if token_hash is None:
|
||||
# The winner's mapping vanished between the unique-constraint
|
||||
|
|
@ -983,6 +998,43 @@ async def _auto_register_jwt_mapping(
|
|||
return auto_registered_key
|
||||
|
||||
|
||||
async def _lookup_jwt_mapping_token_hash(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
virtual_key_claim_field: str,
|
||||
claim_value: str,
|
||||
normalized_issuer: str | None,
|
||||
cache_key: str,
|
||||
ttl: float,
|
||||
) -> str | None:
|
||||
issuer_scoped: Final = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
if issuer_scoped is not None:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=issuer_scoped, ttl=ttl)
|
||||
return issuer_scoped
|
||||
if normalized_issuer is None:
|
||||
return None
|
||||
# Another issuer may have already resolved (and cached) this same
|
||||
# global mapping -- check its cache entry before re-querying the DB.
|
||||
global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, claim_value)
|
||||
cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key)
|
||||
if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__":
|
||||
return cached_global
|
||||
global_row: Final = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=None,
|
||||
)
|
||||
if global_row is not None:
|
||||
await user_api_key_cache.async_set_cache(key=global_cache_key, value=global_row, ttl=ttl)
|
||||
return global_row
|
||||
|
||||
|
||||
async def _resolve_jwt_to_virtual_key(
|
||||
jwt_claims: dict,
|
||||
jwt_handler: JWTHandler,
|
||||
|
|
@ -1041,7 +1093,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
)
|
||||
return None
|
||||
|
||||
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
|
||||
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer)
|
||||
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
|
||||
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
|
||||
cached_mapping: Final = (
|
||||
|
|
@ -1081,6 +1133,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
return None
|
||||
elif cached_mapping is not None:
|
||||
|
|
@ -1094,21 +1147,30 @@ async def _resolve_jwt_to_virtual_key(
|
|||
)
|
||||
|
||||
# Resolve the mapping from DB, or treat prisma_client=None as a definitive
|
||||
# miss (no DB → no mapping can exist → apply no-match policy below).
|
||||
token_hash: str | None = None
|
||||
if prisma_client is not None:
|
||||
token_hash = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=str(claim_value),
|
||||
# miss (no DB → no mapping can exist → apply no-match policy below). An
|
||||
# issuer-scoped row wins; falling back to the global (no-issuer) row keeps
|
||||
# mappings created before issuer scoping existed working for every issuer.
|
||||
# Each tier is cached under ITS OWN key (the global tier under the
|
||||
# issuer-less cache key, not under `cache_key`/this issuer's key) so that
|
||||
# updating or deleting either row invalidates exactly the cache entries it
|
||||
# can affect. Caching a global-row hit under the requesting issuer's key
|
||||
# would leave every OTHER issuer that had fallen back to that same global
|
||||
# mapping serving its stale token until TTL after the row changes.
|
||||
token_hash: Final = (
|
||||
await _lookup_jwt_mapping_token_hash(
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
if token_hash is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=token_hash,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
virtual_key_claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
normalized_issuer=normalized_issuer,
|
||||
cache_key=cache_key,
|
||||
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
|
||||
)
|
||||
if prisma_client is not None
|
||||
else None
|
||||
)
|
||||
|
||||
if token_hash is not None:
|
||||
return IdentityStore.key_from_principal(
|
||||
await IdentityStore(
|
||||
prisma_client,
|
||||
|
|
@ -1149,6 +1211,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
|
||||
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
|
||||
|
|
@ -1641,6 +1704,7 @@ async def _user_api_key_auth_builder(
|
|||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
cache_key=pending_auto_register.cache_key,
|
||||
jwt_issuer=pending_auto_register.jwt_issuer,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
org_id=org_id,
|
||||
|
|
|
|||
|
|
@ -490,7 +490,7 @@ lite codex exec "summarize the repo"
|
|||
|
||||
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
|
||||
|
||||
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
|
||||
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=<path>` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place.
|
||||
|
||||
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
|
||||
|
||||
|
|
|
|||
|
|
@ -4,15 +4,16 @@ import re
|
|||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import click
|
||||
import requests
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
|
||||
from .claude_settings import ClaudeSettingsError, install_statusline_script
|
||||
|
|
@ -65,6 +66,10 @@ _INSTALL_DOCS: Final[dict[str, str]] = {
|
|||
_HIDDEN_AGENTS: Final = frozenset({"pi"})
|
||||
|
||||
CODEX_PROXY_PROVIDER: Final = "litellm"
|
||||
CODEX_HOME_ENV: Final = "CODEX_HOME"
|
||||
CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json"
|
||||
_CODEX_BASE_INSTRUCTIONS_PATH: Final = Path(__file__).with_name("codex_base_instructions.md")
|
||||
_CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0
|
||||
|
||||
|
||||
class AgentRunError(Exception):
|
||||
|
|
@ -252,7 +257,7 @@ def agent_launch_args(command: str, base_url: str) -> list[str]:
|
|||
|
||||
|
||||
class ListedModel(BaseModel):
|
||||
"""The fields of a /v1/models entry that an OpenCode model entry is built from."""
|
||||
"""The fields of a /v1/models entry that an OpenCode or Codex model entry is built from."""
|
||||
|
||||
id: str
|
||||
mode: str | None = None
|
||||
|
|
@ -265,7 +270,7 @@ class _ModelListing(BaseModel):
|
|||
|
||||
|
||||
_MODEL_LISTING: Final = TypeAdapter(_ModelListing)
|
||||
_OPENCODE_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
|
||||
_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
|
||||
_NO_EXTRA_ENV: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
|
|
@ -274,6 +279,40 @@ class ModelSyncSkipped:
|
|||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelSyncArgs:
|
||||
"""CLI args, placed before the user's own, that hand an agent the synced model list."""
|
||||
|
||||
args: tuple[str, ...]
|
||||
|
||||
|
||||
ModelSyncResult: TypeAlias = Mapping[str, str] | ModelSyncArgs | ModelSyncSkipped
|
||||
|
||||
|
||||
def _chat_models(models: Sequence[ListedModel]) -> tuple[ListedModel, ...]:
|
||||
return tuple(m for m in models if m.mode is None or m.mode in _CHAT_MODES)
|
||||
|
||||
|
||||
def _fetch_model_listing(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
*,
|
||||
get: Callable[..., requests.Response],
|
||||
) -> tuple[ListedModel, ...] | ModelSyncSkipped:
|
||||
url: Final = base_url.rstrip("/") + "/v1/models"
|
||||
try:
|
||||
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
|
||||
except requests.RequestException as e:
|
||||
return ModelSyncSkipped(f"could not reach {url}: {e}")
|
||||
if resp.status_code != 200:
|
||||
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
|
||||
try:
|
||||
listing: Final = _MODEL_LISTING.validate_json(resp.content)
|
||||
except ValidationError:
|
||||
return ModelSyncSkipped(f"{url} returned an unexpected body")
|
||||
return listing.data
|
||||
|
||||
|
||||
class _OpenCodeLimit(BaseModel):
|
||||
context: int
|
||||
output: int
|
||||
|
|
@ -317,7 +356,7 @@ def opencode_provider_config(base_url: str, models: Sequence[ListedModel]) -> st
|
|||
it never lands in the config text. OpenCode merges this inline config over
|
||||
the user's own files, leaving unrelated keys and providers untouched.
|
||||
"""
|
||||
chat_models: Final = tuple(m for m in models if m.mode is None or m.mode in _OPENCODE_CHAT_MODES)
|
||||
chat_models: Final = _chat_models(models)
|
||||
provider: Final = _OpenCodeProvider(
|
||||
npm=OPENCODE_PROVIDER_NPM,
|
||||
name=OPENCODE_PROVIDER_NAME,
|
||||
|
|
@ -347,40 +386,269 @@ def opencode_model_sync_env(
|
|||
"""
|
||||
if OPENCODE_CONFIG_CONTENT_ENV in base_env:
|
||||
return ModelSyncSkipped(f"{OPENCODE_CONFIG_CONTENT_ENV} is already set")
|
||||
url: Final = base_url.rstrip("/") + "/v1/models"
|
||||
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
|
||||
if isinstance(listing, ModelSyncSkipped):
|
||||
return listing
|
||||
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing)})
|
||||
|
||||
|
||||
class _CodexTruncationPolicy(BaseModel):
|
||||
mode: Literal["bytes"] = "bytes"
|
||||
limit: int = 10_000
|
||||
|
||||
|
||||
class _CodexModel(BaseModel):
|
||||
"""One `ModelInfo` entry of a Codex model catalog for a model the installed Codex does not know.
|
||||
|
||||
Every field that some Codex release since `model_catalog_json` appeared
|
||||
(0.105.0) deserializes without a default is spelled out here, so one catalog
|
||||
parses on all of them; the values match the fallback metadata Codex uses for
|
||||
a model slug it does not know, so picking such a proxy model behaves the
|
||||
same as `codex -m` did.
|
||||
"""
|
||||
|
||||
slug: str
|
||||
display_name: str
|
||||
description: None = None
|
||||
supported_reasoning_levels: tuple[()] = ()
|
||||
shell_type: Literal["unified_exec"] = "unified_exec"
|
||||
visibility: Literal["list"] = "list"
|
||||
supported_in_api: Literal[True] = True
|
||||
priority: int
|
||||
availability_nux: None = None
|
||||
upgrade: None = None
|
||||
support_verbosity: Literal[False] = False
|
||||
supports_reasoning_summaries: Literal[False] = False
|
||||
supports_parallel_tool_calls: Literal[False] = False
|
||||
default_verbosity: None = None
|
||||
apply_patch_tool_type: None = None
|
||||
truncation_policy: _CodexTruncationPolicy = _CodexTruncationPolicy()
|
||||
experimental_supported_tools: tuple[()] = ()
|
||||
context_window: int | None
|
||||
base_instructions: str
|
||||
|
||||
|
||||
class _StockCodexUpgrade(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
model: str
|
||||
|
||||
|
||||
class _StockCodexModel(BaseModel):
|
||||
"""One `ModelInfo` entry as the installed Codex prints it from `codex debug models`.
|
||||
|
||||
Only the fields the sync rewrites are named; everything else that release
|
||||
knows about the model (its reasoning levels, prompt, tool support) rides
|
||||
along untouched, whatever the release's schema.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
slug: str
|
||||
priority: int
|
||||
visibility: str
|
||||
supported_in_api: bool = True
|
||||
upgrade: _StockCodexUpgrade | None = None
|
||||
|
||||
|
||||
class _StockCodexCatalog(BaseModel):
|
||||
models: tuple[_StockCodexModel, ...]
|
||||
|
||||
|
||||
class _CodexCatalog(BaseModel):
|
||||
models: tuple[_CodexModel | _StockCodexModel, ...]
|
||||
|
||||
|
||||
def _codex_catalog_entry(
|
||||
priority: int,
|
||||
listed: ListedModel,
|
||||
stock: _StockCodexModel | None,
|
||||
served: frozenset[str],
|
||||
instructions: str,
|
||||
) -> _CodexModel | _StockCodexModel:
|
||||
if stock is None:
|
||||
return _CodexModel(
|
||||
slug=listed.id,
|
||||
display_name=listed.id,
|
||||
priority=priority,
|
||||
context_window=listed.max_input_tokens,
|
||||
base_instructions=instructions,
|
||||
)
|
||||
upgrade: Final = stock.upgrade if stock.upgrade is not None and stock.upgrade.model in served else None
|
||||
return stock.model_copy(
|
||||
update={"priority": priority, "visibility": "list", "supported_in_api": True, "upgrade": upgrade}
|
||||
)
|
||||
|
||||
|
||||
def codex_model_catalog(
|
||||
models: Sequence[ListedModel], stock: Sequence[_StockCodexModel], instructions: str
|
||||
) -> str | None:
|
||||
"""The `model_catalog_json` body listing the proxy's chat models, or None if there are none.
|
||||
|
||||
Codex refuses an empty catalog, hence None instead of `{"models": []}`.
|
||||
Passing a catalog replaces Codex's built-in one, so a proxy model the
|
||||
installed Codex knows keeps that Codex's own entry and the proxy only
|
||||
decides its place in the picker: the listing orders it, lists it even when
|
||||
Codex hides it or keeps it off the API, and keeps Codex's upgrade nudge only
|
||||
when the model it points at is served too. A model Codex does not know gets the fallback
|
||||
entry, with the same base instructions Codex itself uses so the agent never
|
||||
runs without a system prompt.
|
||||
"""
|
||||
chat_models: Final = _chat_models(models)
|
||||
if not chat_models:
|
||||
return None
|
||||
served: Final = frozenset(m.id for m in chat_models)
|
||||
known: Final = MappingProxyType({m.slug: m for m in stock})
|
||||
catalog: Final = _CodexCatalog(
|
||||
models=tuple(
|
||||
_codex_catalog_entry(index, m, known.get(m.id), served, instructions) for index, m in enumerate(chat_models)
|
||||
)
|
||||
)
|
||||
return catalog.model_dump_json()
|
||||
|
||||
|
||||
def codex_model_catalog_path(env: Mapping[str, str], *, home: Callable[[], Path] = Path.home) -> Path:
|
||||
override: Final = env.get(CODEX_HOME_ENV)
|
||||
root: Final = Path(override) if override else home() / ".codex"
|
||||
return root / CODEX_MODEL_CATALOG_FILENAME
|
||||
|
||||
|
||||
def _replace_file(path: Path, text: str) -> None:
|
||||
with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as tmp:
|
||||
_ = tmp.write(text)
|
||||
try:
|
||||
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
|
||||
except requests.RequestException as e:
|
||||
return ModelSyncSkipped(f"could not reach {url}: {e}")
|
||||
if resp.status_code != 200:
|
||||
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
|
||||
os.replace(tmp.name, path)
|
||||
except OSError:
|
||||
Path(tmp.name).unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def _codex_debug_models(
|
||||
binary: str,
|
||||
args: Sequence[str],
|
||||
env: Mapping[str, str],
|
||||
*,
|
||||
run: Callable[..., subprocess.CompletedProcess[str]],
|
||||
) -> str | ModelSyncSkipped:
|
||||
"""What `codex debug models` prints with `args` in front, or why the installed Codex could not run it.
|
||||
|
||||
The command prints the catalog Codex would launch with, without touching
|
||||
the network, so it lists the installed Codex's own models and parses a
|
||||
catalog override the way a launch does. Releases before 0.130.0 have no
|
||||
such command and are reported the same way. A batch shim goes through
|
||||
cmd.exe exactly as the launch will.
|
||||
"""
|
||||
name: Final = os.path.basename(binary)
|
||||
command: Final = _windows_command(binary, (binary, *args, "debug", "models"))
|
||||
try:
|
||||
listing: Final = _MODEL_LISTING.validate_json(resp.content)
|
||||
except ValidationError:
|
||||
return ModelSyncSkipped(f"{url} returned an unexpected body")
|
||||
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing.data)})
|
||||
completed: Final = run(
|
||||
command,
|
||||
env=dict(env),
|
||||
stdin=subprocess.DEVNULL,
|
||||
capture_output=True,
|
||||
encoding="utf-8",
|
||||
timeout=_CODEX_PREFLIGHT_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
return ModelSyncSkipped(f"`{name} debug models` failed: {e}")
|
||||
if completed.returncode == 0:
|
||||
return completed.stdout
|
||||
lines: Final = completed.stderr.strip().splitlines()
|
||||
detail: Final = lines[0] if lines else "no output"
|
||||
return ModelSyncSkipped(f"`{name} debug models` exited {completed.returncode}: {detail}")
|
||||
|
||||
|
||||
def _stock_codex_models(
|
||||
binary: str, env: Mapping[str, str], *, run: Callable[..., subprocess.CompletedProcess[str]]
|
||||
) -> tuple[_StockCodexModel, ...] | ModelSyncSkipped:
|
||||
printed: Final = _codex_debug_models(binary, (), env, run=run)
|
||||
if isinstance(printed, ModelSyncSkipped):
|
||||
return printed
|
||||
try:
|
||||
return _StockCodexCatalog.model_validate_json(printed).models
|
||||
except ValidationError as e:
|
||||
name: Final = os.path.basename(binary)
|
||||
return ModelSyncSkipped(f"`{name} debug models` printed no model catalog: {e.errors()[0]['msg']}")
|
||||
|
||||
|
||||
def codex_model_sync_args(
|
||||
base_env: Mapping[str, str],
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
*,
|
||||
binary: str = "codex",
|
||||
get: Callable[..., requests.Response] = requests.get,
|
||||
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
home: Callable[[], Path] = Path.home,
|
||||
instructions_path: Path = _CODEX_BASE_INSTRUCTIONS_PATH,
|
||||
) -> ModelSyncArgs | ModelSyncSkipped:
|
||||
"""`-c model_catalog_json=...` pointing Codex at the proxy's model list, or why it was skipped.
|
||||
|
||||
Codex has no env or inline equivalent of OPENCODE_CONFIG_CONTENT: the catalog
|
||||
must be a file, so it is written under $CODEX_HOME (default ~/.codex) and
|
||||
atomically replaced on every launch. The Codex at `binary` first lists its
|
||||
own models, so the ones the proxy serves keep that Codex's entries, and then
|
||||
reads the file back once before it is handed over. The key never lands in
|
||||
the file. A failed fetch, read, listing, write or read-back is reported
|
||||
rather than raised: Codex still launches with its built-in catalog and takes
|
||||
a proxy model by name via -m, and a rejected file stays on disk to be looked
|
||||
at.
|
||||
"""
|
||||
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
|
||||
if isinstance(listing, ModelSyncSkipped):
|
||||
return listing
|
||||
try:
|
||||
instructions: Final = instructions_path.read_text(encoding="utf-8")
|
||||
except OSError as e:
|
||||
return ModelSyncSkipped(f"could not read {instructions_path}: {e}")
|
||||
path: Final = codex_model_catalog_path(base_env, home=home)
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as e:
|
||||
return ModelSyncSkipped(f"could not write {path}: {e}")
|
||||
stock: Final = _stock_codex_models(binary, base_env, run=run)
|
||||
if isinstance(stock, ModelSyncSkipped):
|
||||
return stock
|
||||
catalog: Final = codex_model_catalog(listing, stock, instructions)
|
||||
if catalog is None:
|
||||
return ModelSyncSkipped(f"{base_url.rstrip('/')}/v1/models lists no chat models")
|
||||
try:
|
||||
_replace_file(path, catalog)
|
||||
except OSError as e:
|
||||
return ModelSyncSkipped(f"could not write {path}: {e}")
|
||||
override: Final = f"model_catalog_json={json.dumps(str(path))}"
|
||||
read_back: Final = _codex_debug_models(binary, ("-c", override), base_env, run=run)
|
||||
if isinstance(read_back, ModelSyncSkipped):
|
||||
return read_back
|
||||
return ModelSyncArgs(("-c", override))
|
||||
|
||||
|
||||
def agent_model_sync_env(
|
||||
command: str,
|
||||
binary: str,
|
||||
base_env: Mapping[str, str],
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
skip_verify: bool,
|
||||
*,
|
||||
get: Callable[..., requests.Response] = requests.get,
|
||||
) -> Mapping[str, str] | ModelSyncSkipped:
|
||||
"""Extra env an agent needs to see the proxy's model list.
|
||||
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
||||
) -> ModelSyncResult:
|
||||
"""Extra env or args an agent needs to see the proxy's model list.
|
||||
|
||||
Only OpenCode needs one: Claude Code discovers models through
|
||||
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY and Codex takes the model by name.
|
||||
skip_verify means the caller wants no pre-launch proxy call at all, so the
|
||||
listing is skipped too rather than hanging on an offline proxy.
|
||||
binary is the resolved path the launch will run (`codex.cmd` on a Windows
|
||||
npm install). OpenCode takes the list as env, Codex as a `-c` override that
|
||||
binary has read back first; Claude Code discovers models itself through
|
||||
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY. skip_verify means the caller
|
||||
wants no pre-launch proxy call at all, so the listing is skipped too rather
|
||||
than hanging on an offline proxy.
|
||||
"""
|
||||
if os.path.basename(command) != "opencode":
|
||||
agent: Final = os.path.splitext(os.path.basename(binary))[0]
|
||||
if agent not in ("opencode", "codex"):
|
||||
return _NO_EXTRA_ENV
|
||||
if skip_verify:
|
||||
return ModelSyncSkipped(f"{_SKIP_VERIFY_FLAG} was passed")
|
||||
if agent == "codex":
|
||||
return codex_model_sync_args(base_env, base_url, api_key, binary=binary, get=get, run=run)
|
||||
return opencode_model_sync_env(base_env, base_url, api_key, get=get)
|
||||
|
||||
|
||||
|
|
@ -508,9 +776,7 @@ def run_agent(
|
|||
base_env: Mapping[str, str] | None = None,
|
||||
which: Callable[[str], str | None] = shutil.which,
|
||||
verify: Callable[[str, str], None] = verify_proxy_key,
|
||||
sync_models: Callable[[str, Mapping[str, str], str, str, bool], Mapping[str, str] | ModelSyncSkipped] = (
|
||||
agent_model_sync_env
|
||||
),
|
||||
sync_models: Callable[[str, Mapping[str, str], str, str, bool], ModelSyncResult] = agent_model_sync_env,
|
||||
warn: Callable[[str], None] = _warn,
|
||||
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
|
||||
reattach_terminal: Callable[[], None] | None = None,
|
||||
|
|
@ -537,7 +803,7 @@ def run_agent(
|
|||
verify(base_url, api_key)
|
||||
|
||||
env_before_sync: Final = base_env if base_env is not None else os.environ
|
||||
synced: Final = sync_models(command[0], env_before_sync, base_url, api_key, skip_verify)
|
||||
synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
|
||||
if isinstance(synced, ModelSyncSkipped):
|
||||
warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
|
||||
|
||||
|
|
@ -547,10 +813,11 @@ def run_agent(
|
|||
env: Final = MappingProxyType(
|
||||
{
|
||||
**build_agent_env(env_before_sync, base_url, api_key, profiles),
|
||||
**(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
|
||||
**(synced if isinstance(synced, Mapping) else _NO_EXTRA_ENV),
|
||||
}
|
||||
)
|
||||
extra_args: Final = (*agent_launch_args(command[0], base_url), *prepared_args)
|
||||
synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else ()
|
||||
extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args)
|
||||
if reattach_terminal is not None:
|
||||
reattach_terminal()
|
||||
launcher(binary, [command[0], *extra_args, *command[1:]], env)
|
||||
|
|
|
|||
275
litellm/proxy/client/cli/commands/codex_base_instructions.md
Normal file
275
litellm/proxy/client/cli/commands/codex_base_instructions.md
Normal file
|
|
@ -0,0 +1,275 @@
|
|||
You are a coding agent running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful.
|
||||
|
||||
Your capabilities:
|
||||
|
||||
- Receive user prompts and other context provided by the harness, such as files in the workspace.
|
||||
- Communicate with the user by streaming thinking & responses, and by making & updating plans.
|
||||
- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section.
|
||||
|
||||
Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI).
|
||||
|
||||
# How you work
|
||||
|
||||
## Personality
|
||||
|
||||
Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work.
|
||||
|
||||
# AGENTS.md spec
|
||||
- Repos often contain AGENTS.md files. These files can appear anywhere within the repository.
|
||||
- These files are a way for humans to give you (the agent) instructions or tips for working within the container.
|
||||
- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code.
|
||||
- Instructions in AGENTS.md files:
|
||||
- The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it.
|
||||
- For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file.
|
||||
- Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise.
|
||||
- More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions.
|
||||
- Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions.
|
||||
- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable.
|
||||
|
||||
## Responsiveness
|
||||
|
||||
### Preamble messages
|
||||
|
||||
Before making tool calls, send a brief preamble to the user explaining what you’re about to do. When sending preamble messages, follow these principles and examples:
|
||||
|
||||
- **Logically group related actions**: if you’re about to run several related commands, describe them together in one preamble rather than sending a separate note for each.
|
||||
- **Keep it concise**: be no more than 1-2 sentences, focused on immediate, tangible next steps. (8–12 words for quick updates).
|
||||
- **Build on prior context**: if this is not your first tool call, use the preamble message to connect the dots with what’s been done so far and create a sense of momentum and clarity for the user to understand your next actions.
|
||||
- **Keep your tone light, friendly and curious**: add small touches of personality in preambles feel collaborative and engaging.
|
||||
- **Exception**: Avoid adding a preamble for every trivial read (e.g., `cat` a single file) unless it’s part of a larger grouped action.
|
||||
|
||||
**Examples:**
|
||||
|
||||
- “I’ve explored the repo; now checking the API route definitions.”
|
||||
- “Next, I’ll patch the config and update the related tests.”
|
||||
- “I’m about to scaffold the CLI commands and helper functions.”
|
||||
- “Ok cool, so I’ve wrapped my head around the repo. Now digging into the API routes.”
|
||||
- “Config’s looking tidy. Next up is patching helpers to keep things in sync.”
|
||||
- “Finished poking at the DB gateway. I will now chase down error handling.”
|
||||
- “Alright, build pipeline order is interesting. Checking how it reports failures.”
|
||||
- “Spotted a clever caching util; now hunting where it gets used.”
|
||||
|
||||
## Planning
|
||||
|
||||
You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go.
|
||||
|
||||
Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately.
|
||||
|
||||
Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step.
|
||||
|
||||
Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so.
|
||||
|
||||
Use a plan when:
|
||||
|
||||
- The task is non-trivial and will require multiple actions over a long time horizon.
|
||||
- There are logical phases or dependencies where sequencing matters.
|
||||
- The work has ambiguity that benefits from outlining high-level goals.
|
||||
- You want intermediate checkpoints for feedback and validation.
|
||||
- When the user asked you to do more than one thing in a single prompt
|
||||
- The user has asked you to use the plan tool (aka "TODOs")
|
||||
- You generate additional steps while working, and plan to do them before yielding to the user
|
||||
|
||||
### Examples
|
||||
|
||||
**High-quality plans**
|
||||
|
||||
Example 1:
|
||||
|
||||
1. Add CLI entry with file args
|
||||
2. Parse Markdown via CommonMark library
|
||||
3. Apply semantic HTML template
|
||||
4. Handle code blocks, images, links
|
||||
5. Add error handling for invalid files
|
||||
|
||||
Example 2:
|
||||
|
||||
1. Define CSS variables for colors
|
||||
2. Add toggle with localStorage state
|
||||
3. Refactor components to use variables
|
||||
4. Verify all views for readability
|
||||
5. Add smooth theme-change transition
|
||||
|
||||
Example 3:
|
||||
|
||||
1. Set up Node.js + WebSocket server
|
||||
2. Add join/leave broadcast events
|
||||
3. Implement messaging with timestamps
|
||||
4. Add usernames + mention highlighting
|
||||
5. Persist messages in lightweight DB
|
||||
6. Add typing indicators + unread count
|
||||
|
||||
**Low-quality plans**
|
||||
|
||||
Example 1:
|
||||
|
||||
1. Create CLI tool
|
||||
2. Add Markdown parser
|
||||
3. Convert to HTML
|
||||
|
||||
Example 2:
|
||||
|
||||
1. Add dark mode toggle
|
||||
2. Save preference
|
||||
3. Make styles look good
|
||||
|
||||
Example 3:
|
||||
|
||||
1. Create single-file HTML game
|
||||
2. Run quick sanity check
|
||||
3. Summarize usage instructions
|
||||
|
||||
If you need to write a plan, only write high quality plans, not low quality ones.
|
||||
|
||||
## Task execution
|
||||
|
||||
You are a coding agent. Please keep going until the query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer.
|
||||
|
||||
You MUST adhere to the following criteria when solving queries:
|
||||
|
||||
- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
|
||||
- Analyzing code for vulnerabilities is allowed.
|
||||
- Showing user code and tool call details is allowed.
|
||||
- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`): {"command":["apply_patch","*** Begin Patch\\n*** Update File: path/to/file.py\\n@@ def example():\\n- pass\\n+ return 123\\n*** End Patch"]}
|
||||
|
||||
If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines:
|
||||
|
||||
- Fix the problem at the root cause rather than applying surface-level patches, when possible.
|
||||
- Avoid unneeded complexity in your solution.
|
||||
- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
|
||||
- Update documentation as necessary.
|
||||
- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
|
||||
- Use `git log` and `git blame` to search the history of the codebase if additional context is required.
|
||||
- NEVER add copyright or license headers unless specifically requested.
|
||||
- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc.
|
||||
- Do not `git commit` your changes or create new git branches unless explicitly requested.
|
||||
- Do not add inline comments within code unless explicitly requested.
|
||||
- Do not use one-letter variable names unless explicitly requested.
|
||||
- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor.
|
||||
|
||||
## Validating your work
|
||||
|
||||
If the codebase has tests or the ability to build or run, consider using them to verify that your work is complete.
|
||||
|
||||
When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests.
|
||||
|
||||
Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one.
|
||||
|
||||
For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
|
||||
|
||||
Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance:
|
||||
|
||||
- When running in the non-interactive approval mode **never**, proactively run tests, lint and do whatever you need to ensure you've completed the task.
|
||||
- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first.
|
||||
- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task.
|
||||
|
||||
## Ambition vs. precision
|
||||
|
||||
For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation.
|
||||
|
||||
If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature.
|
||||
|
||||
You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified.
|
||||
|
||||
## Sharing progress updates
|
||||
|
||||
For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next.
|
||||
|
||||
Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why.
|
||||
|
||||
The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along.
|
||||
|
||||
## Presenting your work and final message
|
||||
|
||||
Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the user’s style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges.
|
||||
|
||||
You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation.
|
||||
|
||||
The user is working on the same computer as you, and has access to your work. As such there's no need to show the full contents of large files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path.
|
||||
|
||||
If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If there’s something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly.
|
||||
|
||||
Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding.
|
||||
|
||||
### Final answer structure and style guidelines
|
||||
|
||||
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
|
||||
|
||||
**Section Headers**
|
||||
|
||||
- Use only when they improve clarity — they are not mandatory for every answer.
|
||||
- Choose descriptive names that fit the content
|
||||
- Keep headers short (1–3 words) and in `**Title Case**`. Always start headers with `**` and end with `**`
|
||||
- Leave no blank line before the first bullet under a header.
|
||||
- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer.
|
||||
|
||||
**Bullets**
|
||||
|
||||
- Use `-` followed by a space for every bullet.
|
||||
- Merge related points when possible; avoid a bullet for every trivial detail.
|
||||
- Keep bullets to one line unless breaking for clarity is unavoidable.
|
||||
- Group into short lists (4–6 bullets) ordered by importance.
|
||||
- Use consistent keyword phrasing and formatting across sections.
|
||||
|
||||
**Monospace**
|
||||
|
||||
- Wrap all commands, file paths, env vars, and code identifiers in backticks (`` `...` ``).
|
||||
- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command.
|
||||
- Never mix monospace and bold markers; choose one based on whether it’s a keyword (`**`) or inline code/path (`` ` ``).
|
||||
|
||||
**File References**
|
||||
When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
|
||||
* Use inline code to make file paths clickable.
|
||||
* Each reference should have a stand alone path. Even if it's the same file.
|
||||
* Accepted: absolute, workspace‑relative, a/ or b/ diff prefixes, or bare filename/suffix.
|
||||
* Line/column (1‑based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
|
||||
* Do not use URIs like file://, vscode://, or https://.
|
||||
* Do not provide range of lines
|
||||
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5
|
||||
|
||||
**Structure**
|
||||
|
||||
- Place related bullets together; don’t mix unrelated concepts in the same section.
|
||||
- Order sections from general → specific → supporting info.
|
||||
- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it.
|
||||
- Match structure to complexity:
|
||||
- Multi-part or detailed results → use clear headers and grouped bullets.
|
||||
- Simple results → minimal headers, possibly just a short list or paragraph.
|
||||
|
||||
**Tone**
|
||||
|
||||
- Keep the voice collaborative and natural, like a coding partner handing off work.
|
||||
- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition
|
||||
- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”).
|
||||
- Keep descriptions self-contained; don’t refer to “above” or “below”.
|
||||
- Use parallel structure in lists for consistency.
|
||||
|
||||
**Don’t**
|
||||
|
||||
- Don’t use literal words “bold” or “monospace” in the content.
|
||||
- Don’t nest bullets or create deep hierarchies.
|
||||
- Don’t output ANSI escape codes directly — the CLI renderer applies them.
|
||||
- Don’t cram unrelated keywords into a single bullet; split for clarity.
|
||||
- Don’t let keyword lists run long — wrap or reformat for scanability.
|
||||
|
||||
Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with what’s needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable.
|
||||
|
||||
For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting.
|
||||
|
||||
# Tool Guidelines
|
||||
|
||||
## Shell commands
|
||||
|
||||
When using the shell, you must adhere to the following guidelines:
|
||||
|
||||
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
|
||||
- Do not use python scripts to attempt to output larger chunks of a file.
|
||||
|
||||
## `update_plan`
|
||||
|
||||
A tool named `update_plan` is available to you. You can use it to keep an up‑to‑date, step‑by‑step plan for the task.
|
||||
|
||||
To create a new plan, call `update_plan` with a short list of 1‑sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`).
|
||||
|
||||
When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call.
|
||||
|
||||
If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`.
|
||||
|
|
@ -189,8 +189,9 @@ async def _read_request_body(request: Request | None) -> dict:
|
|||
|
||||
try:
|
||||
parsed_body = json.loads(body_str)
|
||||
except json.JSONDecodeError:
|
||||
# If both orjson and json.loads fail, throw a proper error
|
||||
json.dumps(parsed_body, ensure_ascii=False).encode("utf-8")
|
||||
except (json.JSONDecodeError, UnicodeEncodeError):
|
||||
# json.loads accepts lone surrogate escapes that no provider can encode
|
||||
verbose_proxy_logger.error("Invalid JSON payload received: %s", e)
|
||||
raise ProxyException(
|
||||
message=f"Invalid JSON payload: {e}",
|
||||
|
|
|
|||
|
|
@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol):
|
|||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def jwt_issuer(self) -> str: ...
|
||||
|
||||
@property
|
||||
def jwt_claim_name(self) -> str: ...
|
||||
|
||||
|
|
@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
|
|||
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
|
||||
return JWTKeyMappingResponse(
|
||||
id=mapping.id,
|
||||
jwt_issuer=mapping.jwt_issuer or None,
|
||||
jwt_claim_name=mapping.jwt_claim_name,
|
||||
jwt_claim_value=mapping.jwt_claim_value,
|
||||
description=mapping.description,
|
||||
|
|
@ -109,6 +113,7 @@ async def create_jwt_key_mapping(
|
|||
try:
|
||||
hashed_key: Final = hash_token(data.key)
|
||||
create_data: Final = {
|
||||
"jwt_issuer": data.jwt_issuer or "",
|
||||
"jwt_claim_name": data.jwt_claim_name,
|
||||
"jwt_claim_value": data.jwt_claim_value,
|
||||
"token": hashed_key,
|
||||
|
|
@ -120,7 +125,7 @@ async def create_jwt_key_mapping(
|
|||
|
||||
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
|
||||
|
||||
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
|
||||
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer)
|
||||
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
return _to_response(new_mapping)
|
||||
|
|
@ -131,7 +136,10 @@ async def create_jwt_key_mapping(
|
|||
if "unique" in error_str or "p2002" in error_str:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
|
||||
detail=(
|
||||
f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' "
|
||||
f"already exists for issuer '{data.jwt_issuer}'."
|
||||
),
|
||||
)
|
||||
if "foreign" in error_str or "p2003" in error_str:
|
||||
raise HTTPException(
|
||||
|
|
@ -161,6 +169,9 @@ async def update_jwt_key_mapping(
|
|||
update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
|
||||
if data.key is not None:
|
||||
update_data["token"] = hash_token(data.key)
|
||||
if "jwt_issuer" in update_data:
|
||||
# DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
|
||||
update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""
|
||||
update_data["updated_by"] = user_api_key_dict.user_id
|
||||
|
||||
try:
|
||||
|
|
@ -178,9 +189,11 @@ async def update_jwt_key_mapping(
|
|||
# Evict only after the write commits: a concurrent request between an
|
||||
# early eviction and the commit would re-cache the old mapping and keep
|
||||
# it authorized until TTL.
|
||||
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
|
||||
old_cache_key: Final = jwt_key_mapping_cache_key(
|
||||
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
|
||||
)
|
||||
new_cache_key: Final = jwt_key_mapping_cache_key(
|
||||
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
|
||||
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer
|
||||
)
|
||||
cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key)
|
||||
await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
|
||||
|
|
@ -227,7 +240,9 @@ async def delete_jwt_key_mapping(
|
|||
|
||||
# Evict only after the row is gone, else a concurrent request can
|
||||
# re-cache the deleted mapping and keep it authorized until TTL.
|
||||
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
|
||||
cache_key: Final = jwt_key_mapping_cache_key(
|
||||
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
|
||||
)
|
||||
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
return {"status": "success"}
|
||||
except HTTPException:
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
|
|||
from litellm.secret_managers.main import get_secret_str, str_to_bool
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
|
||||
|
|
@ -1322,7 +1323,7 @@ def _resolve_vertex_model_from_router(
|
|||
endpoint: str,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
) -> tuple[str, str, str | None, str | None]:
|
||||
) -> tuple[str, str, str | None, str | None, Mapping[str, object] | None]:
|
||||
"""
|
||||
Resolve Vertex AI model configuration from router.
|
||||
|
||||
|
|
@ -1335,18 +1336,21 @@ def _resolve_vertex_model_from_router(
|
|||
vertex_location: Current vertex location (may be from URL)
|
||||
|
||||
Returns:
|
||||
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location)
|
||||
with resolved values from router config
|
||||
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info)
|
||||
with resolved values from router config; deployment_model_info is the resolved
|
||||
deployment's `model_info`, or None when no deployment matched
|
||||
"""
|
||||
if not llm_router:
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
|
||||
|
||||
try:
|
||||
deployment: Final = llm_router.get_available_deployment_for_pass_through(model=model_id)
|
||||
if not deployment:
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
|
||||
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
model_info: Final = deployment.get("model_info")
|
||||
deployment_model_info: Final = model_info if isinstance(model_info, Mapping) else None
|
||||
|
||||
# Always override with router config values (they take precedence over URL values)
|
||||
config_vertex_project: Final = litellm_params.get("vertex_project")
|
||||
|
|
@ -1387,10 +1391,11 @@ def _resolve_vertex_model_from_router(
|
|||
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
|
||||
endpoint = endpoint.replace(model_id, actual_model)
|
||||
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Error resolving vertex model from router for model %s: %s", model_id, e)
|
||||
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
|
||||
|
||||
|
||||
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
|
||||
|
|
@ -2134,6 +2139,7 @@ async def _base_vertex_proxy_route(
|
|||
endpoint,
|
||||
vertex_project,
|
||||
vertex_location,
|
||||
deployment_model_info,
|
||||
) = _resolve_vertex_model_from_router(
|
||||
model_id=model_id,
|
||||
llm_router=llm_router,
|
||||
|
|
@ -2142,6 +2148,8 @@ async def _base_vertex_proxy_route(
|
|||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
if deployment_model_info:
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, deployment_model_info)
|
||||
|
||||
vertex_credentials: Final = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ from litellm.secret_managers.main import get_secret_str
|
|||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
EndpointType,
|
||||
|
|
@ -613,6 +614,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
_metadata.update(
|
||||
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
|
||||
)
|
||||
_request_state: Final = getattr(request, "state", None)
|
||||
deployment_model_info: Final = getattr(
|
||||
_request_state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, None
|
||||
)
|
||||
if isinstance(deployment_model_info, Mapping):
|
||||
_metadata["model_info"] = dict(deployment_model_info)
|
||||
|
||||
kwargs: Final = {
|
||||
"litellm_params": {
|
||||
|
|
@ -2002,6 +2009,8 @@ def create_pass_through_route(
|
|||
delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
|
||||
if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
|
||||
delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
|
||||
if hasattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY):
|
||||
delattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY)
|
||||
|
||||
# The upstream withholds its response headers until its first token, so
|
||||
# the whole time-to-first-token is spent inside _relay with nothing on
|
||||
|
|
|
|||
|
|
@ -17467,6 +17467,16 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie
|
|||
"tab": "prompt_caching",
|
||||
"description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.",
|
||||
},
|
||||
"openai_system_messages_first": {
|
||||
"type": "Boolean",
|
||||
"tab": "prompt_caching",
|
||||
"description": (
|
||||
"Moves system and developer messages to the front of the messages array on OpenAI and "
|
||||
"Azure OpenAI chat completions requests, keeping their relative order. OpenAI's prompt cache "
|
||||
"matches on the exact prefix, so a system message that arrives mid-conversation otherwise "
|
||||
"breaks the cached prefix on every turn."
|
||||
),
|
||||
},
|
||||
"budget_rollover": { # mutable-ok: registry literal, frozen with its siblings below
|
||||
"type": "Boolean",
|
||||
"description": (
|
||||
|
|
|
|||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
|
||||
from litellm.types.integrations.custom_logger import converted_stream_requested
|
||||
from litellm.types.llms.openai import (
|
||||
PART_UNION_TYPES,
|
||||
ResponseAPIUsage,
|
||||
|
|
@ -626,7 +627,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
return
|
||||
|
||||
request_kwargs = getattr(caching_handler, "request_kwargs", None)
|
||||
if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True:
|
||||
if not _is_json_object(request_kwargs):
|
||||
return
|
||||
if request_kwargs.get("stream") is not True and not converted_stream_requested(request_kwargs):
|
||||
return
|
||||
request_kwargs = request_kwargs.copy()
|
||||
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
|
||||
|
|
|
|||
|
|
@ -79,7 +79,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
declared_authenticating_provider,
|
||||
is_registered_custom_provider,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.ptu_pricing import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
|
|
@ -99,6 +102,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
|
|||
mask_sensitive_structure,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.base_llm.passthrough.transformation import replace_path_segment
|
||||
from litellm.llms.base_llm.vector_store.transformation import (
|
||||
RouterVectorStoreEmbeddingExecutor,
|
||||
|
|
@ -3774,7 +3778,16 @@ class Router:
|
|||
self, deployment: dict, kwargs: dict, function_name: str | None = None
|
||||
) -> Deployment:
|
||||
"""
|
||||
Handle clientside credential
|
||||
Build a per-request Deployment carrying the caller-supplied api_key/api_base,
|
||||
with its own stable id for cooldown, logging, and cost-map identity.
|
||||
|
||||
This deployment is deliberately never registered with the router (no
|
||||
upsert_deployment/add_deployment call): doing so used to add it to
|
||||
self.model_list under the shared model_name, which made a request-scoped,
|
||||
caller-supplied provider credential a permanent, load-balanced deployment
|
||||
that every other caller of that model group could be routed onto. Its
|
||||
pricing is still registered directly, so a custom price configured on the
|
||||
underlying deployment still applies to this call.
|
||||
"""
|
||||
model_info: Final = deployment.get("model_info", {}).copy()
|
||||
litellm_params: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -3793,7 +3806,7 @@ class Router:
|
|||
litellm_params=LiteLLM_Params(**dynamic_litellm_params),
|
||||
model_info=model_info,
|
||||
)
|
||||
self.upsert_deployment(deployment=deployment_pydantic_obj) # add new deployment to router
|
||||
Router._register_deployment_pricing(deployment=deployment_pydantic_obj)
|
||||
return deployment_pydantic_obj
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -9554,8 +9567,10 @@ class Router:
|
|||
)
|
||||
# done reading model["litellm_params"]
|
||||
# Check if provider is supported: either in enum or JSON-configured
|
||||
if custom_llm_provider not in litellm.provider_list and not JSONProviderRegistry.exists(
|
||||
custom_llm_provider
|
||||
if (
|
||||
custom_llm_provider not in litellm.provider_list
|
||||
and not JSONProviderRegistry.exists(custom_llm_provider)
|
||||
and not is_registered_custom_provider(custom_llm_provider)
|
||||
):
|
||||
raise Exception(f"Unsupported provider - {custom_llm_provider}")
|
||||
|
||||
|
|
@ -9701,40 +9716,7 @@ class Router:
|
|||
# initialize client
|
||||
self._add_deployment(deployment=deployment)
|
||||
|
||||
_model_info_dict: Final[dict] = deployment.model_info.model_dump(exclude_none=True)
|
||||
for field in CustomPricingLiteLLMParams.model_fields:
|
||||
field_value = deployment.litellm_params.get(field)
|
||||
if field_value is not None:
|
||||
_model_info_dict[field] = field_value
|
||||
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=_model_info_dict,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if _model_info_dict.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=_model_info_dict,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._inherit_builtin_tiered_output_rate(
|
||||
model_info=_model_info_dict,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
|
||||
# Register custom pricing in litellm.model_cost.
|
||||
# Mirrors _create_deployment() logic to ensure dynamically-added deployments
|
||||
# (e.g., loaded from DB) also have their custom pricing registered.
|
||||
# Without this, _is_model_cost_zero() cannot detect explicitly-configured
|
||||
# zero-cost models, causing budget checks to block free models.
|
||||
Router._register_deployment_in_model_cost(
|
||||
model_id=deployment.model_info.id,
|
||||
model_info=_model_info_dict,
|
||||
model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._register_deployment_pricing(deployment=deployment)
|
||||
|
||||
# add to model names
|
||||
self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id)
|
||||
|
|
@ -9996,6 +9978,21 @@ class Router:
|
|||
)
|
||||
return model_info
|
||||
|
||||
@staticmethod
|
||||
def _register_deployment_pricing(deployment: Deployment) -> None:
|
||||
"""Register a deployment's custom/inherited pricing in ``litellm.model_cost``.
|
||||
|
||||
Takes only a ``Deployment``, so it registers pricing for a deployment that
|
||||
is never added to ``self.model_list`` (a per-request client-side-credential
|
||||
deployment) just as readily as one that is.
|
||||
"""
|
||||
Router._register_deployment_in_model_cost(
|
||||
model_id=deployment.model_info.id,
|
||||
model_info=Router._deployment_model_cost_payload(deployment),
|
||||
model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _register_deployment_in_model_cost(
|
||||
*,
|
||||
|
|
@ -10822,6 +10819,7 @@ class Router:
|
|||
"model_group": user_facing_model_group_name,
|
||||
"providers": [llm_provider],
|
||||
**model_info,
|
||||
"supports_fast_mode": True,
|
||||
"supported_reasoning_efforts": None,
|
||||
}
|
||||
)
|
||||
|
|
@ -10900,6 +10898,9 @@ class Router:
|
|||
if model_info.get("rpm", None) is not None and _deployment_rpm is None:
|
||||
_deployment_rpm = model_info.get("rpm")
|
||||
|
||||
model_group_info.supports_fast_mode = model_group_info.supports_fast_mode and (
|
||||
AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider)
|
||||
)
|
||||
deployment_reasoning_efforts = (
|
||||
resolve_supported_reasoning_efforts( # rebind-ok: recalculated per deployment
|
||||
model_info, deployment_is_mapped=deployment_is_mapped
|
||||
|
|
|
|||
|
|
@ -11,6 +11,9 @@ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY: Final = "litellm_pass_through_custom
|
|||
# exact byte/string body, such as AWS SigV4-signed requests.
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY: Final = "litellm_pass_through_raw_body"
|
||||
|
||||
# `model_info` of the router deployment a provider route (e.g. Vertex) resolved for this request.
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: Final = "litellm_pass_through_deployment_model_info"
|
||||
|
||||
# Attribute set on the FastAPI endpoint function of every user-defined pass-through
|
||||
# route. Auth reads it off the dispatched endpoint (``request.scope["endpoint"]``) to
|
||||
# decide whether a request body ``model`` names an upstream model rather than a
|
||||
|
|
|
|||
|
|
@ -722,6 +722,7 @@ class ModelGroupInfo(BaseModel):
|
|||
supports_url_context: bool = Field(default=False)
|
||||
supports_reasoning: bool = Field(default=False)
|
||||
supports_function_calling: bool = Field(default=False)
|
||||
supports_fast_mode: bool = Field(default=False)
|
||||
supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None)
|
||||
supported_openai_params: list[str] | None = Field(default=[])
|
||||
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
|
||||
|
|
|
|||
|
|
@ -846,6 +846,13 @@ def _is_streaming_response_for_correlation(result: object) -> bool:
|
|||
return isinstance(result, CustomStreamWrapper)
|
||||
|
||||
|
||||
def _is_converted_stream_result(result: object) -> bool:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator))
|
||||
|
||||
|
||||
# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
|
||||
def function_setup(
|
||||
original_function: str,
|
||||
|
|
@ -1889,6 +1896,9 @@ def client(original_function):
|
|||
_caching_handler_response.cached_result is not None
|
||||
and _caching_handler_response.final_embedding_cached_response is None
|
||||
):
|
||||
if _is_converted_stream_result(_caching_handler_response.cached_result):
|
||||
logging_obj.stream = True
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
return _caching_handler_response.cached_result
|
||||
|
||||
elif _caching_handler_response.embedding_all_elements_cache_hit is True:
|
||||
|
|
@ -1946,10 +1956,9 @@ def client(original_function):
|
|||
raise
|
||||
end_time = datetime.datetime.now()
|
||||
|
||||
if _is_streaming_request(
|
||||
kwargs=kwargs,
|
||||
call_type=call_type,
|
||||
):
|
||||
if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result):
|
||||
logging_obj.stream = True
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
if "complete_response" in kwargs and kwargs["complete_response"] is True:
|
||||
chunks: Final = []
|
||||
for idx, chunk in enumerate(result):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -716,6 +716,9 @@
|
|||
"supports_embedding_image_input": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_fast_mode": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_forced_tool_use": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.102.0"
|
||||
version = "1.103.0"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
|
|
@ -67,8 +67,8 @@ proxy = [
|
|||
"azure-identity>=1.25.2,<2.0",
|
||||
"azure-storage-blob>=12.28.0,<13.0",
|
||||
"mcp>=1.28.1,<2.0",
|
||||
"litellm-proxy-extras==0.4.97",
|
||||
"litellm-enterprise==0.1.67",
|
||||
"litellm-proxy-extras==0.4.98",
|
||||
"litellm-enterprise==0.1.68",
|
||||
"RestrictedPython>=8.5,<9.0",
|
||||
"rich>=13.9.4,<14.0",
|
||||
"InquirerPy>=0.3.4,<1.0",
|
||||
|
|
@ -290,6 +290,7 @@ editable-profile = "dev"
|
|||
include = [
|
||||
"litellm/proxy/_experimental/out/**",
|
||||
"litellm/router_strategy/complexity_router/artifacts/*.json",
|
||||
"litellm/proxy/client/cli/commands/codex_base_instructions.md",
|
||||
]
|
||||
exclude = [
|
||||
"litellm/proxy/enterprise",
|
||||
|
|
@ -331,7 +332,7 @@ members = ["enterprise", "litellm-proxy-extras"]
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.102.0"
|
||||
version = "1.103.0"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -927,24 +927,14 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
|
|||
|
||||
|
||||
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={"stream": True},
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={"stream": False},
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
_should_defer_streaming_cache_hit_callbacks(
|
||||
kwargs={},
|
||||
)
|
||||
is False
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
stream_replay = CustomStreamWrapper(
|
||||
completion_stream=iter(()), model="gpt-4o", logging_obj=logging_obj
|
||||
)
|
||||
assert _should_defer_streaming_cache_hit_callbacks(cached_result=stream_replay) is True
|
||||
assert _should_defer_streaming_cache_hit_callbacks(cached_result=ModelResponse()) is False
|
||||
assert _should_defer_streaming_cache_hit_callbacks(cached_result={"id": "msg_1"}) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
import sys, os, time
|
||||
import traceback, asyncio
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
|
@ -402,6 +403,10 @@ def test_router_redis_cache():
|
|||
|
||||
|
||||
def test_router_handle_clientside_credential():
|
||||
"""A caller-supplied credential must stay scoped to the current call: it must
|
||||
never be registered as a router deployment, or a later caller with no override
|
||||
of their own can be load-balanced onto it and reach the provider with someone
|
||||
else's credential (see LIT-7811)."""
|
||||
deployment = {
|
||||
"model_name": "gemini/*",
|
||||
"litellm_params": {"model": "gemini/*"},
|
||||
|
|
@ -421,7 +426,67 @@ def test_router_handle_clientside_credential():
|
|||
)
|
||||
|
||||
assert new_deployment.litellm_params.api_key == "123"
|
||||
assert len(router.get_model_list()) == 2
|
||||
assert len(router.get_model_list()) == 1
|
||||
assert router.get_deployment(model_id=new_deployment.model_info.id) is None
|
||||
|
||||
|
||||
async def test_router_clientside_credential_not_reused_by_other_callers(
|
||||
respx_mock, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""End-to-end regression test for LIT-7811.
|
||||
|
||||
One caller's request-scoped api_key must never leak into a later, unrelated
|
||||
caller's request. Before the fix, the router registered the caller-supplied
|
||||
credential as a second, permanent deployment for the shared model group, so
|
||||
plain follow-up calls with no override of their own could be load-balanced
|
||||
onto it and reach the provider with the first caller's key.
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4o",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "shared-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "configured-key"},
|
||||
"model_info": {"id": "configured-deployment"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
await router.acompletion(
|
||||
model="shared-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_key="alternate-tenant-key",
|
||||
)
|
||||
assert route.calls[-1].request.headers["authorization"] == "Bearer alternate-tenant-key"
|
||||
|
||||
# The forwarded credential must never become a routable deployment for the
|
||||
# model group other callers share.
|
||||
assert [d["model_info"]["id"] for d in router.get_model_list(model_name="shared-model")] == [
|
||||
"configured-deployment"
|
||||
]
|
||||
|
||||
for _ in range(20):
|
||||
await router.acompletion(
|
||||
model="shared-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]}
|
||||
assert used_auth_headers == {"Bearer configured-key"}
|
||||
|
||||
|
||||
def test_router_get_async_openai_model_client():
|
||||
|
|
|
|||
|
|
@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution():
|
|||
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
|
||||
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
|
||||
DIFFERENT issuer whose claim value happens to collide, even though both issuers
|
||||
map the same claim field (``sub``) to a virtual key."""
|
||||
issuer_a = "https://issuer-a.example.com"
|
||||
issuer_b = "https://issuer-b.example.com"
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600
|
||||
)
|
||||
|
||||
rows = [
|
||||
{
|
||||
"jwt_issuer": issuer_b,
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "dev-alice",
|
||||
"token": "hashed-issuer-b-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
|
||||
async def fake_find_first(where):
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return MagicMock(**row)
|
||||
return None
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first)
|
||||
|
||||
# Dependency-inject the resolved key via the cache (IdentityStore._resolve_key
|
||||
# reads it from here) instead of monkeypatching IdentityStore itself.
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-issuer-b-key",
|
||||
value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"),
|
||||
)
|
||||
|
||||
# The rightful owner: issuer-b's own claim resolves to its mapping.
|
||||
owner_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(owner_result, UserAPIKeyAuth)
|
||||
assert owner_result.token == "hashed-issuer-b-key"
|
||||
|
||||
# A validly-signed token from issuer-a carrying the SAME claim value must not
|
||||
# inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a
|
||||
# correctly-scoped miss returns None instead of resolving to issuer-b's key.
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert colliding_result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer():
|
||||
"""LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING
|
||||
issuer's key would leave every issuer that falls back to it holding its own
|
||||
stale copy after the row is updated/deleted -- CRUD only evicts the cache key
|
||||
computed from the row's own scope (global), so a copy cached under some other
|
||||
issuer's key would keep resolving to the old token until TTL. Caching it under
|
||||
the global key instead means every issuer shares (and CRUD correctly evicts)
|
||||
the exact same entry."""
|
||||
issuer_a = "https://issuer-a.example.com"
|
||||
issuer_b = "https://issuer-b.example.com"
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
issuers=[
|
||||
{
|
||||
"issuer": issuer_a,
|
||||
"jwks_url": f"{issuer_a}/jwks",
|
||||
"virtual_key_claim_field": "sub",
|
||||
"disable_audience_validation": True,
|
||||
},
|
||||
{
|
||||
"issuer": issuer_b,
|
||||
"jwks_url": f"{issuer_b}/jwks",
|
||||
"virtual_key_claim_field": "sub",
|
||||
"disable_audience_validation": True,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
rows = [
|
||||
{
|
||||
"jwt_issuer": "",
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "legacy-user",
|
||||
"token": "hashed-legacy-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
|
||||
async def fake_find_first(where):
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return MagicMock(**row)
|
||||
return None
|
||||
|
||||
prisma_client = MagicMock()
|
||||
find_first = AsyncMock(side_effect=fake_find_first)
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = find_first
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-legacy-key",
|
||||
value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"),
|
||||
)
|
||||
|
||||
resolved_a = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(resolved_a, UserAPIKeyAuth)
|
||||
assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit
|
||||
|
||||
# issuer-b resolving the SAME global mapping must hit the cache issuer-a's
|
||||
# resolution populated, not issue a fresh DB query for the global row again.
|
||||
resolved_b = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(resolved_b, UserAPIKeyAuth)
|
||||
assert resolved_b.token == "hashed-legacy-key"
|
||||
assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_to_virtual_key_mapping_no_mapping():
|
||||
"""
|
||||
|
|
@ -223,6 +371,7 @@ def test_to_response_excludes_token():
|
|||
now = datetime.now(timezone.utc)
|
||||
mock_mapping = MagicMock()
|
||||
mock_mapping.id = "mapping-1"
|
||||
mock_mapping.jwt_issuer = None
|
||||
mock_mapping.jwt_claim_name = "email"
|
||||
mock_mapping.jwt_claim_value = "user@example.com"
|
||||
mock_mapping.token = "hashed_secret_value"
|
||||
|
|
@ -275,10 +424,12 @@ def _mock_mapping(
|
|||
id="mapping-1",
|
||||
claim_name="email",
|
||||
claim_value="user@example.com",
|
||||
issuer=None,
|
||||
):
|
||||
now = datetime.now(timezone.utc)
|
||||
m = MagicMock()
|
||||
m.id = id
|
||||
m.jwt_issuer = issuer
|
||||
m.jwt_claim_name = claim_name
|
||||
m.jwt_claim_value = claim_value
|
||||
m.token = "hashed_token"
|
||||
|
|
@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token():
|
|||
assert result.jwt_claim_name == "email"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_without_issuer_stores_empty_string_not_null():
|
||||
"""LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL
|
||||
for an unscoped mapping would let Postgres accept unlimited duplicate unscoped
|
||||
rows for the same claim (NULL is never equal to NULL in a unique constraint),
|
||||
so two mappings for the same claim value could point at two different keys with
|
||||
no conflict, and resolution would pick whichever one Postgres returns first."""
|
||||
from litellm.proxy._types import CreateJWTKeyMappingRequest
|
||||
|
||||
mock_prisma = _mock_prisma()
|
||||
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
|
||||
mock_cache = AsyncMock()
|
||||
|
||||
data = CreateJWTKeyMappingRequest(jwt_claim_name="sub", jwt_claim_value="dev-alice", key="sk-test-key")
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
"litellm.proxy.proxy_server.prisma_client", mock_prisma
|
||||
),
|
||||
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
|
||||
),
|
||||
):
|
||||
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
|
||||
|
||||
sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"]
|
||||
assert sent_data["jwt_issuer"] == ""
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# Tests: unregistered_jwt_client_behavior
|
||||
# ──────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -2099,8 +2099,12 @@ def test_handle_clientside_credential_metadata_loading(
|
|||
assert result_deployment.model_info.id != "original-id-123"
|
||||
assert result_deployment.model_info.original_model_id == "original-id-123"
|
||||
|
||||
# Verify the deployment was added to the router
|
||||
assert len(router.model_list) == len(model_list) + 1
|
||||
# The caller-supplied credential must stay scoped to this call: it must never be
|
||||
# registered as a router deployment, or a later caller with no override of their
|
||||
# own could be load-balanced onto it and reach the provider with this credential
|
||||
# (see LIT-7811).
|
||||
assert len(router.model_list) == len(model_list)
|
||||
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
|
||||
|
||||
# Test that the function correctly uses the right metadata key
|
||||
# For acompletion, it should use "metadata"
|
||||
|
|
@ -2260,14 +2264,63 @@ def test_handle_clientside_credential_with_responses_function(model_list):
|
|||
assert result_deployment.model_info.id != "original-id-responses"
|
||||
assert result_deployment.model_info.original_model_id == "original-id-responses"
|
||||
|
||||
# Verify the deployment was added to the router
|
||||
assert len(router.model_list) == len(model_list) + 1
|
||||
# The caller-supplied credential must stay scoped to this call: it must never be
|
||||
# registered as a router deployment (see LIT-7811).
|
||||
assert len(router.model_list) == len(model_list)
|
||||
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
|
||||
|
||||
print(
|
||||
"✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata"
|
||||
)
|
||||
|
||||
|
||||
def test_handle_clientside_credential_still_registers_custom_pricing(model_list):
|
||||
"""A clientside-credential call must still price against the deployment's own
|
||||
custom rate, even though the call's ephemeral deployment is never added to the
|
||||
router (see LIT-7811): losing that registration would silently fall back to
|
||||
public catalog pricing for every clientside-credential call on a deployment
|
||||
with a custom rate configured."""
|
||||
router = Router(model_list=model_list)
|
||||
deployment = {
|
||||
"model_name": "gpt-4.1",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4.1",
|
||||
"api_key": "test_key",
|
||||
"input_cost_per_token": 0.0001234,
|
||||
"output_cost_per_token": 0.0005678,
|
||||
},
|
||||
"model_info": {"id": "original-id-pricing"},
|
||||
}
|
||||
kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}}
|
||||
|
||||
result_deployment = router._handle_clientside_credential(
|
||||
deployment=deployment, kwargs=kwargs, function_name="acompletion"
|
||||
)
|
||||
|
||||
registered = litellm.model_cost.get(result_deployment.model_info.id)
|
||||
assert registered is not None
|
||||
assert registered["input_cost_per_token"] == 0.0001234
|
||||
assert registered["output_cost_per_token"] == 0.0005678
|
||||
|
||||
|
||||
def test_register_deployment_pricing_direct_call():
|
||||
"""Direct-call unit test for the pricing-registration helper `_handle_clientside_credential`
|
||||
relies on, so it prices a deployment that is deliberately never added to `self.model_list`."""
|
||||
deployment = Deployment(
|
||||
model_name="gpt-4.1",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="gpt-4.1",
|
||||
api_key="test_key",
|
||||
input_cost_per_token=0.0009999,
|
||||
),
|
||||
model_info=ModelInfo(id="direct-call-pricing-id"),
|
||||
)
|
||||
|
||||
Router._register_deployment_pricing(deployment=deployment)
|
||||
|
||||
assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999
|
||||
|
||||
|
||||
def test_get_metadata_variable_name_from_kwargs(model_list):
|
||||
"""
|
||||
Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content.
|
||||
|
|
|
|||
|
|
@ -693,3 +693,90 @@ async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monke
|
|||
assert handler.preset_cache_key is not None
|
||||
assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
|
||||
assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
async def aanthropic_messages(**kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
kwargs = {
|
||||
"model": "claude-sonnet-5",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 16,
|
||||
"caching": True,
|
||||
"stream": False,
|
||||
"_websearch_interception_converted_stream": True,
|
||||
}
|
||||
cached_message = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
}
|
||||
await litellm.cache.async_add_cache(cached_message, **kwargs)
|
||||
handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
|
||||
logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
||||
|
||||
hit = await handler._async_get_cache(
|
||||
model="claude-sonnet-5",
|
||||
original_function=aanthropic_messages,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.aanthropic_messages.value,
|
||||
kwargs=kwargs,
|
||||
args=(),
|
||||
)
|
||||
|
||||
assert hit is not None and hit.cached_result == cached_message
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
|
||||
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
async def acompletion(**kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
kwargs = {
|
||||
"model": "gpt-5.6",
|
||||
"messages": [{"role": "user", "content": "run the code"}],
|
||||
"caching": True,
|
||||
"stream": False,
|
||||
"_code_interpreter_interception_converted_stream": True,
|
||||
"_agentic_loop_depth": 1,
|
||||
}
|
||||
await litellm.cache.async_add_cache(
|
||||
litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs
|
||||
)
|
||||
handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
|
||||
logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
||||
|
||||
hit = await handler._async_get_cache(
|
||||
model="gpt-5.6",
|
||||
original_function=acompletion,
|
||||
logging_obj=logging_obj,
|
||||
start_time=datetime.now(),
|
||||
call_type=CallTypes.acompletion.value,
|
||||
kwargs=kwargs,
|
||||
args=(),
|
||||
)
|
||||
|
||||
assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse)
|
||||
assert hit.cached_result.choices[0].message.content == "done"
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
|
||||
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import os
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -7,6 +8,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.integrations.langsmith import LangsmithLogger
|
||||
from litellm.types.integrations.langsmith import LangsmithQueueObject
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -531,3 +533,44 @@ class TestLangsmithRootRunIdConsistency:
|
|||
|
||||
assert data["trace_id"] == "trace-1"
|
||||
assert data["dotted_order"] == dotted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_events_appended_during_flush_are_not_dropped():
|
||||
logger = LangsmithLogger(langsmith_api_key="test-key", langsmith_project="test-project")
|
||||
try:
|
||||
sent_batches: Final[list[list[dict[str, str]]]] = []
|
||||
late_event: Final = LangsmithQueueObject(
|
||||
credentials=logger.default_credentials, data={"id": "late"}
|
||||
)
|
||||
|
||||
async def fake_post(
|
||||
url: str, json: dict[str, list[dict[str, str]]], headers: dict[str, str]
|
||||
) -> MagicMock:
|
||||
if not sent_batches:
|
||||
logger.log_queue.append(late_event)
|
||||
sent_batches.append(json["post"])
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.raise_for_status = MagicMock()
|
||||
return response
|
||||
|
||||
logger.async_httpx_client = MagicMock(post=AsyncMock(side_effect=fake_post))
|
||||
logger.log_queue = [
|
||||
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "a"}),
|
||||
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "b"}),
|
||||
]
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
assert [e["id"] for e in sent_batches[0]] == ["a", "b"]
|
||||
assert logger.log_queue == [late_event]
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
assert [e["id"] for e in sent_batches[1]] == ["late"]
|
||||
assert logger.log_queue == []
|
||||
finally:
|
||||
if logger._flush_task is not None:
|
||||
logger._flush_task.cancel()
|
||||
await asyncio.gather(logger._flush_task, return_exceptions=True)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
responses_reasoning_items_from_thinking_blocks,
|
||||
split_concatenated_json_objects,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
system_messages_first,
|
||||
update_messages_with_model_file_ids,
|
||||
)
|
||||
|
||||
|
|
@ -1107,6 +1108,38 @@ def test_drop_tool_reference_parts_leaves_non_tool_messages_alone():
|
|||
assert result[2]["content"] == ""
|
||||
|
||||
|
||||
class TestSystemMessagesFirst:
|
||||
def test_stable_partition_keeps_order_within_each_group(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "system", "content": "s1"},
|
||||
{"role": "assistant", "content": "a1"},
|
||||
{"role": "developer", "content": "d1"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "t1"},
|
||||
{"role": "system", "content": "s2"},
|
||||
]
|
||||
|
||||
result = system_messages_first(messages)
|
||||
|
||||
assert [m["content"] for m in result] == ["s1", "d1", "s2", "u1", "a1", "t1"]
|
||||
assert [m["content"] for m in messages] == ["u1", "s1", "a1", "d1", "t1", "s2"]
|
||||
assert all(
|
||||
result_message is original for result_message, original in zip(result[3:], messages[::2], strict=True)
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"messages",
|
||||
[
|
||||
[],
|
||||
[{"role": "user", "content": "u1"}, {"role": "assistant", "content": "a1"}],
|
||||
[{"role": "system", "content": "s1"}, {"role": "user", "content": "u1"}],
|
||||
[{"role": "system", "content": "s1"}, {"role": "system", "content": "s2"}],
|
||||
],
|
||||
)
|
||||
def test_already_ordered_messages_come_back_unchanged(self, messages):
|
||||
assert system_messages_first(messages) == messages
|
||||
|
||||
|
||||
class TestFlattenTopLevelSchemaCombinators:
|
||||
def _customer_anyof_schema(self):
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import base64
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
_convert_to_bedrock_tool_call_invoke,
|
||||
_convert_to_bedrock_tool_call_result,
|
||||
anthropic_messages_pt,
|
||||
convert_to_anthropic_tool_result,
|
||||
convert_to_gemini_tool_call_result,
|
||||
make_valid_bedrock_tool_name,
|
||||
ollama_pt,
|
||||
|
|
@ -2208,6 +2210,104 @@ def test_bedrock_tool_call_invoke_empty_arguments():
|
|||
assert result[0]["toolUse"]["input"] == {}
|
||||
|
||||
|
||||
_BEDROCK_TOOL_USE_ID_RE = re.compile(r"^[a-zA-Z0-9_.:-]{1,64}$")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_call_id",
|
||||
[
|
||||
"call_" + "x" * 100,
|
||||
"call|with|pipes",
|
||||
"call_" + "y" * 60 + "|end",
|
||||
"call:ok.dots-and_under",
|
||||
"",
|
||||
],
|
||||
)
|
||||
def test_bedrock_tool_use_id_is_sanitized_consistently_for_invoke_and_result(tool_call_id):
|
||||
"""
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/34239: client-minted
|
||||
tool_call ids longer than 64 chars or with chars outside [a-zA-Z0-9_.:-] made Bedrock
|
||||
return a 400. The invoke and result paths must produce the same valid toolUseId so the
|
||||
toolUse/toolResult pair still correlates.
|
||||
"""
|
||||
invoke = _convert_to_bedrock_tool_call_invoke(
|
||||
[
|
||||
{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"location": "Boston"}'},
|
||||
}
|
||||
]
|
||||
)
|
||||
result = _convert_to_bedrock_tool_call_result(
|
||||
{"tool_call_id": tool_call_id, "role": "tool", "name": "get_weather", "content": "sunny"}
|
||||
)
|
||||
tool_use_id = invoke[0]["toolUse"]["toolUseId"]
|
||||
assert _BEDROCK_TOOL_USE_ID_RE.match(tool_use_id)
|
||||
assert result["toolResult"]["toolUseId"] == tool_use_id
|
||||
|
||||
|
||||
def test_bedrock_tool_use_id_valid_ids_pass_through_unchanged():
|
||||
result = _convert_to_bedrock_tool_call_result(
|
||||
{"tool_call_id": "tooluse_Ab.c:1-2_3", "role": "tool", "name": "f", "content": "ok"}
|
||||
)
|
||||
assert result["toolResult"]["toolUseId"] == "tooluse_Ab.c:1-2_3"
|
||||
|
||||
|
||||
def test_bedrock_tool_use_id_truncation_keeps_distinct_ids_distinct():
|
||||
prefix = "call_" + "z" * 70
|
||||
ids = {
|
||||
_convert_to_bedrock_tool_call_result(
|
||||
{"tool_call_id": f"{prefix}{suffix}", "role": "tool", "name": "f", "content": "ok"}
|
||||
)["toolResult"]["toolUseId"]
|
||||
for suffix in ("a", "b")
|
||||
}
|
||||
assert len(ids) == 2
|
||||
assert all(len(i) == 64 for i in ids)
|
||||
|
||||
|
||||
def test_bedrock_tool_use_id_replaced_chars_do_not_collide_with_existing_ids():
|
||||
ids = {
|
||||
_convert_to_bedrock_tool_call_result({"tool_call_id": i, "role": "tool", "name": "f", "content": "ok"})[
|
||||
"toolResult"
|
||||
]["toolUseId"]
|
||||
for i in ("call|x", "call_x")
|
||||
}
|
||||
assert len(ids) == 2
|
||||
|
||||
|
||||
def test_bedrock_tool_call_invoke_concatenated_json_long_id_stays_within_limit():
|
||||
long_id = "call_" + "q" * 62
|
||||
result = _convert_to_bedrock_tool_call_invoke(
|
||||
[
|
||||
{
|
||||
"id": long_id,
|
||||
"type": "function",
|
||||
"function": {"name": "run", "arguments": '{"cmd":"a"}{"cmd":"b"}'},
|
||||
}
|
||||
]
|
||||
)
|
||||
ids = [block["toolUse"]["toolUseId"] for block in result]
|
||||
assert len(ids) == 2
|
||||
assert len(set(ids)) == 2
|
||||
assert all(_BEDROCK_TOOL_USE_ID_RE.match(i) for i in ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool_call_id", "expected"),
|
||||
[
|
||||
("call|with|pipes", "call_with_pipes"),
|
||||
("call:ok.dots", "call_ok_dots"),
|
||||
("call_" + "x" * 100, "call_" + "x" * 100),
|
||||
("toolu_01AbC-xyz", "toolu_01AbC-xyz"),
|
||||
("", "tool_use_id"),
|
||||
],
|
||||
)
|
||||
def test_anthropic_tool_use_id_keeps_pattern_only_rewrite_with_no_cap_or_hash(tool_call_id, expected):
|
||||
result = convert_to_anthropic_tool_result({"role": "tool", "tool_call_id": tool_call_id, "content": "ok"})
|
||||
assert result["tool_use_id"] == expected
|
||||
|
||||
|
||||
def test_bedrock_tool_call_invoke_concatenated_json():
|
||||
"""
|
||||
Tool call whose arguments contain multiple concatenated JSON objects
|
||||
|
|
|
|||
|
|
@ -0,0 +1,55 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import CustomLLM
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
get_llm_provider,
|
||||
is_registered_custom_provider,
|
||||
)
|
||||
|
||||
CUSTOM_PROVIDER: Final = "test-onprem-llm"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str:
|
||||
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": CUSTOM_PROVIDER, "custom_handler": CustomLLM()}])
|
||||
monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list))
|
||||
monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers))
|
||||
return CUSTOM_PROVIDER
|
||||
|
||||
|
||||
def test_get_llm_provider_resolves_custom_provider_map_prefix_before_first_completion(
|
||||
registered_custom_provider: str,
|
||||
) -> None:
|
||||
assert registered_custom_provider not in litellm.provider_list
|
||||
|
||||
model, provider, dynamic_api_key, api_base = get_llm_provider(model=f"{registered_custom_provider}/my-model")
|
||||
|
||||
assert (model, provider, dynamic_api_key, api_base) == ("my-model", registered_custom_provider, None, None)
|
||||
|
||||
|
||||
def test_get_llm_provider_strips_prefix_when_custom_provider_passed_explicitly(
|
||||
registered_custom_provider: str,
|
||||
) -> None:
|
||||
model, provider, _, api_base = get_llm_provider(
|
||||
model="my-model",
|
||||
custom_llm_provider=registered_custom_provider,
|
||||
api_base="http://onprem.internal:8080",
|
||||
)
|
||||
|
||||
assert (model, provider, api_base) == ("my-model", registered_custom_provider, "http://onprem.internal:8080")
|
||||
|
||||
|
||||
def test_get_llm_provider_still_rejects_unregistered_prefix(registered_custom_provider: str) -> None:
|
||||
with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"):
|
||||
get_llm_provider(model="not-registered-llm/my-model")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("candidate", "expected"),
|
||||
[(CUSTOM_PROVIDER, True), ("not-registered-llm", False), (None, False), ("", False)],
|
||||
)
|
||||
def test_is_registered_custom_provider(registered_custom_provider: str, candidate: str | None, expected: bool) -> None:
|
||||
assert is_registered_custom_provider(candidate) is expected
|
||||
|
|
@ -132,6 +132,32 @@ def test_transform_request_drops_tool_reference_parts():
|
|||
assert request["messages"][2]["content"] == ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"enabled, expected", [(False, ("hi", "sys", "reply", "more")), (True, ("sys", "hi", "reply", "more"))]
|
||||
)
|
||||
def test_transform_request_system_messages_first_follows_global_flag(monkeypatch, enabled, expected):
|
||||
"""Azure OpenAI shares OpenAI's prefix-matched prompt cache, so the same flag moves
|
||||
system messages ahead of the conversation on the Azure request body."""
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", enabled)
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "more"},
|
||||
]
|
||||
|
||||
request = AzureOpenAIConfig().transform_request(
|
||||
model="gpt-4o",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "azure"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert tuple(m["content"] for m in request["messages"]) == expected
|
||||
assert [m["content"] for m in messages] == ["hi", "sys", "reply", "more"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, emitted_key, absent_key",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -68,3 +68,24 @@ def test_azure_o_series_transform_request_flattens_top_level_anyof():
|
|||
assert parameters["required"] == ["id"]
|
||||
assert "anyOf" in tool["function"]["parameters"]
|
||||
assert optional_params["tools"][0] is tool
|
||||
|
||||
|
||||
def test_azure_o_series_transform_request_moves_system_messages_first(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "developer", "content": "dev"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "more"},
|
||||
]
|
||||
|
||||
request = AzureOpenAIO1Config().transform_request(
|
||||
model="o3-mini",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "azure"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert [m["content"] for m in request["messages"]] == ["dev", "hi", "reply", "more"]
|
||||
assert [m["content"] for m in messages] == ["hi", "dev", "reply", "more"]
|
||||
|
|
|
|||
|
|
@ -1124,6 +1124,80 @@ class TestToolReferenceStripping:
|
|||
assert request["messages"][2]["content"] == ""
|
||||
|
||||
|
||||
class TestSystemMessagesFirst:
|
||||
"""With litellm.openai_system_messages_first on, requests bound for OpenAI put system and
|
||||
developer messages ahead of the conversation, keeping each group's order, so the instruction
|
||||
prefix stays byte-stable for OpenAI's prefix-matched prompt cache."""
|
||||
|
||||
MESSAGES: Final = (
|
||||
{"role": "user", "content": "first turn"},
|
||||
{"role": "system", "content": "sys 1"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "developer", "content": "dev"},
|
||||
{"role": "user", "content": "second turn"},
|
||||
{"role": "system", "content": "sys 2"},
|
||||
)
|
||||
ORIGINAL_ORDER: Final = ("first turn", "sys 1", "reply", "dev", "second turn", "sys 2")
|
||||
ORDERED: Final = ("sys 1", "dev", "sys 2", "first turn", "reply", "second turn")
|
||||
|
||||
def setup_method(self):
|
||||
self.config = OpenAIGPTConfig()
|
||||
|
||||
def _messages(self):
|
||||
return [dict(m) for m in self.MESSAGES]
|
||||
|
||||
def _transform(self, provider):
|
||||
return self.config.transform_request(
|
||||
model="gpt-4.1",
|
||||
messages=self._messages(),
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": provider},
|
||||
headers={},
|
||||
)
|
||||
|
||||
def test_default_off_keeps_caller_order(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
|
||||
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORIGINAL_ORDER
|
||||
|
||||
def test_moves_system_and_developer_messages_first_for_openai(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
|
||||
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORDERED
|
||||
|
||||
def test_leaves_openai_compatible_providers_alone(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
|
||||
assert tuple(m["content"] for m in self._transform("deepseek")["messages"]) == self.ORIGINAL_ORDER
|
||||
|
||||
def test_does_not_mutate_caller_messages(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
|
||||
messages = self._messages()
|
||||
self.config.transform_request(
|
||||
model="gpt-4.1",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "openai"},
|
||||
headers={},
|
||||
)
|
||||
assert tuple(m["content"] for m in messages) == self.ORIGINAL_ORDER
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_transform_request_moves_system_messages_first(self, monkeypatch):
|
||||
class UninstantiatedOpenAIGPTConfig(OpenAIGPTConfig):
|
||||
_is_base_class = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
|
||||
request = await UninstantiatedOpenAIGPTConfig().async_transform_request(
|
||||
model="gpt-4.1",
|
||||
messages=self._messages(),
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "openai"},
|
||||
headers={},
|
||||
)
|
||||
assert tuple(m["content"] for m in request["messages"]) == self.ORDERED
|
||||
|
||||
|
||||
class TestOpenAIPromptCacheBreakpointChatPath:
|
||||
"""Chat-path shape for OpenAI explicit prompt caching (#37509)."""
|
||||
|
||||
|
|
|
|||
|
|
@ -32,7 +32,13 @@ from litellm.proxy._types import (
|
|||
JWTRoutingOverride,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
TeamNotFoundError,
|
||||
UserNotFoundError,
|
||||
get_key_object,
|
||||
_cache_key_object,
|
||||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
|
|
@ -7948,13 +7954,38 @@ def _per_issuer_virtual_key_jwt_handler(
|
|||
|
||||
|
||||
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
|
||||
"""Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way."""
|
||||
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
|
||||
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
|
||||
return prisma_client, find_first
|
||||
|
||||
|
||||
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
|
||||
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
|
||||
def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]:
|
||||
"""A ``find_first`` whose result depends on the ``where`` clause, like a real table.
|
||||
|
||||
Matches a row when every key present in ``where`` equals that key on the row --
|
||||
a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never
|
||||
sending ``jwt_issuer``) does not constrain the match, exactly like Prisma.
|
||||
"""
|
||||
|
||||
async def _find_first(where: dict[str, object]) -> SimpleNamespace | None:
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return SimpleNamespace(**row)
|
||||
return None
|
||||
|
||||
find_first = AsyncMock(side_effect=_find_first)
|
||||
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
|
||||
return prisma_client, find_first
|
||||
|
||||
|
||||
def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]:
|
||||
return {
|
||||
"jwt_claim_name": claim_name,
|
||||
"jwt_claim_value": claim_value,
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"is_active": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -7978,11 +8009,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for
|
|||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
|
||||
# Issuer-scoped lookup hits on the first query, so no global fallback query runs.
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO))
|
||||
assert isinstance(resolved, UserAPIKeyAuth)
|
||||
assert resolved.token == "hashed-mapped-key"
|
||||
assert resolved.team_id == "svc-team"
|
||||
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
|
||||
cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO)
|
||||
assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8015,7 +8048,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
|
|||
|
||||
assert exc.value.status_code == 403
|
||||
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
|
||||
# REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row.
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("sub", "unknown-svc", ISSUER_TWO),
|
||||
_mapping_where("sub", "unknown-svc", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8025,7 +8062,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
|
|||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
|
||||
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
|
||||
# Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup.
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL
|
||||
)
|
||||
|
||||
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
|
||||
|
|
@ -8050,7 +8090,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
|
|||
|
||||
assert exc.value.status_code == 403
|
||||
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("sub", "admin-7", ISSUER_TWO),
|
||||
_mapping_where("sub", "admin-7", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8079,7 +8122,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f
|
|||
|
||||
assert with_claim is None
|
||||
assert without_claim is None
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
|
||||
# without_claim has no claim value and returns before ever reaching the DB.
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("client_id", "app-9", ISSUER_ONE),
|
||||
_mapping_where("client_id", "app-9", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
|
||||
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
|
||||
DIFFERENT issuer whose claim value happens to collide, even though both issuers
|
||||
use the same claim field (``sub``) for their virtual-key mapping."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table(
|
||||
[
|
||||
{
|
||||
"jwt_issuer": ISSUER_TWO,
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "dev-alice",
|
||||
"token": "hashed-issuer-b-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-issuer-b-key",
|
||||
value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"),
|
||||
)
|
||||
|
||||
owner_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert isinstance(owner_result, UserAPIKeyAuth)
|
||||
assert owner_result.token == "hashed-issuer-b-key"
|
||||
|
||||
# issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must
|
||||
# return None (fall through to team-based JWT auth), never issuer-two's key.
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert colliding_result is None
|
||||
assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer():
|
||||
"""A cached positive resolution must be keyed by issuer too, or a colliding
|
||||
claim value from another issuer could be served straight from cache without
|
||||
ever reaching the (correctly issuer-scoped) DB lookup."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([])
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key"
|
||||
)
|
||||
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert colliding_result is None
|
||||
# Must have gone to the DB rather than serving issuer-two's cached token.
|
||||
assert find_first.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_issuer_agnostic_mapping_matches_every_issuer():
|
||||
"""A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps
|
||||
matching any issuer, so existing global mappings are not broken by this fix."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table(
|
||||
[
|
||||
{
|
||||
"jwt_issuer": "",
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "legacy-user",
|
||||
"token": "hashed-legacy-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-legacy-key",
|
||||
value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"),
|
||||
)
|
||||
|
||||
for issuer in (ISSUER_ONE, ISSUER_TWO):
|
||||
resolved = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert isinstance(resolved, UserAPIKeyAuth)
|
||||
assert resolved.token == "hashed-legacy-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import inspect
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import click
|
||||
|
|
@ -9,10 +11,9 @@ import pytest
|
|||
import requests
|
||||
from click.testing import CliRunner
|
||||
|
||||
|
||||
|
||||
from litellm.proxy.client.cli.commands.agents import (
|
||||
AgentRunError,
|
||||
ModelSyncArgs,
|
||||
ModelSyncSkipped,
|
||||
_hand_off,
|
||||
_replace_process,
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.proxy.client.cli.commands.agents import (
|
|||
agent_model_sync_env,
|
||||
agent_profile,
|
||||
build_agent_env,
|
||||
codex_model_sync_args,
|
||||
opencode_model_sync_env,
|
||||
run_agent,
|
||||
verify_proxy_key,
|
||||
|
|
@ -55,6 +57,112 @@ class _Recorder:
|
|||
return self.returns
|
||||
|
||||
|
||||
_STOCK_REASONING_LEVELS = [
|
||||
{"effort": "low", "description": "Fast responses with lighter reasoning"},
|
||||
{"effort": "medium", "description": "Balances speed and reasoning depth for everyday tasks"},
|
||||
{"effort": "high", "description": "Greater reasoning depth for complex problems"},
|
||||
]
|
||||
|
||||
_STOCK_MODELS = {
|
||||
"gpt-5.6-terra": {
|
||||
"slug": "gpt-5.6-terra",
|
||||
"display_name": "GPT-5.6 Terra",
|
||||
"description": "Balanced agentic coding model for everyday work.",
|
||||
"default_reasoning_level": "medium",
|
||||
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
|
||||
"shell_type": "unified_exec",
|
||||
"visibility": "list",
|
||||
"supported_in_api": True,
|
||||
"priority": 7,
|
||||
"availability_nux": None,
|
||||
"upgrade": None,
|
||||
"base_instructions": "You are Codex, a coding agent based on GPT-5.6.",
|
||||
"apply_patch_tool_type": "freeform",
|
||||
"supports_parallel_tool_calls": True,
|
||||
"context_window": 272000,
|
||||
"comp_hash": "terra-hash",
|
||||
},
|
||||
"gpt-5.5": {
|
||||
"slug": "gpt-5.5",
|
||||
"display_name": "GPT-5.5",
|
||||
"description": "Frontier model for complex coding, research, and real-world work.",
|
||||
"default_reasoning_level": "medium",
|
||||
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
|
||||
"shell_type": "unified_exec",
|
||||
"visibility": "list",
|
||||
"supported_in_api": True,
|
||||
"priority": 12,
|
||||
"availability_nux": None,
|
||||
"upgrade": None,
|
||||
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
|
||||
"apply_patch_tool_type": "freeform",
|
||||
"supports_parallel_tool_calls": True,
|
||||
"context_window": 272000,
|
||||
"comp_hash": "gpt-5.5-hash",
|
||||
},
|
||||
"gpt-5.4": {
|
||||
"slug": "gpt-5.4",
|
||||
"display_name": "GPT-5.4",
|
||||
"description": "Strong model for everyday coding.",
|
||||
"default_reasoning_level": "medium",
|
||||
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
|
||||
"shell_type": "unified_exec",
|
||||
"visibility": "hide",
|
||||
"supported_in_api": True,
|
||||
"priority": 16,
|
||||
"availability_nux": None,
|
||||
"upgrade": {
|
||||
"model": "gpt-5.6-terra",
|
||||
"migration_markdown": "GPT-5.4 is no longer available. Switch to GPT-5.6 Terra to continue.",
|
||||
"retirement_at": "2026-08-31T19:00:00Z",
|
||||
},
|
||||
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
|
||||
"apply_patch_tool_type": "freeform",
|
||||
"supports_parallel_tool_calls": True,
|
||||
"context_window": 272000,
|
||||
"comp_hash": "gpt-5.4-hash",
|
||||
},
|
||||
"codex-auto-review": {
|
||||
"slug": "codex-auto-review",
|
||||
"display_name": "Codex Auto Review",
|
||||
"description": None,
|
||||
"supported_reasoning_levels": [],
|
||||
"shell_type": "unified_exec",
|
||||
"visibility": "hide",
|
||||
"supported_in_api": False,
|
||||
"priority": 43,
|
||||
"availability_nux": None,
|
||||
"upgrade": None,
|
||||
"base_instructions": "You are Codex, reviewing a change.",
|
||||
"apply_patch_tool_type": None,
|
||||
"supports_parallel_tool_calls": True,
|
||||
"context_window": 272000,
|
||||
"comp_hash": "review-hash",
|
||||
},
|
||||
}
|
||||
|
||||
_STOCK_CATALOG = json.dumps({"models": list(_STOCK_MODELS.values())})
|
||||
|
||||
|
||||
class _FakeRun:
|
||||
"""A `codex` that prints `stock` from a bare `debug models` and answers a catalog override with `returncode`.
|
||||
|
||||
`stock=None` is a Codex with no `debug models` at all: every call answers with `returncode` and `stderr`.
|
||||
"""
|
||||
|
||||
def __init__(self, returncode=0, stderr="", stock=_STOCK_CATALOG):
|
||||
self.returncode = returncode
|
||||
self.stderr = stderr
|
||||
self.stock = stock
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, args, **kwargs):
|
||||
self.calls.append((args, kwargs))
|
||||
if self.stock is not None and "model_catalog_json=" not in str(args):
|
||||
return subprocess.CompletedProcess(args, 0, self.stock, "")
|
||||
return subprocess.CompletedProcess(args, self.returncode, "", self.stderr)
|
||||
|
||||
|
||||
class _FakeJsonResponse:
|
||||
def __init__(self, status_code, payload=None):
|
||||
self.status_code = status_code
|
||||
|
|
@ -90,9 +198,7 @@ class TestAgentProfile:
|
|||
|
||||
class TestBuildAgentEnv:
|
||||
def test_anthropic_profile_uses_bare_root_and_bearer(self):
|
||||
env = build_agent_env(
|
||||
{}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"})
|
||||
)
|
||||
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}))
|
||||
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
|
||||
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
|
||||
assert env["ENABLE_TOOL_SEARCH"] == "true"
|
||||
|
|
@ -128,9 +234,7 @@ class TestBuildAgentEnv:
|
|||
assert "ANTHROPIC_API_KEY" not in env
|
||||
|
||||
def test_openai_profile_appends_v1(self):
|
||||
env = build_agent_env(
|
||||
{}, "http://localhost:4000/", "sk-key", frozenset({"openai"})
|
||||
)
|
||||
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}))
|
||||
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
|
||||
assert env["OPENAI_API_KEY"] == "sk-key"
|
||||
assert "ANTHROPIC_BASE_URL" not in env
|
||||
|
|
@ -138,9 +242,7 @@ class TestBuildAgentEnv:
|
|||
assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env
|
||||
|
||||
def test_both_profiles_set_everything(self):
|
||||
env = build_agent_env(
|
||||
{}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"})
|
||||
)
|
||||
env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}))
|
||||
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
|
||||
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
|
||||
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
|
||||
|
|
@ -148,9 +250,7 @@ class TestBuildAgentEnv:
|
|||
assert env["ENABLE_TOOL_SEARCH"] == "true"
|
||||
|
||||
def test_litellm_profile_exports_only_the_proxy_key(self):
|
||||
env = build_agent_env(
|
||||
{}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
|
||||
)
|
||||
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"litellm"}))
|
||||
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
|
||||
assert "ANTHROPIC_BASE_URL" not in env
|
||||
assert "OPENAI_BASE_URL" not in env
|
||||
|
|
@ -158,9 +258,7 @@ class TestBuildAgentEnv:
|
|||
|
||||
def test_preserves_unrelated_env_and_does_not_mutate_input(self):
|
||||
base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
|
||||
env = build_agent_env(
|
||||
base, "http://localhost:4000", "sk-key", frozenset({"anthropic"})
|
||||
)
|
||||
env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}))
|
||||
assert env["PATH"] == "/usr/bin"
|
||||
assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
|
||||
|
||||
|
|
@ -322,9 +420,7 @@ class TestOpencodeModelSync:
|
|||
assert "refused" in result.reason
|
||||
|
||||
def test_non_200_is_reported(self):
|
||||
result = opencode_model_sync_env(
|
||||
{}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
|
||||
)
|
||||
result = opencode_model_sync_env({}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500))
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "HTTP 500" in result.reason
|
||||
|
||||
|
|
@ -335,10 +431,10 @@ class TestOpencodeModelSync:
|
|||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "unexpected body" in result.reason
|
||||
|
||||
@pytest.mark.parametrize("command", ["claude", "codex", "/usr/bin/claude"])
|
||||
def test_only_opencode_syncs(self, command):
|
||||
@pytest.mark.parametrize("command", ["claude", "pi", "/usr/bin/claude"])
|
||||
def test_only_opencode_and_codex_sync(self, command):
|
||||
def boom(*a, **k):
|
||||
raise AssertionError("no agent other than opencode should call the proxy")
|
||||
raise AssertionError("no agent other than opencode or codex should call the proxy")
|
||||
|
||||
assert agent_model_sync_env(command, {}, "http://localhost:4000", "sk-key", False, get=boom) == {}
|
||||
|
||||
|
|
@ -367,7 +463,369 @@ class TestOpencodeModelSync:
|
|||
assert _default_of(opencode_model_sync_env, "get") is requests.get
|
||||
|
||||
|
||||
class TestCodexModelSync:
|
||||
@staticmethod
|
||||
def _listing(*models):
|
||||
return {"object": "list", "data": list(models)}
|
||||
|
||||
@staticmethod
|
||||
def _row(model_id, **extra):
|
||||
return {"id": model_id, "object": "model", "created": 1, "owned_by": "openai", **extra}
|
||||
|
||||
def _sync(self, listing, codex_home, base_url="http://localhost:4000/", run=None):
|
||||
captured = {}
|
||||
|
||||
def fake_get(url, headers, timeout):
|
||||
captured["url"] = url
|
||||
captured["headers"] = headers
|
||||
return _FakeResponse(200, listing)
|
||||
|
||||
result = codex_model_sync_args(
|
||||
{"CODEX_HOME": str(codex_home)},
|
||||
base_url,
|
||||
"sk-key",
|
||||
get=fake_get,
|
||||
run=_FakeRun() if run is None else run,
|
||||
)
|
||||
return captured, result
|
||||
|
||||
@staticmethod
|
||||
def _catalog_path(result):
|
||||
assert isinstance(result, ModelSyncArgs)
|
||||
flag, override = result.args
|
||||
assert flag == "-c"
|
||||
key, _, value = override.partition("=")
|
||||
assert key == "model_catalog_json"
|
||||
return json.loads(value)
|
||||
|
||||
def test_writes_catalog_under_codex_home_and_points_codex_at_it(self, tmp_path):
|
||||
listing = self._listing(self._row("gpt-5.5", mode="chat"), self._row("claude-opus-4-7"))
|
||||
captured, result = self._sync(listing, tmp_path / "codex")
|
||||
|
||||
assert captured["url"] == "http://localhost:4000/v1/models"
|
||||
assert captured["headers"] == {"Authorization": "Bearer sk-key"}
|
||||
path = self._catalog_path(result)
|
||||
assert path == str(tmp_path / "codex" / "litellm-models.json")
|
||||
text = (tmp_path / "codex" / "litellm-models.json").read_text()
|
||||
assert "sk-key" not in text
|
||||
catalog = json.loads(text)
|
||||
assert [m["slug"] for m in catalog["models"]] == ["gpt-5.5", "claude-opus-4-7"]
|
||||
assert [m["display_name"] for m in catalog["models"]] == ["GPT-5.5", "claude-opus-4-7"]
|
||||
assert [m["priority"] for m in catalog["models"]] == [0, 1]
|
||||
|
||||
def _entries(self, codex_home):
|
||||
return {m["slug"]: m for m in json.loads((codex_home / "litellm-models.json").read_text())["models"]}
|
||||
|
||||
def test_known_model_keeps_the_installed_codex_entry(self, tmp_path):
|
||||
self._sync(self._listing(self._row("gpt-5.5", mode="chat")), tmp_path)
|
||||
assert self._entries(tmp_path)["gpt-5.5"] == {**_STOCK_MODELS["gpt-5.5"], "priority": 0}
|
||||
|
||||
def test_hidden_stock_model_is_listed_when_the_proxy_serves_it(self, tmp_path):
|
||||
self._sync(self._listing(self._row("gpt-5.4")), tmp_path)
|
||||
entry = self._entries(tmp_path)["gpt-5.4"]
|
||||
assert entry["visibility"] == "list"
|
||||
assert entry["upgrade"] is None
|
||||
assert entry["supported_reasoning_levels"] == _STOCK_REASONING_LEVELS
|
||||
|
||||
def test_api_disabled_stock_model_is_selectable_when_the_proxy_serves_it(self, tmp_path):
|
||||
self._sync(self._listing(self._row("codex-auto-review")), tmp_path)
|
||||
entry = self._entries(tmp_path)["codex-auto-review"]
|
||||
assert entry["supported_in_api"] is True
|
||||
assert entry["visibility"] == "list"
|
||||
assert entry["base_instructions"] == _STOCK_MODELS["codex-auto-review"]["base_instructions"]
|
||||
|
||||
def test_stock_upgrade_nudge_survives_when_its_target_is_listed(self, tmp_path):
|
||||
self._sync(self._listing(self._row("gpt-5.4"), self._row("gpt-5.6-terra")), tmp_path)
|
||||
entries = self._entries(tmp_path)
|
||||
assert entries["gpt-5.4"]["upgrade"] == _STOCK_MODELS["gpt-5.4"]["upgrade"]
|
||||
assert [entries["gpt-5.4"]["priority"], entries["gpt-5.6-terra"]["priority"]] == [0, 1]
|
||||
|
||||
def test_stock_catalog_is_decoded_as_utf8_regardless_of_locale(self, tmp_path):
|
||||
description = "Modelo equilibrado para el trabajo diario, con acentos y ñ."
|
||||
catalog = {"models": [{**_STOCK_MODELS["gpt-5.5"], "description": description}]}
|
||||
stock = json.dumps(catalog, ensure_ascii=False).encode("utf-8")
|
||||
|
||||
def locale_bound_run(args, **kwargs):
|
||||
if "model_catalog_json=" in str(args):
|
||||
return subprocess.CompletedProcess(args, 0, "", "")
|
||||
return subprocess.CompletedProcess(args, 0, stock.decode(kwargs.get("encoding") or "ascii"), "")
|
||||
|
||||
_, result = self._sync(self._listing(self._row("gpt-5.5")), tmp_path, run=locale_bound_run)
|
||||
|
||||
assert isinstance(result, ModelSyncArgs)
|
||||
written = json.loads((tmp_path / "litellm-models.json").read_text(encoding="utf-8"))["models"]
|
||||
assert [m["description"] for m in written] == [description]
|
||||
|
||||
def test_unparseable_stock_catalog_is_reported(self, tmp_path):
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(stock="not json"))
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert result.reason.startswith("`codex debug models` printed no model catalog: ")
|
||||
assert not (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
def test_unknown_model_gets_the_fields_codex_requires(self, tmp_path):
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path)
|
||||
entry = json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]
|
||||
|
||||
assert entry["visibility"] == "list"
|
||||
assert entry["supported_in_api"] is True
|
||||
assert entry["shell_type"] == "unified_exec"
|
||||
assert entry["supported_reasoning_levels"] == []
|
||||
assert entry["truncation_policy"] == {"mode": "bytes", "limit": 10000}
|
||||
assert entry["experimental_supported_tools"] == []
|
||||
assert entry["support_verbosity"] is False
|
||||
assert entry["supports_reasoning_summaries"] is False
|
||||
assert entry["supports_parallel_tool_calls"] is False
|
||||
for nullable in ("description", "availability_nux", "upgrade", "default_verbosity", "apply_patch_tool_type"):
|
||||
assert nullable in entry and entry[nullable] is None
|
||||
assert entry["base_instructions"].startswith("You are a coding agent running in the Codex CLI")
|
||||
|
||||
def test_context_window_comes_from_max_input_tokens_for_unknown_models_only(self, tmp_path):
|
||||
listing = self._listing(
|
||||
self._row("big", max_input_tokens=400000),
|
||||
self._row("unknown"),
|
||||
self._row("gpt-5.5", max_input_tokens=400000),
|
||||
)
|
||||
self._sync(listing, tmp_path)
|
||||
models = self._entries(tmp_path)
|
||||
assert models["big"]["context_window"] == 400000
|
||||
assert models["unknown"]["context_window"] is None
|
||||
assert models["gpt-5.5"]["context_window"] == 272000
|
||||
|
||||
def test_non_chat_models_are_left_out(self, tmp_path):
|
||||
listing = self._listing(
|
||||
self._row("chat", mode="chat"),
|
||||
self._row("resp", mode="responses"),
|
||||
self._row("embed", mode="embedding"),
|
||||
self._row("img", mode="image_generation"),
|
||||
)
|
||||
self._sync(listing, tmp_path)
|
||||
slugs = {m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]}
|
||||
assert slugs == {"chat", "resp"}
|
||||
|
||||
def test_listing_without_chat_models_is_skipped_and_writes_nothing(self, tmp_path):
|
||||
_, result = self._sync(self._listing(self._row("embed", mode="embedding")), tmp_path)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "no chat models" in result.reason
|
||||
assert not (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
def test_catalog_is_rewritten_on_every_launch(self, tmp_path):
|
||||
self._sync(self._listing(self._row("old")), tmp_path)
|
||||
self._sync(self._listing(self._row("new")), tmp_path)
|
||||
slugs = [m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]]
|
||||
assert slugs == ["new"]
|
||||
|
||||
def test_catalog_is_replaced_whole_and_leaves_no_temp_files(self, tmp_path):
|
||||
self._sync(self._listing(*(self._row(f"m{i}") for i in range(50))), tmp_path)
|
||||
self._sync(self._listing(self._row("new")), tmp_path)
|
||||
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
|
||||
assert json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]["slug"] == "new"
|
||||
|
||||
def test_defaults_to_dot_codex_in_home(self, tmp_path):
|
||||
result = codex_model_sync_args(
|
||||
{},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
|
||||
run=_FakeRun(),
|
||||
home=lambda: tmp_path,
|
||||
)
|
||||
assert self._catalog_path(result) == str(tmp_path / ".codex" / "litellm-models.json")
|
||||
|
||||
def test_default_home_is_the_users(self):
|
||||
assert _default_of(codex_model_sync_args, "home") == Path.home
|
||||
|
||||
def test_missing_base_instructions_is_reported_not_raised(self, tmp_path):
|
||||
result = codex_model_sync_args(
|
||||
{"CODEX_HOME": str(tmp_path)},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
|
||||
instructions_path=tmp_path / "missing.md",
|
||||
)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "could not read" in result.reason
|
||||
assert not (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
def test_unwritable_catalog_path_is_reported_not_raised(self, tmp_path):
|
||||
blocker = tmp_path / "file"
|
||||
blocker.write_text("")
|
||||
_, result = self._sync(self._listing(self._row("m")), blocker / "codex")
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "could not write" in result.reason
|
||||
|
||||
def test_failed_replace_is_reported_and_leaves_no_temp_file(self, tmp_path):
|
||||
(tmp_path / "litellm-models.json").mkdir()
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "could not write" in result.reason
|
||||
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
|
||||
|
||||
def test_unreachable_proxy_is_reported_not_raised(self, tmp_path):
|
||||
def boom(*a, **k):
|
||||
raise requests.ConnectionError("refused")
|
||||
|
||||
result = codex_model_sync_args({"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=boom)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "refused" in result.reason
|
||||
assert not (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("response", "reason"),
|
||||
[(_FakeResponse(500), "HTTP 500"), (_FakeResponse(200, {"data": "nope"}), "unexpected body")],
|
||||
)
|
||||
def test_bad_response_is_reported(self, tmp_path, response, reason):
|
||||
result = codex_model_sync_args(
|
||||
{"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=lambda *a, **k: response
|
||||
)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert reason in result.reason
|
||||
|
||||
@pytest.mark.parametrize("binary", ["codex", "/opt/bin/codex", "codex.cmd", "/c/npm/codex.CMD"])
|
||||
def test_codex_syncs_through_the_agent_dispatch_with_the_binary_it_will_run(self, tmp_path, binary):
|
||||
run = _FakeRun()
|
||||
result = agent_model_sync_env(
|
||||
binary,
|
||||
{"CODEX_HOME": str(tmp_path)},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
False,
|
||||
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
|
||||
run=run,
|
||||
)
|
||||
assert self._catalog_path(result) == str(tmp_path / "litellm-models.json")
|
||||
assert len(run.calls) == 2
|
||||
assert all(binary in command for command, _ in run.calls)
|
||||
|
||||
def test_opencode_dispatch_never_runs_codex(self):
|
||||
def boom(*a, **k):
|
||||
raise AssertionError("only the Codex sync reads its catalog back")
|
||||
|
||||
result = agent_model_sync_env(
|
||||
"opencode",
|
||||
{},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
False,
|
||||
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
|
||||
run=boom,
|
||||
)
|
||||
assert "OPENCODE_CONFIG_CONTENT" in result
|
||||
|
||||
def test_codex_lists_its_own_models_then_reads_the_catalog_back_before_launch(self, tmp_path):
|
||||
run = _FakeRun()
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
|
||||
path = self._catalog_path(result)
|
||||
|
||||
assert [command for command, _ in run.calls] == [
|
||||
("codex", "debug", "models"),
|
||||
("codex", "-c", f"model_catalog_json={json.dumps(path)}", "debug", "models"),
|
||||
]
|
||||
for _, options in run.calls:
|
||||
assert options["env"] == {"CODEX_HOME": str(tmp_path)}
|
||||
assert options["stdin"] is subprocess.DEVNULL
|
||||
assert options["capture_output"] is True
|
||||
assert options["encoding"] == "utf-8"
|
||||
assert options["timeout"] == 10
|
||||
|
||||
def test_codex_rejecting_the_catalog_skips_the_sync_and_keeps_the_file(self, tmp_path):
|
||||
stderr = (
|
||||
"Error: failed to parse model_catalog_json path `/home/me/.codex/litellm-models.json` as JSON: "
|
||||
"missing field `supports_parallel_tool_calls` at line 1 column 21648\n"
|
||||
)
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1, stderr))
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert result.reason == (
|
||||
"`codex debug models` exited 1: Error: failed to parse model_catalog_json path "
|
||||
"`/home/me/.codex/litellm-models.json` as JSON: missing field `supports_parallel_tool_calls` "
|
||||
"at line 1 column 21648"
|
||||
)
|
||||
assert (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
def test_codex_without_debug_models_skips_the_sync(self, tmp_path):
|
||||
stderr = "error: unrecognized subcommand 'models'\n\nUsage: codex debug [OPTIONS] <COMMAND>\n"
|
||||
run = _FakeRun(2, stderr, stock=None)
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert result.reason == "`codex debug models` exited 2: error: unrecognized subcommand 'models'"
|
||||
assert len(run.calls) == 1
|
||||
assert not (tmp_path / "litellm-models.json").exists()
|
||||
|
||||
def test_codex_failing_silently_is_reported(self, tmp_path):
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1))
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert result.reason == "`codex debug models` exited 1: no output"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error", [OSError("codex vanished"), subprocess.TimeoutExpired("codex", 10)], ids=["oserror", "timeout"]
|
||||
)
|
||||
def test_unrunnable_preflight_is_reported_not_raised(self, tmp_path, error):
|
||||
def failing_run(*a, **k):
|
||||
raise error
|
||||
|
||||
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=failing_run)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert result.reason.startswith("`codex debug models` failed: ")
|
||||
assert str(error) in result.reason
|
||||
|
||||
def test_windows_shim_preflight_goes_through_cmd_exe(self, tmp_path):
|
||||
shim = _WINDOWS_CLAUDE_CMD.replace("claude", "codex")
|
||||
run = _FakeRun()
|
||||
result = codex_model_sync_args(
|
||||
{"CODEX_HOME": str(tmp_path)},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
binary=shim,
|
||||
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
|
||||
run=run,
|
||||
)
|
||||
override = f"model_catalog_json={json.dumps(self._catalog_path(result))}"
|
||||
doubled = override.replace('"', '""')
|
||||
assert [command for command, _ in run.calls] == [
|
||||
f'{_CMD_PREFIX}""{shim}" "debug" "models""',
|
||||
f'{_CMD_PREFIX}""{shim}" "-c" "{doubled}" "debug" "models""',
|
||||
]
|
||||
|
||||
def test_default_binary_is_codex_on_path(self):
|
||||
assert _default_of(codex_model_sync_args, "binary") == "codex"
|
||||
|
||||
def test_default_runner_is_subprocess_run(self):
|
||||
assert _default_of(codex_model_sync_args, "run") is subprocess.run
|
||||
assert _default_of(agent_model_sync_env, "run") is subprocess.run
|
||||
|
||||
def test_skip_verify_keeps_the_launch_offline(self):
|
||||
def boom(*a, **k):
|
||||
raise AssertionError("--skip-verify must not touch the proxy")
|
||||
|
||||
result = agent_model_sync_env("codex", {}, "http://localhost:4000", "sk-key", True, get=boom)
|
||||
assert isinstance(result, ModelSyncSkipped)
|
||||
assert "--skip-verify" in result.reason
|
||||
|
||||
def test_default_http_client_is_requests_get(self):
|
||||
assert _default_of(codex_model_sync_args, "get") is requests.get
|
||||
|
||||
|
||||
class TestRunAgent:
|
||||
def test_synced_args_precede_user_args_and_follow_provider_overrides(self):
|
||||
calls = {}
|
||||
run_agent(
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
["codex", "exec", "hi"],
|
||||
base_env={},
|
||||
sync_models=lambda *a: ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')),
|
||||
which=lambda name: "/usr/local/bin/codex",
|
||||
verify=lambda *a: None,
|
||||
launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
|
||||
)
|
||||
args = calls["args"]
|
||||
assert args[-2:] == ("exec", "hi")
|
||||
assert args[args.index('model_catalog_json="/tmp/c.json"') - 1] == "-c"
|
||||
assert (
|
||||
args.index('model_provider="litellm"') < args.index('model_catalog_json="/tmp/c.json"') < args.index("exec")
|
||||
)
|
||||
assert calls["env"]["OPENAI_API_KEY"] == "sk-key"
|
||||
assert "model_catalog_json" not in json.dumps(calls["env"])
|
||||
|
||||
def test_synced_model_config_reaches_the_agent_alongside_profile_env(self):
|
||||
calls = {}
|
||||
run_agent(
|
||||
|
|
@ -405,7 +863,13 @@ class TestRunAgent:
|
|||
launcher=lambda p, a, e: order.append("launch"),
|
||||
)
|
||||
assert order == ["verify", "sync", "launch"]
|
||||
assert calls["args"] == ("opencode", {"HOME": "/home/me"}, "http://localhost:4000", "sk-key", False)
|
||||
assert calls["args"] == (
|
||||
"/usr/local/bin/opencode",
|
||||
{"HOME": "/home/me"},
|
||||
"http://localhost:4000",
|
||||
"sk-key",
|
||||
False,
|
||||
)
|
||||
|
||||
def test_unreachable_proxy_is_not_asked_for_models(self):
|
||||
def failing_verify(*a):
|
||||
|
|
@ -1050,10 +1514,7 @@ class TestAgentCommands:
|
|||
assert captured["api_key"] == "sk-key"
|
||||
assert captured["command"] == ["claude", "--resume", "-p", "hi"]
|
||||
assert captured["skip_verify"] is False
|
||||
assert (
|
||||
"routing Claude Code through proxy at http://localhost:4000"
|
||||
in result.output
|
||||
)
|
||||
assert "routing Claude Code through proxy at http://localhost:4000" in result.output
|
||||
|
||||
def test_codex_shows_friendly_name(self):
|
||||
captured = {}
|
||||
|
|
@ -1126,14 +1587,10 @@ class TestAgentCommands:
|
|||
with (
|
||||
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
|
||||
patch(f"{AGENTS_MODULE}.login", fake_login),
|
||||
patch(
|
||||
f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login"
|
||||
) as mock_get,
|
||||
patch(f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login") as mock_get,
|
||||
patch(
|
||||
f"{AGENTS_MODULE}.run_agent",
|
||||
side_effect=lambda base_url, api_key, command, **k: captured.update(
|
||||
api_key=api_key
|
||||
),
|
||||
side_effect=lambda base_url, api_key, command, **k: captured.update(api_key=api_key),
|
||||
),
|
||||
):
|
||||
result = self.runner.invoke(
|
||||
|
|
|
|||
|
|
@ -512,8 +512,8 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
|
|||
the repair must be skipped and the existing 400 raised immediately, while bodies
|
||||
at or below the limit still get repaired.
|
||||
|
||||
`\\ud83d` is a lone high-surrogate escape: orjson rejects it, the json fallback
|
||||
accepts it, so a body containing it is only salvaged when the repair path runs.
|
||||
`NaN` is rejected by orjson and accepted by the json fallback, so a body containing
|
||||
it is only salvaged when the repair path runs.
|
||||
"""
|
||||
import litellm.proxy.common_utils.http_parsing_utils as http_parsing_utils
|
||||
|
||||
|
|
@ -522,14 +522,14 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
|
|||
http_parsing_utils, "MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 100 / (1024 * 1024)
|
||||
)
|
||||
|
||||
small_body = b'{"model":"gpt-4o","x":"\\ud83d"}'
|
||||
small_body = b'{"model":"gpt-4o","x":NaN}'
|
||||
assert len(small_body) <= 100
|
||||
repaired = await _read_request_body(_make_json_request(small_body))
|
||||
assert repaired["model"] == "gpt-4o"
|
||||
|
||||
padding = "a" * 200
|
||||
large_body = (
|
||||
b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":"\\ud83d"}'
|
||||
b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":NaN}'
|
||||
)
|
||||
assert len(large_body) > 100
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
|
|
@ -546,6 +546,33 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
|
|||
assert repaired_large["model"] == "gpt-4o"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[
|
||||
pytest.param(b"say ok \\ud83d", id="lone-high-surrogate"),
|
||||
pytest.param(b"say ok \\ude00", id="lone-low-surrogate"),
|
||||
pytest.param(b"\\ud83d\\ud83d\\ude00", id="lone-high-before-valid-pair"),
|
||||
],
|
||||
)
|
||||
async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes):
|
||||
"""
|
||||
orjson rejects a lone surrogate escape, and the json fallback accepts it, so the
|
||||
parsed body used to carry a code point no provider request can UTF-8 encode. That
|
||||
surfaced as a 500 from the provider handler instead of a 400 for the bad input.
|
||||
"""
|
||||
body = b'{"model":"gpt-4o","messages":[{"role":"user","content":"' + content + b'"}]}'
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _read_request_body(_make_json_request(body))
|
||||
assert exc_info.value.code == "400"
|
||||
assert exc_info.value.type == "invalid_request_error"
|
||||
assert "Invalid JSON payload" in exc_info.value.message
|
||||
|
||||
paired = body.replace(content, b"say ok \\ud83d\\ude00")
|
||||
parsed = await _read_request_body(_make_json_request(paired))
|
||||
assert parsed["messages"][0]["content"] == "say ok \U0001F600"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_form_data():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -36,7 +36,11 @@ from litellm.proxy._types import (
|
|||
UpdateKeyRequest,
|
||||
)
|
||||
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_delete_cache_key_object,
|
||||
_project_cache_key,
|
||||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
|
|
@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
|
|||
|
||||
|
||||
class _JWTMappingRow:
|
||||
def __init__(self, token, jwt_claim_name, jwt_claim_value):
|
||||
def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None):
|
||||
self.token = token
|
||||
self.jwt_claim_name = jwt_claim_name
|
||||
self.jwt_claim_value = jwt_claim_value
|
||||
self.jwt_issuer = jwt_issuer
|
||||
|
||||
|
||||
class _CascadingJWTMappingTable:
|
||||
|
|
@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
|
|||
),
|
||||
)
|
||||
|
||||
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
|
||||
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new
|
|||
_execute_virtual_key_regeneration,
|
||||
)
|
||||
|
||||
stale_cache_key = "jwt_key_mapping:sub:user1"
|
||||
stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None)
|
||||
existing_key = _make_regenerate_existing_key()
|
||||
mock_prisma_client = _make_regenerate_mock_prisma()
|
||||
mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(
|
||||
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")]
|
||||
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)]
|
||||
)
|
||||
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
|
||||
return_value=MagicMock(token="new-hashed-token")
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
|
|
@ -5934,6 +5935,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"])
|
||||
def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str):
|
||||
"""A provider route that resolved a router deployment stashes its model_info on request.state. That
|
||||
deployment, not a model_info the client put in its own body, is what spend logs and metrics attribute
|
||||
the call to (LIT-1761: passthrough successes carried model_id="")."""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
|
||||
mock_request.headers = Headers({})
|
||||
mock_request.scope = {}
|
||||
mock_request.state = SimpleNamespace(
|
||||
**{LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: {"id": "vertex-gemini-38-flash-dep"}}
|
||||
)
|
||||
|
||||
kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=mock_request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
passthrough_logging_payload=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
_parsed_body={client_metadata_key: {"model_info": {"id": "client-forged-id"}}},
|
||||
litellm_call_id="lit-1761-call-id",
|
||||
)
|
||||
|
||||
assert kwargs["litellm_params"]["metadata"]["model_info"] == {"id": "vertex-gemini-38-flash-dep"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -1,13 +1,19 @@
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import Headers, State
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
VertexAIPassThroughHandler,
|
||||
_base_vertex_proxy_route,
|
||||
_resolve_vertex_model_from_router,
|
||||
_upstream_headers_for_vertex_route,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
HttpPassThroughEndpointHelpers,
|
||||
)
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
|
||||
|
|
@ -758,3 +764,115 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
|
|||
assert (
|
||||
"gemini-3-pro" in target_url
|
||||
), f"Actual Vertex AI model name should be in target URL. Got: {target_url}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment():
|
||||
"""The router deployment that rewrote the upstream URL is the one the logging kwargs must name, so
|
||||
the Prometheus model_id label (and SpendLogs.model_id) on a Vertex passthrough success reads the
|
||||
deployment's id instead of "" (LIT-1761)."""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
|
||||
mock_request.headers = Headers({})
|
||||
mock_request.scope = {}
|
||||
mock_request.state = State()
|
||||
mock_handler = MagicMock()
|
||||
mock_handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_available_deployment_for_pass_through.return_value = {
|
||||
"model_name": "gemini-3.8-flash",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-3.8-flash",
|
||||
"vertex_project": "p",
|
||||
"vertex_location": "global",
|
||||
"use_in_pass_through": True,
|
||||
},
|
||||
"model_info": {"id": "vertex-gemini-38-flash-dep"},
|
||||
}
|
||||
|
||||
async def relay_returning_logging_kwargs(
|
||||
request: Request, fastapi_response: object, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> dict:
|
||||
return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
passthrough_logging_payload=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
_parsed_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
|
||||
litellm_call_id="lit-1761-call-id",
|
||||
)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
|
||||
"litellm.proxy.proxy_server.llm_router", mock_router
|
||||
),
|
||||
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch( # test-quality-ok: the route offers no injection point for its header preparation
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=({}, False, "p", "global"),
|
||||
),
|
||||
patch( # test-quality-ok: the relay is captured here to read the logging kwargs, the route offers no seam
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
return_value=relay_returning_logging_kwargs,
|
||||
),
|
||||
patch( # test-quality-ok: the route calls auth directly rather than through Depends
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
return_value=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
|
||||
logging_kwargs = await _base_vertex_proxy_route(
|
||||
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
request=mock_request,
|
||||
fastapi_response=MagicMock(),
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
)
|
||||
|
||||
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"
|
||||
|
||||
|
||||
def _router_without_deployment() -> MagicMock:
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = None
|
||||
return router
|
||||
|
||||
|
||||
def _router_raising_on_lookup() -> MagicMock:
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.side_effect = ValueError("no healthy deployment")
|
||||
return router
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"llm_router",
|
||||
[None, _router_without_deployment(), _router_raising_on_lookup()],
|
||||
ids=["no-router", "no-matching-deployment", "lookup-raises"],
|
||||
)
|
||||
def test_vertex_passthrough_without_a_resolved_deployment_keeps_the_url_and_reports_no_model_info(
|
||||
llm_router: MagicMock | None,
|
||||
):
|
||||
"""A Vertex passthrough call that no router deployment serves must keep the URL-derived values and carry no
|
||||
deployment model_info, so logging cannot attribute it to a deployment that never handled it."""
|
||||
resolved = _resolve_vertex_model_from_router(
|
||||
model_id="gemini-3.8-flash",
|
||||
llm_router=llm_router,
|
||||
encoded_endpoint="/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
vertex_project="url-project",
|
||||
vertex_location="url-location",
|
||||
)
|
||||
|
||||
assert resolved == (
|
||||
"/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
"v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
"url-project",
|
||||
"url-location",
|
||||
None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10750,6 +10750,7 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
|
|||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
|
||||
monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h")
|
||||
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
|
@ -10771,6 +10772,10 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
|
|||
assert fields["enable_anthropic_prompt_caching"]["field_tab"] == "prompt_caching"
|
||||
assert fields["anthropic_prompt_caching_ttl"]["field_tab"] == "prompt_caching"
|
||||
assert fields["budget_exceeded_throttle_percentage"]["field_tab"] is None
|
||||
|
||||
assert fields["openai_system_messages_first"]["field_type"] == "Boolean"
|
||||
assert fields["openai_system_messages_first"]["field_value"] is False
|
||||
assert fields["openai_system_messages_first"]["field_tab"] == "prompt_caching"
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
|
@ -10975,6 +10980,7 @@ def test_general_settings_ui_defaults_unchanged_for_existing_fields():
|
|||
[
|
||||
("enable_anthropic_prompt_caching", True),
|
||||
("anthropic_prompt_caching_ttl", "1h"),
|
||||
("openai_system_messages_first", True),
|
||||
],
|
||||
)
|
||||
def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_name, db_value):
|
||||
|
|
@ -11033,6 +11039,8 @@ def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypa
|
|||
("enable_anthropic_prompt_caching", False),
|
||||
("anthropic_prompt_caching_ttl", "5m"),
|
||||
("anthropic_prompt_caching_ttl", "1h"),
|
||||
("openai_system_messages_first", True),
|
||||
("openai_system_messages_first", False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -11081,6 +11089,8 @@ async def test_update_config_field_prompt_caching_persists_to_litellm_settings(m
|
|||
("anthropic_prompt_caching_ttl", "10m"),
|
||||
("anthropic_prompt_caching_ttl", "1H"),
|
||||
("anthropic_prompt_caching_ttl", 3600),
|
||||
("openai_system_messages_first", "yes"),
|
||||
("openai_system_messages_first", 1),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -11120,6 +11130,7 @@ async def test_update_config_field_prompt_caching_rejects_invalid(monkeypatch, f
|
|||
[
|
||||
("enable_anthropic_prompt_caching", False),
|
||||
("anthropic_prompt_caching_ttl", None),
|
||||
("openai_system_messages_first", False),
|
||||
("budget_exceeded_throttle_percentage", None),
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1360,6 +1360,56 @@ def test_add_invalid_provider_to_router():
|
|||
assert router.pattern_router.patterns == {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str:
|
||||
from litellm import CustomLLM
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
class OnPremLLM(CustomLLM):
|
||||
def completion(self, *args, **kwargs) -> ModelResponse:
|
||||
return litellm.completion(
|
||||
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler"
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}])
|
||||
monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list))
|
||||
monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers))
|
||||
return "test-onprem-llm"
|
||||
|
||||
|
||||
def test_router_init_accepts_custom_provider_map_prefix_before_first_completion(registered_custom_provider: str):
|
||||
assert registered_custom_provider not in litellm.provider_list
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "onprem", "litellm_params": {"model": f"{registered_custom_provider}/my-model"}},
|
||||
],
|
||||
)
|
||||
|
||||
assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == (
|
||||
f"{registered_custom_provider}/my-model"
|
||||
)
|
||||
response = router.completion(model="onprem", messages=[{"role": "user", "content": "hi"}])
|
||||
assert response.choices[0].message.content == "served by onprem handler"
|
||||
|
||||
|
||||
def test_router_add_deployment_accepts_explicit_custom_provider_from_custom_provider_map(
|
||||
registered_custom_provider: str,
|
||||
):
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
router = litellm.Router(model_list=[])
|
||||
|
||||
router.add_deployment(
|
||||
Deployment(
|
||||
model_name="onprem",
|
||||
litellm_params={"model": "my-model", "custom_llm_provider": registered_custom_provider},
|
||||
)
|
||||
)
|
||||
|
||||
assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == "my-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_ageneric_api_call_with_fallbacks_helper():
|
||||
"""
|
||||
|
|
@ -12156,6 +12206,83 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o
|
|||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,provider,expected",
|
||||
[
|
||||
("anthropic/claude-opus-5", None, True),
|
||||
("claude-opus-4-8", None, True),
|
||||
("anthropic/claude-opus-4-7", None, False),
|
||||
("anthropic/claude-opus-4-6", None, False),
|
||||
("anthropic/claude-sonnet-5", None, False),
|
||||
("anthropic/off-map-opus", None, False),
|
||||
("vertex_ai/claude-opus-5", None, False),
|
||||
("bedrock/claude-opus-5", None, False),
|
||||
("claude-opus-5", "vertex_ai", False),
|
||||
("claude-opus-5", "bedrock", False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("operator_flag", [True, False])
|
||||
def test_model_group_info_fast_mode_uses_exact_provider_catalog(
|
||||
local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool
|
||||
) -> None:
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "fast-group",
|
||||
"litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"},
|
||||
"model_info": {"supports_fast_mode": operator_flag},
|
||||
}])
|
||||
|
||||
result: Final = router.get_model_group_info("fast-group")
|
||||
|
||||
assert result is not None
|
||||
assert result.supports_fast_mode is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flag", [None, False, "true", 1])
|
||||
def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean(
|
||||
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object
|
||||
) -> None:
|
||||
entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items()
|
||||
if key != "supports_fast_mode"}
|
||||
if flag is not None:
|
||||
entry["supports_fast_mode"] = flag
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry)
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "fast-group",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"},
|
||||
"model_info": {"supports_fast_mode": True},
|
||||
}])
|
||||
|
||||
result: Final = router.get_model_group_info("fast-group")
|
||||
|
||||
assert result is not None
|
||||
assert result.supports_fast_mode is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("other_model,expected", [
|
||||
("anthropic/claude-opus-4-8", True),
|
||||
("anthropic/claude-opus-4-7", False),
|
||||
("anthropic/off-map-opus", False),
|
||||
("vertex_ai/claude-opus-5", False),
|
||||
("bedrock/claude-opus-5", False),
|
||||
])
|
||||
@pytest.mark.parametrize("reverse", [True, False])
|
||||
def test_model_group_info_fast_mode_requires_every_deployment(
|
||||
local_model_cost_map: None, other_model: str, expected: bool, reverse: bool
|
||||
) -> None:
|
||||
models: Final = (other_model, "anthropic/claude-opus-5") if reverse else (
|
||||
"anthropic/claude-opus-5", other_model
|
||||
)
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "fast-group",
|
||||
"litellm_params": {"model": model, "api_key": "fake-key"},
|
||||
} for model in models])
|
||||
|
||||
result: Final = router.get_model_group_info("fast-group")
|
||||
|
||||
assert result is not None
|
||||
assert result.supports_fast_mode is expected
|
||||
|
||||
|
||||
def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map):
|
||||
"""``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose
|
||||
registry entry declares parallel function calling must flip the group to True instead of False."""
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from jsonschema import validate
|
|||
|
||||
import litellm
|
||||
from litellm._internal_context import is_internal_call
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
|
||||
from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
|
||||
from litellm._logging import (
|
||||
CorrelationContextFilter,
|
||||
|
|
@ -32,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor
|
||||
from litellm.proxy.utils import is_valid_api_key
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
Delta,
|
||||
|
|
@ -44,6 +47,7 @@ from litellm.types.utils import (
|
|||
from litellm.types.utils import all_litellm_params, bedrock_batch_litellm_params
|
||||
from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams
|
||||
from litellm.utils import (
|
||||
CustomStreamWrapper,
|
||||
ProviderConfigManager,
|
||||
TextCompletionStreamWrapper,
|
||||
_check_provider_match,
|
||||
|
|
@ -1099,6 +1103,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"supports_sampling_params": {"type": "boolean"},
|
||||
"supports_output_config": {"type": "boolean"},
|
||||
"supports_speed": {"type": "boolean"},
|
||||
"supports_fast_mode": {"type": "boolean"},
|
||||
"supported_audio_formats": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
|
|
@ -4953,6 +4958,208 @@ async def test_wrapper_async_restores_originating_task_context_after_success(mon
|
|||
session_id_var.set("")
|
||||
|
||||
|
||||
class _ConvertStreamDeploymentHook(CustomLogger):
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, object], call_type: CallTypes | None
|
||||
) -> dict[str, object] | None:
|
||||
if not kwargs.get("stream"):
|
||||
return None
|
||||
return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True}
|
||||
|
||||
|
||||
class _SuccessKwargsCapture(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.success_kwargs: list[dict[str, object]] = []
|
||||
self.stream_event_responses: list[object] = []
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.success_kwargs.append(kwargs)
|
||||
|
||||
async def async_log_stream_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.stream_event_responses.append(response_obj)
|
||||
|
||||
|
||||
def _install_converted_stream_callbacks(monkeypatch: pytest.MonkeyPatch) -> _SuccessKwargsCapture:
|
||||
capture: Final = _SuccessKwargsCapture()
|
||||
monkeypatch.setattr(litellm, "callbacks", [_ConvertStreamDeploymentHook(), capture])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [])
|
||||
return capture
|
||||
|
||||
|
||||
async def _wait_for_success_kwargs(capture: _SuccessKwargsCapture, count: int = 1) -> dict[str, object]:
|
||||
for _ in range(50):
|
||||
if len(capture.success_kwargs) >= count and not _PENDING_CACHE_WRITES:
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
await asyncio.sleep(0.2)
|
||||
assert len(capture.success_kwargs) == count
|
||||
return capture.success_kwargs[-1]
|
||||
|
||||
|
||||
def _assert_cache_hit_logged_as_stream(capture: _SuccessKwargsCapture, success_kwargs: dict[str, object]) -> None:
|
||||
standard_logging_object: Final = success_kwargs["standard_logging_object"]
|
||||
assert isinstance(standard_logging_object, dict)
|
||||
assert standard_logging_object["cache_hit"] is True
|
||||
assert standard_logging_object["stream"] is True
|
||||
assert success_kwargs["stream"] is True
|
||||
assert capture.stream_event_responses == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wrapper_async_logs_converted_chat_stream_with_standard_logging_object(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
capture: Final = _install_converted_stream_callbacks(monkeypatch)
|
||||
|
||||
response: Final = await litellm.acompletion(
|
||||
model="gpt-5.6",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
mock_response="converted stream body",
|
||||
num_retries=0,
|
||||
)
|
||||
assert isinstance(response, CustomStreamWrapper)
|
||||
chunks: Final = [chunk async for chunk in response]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "converted stream body"
|
||||
|
||||
success_kwargs: Final = await _wait_for_success_kwargs(capture)
|
||||
standard_logging_object: Final = success_kwargs["standard_logging_object"]
|
||||
assert isinstance(standard_logging_object, dict)
|
||||
assert standard_logging_object["response_cost"] > 0
|
||||
assert standard_logging_object["stream"] is True
|
||||
assert success_kwargs["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_wrapper_async_logs_converted_responses_stream_with_standard_logging_object(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
capture: Final = _install_converted_stream_callbacks(monkeypatch)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
respx.post("https://api.openai.com/v1/responses").respond(
|
||||
json={
|
||||
"id": "resp_converted",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5.6",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_converted",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
|
||||
}
|
||||
)
|
||||
|
||||
response: Final = await litellm.aresponses(
|
||||
model="openai/gpt-5.6", input="hi", stream=True, api_key="sk-test", num_retries=0
|
||||
)
|
||||
assert isinstance(response, BaseResponsesAPIStreamingIterator)
|
||||
events: Final = [event async for event in response]
|
||||
assert events[-1].type == "response.completed"
|
||||
|
||||
success_kwargs: Final = await _wait_for_success_kwargs(capture)
|
||||
standard_logging_object: Final = success_kwargs["standard_logging_object"]
|
||||
assert isinstance(standard_logging_object, dict)
|
||||
assert standard_logging_object["response_cost"] > 0
|
||||
assert standard_logging_object["stream"] is True
|
||||
assert success_kwargs["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wrapper_async_replays_cached_converted_chat_stream_as_stream(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
capture: Final = _install_converted_stream_callbacks(monkeypatch)
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
request: Final = {
|
||||
"model": "gpt-5.6",
|
||||
"messages": [{"role": "user", "content": "replay me from cache"}],
|
||||
"stream": True,
|
||||
"mock_response": "converted stream body",
|
||||
"num_retries": 0,
|
||||
}
|
||||
|
||||
first: Final = await litellm.acompletion(**request)
|
||||
first_chunks: Final = [chunk async for chunk in first]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in first_chunks) == "converted stream body"
|
||||
await _wait_for_success_kwargs(capture)
|
||||
|
||||
replay: Final = await litellm.acompletion(**request)
|
||||
assert isinstance(replay, CustomStreamWrapper)
|
||||
replay_chunks: Final = [chunk async for chunk in replay]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in replay_chunks) == "converted stream body"
|
||||
|
||||
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_wrapper_async_replays_cached_converted_responses_stream_as_stream(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
capture: Final = _install_converted_stream_callbacks(monkeypatch)
|
||||
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
route: Final = respx.post("https://api.openai.com/v1/responses").respond(
|
||||
json={
|
||||
"id": "resp_cached_converted",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5.6",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_cached_converted",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
|
||||
}
|
||||
)
|
||||
request: Final = {
|
||||
"model": "openai/gpt-5.6",
|
||||
"input": "replay me from cache",
|
||||
"stream": True,
|
||||
"api_key": "sk-test",
|
||||
"num_retries": 0,
|
||||
}
|
||||
|
||||
first: Final = await litellm.aresponses(**request)
|
||||
assert [event async for event in first][-1].type == "response.completed"
|
||||
await _wait_for_success_kwargs(capture)
|
||||
|
||||
replay: Final = await litellm.aresponses(**request)
|
||||
assert isinstance(replay, BaseResponsesAPIStreamingIterator)
|
||||
assert [event async for event in replay][-1].type == "response.completed"
|
||||
assert route.call_count == 1
|
||||
|
||||
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
|
||||
|
||||
|
||||
def test_function_setup_failure_after_logging_construction_restores_context(monkeypatch):
|
||||
"""If function_setup() constructs Logging() (which already mutated
|
||||
trace_id_var/session_id_var in __init__) but then raises before returning,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
|
||||
import { fireEvent, screen, waitFor, within } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
|
||||
import { beforeEach, describe, expect, it, Mock, vi } from "vitest";
|
||||
|
||||
import { renderWithProviders } from "../../../../../tests/test-utils";
|
||||
import AllModelsTab from "./AllModelsTab";
|
||||
import { STATUS_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
|
||||
|
||||
|
|
@ -111,6 +113,9 @@ const setModelsInfo = (rows: Record<string, unknown>[], totalCount = rows.length
|
|||
|
||||
const lastModelsInfoCall = (): ModelsInfoArgs => modelsInfoCalls[modelsInfoCalls.length - 1];
|
||||
|
||||
const lastUrlParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>): URLSearchParams | undefined =>
|
||||
onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
|
||||
|
||||
const SEARCH_SETTLE_MS = 400;
|
||||
|
||||
const MOCK_AUTHORIZED = {
|
||||
|
|
@ -121,6 +126,8 @@ const MOCK_AUTHORIZED = {
|
|||
userId: "user-123",
|
||||
userEmail: "test@example.com",
|
||||
userRole: "Admin",
|
||||
userRoleLabel: "Admin",
|
||||
isViewOnly: false,
|
||||
premiumUser: true,
|
||||
disabledPersonalKeyCreation: false,
|
||||
showSSOBanner: false,
|
||||
|
|
@ -149,14 +156,14 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("renders the fetched models and the server row count", async () => {
|
||||
setModelsInfo([makeRow()], 137);
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(await screen.findByText("gpt-4")).toBeInTheDocument();
|
||||
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 137");
|
||||
});
|
||||
|
||||
it("does not re-query after the mount-time debounced search settles unchanged", async () => {
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
const callsAfterMount = modelsInfoCalls.length;
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, SEARCH_SETTLE_MS));
|
||||
|
|
@ -166,14 +173,14 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("shows the empty state when the proxy returns no models", () => {
|
||||
setModelsInfo([], 0);
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(screen.getByText("No models found")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the loading skeleton while the first page is in flight", () => {
|
||||
setModelsInfo([], 0, true);
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
|
||||
expect(screen.queryByText("No models found")).not.toBeInTheDocument();
|
||||
|
|
@ -197,7 +204,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it.each(cases)("sorts %s using the server field %s", async (_label, columnId, serverField, firstDirection) => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(sortHeader(columnId));
|
||||
await expectIndicator(columnId, firstDirection);
|
||||
|
|
@ -212,7 +219,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("cycles a sorted column back to unsorted", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(sortHeader("model_info_updated_at"));
|
||||
await expectIndicator("model_info_updated_at", "asc");
|
||||
|
|
@ -230,7 +237,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("queries the selected team and resets to the first page", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(lastModelsInfoCall().teamId).toBeUndefined();
|
||||
|
||||
|
|
@ -244,8 +251,7 @@ describe("AllModelsTab", () => {
|
|||
});
|
||||
|
||||
it("debounces the model name search into the server query", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
|
||||
|
||||
|
|
@ -254,9 +260,138 @@ describe("AllModelsTab", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("URL persistence", () => {
|
||||
it("writes the typed search to the URL and drops the page so a reload keeps the search", async () => {
|
||||
setModelsInfo([makeRow()], 200);
|
||||
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "3" }, onUrlUpdate });
|
||||
expect(lastModelsInfoCall().page).toBe(3);
|
||||
|
||||
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("model_search")).toBe("claude");
|
||||
});
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
|
||||
await waitFor(() => {
|
||||
expect(lastModelsInfoCall().page).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
it("restores the search box and server query from ?model_search= on mount", () => {
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { model_search: "haiku" } });
|
||||
|
||||
expect(screen.getByTestId("datatable-search")).toHaveValue("haiku");
|
||||
expect(lastModelsInfoCall().search).toBe("haiku");
|
||||
});
|
||||
|
||||
it("restores team, sort, page and page size from the URL into the server query", () => {
|
||||
setModelsInfo([makeRow()], 200);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, {
|
||||
searchParams: {
|
||||
filter_team: "team-1",
|
||||
sort_by: "model_info_updated_at",
|
||||
sort_order: "desc",
|
||||
page: "2",
|
||||
page_size: "25",
|
||||
},
|
||||
});
|
||||
|
||||
const expectedQuery: ModelsInfoArgs = {
|
||||
teamId: "team-1",
|
||||
sortBy: "updated_at",
|
||||
sortOrder: "desc",
|
||||
page: 2,
|
||||
size: 25,
|
||||
};
|
||||
expect(lastModelsInfoCall()).toMatchObject(expectedQuery);
|
||||
expect(screen.getByTestId("models-team-select")).toHaveTextContent("Engineering");
|
||||
});
|
||||
|
||||
it("restores the access group and view mode from the URL", () => {
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, {
|
||||
searchParams: { access_group: "sales-team", view_mode: "all" },
|
||||
});
|
||||
|
||||
expect(lastModelsInfoCall().accessGroup).toBe("sales-team");
|
||||
expect(screen.queryByText(/To access these models/)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("clamps a hand-edited page and page size into the range the table supports", () => {
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "0", page_size: "5000" } });
|
||||
|
||||
expect(lastModelsInfoCall().page).toBe(1);
|
||||
expect(lastModelsInfoCall().size).toBe(100);
|
||||
});
|
||||
|
||||
it("keeps the default page size when the URL value is not a number", () => {
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page_size: "lots" } });
|
||||
|
||||
expect(lastModelsInfoCall().size).toBe(50);
|
||||
});
|
||||
|
||||
it("ignores a sort_by the table cannot sort by instead of forwarding it to the server", () => {
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, {
|
||||
searchParams: { sort_by: "litellm_credential_name", sort_order: "desc" },
|
||||
});
|
||||
|
||||
expect(lastModelsInfoCall().sortBy).toBeUndefined();
|
||||
expect(lastModelsInfoCall().sortOrder).toBeUndefined();
|
||||
});
|
||||
|
||||
it("writes sort changes to the URL with the page cleared", async () => {
|
||||
setModelsInfo([makeRow()], 200);
|
||||
const user = userEvent.setup();
|
||||
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "2" }, onUrlUpdate });
|
||||
|
||||
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("sort_by")).toBe("model_info_updated_at");
|
||||
});
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBeNull();
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
|
||||
|
||||
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBe("desc");
|
||||
});
|
||||
});
|
||||
|
||||
it("clears every table param from the URL on drawer reset", async () => {
|
||||
setModelsInfo([makeRow()], 200);
|
||||
const user = userEvent.setup();
|
||||
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />, {
|
||||
searchParams: {
|
||||
model_search: "haiku",
|
||||
filter_team: "team-1",
|
||||
sort_by: "model_name",
|
||||
page: "2",
|
||||
view_mode: "all",
|
||||
},
|
||||
onUrlUpdate,
|
||||
});
|
||||
|
||||
await user.click(screen.getByTestId("datatable-filters-trigger"));
|
||||
await user.click(await screen.findByTestId("filter-drawer-reset"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(lastUrlParams(onUrlUpdate)?.toString()).toBe("");
|
||||
});
|
||||
expect(screen.getByTestId("datatable-search")).toHaveValue("");
|
||||
const defaultQuery: ModelsInfoArgs = { search: undefined, teamId: undefined, sortBy: undefined, page: 1 };
|
||||
await waitFor(() => {
|
||||
expect(lastModelsInfoCall()).toMatchObject(defaultQuery);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
it("applies a public model name filter through the drawer", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(screen.getByTestId("datatable-filters-trigger"));
|
||||
await user.click(await screen.findByPlaceholderText("Filter by Public Model Name"));
|
||||
|
|
@ -270,7 +405,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("renders every row the server returned for the selected model group so rows match the footer total", () => {
|
||||
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "claude-opus" }], 2);
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
|
||||
const table = screen.getByRole("table");
|
||||
expect(within(table).getByText("claude-opus")).toBeInTheDocument();
|
||||
|
|
@ -280,7 +415,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("asks the server for wildcard deployments instead of hiding rows client-side", () => {
|
||||
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "openai/*" }], 2);
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
|
||||
|
||||
expect(lastModelsInfoCall().wildcardOnly).toBe(true);
|
||||
expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument();
|
||||
|
|
@ -289,7 +424,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("asks the server for the selected access group instead of hiding rows client-side", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
expect(lastModelsInfoCall().wildcardOnly).toBe(false);
|
||||
|
||||
await user.click(screen.getByTestId("datatable-filters-trigger"));
|
||||
|
|
@ -303,20 +438,20 @@ describe("AllModelsTab", () => {
|
|||
});
|
||||
|
||||
it("asks the server for the exact selected model group so deployments beyond the first page are found", () => {
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
|
||||
expect(lastModelsInfoCall().modelName).toBe("claude-opus");
|
||||
expect(lastModelsInfoCall().search).toBeUndefined();
|
||||
});
|
||||
|
||||
it.each(["all", "wildcard"])("sends no exact model name for the %s pseudo group", (group) => {
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
|
||||
|
||||
expect(lastModelsInfoCall().modelName).toBeUndefined();
|
||||
});
|
||||
|
||||
it("keeps the exact model group alongside a typed search", async () => {
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
|
||||
|
||||
fireEvent.change(screen.getByPlaceholderText("Search model names…"), { target: { value: "opus" } });
|
||||
|
||||
|
|
@ -326,7 +461,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("resets search, filters, team and sorting from the drawer reset button", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
|
||||
|
||||
await user.click(screen.getByTestId("models-team-select"));
|
||||
await user.click(await screen.findByRole("option", { name: "Engineering" }));
|
||||
|
|
@ -343,7 +478,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("opens the delete modal from the row and deletes the model", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(await screen.findByTestId("model-delete-model-1"));
|
||||
expect(await screen.findByText("Delete Model")).toBeInTheDocument();
|
||||
|
|
@ -357,7 +492,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("pauses a model through the row toggle", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(await screen.findByTestId("model-pause-toggle-model-1"));
|
||||
|
||||
|
|
@ -368,7 +503,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("opens the model settings modal from the toolbar", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(screen.queryByTestId("model-settings-modal")).not.toBeInTheDocument();
|
||||
await user.click(screen.getByTestId("models-settings-trigger"));
|
||||
|
|
@ -377,7 +512,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("opens the model detail view from the model ID cell", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(await screen.findByTestId("model-id-model-1"));
|
||||
|
||||
|
|
@ -386,7 +521,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("opens the team detail view from the team ID cell", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(await screen.findByTestId("model-team-id-model-1"));
|
||||
|
||||
|
|
@ -395,20 +530,20 @@ describe("AllModelsTab", () => {
|
|||
|
||||
describe("virtual key hint", () => {
|
||||
it("explains personal key creation while viewing current team models", () => {
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(screen.getByText(/create a Virtual Key without selecting a team/i)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("links the Virtual Keys page through the migrated /ui route", () => {
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
expect(screen.getByRole("link", { name: "Virtual Keys page" })).toHaveAttribute("href", "/ui/api-keys");
|
||||
});
|
||||
|
||||
it("links the team hint's Virtual Keys page through the migrated /ui route", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(screen.getByTestId("models-team-select"));
|
||||
await user.click(await screen.findByRole("option", { name: "Engineering" }));
|
||||
|
|
@ -419,7 +554,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("names the selected team in the hint", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(screen.getByTestId("models-team-select"));
|
||||
await user.click(await screen.findByRole("option", { name: "Engineering" }));
|
||||
|
|
@ -429,7 +564,7 @@ describe("AllModelsTab", () => {
|
|||
|
||||
it("hides the hint when viewing all available models", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<AllModelsTab {...defaultProps} />);
|
||||
renderWithProviders(<AllModelsTab {...defaultProps} />);
|
||||
|
||||
await user.click(screen.getByTestId("models-view-select"));
|
||||
await user.click(await screen.findByRole("option", { name: "All Available Models" }));
|
||||
|
|
|
|||
|
|
@ -10,10 +10,11 @@ import { toast } from "@/lib/toast";
|
|||
import { uiHref } from "@/utils/uiHref";
|
||||
import { modelDeleteCall, modelPatchUpdateCall } from "@/components/networking";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
|
||||
import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
|
||||
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
|
||||
import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
|
||||
import { Info } from "lucide-react";
|
||||
import { useCallback, useEffect, useMemo, useState } from "react";
|
||||
import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs";
|
||||
import { useCallback, useMemo, useState } from "react";
|
||||
|
||||
import { useModelsInfo } from "../../hooks/models/useModels";
|
||||
import { transformModelData } from "../utils/modelDataTransformer";
|
||||
|
|
@ -24,11 +25,40 @@ import {
|
|||
PERSONAL_TEAM_VALUE,
|
||||
WILDCARD_MODEL_GROUP_VALUE,
|
||||
} from "./AllModelsTable";
|
||||
import { ACCESS_GROUPS_COLUMN_ID, MODEL_NAME_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
|
||||
import {
|
||||
ACCESS_GROUPS_COLUMN_ID,
|
||||
isModelTableSortColumnId,
|
||||
MODEL_NAME_COLUMN_ID,
|
||||
MODEL_TABLE_SORT_COLUMN_IDS,
|
||||
toServerSortField,
|
||||
} from "./ModelsTableColumns";
|
||||
|
||||
const SEARCH_DEBOUNCE_WAIT_MS = 200;
|
||||
const DEFAULT_PAGE_SIZE = 50;
|
||||
const DEFAULT_PAGINATION: PaginationState = { pageIndex: 0, pageSize: DEFAULT_PAGE_SIZE };
|
||||
const MAX_PAGE_SIZE = 100;
|
||||
const MAX_PAGE = 100_000;
|
||||
|
||||
const MODEL_VIEW_MODES = ["current_team", "all"] as const satisfies readonly ModelViewMode[];
|
||||
|
||||
const boundedInteger = (min: number, max: number, fallback: number) =>
|
||||
createParser({
|
||||
parse: (value: string) => {
|
||||
const parsed = parseAsInteger.parse(value);
|
||||
return parsed === null ? null : Math.min(Math.max(parsed, min), max);
|
||||
},
|
||||
serialize: String,
|
||||
}).withDefault(fallback);
|
||||
|
||||
const TABLE_STATE = {
|
||||
model_search: parseAsString.withDefault(""),
|
||||
view_mode: parseAsStringLiteral(MODEL_VIEW_MODES).withDefault("current_team"),
|
||||
filter_team: parseAsString.withDefault(PERSONAL_TEAM_VALUE),
|
||||
access_group: parseAsString.withDefault(""),
|
||||
sort_by: parseAsStringLiteral(MODEL_TABLE_SORT_COLUMN_IDS),
|
||||
sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault("asc"),
|
||||
page: boundedInteger(1, MAX_PAGE, 1),
|
||||
page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE),
|
||||
};
|
||||
|
||||
interface AllModelsTabProps {
|
||||
selectedModelGroup: string | null;
|
||||
|
|
@ -52,34 +82,25 @@ const AllModelsTab = ({
|
|||
const { data: teams, isLoading: isLoadingTeams } = useTeams();
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
const [modelNameSearch, setModelNameSearch] = useState<string>("");
|
||||
const [debouncedSearch, setDebouncedSearch] = useState<string>("");
|
||||
const [modelViewMode, setModelViewMode] = useState<ModelViewMode>("current_team");
|
||||
const [selectedTeamValue, setSelectedTeamValue] = useState<string>(PERSONAL_TEAM_VALUE);
|
||||
const [selectedModelAccessGroupFilter, setSelectedModelAccessGroupFilter] = useState<string | null>(null);
|
||||
const [pagination, setPagination] = useState<PaginationState>(DEFAULT_PAGINATION);
|
||||
const [sorting, setSorting] = useState<SortingState>([]);
|
||||
const [tableState, setTableState] = useQueryStates(TABLE_STATE);
|
||||
const modelNameSearch = tableState.model_search;
|
||||
const [debouncedSearch] = useDebouncedValue(modelNameSearch, { wait: SEARCH_DEBOUNCE_WAIT_MS });
|
||||
const modelViewMode = tableState.view_mode;
|
||||
const selectedTeamValue = tableState.filter_team;
|
||||
const selectedModelAccessGroupFilter = tableState.access_group || null;
|
||||
const pagination = useMemo<PaginationState>(
|
||||
() => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }),
|
||||
[tableState.page, tableState.page_size],
|
||||
);
|
||||
const sorting = useMemo<SortingState>(
|
||||
() => (tableState.sort_by ? [{ id: tableState.sort_by, desc: tableState.sort_order === "desc" }] : []),
|
||||
[tableState.sort_by, tableState.sort_order],
|
||||
);
|
||||
const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false);
|
||||
const [deleteModalModelId, setDeleteModalModelId] = useState<string | null>(null);
|
||||
const [deleteLoading, setDeleteLoading] = useState(false);
|
||||
const [pausingModelId, setPausingModelId] = useState<string | null>(null);
|
||||
|
||||
const resetToFirstPage = useCallback(() => {
|
||||
setPagination((previous) => (previous.pageIndex === 0 ? previous : { ...previous, pageIndex: 0 }));
|
||||
}, []);
|
||||
|
||||
const debouncedUpdateSearch = useDebouncedCallback(
|
||||
(value: string) => {
|
||||
setDebouncedSearch(value);
|
||||
resetToFirstPage();
|
||||
},
|
||||
{ wait: SEARCH_DEBOUNCE_WAIT_MS },
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
debouncedUpdateSearch(modelNameSearch);
|
||||
}, [modelNameSearch, debouncedUpdateSearch]);
|
||||
|
||||
const teamIdForQuery = selectedTeamValue === PERSONAL_TEAM_VALUE ? undefined : selectedTeamValue;
|
||||
const isConcreteModelGroup =
|
||||
Boolean(selectedModelGroup) &&
|
||||
|
|
@ -152,33 +173,49 @@ const AllModelsTab = ({
|
|||
[selectedModelGroup, selectedModelAccessGroupFilter],
|
||||
);
|
||||
|
||||
const handleSearchChange = useCallback(
|
||||
(value: string) => {
|
||||
void setTableState({ model_search: value || null, page: null });
|
||||
},
|
||||
[setTableState],
|
||||
);
|
||||
|
||||
const handleColumnFiltersChange: OnChangeFn<ColumnFiltersState> = (updater) => {
|
||||
const next = typeof updater === "function" ? updater(columnFilters) : updater;
|
||||
const next = functionalUpdate(updater, columnFilters);
|
||||
const modelGroup = next.find((entry) => entry.id === MODEL_NAME_COLUMN_ID)?.value;
|
||||
const accessGroup = next.find((entry) => entry.id === ACCESS_GROUPS_COLUMN_ID)?.value;
|
||||
setSelectedModelGroup(typeof modelGroup === "string" ? modelGroup : ALL_MODEL_GROUPS_VALUE);
|
||||
setSelectedModelAccessGroupFilter(typeof accessGroup === "string" ? accessGroup : null);
|
||||
resetToFirstPage();
|
||||
void setTableState({ access_group: typeof accessGroup === "string" ? accessGroup : null, page: null });
|
||||
};
|
||||
|
||||
const handleSortingChange: OnChangeFn<SortingState> = (updater) => {
|
||||
setSorting(typeof updater === "function" ? updater(sorting) : updater);
|
||||
resetToFirstPage();
|
||||
const active = functionalUpdate(updater, sorting)[0];
|
||||
void setTableState({
|
||||
sort_by: active && isModelTableSortColumnId(active.id) ? active.id : null,
|
||||
sort_order: active?.desc ? "desc" : null,
|
||||
page: null,
|
||||
});
|
||||
};
|
||||
|
||||
const handlePaginationChange = useCallback<OnChangeFn<PaginationState>>(
|
||||
(updater) => {
|
||||
const next = functionalUpdate(updater, pagination);
|
||||
void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize });
|
||||
},
|
||||
[pagination, setTableState],
|
||||
);
|
||||
|
||||
const handleTeamChange = (value: string) => {
|
||||
setSelectedTeamValue(value);
|
||||
resetToFirstPage();
|
||||
void setTableState({ filter_team: value, page: null });
|
||||
};
|
||||
|
||||
const handleViewModeChange = (value: ModelViewMode) => {
|
||||
void setTableState({ view_mode: value });
|
||||
};
|
||||
|
||||
const resetFilters = () => {
|
||||
setModelNameSearch("");
|
||||
setSelectedModelGroup(ALL_MODEL_GROUPS_VALUE);
|
||||
setSelectedModelAccessGroupFilter(null);
|
||||
setSelectedTeamValue(PERSONAL_TEAM_VALUE);
|
||||
setModelViewMode("current_team");
|
||||
setPagination(DEFAULT_PAGINATION);
|
||||
setSorting([]);
|
||||
void setTableState(null);
|
||||
};
|
||||
|
||||
const teamOptions = useMemo(
|
||||
|
|
@ -264,18 +301,18 @@ const AllModelsTab = ({
|
|||
sorting={sorting}
|
||||
onSortingChange={handleSortingChange}
|
||||
pagination={pagination}
|
||||
onPaginationChange={setPagination}
|
||||
onPaginationChange={handlePaginationChange}
|
||||
columnFilters={columnFilters}
|
||||
onColumnFiltersChange={handleColumnFiltersChange}
|
||||
onResetFilters={resetFilters}
|
||||
searchValue={modelNameSearch}
|
||||
onSearchChange={setModelNameSearch}
|
||||
onSearchChange={handleSearchChange}
|
||||
teamOptions={teamOptions}
|
||||
selectedTeamValue={selectedTeamValue}
|
||||
onTeamChange={handleTeamChange}
|
||||
isLoadingTeams={isLoadingTeams}
|
||||
viewMode={modelViewMode}
|
||||
onViewModeChange={setModelViewMode}
|
||||
onViewModeChange={handleViewModeChange}
|
||||
onOpenModelSettings={handleOpenModelSettings}
|
||||
availableModelGroups={availableModelGroups}
|
||||
availableModelAccessGroups={availableModelAccessGroups}
|
||||
|
|
|
|||
|
|
@ -24,6 +24,19 @@ export const TEAM_ID_COLUMN_ID = "model_info_team_id";
|
|||
export const ACCESS_GROUPS_COLUMN_ID = "model_info_access_groups";
|
||||
export const STATUS_COLUMN_ID = "model_info_db_model";
|
||||
|
||||
export const MODEL_TABLE_SORT_COLUMN_IDS = [
|
||||
MODEL_NAME_COLUMN_ID,
|
||||
CREATED_BY_COLUMN_ID,
|
||||
UPDATED_AT_COLUMN_ID,
|
||||
COSTS_COLUMN_ID,
|
||||
STATUS_COLUMN_ID,
|
||||
] as const;
|
||||
|
||||
export type ModelTableSortColumnId = (typeof MODEL_TABLE_SORT_COLUMN_IDS)[number];
|
||||
|
||||
export const isModelTableSortColumnId = (columnId: string): columnId is ModelTableSortColumnId =>
|
||||
(MODEL_TABLE_SORT_COLUMN_IDS as readonly string[]).includes(columnId);
|
||||
|
||||
const COLUMN_ID_TO_SERVER_SORT_FIELD: Record<string, string> = {
|
||||
[COSTS_COLUMN_ID]: "costs",
|
||||
[STATUS_COLUMN_ID]: "status",
|
||||
|
|
|
|||
|
|
@ -35,10 +35,12 @@ beforeEach(() => {
|
|||
Element.prototype.scrollIntoView = () => {};
|
||||
});
|
||||
|
||||
const CHAT_REQUEST_ARG_COUNT = 26;
|
||||
const CHAT_REQUEST_ARG_COUNT = 27;
|
||||
const STREAMING_ENABLED_ARG_INDEX = 25;
|
||||
const MESSAGES_REQUEST_ARG_COUNT = 19;
|
||||
const CHAT_CUSTOM_HEADERS_ARG_INDEX = 26;
|
||||
const MESSAGES_REQUEST_ARG_COUNT = 20;
|
||||
const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
|
||||
const MESSAGES_CUSTOM_HEADERS_ARG_INDEX = 19;
|
||||
|
||||
async function openComboboxByPlaceholder(placeholder: string) {
|
||||
const user = userEvent.setup();
|
||||
|
|
@ -447,6 +449,63 @@ describe("ChatUI", () => {
|
|||
expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
|
||||
});
|
||||
|
||||
it("should send custom headers entered in the sidebar with /v1/chat/completions and /v1/messages requests", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(
|
||||
<ChatUI
|
||||
accessToken="1234567890"
|
||||
token="1234567890"
|
||||
userRole="user"
|
||||
userID="1234567890"
|
||||
disabledPersonalKeyCreation={false}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Test Key")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
await selectComboboxOption("Select a Model", "Model 1");
|
||||
await user.click(screen.getByRole("button", { name: "Add Header" }));
|
||||
await user.click(screen.getByRole("button", { name: "Add Header" }));
|
||||
const [firstName] = screen.getAllByPlaceholderText("Header Name");
|
||||
const [firstValue, secondValue] = screen.getAllByPlaceholderText("Header Value");
|
||||
fireEvent.change(firstName, { target: { value: "anthropic-beta" } });
|
||||
fireEvent.change(firstValue, { target: { value: "context-1m-2025-08-07" } });
|
||||
fireEvent.change(secondValue, { target: { value: "ignored because the name is blank" } });
|
||||
|
||||
const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
|
||||
await act(async () => {
|
||||
fireEvent.change(messageInput, { target: { value: "hello" } });
|
||||
});
|
||||
await act(async () => {
|
||||
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
const chatArgs = vi.mocked(makeOpenAIChatCompletionRequest).mock.calls[0];
|
||||
expect(chatArgs).toHaveLength(CHAT_REQUEST_ARG_COUNT);
|
||||
expect(chatArgs[CHAT_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
|
||||
|
||||
await selectComboboxOption("Select an endpoint", "/v1/messages");
|
||||
await selectComboboxOption("Select a Model", "Model 1");
|
||||
await act(async () => {
|
||||
fireEvent.change(messageInput, { target: { value: "hello again" } });
|
||||
});
|
||||
await act(async () => {
|
||||
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
const messagesArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
|
||||
expect(messagesArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
|
||||
expect(messagesArgs[MESSAGES_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
|
||||
});
|
||||
|
||||
it("should force streaming in simplified mode even when the playground setting is off", async () => {
|
||||
sessionStorage.setItem("streamingEnabled", "false");
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import {
|
|||
Info,
|
||||
Key,
|
||||
Link2,
|
||||
ListPlus,
|
||||
Loader2,
|
||||
Settings,
|
||||
Shield,
|
||||
|
|
@ -40,6 +41,8 @@ import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages
|
|||
import { makeOpenAIAudioSpeechRequest } from "../../llm_calls/audio_speech";
|
||||
import { makeOpenAIAudioTranscriptionRequest } from "../../llm_calls/audio_transcriptions";
|
||||
import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion";
|
||||
import { customHeadersFromPairs, parseStoredHeaderPairs } from "@/components/llm_calls/request_headers";
|
||||
import KeyValueInput, { type KeyValuePair } from "@/components/key_value_input";
|
||||
import { makeOpenAIEmbeddingsRequest } from "../../llm_calls/embeddings_api";
|
||||
import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents";
|
||||
import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models";
|
||||
|
|
@ -220,6 +223,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
return [];
|
||||
}
|
||||
});
|
||||
const [customHeaderPairs, setCustomHeaderPairs] = useState<readonly KeyValuePair[]>(() =>
|
||||
parseStoredHeaderPairs(getSecureItem("customHeaders")),
|
||||
);
|
||||
const customHeaders = useMemo(() => customHeadersFromPairs(customHeaderPairs), [customHeaderPairs]);
|
||||
const [selectedVoice, setSelectedVoice] = useState<OpenAIVoice>(() => {
|
||||
const saved = sessionStorage.getItem("selectedVoice");
|
||||
if (!saved) return "alloy";
|
||||
|
|
@ -346,6 +353,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
selectedSdk,
|
||||
selectedVoice,
|
||||
proxySettings,
|
||||
customHeaders,
|
||||
});
|
||||
setGeneratedCode(code);
|
||||
}
|
||||
|
|
@ -367,12 +375,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
endpointType,
|
||||
selectedModel,
|
||||
proxySettings,
|
||||
customHeaders,
|
||||
]);
|
||||
|
||||
useEffect(() => {
|
||||
try {
|
||||
setSecureItem("apiKeySource", JSON.stringify(apiKeySource));
|
||||
setSecureItem("apiKey", apiKey);
|
||||
setSecureItem("customHeaders", JSON.stringify(customHeaderPairs));
|
||||
} catch {
|
||||
// Storage full or unavailable — non-critical, skip persisting.
|
||||
}
|
||||
|
|
@ -410,6 +420,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
mcpServerToolRestrictions,
|
||||
selectedVoice,
|
||||
streamingEnabled,
|
||||
customHeaderPairs,
|
||||
]);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -921,6 +932,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
mockTestFallbacks,
|
||||
mcpToolsets,
|
||||
streamingEnabled,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.IMAGE) {
|
||||
// For image generation
|
||||
|
|
@ -932,6 +944,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
selectedTags,
|
||||
signal,
|
||||
customProxyBaseUrl || undefined,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.SPEECH) {
|
||||
// For audio speech
|
||||
|
|
@ -946,6 +959,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
undefined, // responseFormat
|
||||
undefined, // speed
|
||||
customProxyBaseUrl || undefined,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.IMAGE_EDITS) {
|
||||
// For image edits
|
||||
|
|
@ -959,6 +973,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
selectedTags,
|
||||
signal,
|
||||
customProxyBaseUrl || undefined,
|
||||
customHeaders,
|
||||
);
|
||||
}
|
||||
} else if (endpointType === EndpointType.RESPONSES) {
|
||||
|
|
@ -1004,6 +1019,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
mcpToolsets,
|
||||
streamingEnabled,
|
||||
updateTotalLatency,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) {
|
||||
const apiChatHistory = [
|
||||
|
|
@ -1033,6 +1049,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
mcpServerToolRestrictions,
|
||||
mcpToolsets,
|
||||
streamingEnabled,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.EMBEDDINGS) {
|
||||
await makeOpenAIEmbeddingsRequest(
|
||||
|
|
@ -1042,6 +1059,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
effectiveApiKey,
|
||||
selectedTags,
|
||||
customProxyBaseUrl || undefined,
|
||||
customHeaders,
|
||||
);
|
||||
} else if (endpointType === EndpointType.TRANSCRIPTION) {
|
||||
// For audio transcriptions
|
||||
|
|
@ -1058,6 +1076,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
undefined, // responseFormat
|
||||
undefined, // temperature
|
||||
customProxyBaseUrl || undefined,
|
||||
customHeaders,
|
||||
);
|
||||
}
|
||||
} else if (endpointType === EndpointType.INTERACTIONS) {
|
||||
|
|
@ -1069,6 +1088,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
selectedTags,
|
||||
signal,
|
||||
customProxyBaseUrl || undefined,
|
||||
undefined,
|
||||
customHeaders,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1086,13 +1107,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
resolvedServerId = toolEntry?.server_id ?? rawSelected;
|
||||
}
|
||||
if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) {
|
||||
const result = await callMCPTool(
|
||||
effectiveApiKey,
|
||||
resolvedServerId,
|
||||
selectedMCPDirectTool,
|
||||
mcpToolArguments,
|
||||
selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined,
|
||||
);
|
||||
const result = await callMCPTool(effectiveApiKey, resolvedServerId, selectedMCPDirectTool, mcpToolArguments, {
|
||||
...(selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : {}),
|
||||
customHeaders,
|
||||
});
|
||||
const resultText =
|
||||
result?.content?.length > 0
|
||||
? JSON.stringify(
|
||||
|
|
@ -1118,6 +1136,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
updateA2AMetadata,
|
||||
customProxyBaseUrl || undefined,
|
||||
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
|
||||
customHeaders,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
|
|
@ -1485,6 +1504,18 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
|||
/>
|
||||
</div>
|
||||
|
||||
{endpointType !== EndpointType.REALTIME && (
|
||||
<div>
|
||||
<label className="mb-2 flex items-center text-sm font-medium text-foreground">
|
||||
<ListPlus className="mr-2 size-4" aria-hidden="true" /> Custom Headers
|
||||
</label>
|
||||
<KeyValueInput value={customHeaderPairs} onChange={setCustomHeaderPairs} />
|
||||
<p className="mt-2 text-xs text-muted-foreground">
|
||||
Sent with every playground request, e.g. provider-specific headers like anthropic-beta.
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div>
|
||||
<div className="mb-2 flex items-center gap-1 text-sm font-medium text-foreground">
|
||||
<Wrench className="mr-1 size-4" aria-hidden="true" />
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
|
||||
import { type CustomHeaders, withRequiredHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { A2ATaskMetadata } from "@/components/chat_ui/types";
|
||||
|
||||
interface A2AMessagePart {
|
||||
|
|
@ -116,6 +117,7 @@ export const makeA2ASendMessageRequest = async (
|
|||
onA2AMetadata?: (metadata: A2ATaskMetadata) => void,
|
||||
customBaseUrl?: string,
|
||||
guardrails?: string[],
|
||||
customHeaders?: CustomHeaders,
|
||||
): Promise<void> => {
|
||||
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/a2a/${agentId}/message/send` : `/a2a/${agentId}/message/send`;
|
||||
|
|
@ -146,10 +148,10 @@ export const makeA2ASendMessageRequest = async (
|
|||
try {
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
headers: withRequiredHeaders(customHeaders ?? {}, {
|
||||
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
}),
|
||||
body: JSON.stringify(jsonRpcRequest),
|
||||
signal,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import Anthropic from "@anthropic-ai/sdk";
|
||||
import { makeAnthropicMessagesRequest } from "./anthropic_messages";
|
||||
import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
|
||||
|
||||
|
|
@ -122,4 +123,26 @@ describe("anthropic_messages non-streaming", () => {
|
|||
expect(mockMessagesCreate).not.toHaveBeenCalled();
|
||||
expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true });
|
||||
});
|
||||
|
||||
it("sends custom headers alongside the tags header on the Anthropic client", async () => {
|
||||
mockMessagesCreate.mockResolvedValue({ content: [{ type: "text", text: "OK" }], usage: {} });
|
||||
|
||||
await makeAnthropicMessagesRequest(
|
||||
[{ role: "user", content: "Hello" }],
|
||||
vi.fn(),
|
||||
"claude-haiku-4-5",
|
||||
"test-token",
|
||||
["team-a"],
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
...NON_STREAMING_ARGS,
|
||||
{ "anthropic-beta": "context-1m-2025-08-07" },
|
||||
);
|
||||
|
||||
expect(vi.mocked(Anthropic).mock.calls[0][0]).toMatchObject({
|
||||
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import Anthropic from "@anthropic-ai/sdk";
|
|||
import { MessageType } from "@/components/chat_ui/types";
|
||||
import { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
|
||||
import { buildMcpToolBlocks } from "@/components/llm_calls/mcp_tool_blocks";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
|
@ -34,6 +35,7 @@ export async function makeAnthropicMessagesRequest(
|
|||
mcpServerToolRestrictions?: Record<string, string[]>,
|
||||
mcpToolsets?: MCPToolset[],
|
||||
streamingEnabled: boolean = true,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
if (!accessToken) {
|
||||
throw new Error("Virtual Key is required");
|
||||
|
|
@ -46,11 +48,7 @@ export async function makeAnthropicMessagesRequest(
|
|||
|
||||
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
|
||||
|
||||
// Prepare headers with tags and trace ID
|
||||
const headers: Record<string, string> = {};
|
||||
if (tags && tags.length > 0) {
|
||||
headers["x-litellm-tags"] = tags.join(",");
|
||||
}
|
||||
const headers = buildPlaygroundHeaders(tags, customHeaders);
|
||||
|
||||
const client = new Anthropic({
|
||||
apiKey: accessToken,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import openai from "openai";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { toast } from "@/lib/toast";
|
||||
import type { OpenAIVoice } from "../components/chat_ui/chatConstants";
|
||||
|
||||
|
|
@ -14,6 +15,7 @@ export async function makeOpenAIAudioSpeechRequest(
|
|||
responseFormat?: string,
|
||||
speed?: number,
|
||||
customBaseUrl?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
// base url should be the current base_url
|
||||
const isLocal = process.env.NODE_ENV === "development";
|
||||
|
|
@ -25,7 +27,7 @@ export async function makeOpenAIAudioSpeechRequest(
|
|||
apiKey: accessToken,
|
||||
baseURL: proxyBaseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
|
||||
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
|
||||
});
|
||||
|
||||
try {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import openai from "openai";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
export async function makeOpenAIAudioTranscriptionRequest(
|
||||
|
|
@ -14,6 +15,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
|
|||
responseFormat?: string,
|
||||
temperature?: number,
|
||||
customBaseUrl?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
// base url should be the current base_url
|
||||
const isLocal = process.env.NODE_ENV === "development";
|
||||
|
|
@ -26,7 +28,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
|
|||
apiKey: accessToken,
|
||||
baseURL: proxyBaseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
|
||||
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
|
||||
});
|
||||
|
||||
try {
|
||||
|
|
|
|||
|
|
@ -76,4 +76,42 @@ describe("embeddings_api", () => {
|
|||
input: "Sample text",
|
||||
});
|
||||
});
|
||||
|
||||
it("sends custom headers on the fetch request, letting them override the tags header", async () => {
|
||||
await makeOpenAIEmbeddingsRequest(
|
||||
"Sample text",
|
||||
mockUpdateEmbeddingsUI,
|
||||
"text-embedding-3-small",
|
||||
"abcdef",
|
||||
["team-a"],
|
||||
undefined,
|
||||
{ "x-litellm-tags": "team-b", "x-request-source": "playground" },
|
||||
);
|
||||
|
||||
expect(mockFetch.mock.calls[0][1]).toMatchObject({
|
||||
headers: {
|
||||
Authorization: "Bearer abcdef",
|
||||
"x-litellm-tags": "team-b",
|
||||
"x-request-source": "playground",
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("does not let custom headers replace the gateway auth or content-type headers", async () => {
|
||||
await makeOpenAIEmbeddingsRequest(
|
||||
"Sample text",
|
||||
mockUpdateEmbeddingsUI,
|
||||
"text-embedding-3-small",
|
||||
"abcdef",
|
||||
undefined,
|
||||
undefined,
|
||||
{ authorization: "Bearer stolen", "Content-Type": "text/plain", "x-request-source": "playground" },
|
||||
);
|
||||
|
||||
expect(mockFetch.mock.calls[0][1].headers).toEqual({
|
||||
Authorization: "Bearer abcdef",
|
||||
"Content-Type": "application/json",
|
||||
"x-request-source": "playground",
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
import { toast } from "@/lib/toast";
|
||||
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
|
||||
import {
|
||||
buildPlaygroundHeaders,
|
||||
type CustomHeaders,
|
||||
withRequiredHeaders,
|
||||
} from "@/components/llm_calls/request_headers";
|
||||
|
||||
export async function makeOpenAIEmbeddingsRequest(
|
||||
input: string,
|
||||
|
|
@ -8,6 +13,7 @@ export async function makeOpenAIEmbeddingsRequest(
|
|||
accessToken: string,
|
||||
tags?: string[],
|
||||
customBaseUrl?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
if (!accessToken) {
|
||||
throw new Error("Virtual Key is required");
|
||||
|
|
@ -20,11 +26,10 @@ export async function makeOpenAIEmbeddingsRequest(
|
|||
}
|
||||
|
||||
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
|
||||
// Prepare headers with tags and trace ID
|
||||
const headers: Record<string, string> = {};
|
||||
if (tags && tags.length > 0) {
|
||||
headers["x-litellm-tags"] = tags.join(",");
|
||||
}
|
||||
const headers = withRequiredHeaders(buildPlaygroundHeaders(tags, customHeaders), {
|
||||
"Content-Type": "application/json",
|
||||
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
|
||||
});
|
||||
|
||||
try {
|
||||
const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl;
|
||||
|
|
@ -32,11 +37,7 @@ export async function makeOpenAIEmbeddingsRequest(
|
|||
|
||||
const response = await fetch(requestUrl, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
|
||||
...headers,
|
||||
},
|
||||
headers,
|
||||
body: JSON.stringify({
|
||||
model: selectedModel,
|
||||
input,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import openai from "openai";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
export async function makeOpenAIImageEditsRequest(
|
||||
|
|
@ -11,6 +12,7 @@ export async function makeOpenAIImageEditsRequest(
|
|||
tags?: string[],
|
||||
signal?: AbortSignal,
|
||||
customBaseUrl?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
// base url should be the current base_url
|
||||
const isLocal = process.env.NODE_ENV === "development";
|
||||
|
|
@ -23,7 +25,7 @@ export async function makeOpenAIImageEditsRequest(
|
|||
apiKey: accessToken,
|
||||
baseURL: proxyBaseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
|
||||
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
|
||||
});
|
||||
|
||||
try {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import openai from "openai";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
export async function makeOpenAIImageGenerationRequest(
|
||||
|
|
@ -10,6 +11,7 @@ export async function makeOpenAIImageGenerationRequest(
|
|||
tags?: string[],
|
||||
signal?: AbortSignal,
|
||||
customBaseUrl?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
// base url should be the current base_url
|
||||
const isLocal = process.env.NODE_ENV === "development";
|
||||
|
|
@ -21,7 +23,7 @@ export async function makeOpenAIImageGenerationRequest(
|
|||
apiKey: accessToken,
|
||||
baseURL: proxyBaseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
|
||||
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
|
||||
});
|
||||
|
||||
try {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
import { toast } from "@/lib/toast";
|
||||
import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking";
|
||||
import {
|
||||
buildPlaygroundHeaders,
|
||||
type CustomHeaders,
|
||||
withRequiredHeaders,
|
||||
} from "@/components/llm_calls/request_headers";
|
||||
|
||||
export async function makeInteractionsRequest(
|
||||
input: string,
|
||||
|
|
@ -10,6 +15,7 @@ export async function makeInteractionsRequest(
|
|||
signal?: AbortSignal,
|
||||
customBaseUrl?: string,
|
||||
previousInteractionId?: string,
|
||||
customHeaders?: CustomHeaders,
|
||||
): Promise<void> {
|
||||
if (!accessToken) {
|
||||
throw new Error("Virtual Key is required");
|
||||
|
|
@ -24,13 +30,10 @@ export async function makeInteractionsRequest(
|
|||
const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl;
|
||||
const requestUrl = `${normalizedBaseUrl}/v1beta/interactions`;
|
||||
|
||||
const headers: Record<string, string> = {
|
||||
const headers: Record<string, string> = withRequiredHeaders(buildPlaygroundHeaders(tags, customHeaders), {
|
||||
"Content-Type": "application/json",
|
||||
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
|
||||
};
|
||||
if (tags && tags.length > 0) {
|
||||
headers["x-litellm-tags"] = tags.join(",");
|
||||
}
|
||||
});
|
||||
|
||||
const body: Record<string, unknown> = {
|
||||
model: selectedModel,
|
||||
|
|
|
|||
|
|
@ -45,6 +45,15 @@ const SETTINGS_FIXTURE = [
|
|||
field_tab: "prompt_caching",
|
||||
field_default_value: null,
|
||||
},
|
||||
{
|
||||
field_name: "openai_system_messages_first",
|
||||
field_type: "Boolean",
|
||||
field_value: false,
|
||||
field_description: "openai system first toggle",
|
||||
stored_in_db: null,
|
||||
field_tab: "prompt_caching",
|
||||
field_default_value: false,
|
||||
},
|
||||
{
|
||||
field_name: "max_ui_session_budget",
|
||||
field_type: "Dollar",
|
||||
|
|
@ -157,6 +166,39 @@ describe("GeneralSettings General tab", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("GeneralSettings Prompt Caching tab", () => {
|
||||
beforeEach(() => {
|
||||
vi.mocked(getGeneralSettingsCall).mockResolvedValue([...SETTINGS_FIXTURE.map((s) => ({ ...s }))]);
|
||||
vi.mocked(updateConfigFieldSetting).mockClear();
|
||||
vi.mocked(deleteConfigFieldSetting).mockClear();
|
||||
});
|
||||
|
||||
it("persists openai_system_messages_first when its switch is turned on", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
|
||||
|
||||
await user.click(await screen.findByRole("tab", { name: "Prompt Caching" }));
|
||||
const toggle = await screen.findByRole("switch", { name: "System messages first for OpenAI" });
|
||||
expect(toggle).not.toBeChecked();
|
||||
|
||||
await user.click(toggle);
|
||||
|
||||
expect(toggle).toBeChecked();
|
||||
expect(updateConfigFieldSetting).toHaveBeenCalledWith("token", "openai_system_messages_first", true);
|
||||
expect(deleteConfigFieldSetting).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the prompt caching rows off the General tab table", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
|
||||
|
||||
await user.click(screen.getByText("General"));
|
||||
await settingsRow("max_ui_session_budget");
|
||||
|
||||
expect(screen.queryByText("openai_system_messages_first")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
// The five tabs here are proxy-wide settings. Auto-routers moved to Models + Endpoints.
|
||||
describe("GeneralSettings tabs", () => {
|
||||
beforeEach(() => {
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ import RoutingGroups from "@/components/routing_groups";
|
|||
const PROMPT_CACHING_TAB = "prompt_caching";
|
||||
const ENABLE_ANTHROPIC_PROMPT_CACHING = "enable_anthropic_prompt_caching";
|
||||
const ANTHROPIC_PROMPT_CACHING_TTL = "anthropic_prompt_caching_ttl";
|
||||
const OPENAI_SYSTEM_MESSAGES_FIRST = "openai_system_messages_first";
|
||||
|
||||
const isOn = (value: unknown) => value === true || value === "true";
|
||||
|
||||
interface GeneralSettingsPageProps {
|
||||
accessToken: string | null;
|
||||
|
|
@ -129,14 +132,15 @@ export const PromptCachingPanel: React.FC<{
|
|||
}> = ({ accessToken, settings, onChange }) => {
|
||||
const enableSetting = settings.find((s) => s.field_name === ENABLE_ANTHROPIC_PROMPT_CACHING);
|
||||
const ttlSetting = settings.find((s) => s.field_name === ANTHROPIC_PROMPT_CACHING_TTL);
|
||||
const systemFirstSetting = settings.find((s) => s.field_name === OPENAI_SYSTEM_MESSAGES_FIRST);
|
||||
|
||||
// The two rows come from the same registry the General tab reads; if they
|
||||
// The rows come from the same registry the General tab reads; if they
|
||||
// are not loaded yet there is nothing to render.
|
||||
if (!enableSetting) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const enabled = enableSetting.field_value === true || enableSetting.field_value === "true";
|
||||
const enabled = isOn(enableSetting.field_value);
|
||||
|
||||
// Apply immediately: a toggle and a dropdown are direct controls, so there is
|
||||
// no separate Update button. Clearing the ttl resets it to the provider default.
|
||||
|
|
@ -187,6 +191,20 @@ export const PromptCachingPanel: React.FC<{
|
|||
</Select>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{systemFirstSetting && (
|
||||
<div className="mt-6 flex items-start justify-between gap-8">
|
||||
<div className="min-w-0 max-w-2xl">
|
||||
<p className="font-medium">System messages first for OpenAI</p>
|
||||
<p className="mt-1 break-words text-xs text-muted-foreground">{systemFirstSetting.field_description}</p>
|
||||
</div>
|
||||
<Switch
|
||||
aria-label="System messages first for OpenAI"
|
||||
checked={isOn(systemFirstSetting.field_value)}
|
||||
onCheckedChange={(checked) => persist(OPENAI_SYSTEM_MESSAGES_FIRST, checked)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -42,9 +42,10 @@ import { Restricted, restrictedBy } from "./TierRestrictions";
|
|||
import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions";
|
||||
import {
|
||||
ReasoningEffort,
|
||||
TierModelParamChange,
|
||||
TierModelParamsByTier,
|
||||
classifierEffortOptionsForModels,
|
||||
setTierModelReasoningEffort,
|
||||
setTierModelParam,
|
||||
tierEffortOptionsForModels,
|
||||
tierRowLabel,
|
||||
} from "./complexity_router_tiers";
|
||||
|
|
@ -621,6 +622,9 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
const exitToBuiltInTiers = () => dispatch({ kind: "restore" });
|
||||
|
||||
const tierEffortOptionsByModel = tierEffortOptionsForModels(modelInfo);
|
||||
const fastModeByModel = Object.fromEntries(
|
||||
modelInfo.map((model) => [model.model_group, model.supports_fast_mode === true]),
|
||||
);
|
||||
const classifierEffortOptionsByModel = classifierEffortOptionsForModels(modelInfo);
|
||||
|
||||
// Embedding models can't serve a chat-completion role, so they're excluded here.
|
||||
|
|
@ -631,12 +635,11 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
label: model.model_group,
|
||||
}));
|
||||
|
||||
const handleTierModelEffortChange = (tier: string, model: string, effort: ReasoningEffort | undefined) => {
|
||||
const handleTierModelParamChange = (tier: string, model: string, change: TierModelParamChange) =>
|
||||
onChange({
|
||||
...value,
|
||||
tier_model_params: setTierModelReasoningEffort(value.tier_model_params, tier, model, effort),
|
||||
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, change),
|
||||
});
|
||||
};
|
||||
|
||||
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
|
||||
// "track the tiers" everywhere downstream instead of as a blank model name.
|
||||
|
|
@ -734,7 +737,13 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
models={row.models}
|
||||
effortOptionsByModel={tierEffortOptionsByModel}
|
||||
paramsByModel={row.params}
|
||||
onEffortChange={(model, effort) => handleTierModelEffortChange(row.id, model, effort)}
|
||||
fastModeByModel={fastModeByModel}
|
||||
onEffortChange={(model, effort) =>
|
||||
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
|
||||
}
|
||||
onFastModeChange={(model, enabled) =>
|
||||
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
|
||||
}
|
||||
/>
|
||||
{row.models.length > 1 && (
|
||||
<span className="text-xs text-muted-foreground">
|
||||
|
|
|
|||
|
|
@ -0,0 +1,143 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders, screen } from "../../../tests/test-utils";
|
||||
import {
|
||||
buildUpdatedComplexityRouterConfig,
|
||||
hydrateComplexityRouterConfig,
|
||||
} from "../edit_auto_router/edit_auto_router_modal";
|
||||
import type { ModelGroup } from "../llm_calls/fetch_models";
|
||||
import ComplexityRouterConfig, { type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
|
||||
|
||||
const modelInfo: ModelGroup[] = [
|
||||
{ model_group: "primary", supported_reasoning_efforts: ["low", "high"], supports_fast_mode: true },
|
||||
{ model_group: "secondary", supports_fast_mode: true },
|
||||
{ model_group: "blocked", supported_reasoning_efforts: ["low"], supports_fast_mode: false },
|
||||
{ model_group: "missing", supported_reasoning_efforts: ["low"] },
|
||||
];
|
||||
|
||||
it.each([false, true])("edits and round-trips independent model settings with custom tiers=%s", async (custom) => {
|
||||
const user = userEvent.setup();
|
||||
const tier = custom ? "custom-a" : "COMPLEX";
|
||||
const otherTier = custom ? "custom-b" : "REASONING";
|
||||
const label = custom ? "Interactive" : "Complex";
|
||||
const models = ["primary", "secondary", "blocked", "missing"];
|
||||
const initial: ComplexityRouterConfigValue = {
|
||||
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: models, REASONING: ["primary"] },
|
||||
classifier_type: "heuristic",
|
||||
...(custom && {
|
||||
custom_tier_set: {
|
||||
tiers: [
|
||||
{ id: tier, name: label, definition: "Interactive requests", models },
|
||||
{ id: otherTier, name: "Deliberate", definition: "Careful requests", models: ["primary"] },
|
||||
],
|
||||
fallback_tier_id: tier,
|
||||
},
|
||||
}),
|
||||
tier_model_params: {
|
||||
[tier]: {
|
||||
primary: { reasoning_effort: "high", max_tokens: 1024 },
|
||||
secondary: { speed: "fast" },
|
||||
blocked: { speed: "fast" },
|
||||
},
|
||||
[otherTier]: { primary: { speed: "fast", reasoning_effort: "low" } },
|
||||
},
|
||||
};
|
||||
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
|
||||
const editor = (value: ComplexityRouterConfigValue) => (
|
||||
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
|
||||
);
|
||||
const view = renderWithProviders(editor(initial));
|
||||
const fast = () => screen.getByRole("switch", { name: `Fast mode for primary in the ${label} tier` });
|
||||
|
||||
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(3);
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for (blocked|missing)/ })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("combobox", { name: /^Reasoning effort for secondary/ })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("switch", { name: `Fast mode for secondary in the ${label} tier` })).toBeChecked();
|
||||
expect(fast()).not.toBeChecked();
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
|
||||
await user.click(fast());
|
||||
const enabled = onChange.mock.lastCall![0];
|
||||
expect(enabled.tier_model_params).toEqual({
|
||||
...initial.tier_model_params,
|
||||
[tier]: {
|
||||
...initial.tier_model_params![tier],
|
||||
primary: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" },
|
||||
},
|
||||
});
|
||||
const saved = buildUpdatedComplexityRouterConfig({}, enabled);
|
||||
expect(saved.tier_model_configs).toEqual({
|
||||
[custom ? label : tier]: [
|
||||
{ model_name: "primary", litellm_params: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" } },
|
||||
{ model_name: "secondary", litellm_params: { speed: "fast" } },
|
||||
{ model_name: "blocked", litellm_params: { speed: "fast" } },
|
||||
],
|
||||
[custom ? "Deliberate" : otherTier]: [
|
||||
{ model_name: "primary", litellm_params: { speed: "fast", reasoning_effort: "low" } },
|
||||
],
|
||||
});
|
||||
const reopened = hydrateComplexityRouterConfig(saved, undefined);
|
||||
const reopenedTier = custom ? reopened.custom_tier_set!.tiers[0].id : tier;
|
||||
view.rerender(editor(reopened));
|
||||
expect(fast()).toBeChecked();
|
||||
|
||||
await user.click(screen.getByRole("combobox", { name: `Reasoning effort for primary in the ${label} tier` }));
|
||||
await user.click(await screen.findByRole("option", { name: "low" }));
|
||||
const effortChanged = onChange.mock.lastCall![0];
|
||||
expect(effortChanged.tier_model_params?.[reopenedTier].primary).toEqual({
|
||||
reasoning_effort: "low",
|
||||
max_tokens: 1024,
|
||||
speed: "fast",
|
||||
});
|
||||
view.rerender(editor(effortChanged));
|
||||
await user.click(fast());
|
||||
const disabled = onChange.mock.lastCall![0];
|
||||
expect(disabled.tier_model_params).toEqual({
|
||||
...effortChanged.tier_model_params,
|
||||
[reopenedTier]: {
|
||||
...effortChanged.tier_model_params![reopenedTier],
|
||||
primary: { reasoning_effort: "low", max_tokens: 1024 },
|
||||
},
|
||||
});
|
||||
view.rerender(editor(disabled));
|
||||
expect(fast()).not.toBeChecked();
|
||||
|
||||
const picker = () => screen.getByRole("combobox", { name: `Select model(s) for ${label.toLowerCase()} queries` });
|
||||
await user.click(picker());
|
||||
await user.click(await screen.findByRole("option", { name: "primary" }));
|
||||
await user.keyboard("{Escape}");
|
||||
const deselected = onChange.mock.lastCall![0];
|
||||
expect(deselected.tier_model_params?.[reopenedTier]).toEqual({
|
||||
secondary: { speed: "fast" },
|
||||
blocked: { speed: "fast" },
|
||||
});
|
||||
view.rerender(editor(deselected));
|
||||
expect(screen.queryByRole("switch", { name: `Fast mode for primary in the ${label} tier` })).not.toBeInTheDocument();
|
||||
await user.click(picker());
|
||||
await user.click(await screen.findByRole("option", { name: "primary" }));
|
||||
await user.keyboard("{Escape}");
|
||||
const reselected = onChange.mock.lastCall![0];
|
||||
view.rerender(editor(reselected));
|
||||
expect(fast()).not.toBeChecked();
|
||||
expect(screen.getByRole("combobox", { name: `Reasoning effort for primary in the ${label} tier` })).toHaveTextContent(
|
||||
"Default",
|
||||
);
|
||||
});
|
||||
|
||||
describe("Fast mode metadata", () => {
|
||||
it("offers nothing before model capabilities load and leaves stored speed untouched", () => {
|
||||
const value: ComplexityRouterConfigValue = {
|
||||
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: [] },
|
||||
classifier_type: "heuristic",
|
||||
tier_model_params: { SIMPLE: { primary: { speed: "fast" } } },
|
||||
};
|
||||
const onChange = vi.fn();
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={[]} value={value} onChange={onChange} />);
|
||||
expect(screen.queryByRole("switch", { name: /^Fast mode for/ })).not.toBeInTheDocument();
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(buildUpdatedComplexityRouterConfig({}, value).tier_model_configs).toEqual({
|
||||
SIMPLE: [{ model_name: "primary", litellm_params: { speed: "fast" } }],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { Info } from "lucide-react";
|
||||
import React from "react";
|
||||
import { ReasoningEffort, TierModelParams } from "./complexity_router_tiers";
|
||||
|
|
@ -18,6 +19,8 @@ interface TierModelEffortRowsProps {
|
|||
effortOptionsByModel: Record<string, string[]>;
|
||||
paramsByModel: Record<string, TierModelParams> | undefined;
|
||||
onEffortChange: (model: string, effort: ReasoningEffort | undefined) => void;
|
||||
fastModeByModel?: Record<string, boolean>;
|
||||
onFastModeChange: (model: string, enabled: boolean) => void;
|
||||
}
|
||||
|
||||
export interface TierEffortRow {
|
||||
|
|
@ -29,13 +32,16 @@ export interface TierEffortRow {
|
|||
/**
|
||||
* A stored effort outside the model's supported set (hand-authored, or capabilities changed since
|
||||
* it was saved) is listed anyway, so the row renders with its value selected and can be cleared.
|
||||
* Only a model with no supported level and nothing stored drops out.
|
||||
*/
|
||||
export const tierEffortRows = ({
|
||||
models,
|
||||
effortOptionsByModel,
|
||||
paramsByModel,
|
||||
}: Pick<TierModelEffortRowsProps, "models" | "effortOptionsByModel" | "paramsByModel">): TierEffortRow[] =>
|
||||
fastModeByModel,
|
||||
}: Pick<
|
||||
TierModelEffortRowsProps,
|
||||
"models" | "effortOptionsByModel" | "paramsByModel" | "fastModeByModel"
|
||||
>): TierEffortRow[] =>
|
||||
models
|
||||
.map((model) => {
|
||||
const effort = storedEffort(paramsByModel?.[model]);
|
||||
|
|
@ -43,56 +49,74 @@ export const tierEffortRows = ({
|
|||
const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported;
|
||||
return { model, effort, options: Array.from(new Set(listed)) };
|
||||
})
|
||||
.filter(({ options }) => options.length > 0);
|
||||
.filter(({ model, options }) => options.length > 0 || fastModeByModel?.[model] === true);
|
||||
|
||||
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = ({
|
||||
tierLabel,
|
||||
models,
|
||||
effortOptionsByModel,
|
||||
paramsByModel,
|
||||
onEffortChange,
|
||||
}) => {
|
||||
const rows = tierEffortRows({ models, effortOptionsByModel, paramsByModel });
|
||||
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
|
||||
const { tierLabel, paramsByModel, onEffortChange, fastModeByModel, onFastModeChange } = props;
|
||||
const rows = tierEffortRows(props);
|
||||
if (rows.length === 0) return null;
|
||||
return (
|
||||
<div className="mt-2 space-y-1">
|
||||
<div className="flex items-center gap-1">
|
||||
<span className="text-xs font-medium text-muted-foreground">Reasoning effort</span>
|
||||
<SimpleTooltip
|
||||
content={`Sent as reasoning_effort on requests this tier routes to the model, overriding the caller's value. Default leaves the request untouched.`}
|
||||
>
|
||||
<Info className="size-3 text-muted-foreground/70" />
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
{rows.map(({ model, effort, options }) => (
|
||||
<div key={model} className="flex items-center justify-between gap-2">
|
||||
<span className="truncate text-xs">{model}</span>
|
||||
<Select
|
||||
items={[
|
||||
{ value: PROVIDER_DEFAULT, label: "Default" },
|
||||
...options.map((option) => ({ value: option, label: option })),
|
||||
]}
|
||||
value={effort ?? PROVIDER_DEFAULT}
|
||||
onValueChange={(selected: string | null) =>
|
||||
selected !== null && onEffortChange(model, selected === PROVIDER_DEFAULT ? undefined : selected)
|
||||
}
|
||||
{rows.some(({ options }) => options.length > 0) && (
|
||||
<div className="flex items-center gap-1">
|
||||
<span className="text-xs font-medium text-muted-foreground">Reasoning effort</span>
|
||||
<SimpleTooltip
|
||||
content={`Sent as reasoning_effort on requests this tier routes to the model, overriding the caller's value. Default leaves the request untouched.`}
|
||||
>
|
||||
<SelectTrigger
|
||||
size="sm"
|
||||
className="w-36"
|
||||
aria-label={`Reasoning effort for ${model} in the ${tierLabel} tier`}
|
||||
>
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value={PROVIDER_DEFAULT}>Default</SelectItem>
|
||||
{options.map((option) => (
|
||||
<SelectItem key={option} value={option}>
|
||||
{option}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<Info className="size-3 text-muted-foreground/70" />
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
)}
|
||||
{rows.map(({ model, effort, options }) => (
|
||||
<div key={model} className="flex flex-wrap items-center justify-between gap-2">
|
||||
<span className="min-w-0 flex-1 basis-32 truncate text-xs" title={model}>
|
||||
{model}
|
||||
</span>
|
||||
<div className="flex flex-wrap items-center gap-3">
|
||||
{options.length > 0 && (
|
||||
<Select
|
||||
items={[
|
||||
{ value: PROVIDER_DEFAULT, label: "Default" },
|
||||
...options.map((option) => ({ value: option, label: option })),
|
||||
]}
|
||||
value={effort ?? PROVIDER_DEFAULT}
|
||||
onValueChange={(selected: string | null) =>
|
||||
selected !== null && onEffortChange(model, selected === PROVIDER_DEFAULT ? undefined : selected)
|
||||
}
|
||||
>
|
||||
<SelectTrigger
|
||||
size="sm"
|
||||
className="w-36"
|
||||
aria-label={`Reasoning effort for ${model} in the ${tierLabel} tier`}
|
||||
>
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value={PROVIDER_DEFAULT}>Default</SelectItem>
|
||||
{options.map((option) => (
|
||||
<SelectItem key={option} value={option}>
|
||||
{option}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
)}
|
||||
{fastModeByModel?.[model] === true && (
|
||||
<SimpleTooltip content="Fast mode has higher pricing and requires an eligible provider account. Off removes this tier's speed override and inherits the request or provider default">
|
||||
<label
|
||||
className="flex items-center gap-2 text-xs"
|
||||
aria-label={`Fast mode for ${model} in the ${tierLabel} tier`}
|
||||
>
|
||||
<Switch
|
||||
size="sm"
|
||||
checked={paramsByModel?.[model]?.speed === "fast"}
|
||||
onCheckedChange={(enabled) => onFastModeChange(model, enabled)}
|
||||
/>
|
||||
Fast mode
|
||||
</label>
|
||||
</SimpleTooltip>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -48,6 +48,19 @@ const baseParams: BuildComplexityRouterConfigParams = {
|
|||
};
|
||||
|
||||
describe("buildComplexityRouterConfig", () => {
|
||||
it("carries Fast and reasoning overrides independently into a new router payload", () => {
|
||||
const params = { speed: "fast", reasoning_effort: "high", max_tokens: 1024 };
|
||||
const config = buildComplexityRouterConfig({
|
||||
...baseParams,
|
||||
tiers: { ...tiers, COMPLEX: ["primary"], REASONING: ["secondary"] },
|
||||
tierModelParams: { COMPLEX: { primary: params }, REASONING: { secondary: { speed: "fast" } } },
|
||||
});
|
||||
expect(config.tier_model_configs).toEqual({
|
||||
COMPLEX: [{ model_name: "primary", litellm_params: params }],
|
||||
REASONING: [{ model_name: "secondary", litellm_params: { speed: "fast" } }],
|
||||
});
|
||||
});
|
||||
|
||||
it("emits tiers, classifier_type, and escalation_keywords when nothing else is configured", () => {
|
||||
const config = buildComplexityRouterConfig(baseParams);
|
||||
const expected = {
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
serializeTierModelConfigs,
|
||||
tierRowLabel,
|
||||
setTierModelReasoningEffort,
|
||||
setTierModelParam,
|
||||
} from "./complexity_router_tiers";
|
||||
import { resolveComplexityDefaultModel } from "./tier_rows";
|
||||
|
||||
|
|
@ -218,6 +219,28 @@ describe("setTierModelReasoningEffort", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("setTierModelParam", () => {
|
||||
it.each(["reasoning_effort", "speed"] as const)("clears only %s and preserves the input", (key) => {
|
||||
const params = { reasoning_effort: "high", speed: "fast", max_tokens: 512 };
|
||||
const current = { COMPLEX: { primary: params, secondary: { speed: "fast" } }, REASONING: { primary: params } };
|
||||
const cleared = setTierModelParam(current, "COMPLEX", "primary", [key, undefined]);
|
||||
expect(cleared).toEqual({
|
||||
...current,
|
||||
COMPLEX: {
|
||||
...current.COMPLEX,
|
||||
primary: key === "speed" ? { reasoning_effort: "high", max_tokens: 512 } : { speed: "fast", max_tokens: 512 },
|
||||
},
|
||||
});
|
||||
expect(current.COMPLEX.primary).toEqual({ reasoning_effort: "high", speed: "fast", max_tokens: 512 });
|
||||
});
|
||||
|
||||
it("removes empty records when the only override is Fast", () => {
|
||||
const enabled = setTierModelParam(undefined, "COMPLEX", "primary", ["speed", "fast"]);
|
||||
expect(enabled).toEqual({ COMPLEX: { primary: { speed: "fast" } } });
|
||||
expect(setTierModelParam(enabled, "COMPLEX", "primary", ["speed", undefined])).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("pruneTierModelParams", () => {
|
||||
it("drops params for models deselected from the tier", () => {
|
||||
expect(
|
||||
|
|
|
|||
|
|
@ -114,14 +114,16 @@ export const serializeTierModelConfigs = (
|
|||
return serialized.length > 0 ? Object.fromEntries(serialized) : undefined;
|
||||
};
|
||||
|
||||
export const setTierModelReasoningEffort = (
|
||||
export type TierModelParamChange = ["reasoning_effort", ReasoningEffort | undefined] | ["speed", "fast" | undefined];
|
||||
|
||||
export const setTierModelParam = (
|
||||
current: TierModelParamsByTier | undefined,
|
||||
tier: string,
|
||||
model: string,
|
||||
effort: ReasoningEffort | undefined,
|
||||
[key, value]: TierModelParamChange,
|
||||
): TierModelParamsByTier | undefined => {
|
||||
const { reasoning_effort: _dropped, ...rest } = current?.[tier]?.[model] ?? {};
|
||||
const params = effort === undefined ? rest : { ...rest, reasoning_effort: effort };
|
||||
const { [key]: _dropped, ...rest } = current?.[tier]?.[model] ?? {};
|
||||
const params = value === undefined ? rest : { ...rest, [key]: value };
|
||||
const byModel = Object.fromEntries(
|
||||
Object.entries({ ...current?.[tier], [model]: params }).filter(([, value]) => Object.keys(value).length > 0),
|
||||
);
|
||||
|
|
@ -131,6 +133,13 @@ export const setTierModelReasoningEffort = (
|
|||
return Object.keys(next).length > 0 ? next : undefined;
|
||||
};
|
||||
|
||||
export const setTierModelReasoningEffort = (
|
||||
current: TierModelParamsByTier | undefined,
|
||||
tier: string,
|
||||
model: string,
|
||||
effort: ReasoningEffort | undefined,
|
||||
): TierModelParamsByTier | undefined => setTierModelParam(current, tier, model, ["reasoning_effort", effort]);
|
||||
|
||||
export const pruneTierModelParams = (
|
||||
current: TierModelParamsByTier | undefined,
|
||||
tier: string,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,28 @@ describe("CodeSnippets", () => {
|
|||
expect(code).toContain("print(response.data[0].embedding)");
|
||||
});
|
||||
|
||||
describe("custom headers", () => {
|
||||
const customHeaders = { "anthropic-beta": "context-1m-2025-08-07", "x-request-source": "playground" };
|
||||
|
||||
it("passes configured headers as default_headers on the OpenAI client", () => {
|
||||
const code = generateCodeSnippet({ ...baseParams, endpointType: EndpointType.CHAT, customHeaders });
|
||||
expect(code).toContain('base_url="http://localhost:4000",\n\tdefault_headers={');
|
||||
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
|
||||
expect(code).toContain('"x-request-source": "playground"');
|
||||
});
|
||||
|
||||
it("passes configured headers as default_headers on the Azure client", () => {
|
||||
const code = generateCodeSnippet({ ...baseParams, selectedSdk: "azure", customHeaders });
|
||||
expect(code).toContain('api_version="2024-02-01",\n\tdefault_headers={');
|
||||
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
|
||||
});
|
||||
|
||||
it("omits default_headers when no custom headers are configured", () => {
|
||||
expect(generateCodeSnippet(baseParams)).not.toContain("default_headers");
|
||||
expect(generateCodeSnippet({ ...baseParams, customHeaders: {} })).not.toContain("default_headers");
|
||||
});
|
||||
});
|
||||
|
||||
describe("base URL selection", () => {
|
||||
it("should use LITELLM_UI_API_DOC_BASE_URL when provided", () => {
|
||||
const customBaseUrl = "https://custom-doc.example.com";
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import { MessageType } from "./types";
|
||||
import { EndpointType } from "./mode_endpoint_mapping";
|
||||
import { MCPServer } from "@/components/mcp_tools/types";
|
||||
import type { CustomHeaders } from "@/components/llm_calls/request_headers";
|
||||
|
||||
interface CodeGenMetadata {
|
||||
tags?: string[];
|
||||
|
|
@ -30,6 +31,7 @@ interface GenerateCodeParams {
|
|||
PROXY_BASE_URL?: string;
|
||||
LITELLM_UI_API_DOC_BASE_URL?: string | null;
|
||||
};
|
||||
customHeaders?: CustomHeaders;
|
||||
}
|
||||
|
||||
export const generateCodeSnippet = (params: GenerateCodeParams): string => {
|
||||
|
|
@ -48,6 +50,7 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
|
|||
selectedModel,
|
||||
selectedSdk,
|
||||
proxySettings,
|
||||
customHeaders,
|
||||
} = params;
|
||||
const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey;
|
||||
|
||||
|
|
@ -76,6 +79,11 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
|
|||
|
||||
const modelNameForCode = selectedModel || "your-model-name";
|
||||
|
||||
const defaultHeadersCode =
|
||||
customHeaders && Object.keys(customHeaders).length > 0
|
||||
? `,\n\tdefault_headers=${JSON.stringify(customHeaders, null, 2).replace(/\n/g, "\n\t")}`
|
||||
: "";
|
||||
|
||||
const clientInitialization =
|
||||
selectedSdk === "azure"
|
||||
? `import openai
|
||||
|
|
@ -83,13 +91,13 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
|
|||
client = openai.AzureOpenAI(
|
||||
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
|
||||
azure_endpoint="${apiBase}",
|
||||
api_version="2024-02-01"
|
||||
api_version="2024-02-01"${defaultHeadersCode}
|
||||
)`
|
||||
: `import openai
|
||||
|
||||
client = openai.OpenAI(
|
||||
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
|
||||
base_url="${apiBase}"
|
||||
base_url="${apiBase}"${defaultHeadersCode}
|
||||
)`;
|
||||
|
||||
let endpointSpecificCode;
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import openai from "openai";
|
||||
import { makeOpenAIChatCompletionRequest } from "./chat_completion";
|
||||
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
|
||||
|
||||
|
|
@ -615,3 +616,47 @@ describe("chat_completion response cache", () => {
|
|||
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }));
|
||||
});
|
||||
});
|
||||
|
||||
describe("chat_completion custom headers", () => {
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
|
||||
mockCreate.mockReturnValueOnce(nonStreamingResponse({ choices: [{ message: { content: "Hi" } }] }));
|
||||
|
||||
await makeOpenAIChatCompletionRequest(
|
||||
[{ role: "user", content: "Hello" }],
|
||||
vi.fn(),
|
||||
"gpt-4",
|
||||
"test-token",
|
||||
["team-a"],
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
false,
|
||||
{ "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
|
||||
);
|
||||
|
||||
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
|
||||
defaultHeaders: { "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { getProxyBaseUrl } from "@/components/networking";
|
|||
import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types";
|
||||
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
|
||||
import { parseUsageCost } from "./usage_cost";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
|
||||
|
||||
const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk =>
|
||||
({
|
||||
|
|
@ -50,6 +51,7 @@ export async function makeOpenAIChatCompletionRequest(
|
|||
mockTestFallbacks?: boolean,
|
||||
mcpToolsets?: MCPToolset[],
|
||||
streamingEnabled: boolean = true,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
// base url should be the current base_url
|
||||
const isLocal = process.env.NODE_ENV === "development";
|
||||
|
|
@ -57,11 +59,7 @@ export async function makeOpenAIChatCompletionRequest(
|
|||
console.log = function () {};
|
||||
}
|
||||
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
|
||||
// Prepare headers with tags and trace ID
|
||||
const headers: Record<string, string> = {};
|
||||
if (tags && tags.length > 0) {
|
||||
headers["x-litellm-tags"] = tags.join(",");
|
||||
}
|
||||
const headers = buildPlaygroundHeaders(tags, customHeaders);
|
||||
|
||||
const client = new openai.OpenAI({
|
||||
apiKey: accessToken,
|
||||
|
|
|
|||
|
|
@ -52,6 +52,23 @@ describe("fetchAvailableModels", () => {
|
|||
]);
|
||||
});
|
||||
|
||||
it("carries only explicitly supported Fast capabilities, not accepted speed parameters", async () => {
|
||||
modelHubCallMock.mockResolvedValue({
|
||||
data: [
|
||||
{ model_group: "fast", supports_fast_mode: true },
|
||||
{ model_group: "blocked", supports_fast_mode: false },
|
||||
{ model_group: "missing", supports_speed: true },
|
||||
{ model_group: "unknown", supports_fast_mode: null },
|
||||
],
|
||||
});
|
||||
expect(await fetchAvailableModels("token")).toEqual([
|
||||
{ model_group: "blocked" },
|
||||
{ model_group: "fast", supports_fast_mode: true },
|
||||
{ model_group: "missing" },
|
||||
{ model_group: "unknown" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("preserves absent, unknown, empty, and explicit effort capability states", async () => {
|
||||
modelHubCallMock.mockResolvedValue({
|
||||
data: [
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ export interface ModelGroup {
|
|||
model_group: string;
|
||||
mode?: string;
|
||||
supports_reasoning?: boolean;
|
||||
supports_fast_mode?: boolean;
|
||||
supported_reasoning_efforts?: string[] | null;
|
||||
}
|
||||
|
||||
|
|
@ -16,6 +17,7 @@ interface AvailableModel {
|
|||
id?: string | null;
|
||||
mode?: string | null;
|
||||
supports_reasoning?: boolean | null;
|
||||
supports_fast_mode?: boolean | null;
|
||||
supported_reasoning_efforts?: string[] | null;
|
||||
}
|
||||
|
||||
|
|
@ -25,6 +27,7 @@ const toModelGroup = (item: AvailableModel): ModelGroup => {
|
|||
model_group: groupName,
|
||||
...(item.mode && { mode: item.mode }),
|
||||
...(item.supports_reasoning === true && { supports_reasoning: true }),
|
||||
...(item.supports_fast_mode === true && { supports_fast_mode: true }),
|
||||
...(item.supported_reasoning_efforts !== undefined && {
|
||||
supported_reasoning_efforts: item.supported_reasoning_efforts,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,60 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
buildPlaygroundHeaders,
|
||||
customHeadersFromPairs,
|
||||
parseStoredHeaderPairs,
|
||||
withRequiredHeaders,
|
||||
} from "./request_headers";
|
||||
|
||||
describe("customHeadersFromPairs", () => {
|
||||
it("trims header names and drops rows without a name", () => {
|
||||
expect(
|
||||
customHeadersFromPairs([
|
||||
[" anthropic-beta ", "context-1m-2025-08-07"],
|
||||
["", "orphan value"],
|
||||
[" ", "whitespace name"],
|
||||
["x-empty", ""],
|
||||
]),
|
||||
).toEqual({ "anthropic-beta": "context-1m-2025-08-07", "x-empty": "" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("parseStoredHeaderPairs", () => {
|
||||
it("round-trips pairs persisted as JSON", () => {
|
||||
const pairs = [["anthropic-beta", "context-1m-2025-08-07"]] as const;
|
||||
expect(parseStoredHeaderPairs(JSON.stringify(pairs))).toEqual(pairs);
|
||||
});
|
||||
|
||||
it("returns no pairs for missing, malformed, or wrongly shaped storage", () => {
|
||||
expect(parseStoredHeaderPairs(null)).toEqual([]);
|
||||
expect(parseStoredHeaderPairs("not json")).toEqual([]);
|
||||
expect(parseStoredHeaderPairs(JSON.stringify({ "anthropic-beta": "x" }))).toEqual([]);
|
||||
expect(parseStoredHeaderPairs(JSON.stringify([["ok", "pair"], ["one"], [1, 2], "str"]))).toEqual([["ok", "pair"]]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildPlaygroundHeaders", () => {
|
||||
it("joins tags into x-litellm-tags and lets custom headers override it", () => {
|
||||
expect(buildPlaygroundHeaders(["a", "b"], { "x-custom": "1" })).toEqual({
|
||||
"x-litellm-tags": "a,b",
|
||||
"x-custom": "1",
|
||||
});
|
||||
expect(buildPlaygroundHeaders(["a"], { "x-litellm-tags": "b" })).toEqual({ "x-litellm-tags": "b" });
|
||||
});
|
||||
|
||||
it("omits x-litellm-tags when there are no tags", () => {
|
||||
expect(buildPlaygroundHeaders([], { "x-custom": "1" })).toEqual({ "x-custom": "1" });
|
||||
expect(buildPlaygroundHeaders(undefined, undefined)).toEqual({});
|
||||
});
|
||||
});
|
||||
|
||||
describe("withRequiredHeaders", () => {
|
||||
it("keeps required headers regardless of custom header name casing", () => {
|
||||
expect(
|
||||
withRequiredHeaders(
|
||||
{ authorization: "Bearer stolen", "content-type": "text/plain", "x-custom": "1" },
|
||||
{ Authorization: "Bearer real", "Content-Type": "application/json" },
|
||||
),
|
||||
).toEqual({ Authorization: "Bearer real", "Content-Type": "application/json", "x-custom": "1" });
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
import type { KeyValuePair } from "@/components/key_value_input";
|
||||
|
||||
export type CustomHeaders = Readonly<Record<string, string>>;
|
||||
|
||||
export const customHeadersFromPairs = (pairs: readonly KeyValuePair[]): CustomHeaders =>
|
||||
Object.fromEntries(pairs.map(([name, value]) => [name.trim(), value]).filter(([name]) => name !== ""));
|
||||
|
||||
const isHeaderPair = (entry: unknown): entry is KeyValuePair =>
|
||||
Array.isArray(entry) && entry.length === 2 && entry.every((part) => typeof part === "string");
|
||||
|
||||
export const parseStoredHeaderPairs = (raw: string | null): readonly KeyValuePair[] => {
|
||||
if (!raw) return [];
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(raw);
|
||||
return Array.isArray(parsed) ? parsed.filter(isHeaderPair) : [];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
};
|
||||
|
||||
export const buildPlaygroundHeaders = (
|
||||
tags?: readonly string[],
|
||||
customHeaders?: CustomHeaders,
|
||||
): Record<string, string> => ({
|
||||
...(tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : {}),
|
||||
...customHeaders,
|
||||
});
|
||||
|
||||
export const withRequiredHeaders = (
|
||||
headers: Readonly<Record<string, string>>,
|
||||
required: Readonly<Record<string, string>>,
|
||||
): Record<string, string> => {
|
||||
const reserved = new Set(Object.keys(required).map((name) => name.toLowerCase()));
|
||||
return {
|
||||
...Object.fromEntries(Object.entries(headers).filter(([name]) => !reserved.has(name.toLowerCase()))),
|
||||
...required,
|
||||
};
|
||||
};
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import openai from "openai";
|
||||
import { makeOpenAIResponsesRequest } from "./responses_api";
|
||||
import { MessageType } from "../chat_ui/types";
|
||||
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
|
||||
|
|
@ -611,3 +612,46 @@ describe("responses_api response cache", () => {
|
|||
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }), "");
|
||||
});
|
||||
});
|
||||
|
||||
describe("responses_api custom headers", () => {
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
|
||||
mockResponsesCreate.mockReturnValueOnce(nonStreamingResponse({ id: "resp_1", output: [] }));
|
||||
|
||||
await makeOpenAIResponsesRequest(
|
||||
[{ role: "user", content: "Hello" }],
|
||||
vi.fn(),
|
||||
"gpt-4",
|
||||
"test-token",
|
||||
["team-a"],
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
false,
|
||||
undefined,
|
||||
{ "anthropic-beta": "context-1m-2025-08-07" },
|
||||
);
|
||||
|
||||
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
|
||||
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import { getProxyBaseUrl } from "@/components/networking";
|
|||
import { toast } from "@/lib/toast";
|
||||
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
|
||||
import { parseUsageCost } from "./usage_cost";
|
||||
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
|
||||
import type { MCPEvent } from "@/components/mcp_tools/types";
|
||||
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
|
||||
import {
|
||||
|
|
@ -85,6 +86,7 @@ export async function makeOpenAIResponsesRequest(
|
|||
mcpToolsets?: MCPToolset[],
|
||||
streamingEnabled: boolean = true,
|
||||
onTotalLatency?: (latency: number) => void,
|
||||
customHeaders?: CustomHeaders,
|
||||
) {
|
||||
if (!accessToken) {
|
||||
throw new Error("Virtual Key is required");
|
||||
|
|
@ -101,11 +103,7 @@ export async function makeOpenAIResponsesRequest(
|
|||
}
|
||||
|
||||
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
|
||||
// Prepare headers with tags and trace ID
|
||||
const headers: Record<string, string> = {};
|
||||
if (tags && tags.length > 0) {
|
||||
headers["x-litellm-tags"] = tags.join(",");
|
||||
}
|
||||
const headers = buildPlaygroundHeaders(tags, customHeaders);
|
||||
|
||||
const client = new openai.OpenAI({
|
||||
apiKey: accessToken,
|
||||
|
|
|
|||
11
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
11
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27297,6 +27297,8 @@ export interface components {
|
|||
jwt_claim_name: string;
|
||||
/** Jwt Claim Value */
|
||||
jwt_claim_value: string;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/** Key */
|
||||
key: string;
|
||||
};
|
||||
|
|
@ -28910,6 +28912,8 @@ export interface components {
|
|||
jwt_claim_name: string;
|
||||
/** Jwt Claim Value */
|
||||
jwt_claim_value: string;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/**
|
||||
* Updated At
|
||||
* Format: date-time
|
||||
|
|
@ -32704,6 +32708,11 @@ export interface components {
|
|||
supported_openai_params: string[] | null;
|
||||
/** Supported Reasoning Efforts */
|
||||
supported_reasoning_efforts?: string[] | null;
|
||||
/**
|
||||
* Supports Fast Mode
|
||||
* @default false
|
||||
*/
|
||||
supports_fast_mode: boolean;
|
||||
/**
|
||||
* Supports Function Calling
|
||||
* @default false
|
||||
|
|
@ -38643,6 +38652,8 @@ export interface components {
|
|||
id: string;
|
||||
/** Is Active */
|
||||
is_active?: boolean | null;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/** Key */
|
||||
key?: string | null;
|
||||
};
|
||||
|
|
|
|||
8
uv.lock
generated
8
uv.lock
generated
|
|
@ -10,7 +10,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-09-12T03:51:49.261499386Z"
|
||||
exclude-newer = "2026-09-12T22:48:38.53978Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -4457,7 +4457,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.102.0"
|
||||
version = "1.103.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
|
@ -4874,12 +4874,12 @@ proxy-dev = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.67"
|
||||
version = "0.1.68"
|
||||
source = { editable = "enterprise" }
|
||||
|
||||
[[package]]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.97"
|
||||
version = "0.4.98"
|
||||
source = { editable = "litellm-proxy-extras" }
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue