scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

Merge remote-tracking branch 'origin/main' into litellm_drop_non_finite_spend_cost

Browse source
This commit is contained in:
mateo-berri 2026-10-29 00:33:07 -07:00
commit 0e44f74b5a
42 changed files with 2078 additions and 467 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -84,9 +84,10 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
## Screenshots / Proof of Fix
<!-- Include screenshots, screen recordings, or command (e.g., curl) + output demonstrating that your changes work as expected
The proof must be completely e2e with no mocks, using actual LLM calls costing real $$$ if applicable. `pytest` commands are not enough
Show ONLY the latest run: capture Before at the merge base and After at the PR's current tip, and when new commits change behavior, replace this whole section with the fresh run instead of stacking it on top of older ones. The run must be up to date. As soon as a new commit is made and it makes this PR description's after sha stale (it's no longer tip of PR), you must re-run the QA
Structure the section exactly as below: Before and After one heading level below this section, each naming the commit hash it was captured at, one lower-level heading per case inside each, the same case names in the same order on both sides, and numbered steps (command, observed output) under every case, never loose prose; shared setup (config, payloads) goes above Before, and with a single case, drop the case headings and number the steps directly
Simple PR exception (2026-10-28, #49030). A mechanical change fully proven by focused checks may skip before/after e2e. Follow AGENTS.md's eligibility rule, trace callers, and verify bug-fix regression sensitivity. Record E2E exception at <tip SHA>, scope, rationale, commands and results here. Tests/static checks suffice only for this exception. Keep CI and bot gates, reassess changed diffs, and honor explicit e2e requests. Omit the Before/After scaffold for an exception
Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof
Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa
Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings
### Before (<hash>)

4
AGENTS.md
View file

@ -41,7 +41,9 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`. The Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into. I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis

7
litellm/google_genai/adapters/transformation.py
View file

@ -198,6 +198,11 @@ def _genai_message_part(part: object) -> _GenAIMessagePart | None:
return ChatCompletionTextObject(type="text", text=text) if isinstance(text, str) else _genai_media_part(fields)
def _genai_content_parts(value: object) -> Sequence[object]:
single_part: Final = _validated(_CONFIG_FIELDS, value)
return (single_part,) if single_part is not None else _validated(_GENAI_CONTENT_PARTS, value) or ()
def _genai_message_parts(parts: object) -> tuple[_GenAIMessagePart, ...]:
items: Final = _validated(_GENAI_CONTENT_PARTS, parts) or ()
return tuple(converted for part in items if (converted := _genai_message_part(part)) is not None)
@ -248,7 +253,7 @@ def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]:
fields: Final = _validated(_CONFIG_FIELDS, content)
if fields is None:
return ()
parts: Final = _validated(_GENAI_CONTENT_PARTS, fields.get("parts", ())) or ()
parts: Final = _genai_content_parts(fields.get("parts", ()))
converted: Final = _genai_message_parts(parts)
role: Final = fields.get("role", "user")
if role == "user":

34
litellm/litellm_core_utils/litellm_logging.py
View file

@ -5063,6 +5063,23 @@ class Logging(LiteLLMLoggingBaseClass):
return result_copy
_SENSITIVE_KEYWORDS: Final = (
"authorization",
"token",
"key",
"secret",
"vertex_credentials",
"credentials",
"password",
"passwd",
)
def has_sensitive_keyword(key: str) -> bool:
lowered: Final = key.lower()
return any(keyword in lowered for keyword in _SENSITIVE_KEYWORDS)
def get_masked_values(
sensitive_object: Mapping[str, object],
ignore_sensitive_values: bool = False,
@ -5081,16 +5098,6 @@ def get_masked_values(
masked_length: Optional length for the masked portion (number of *). If set, will use exactly this many *
regardless of original string length. The total length will be unmasked_length + masked_length.
"""
sensitive_keywords: Final = (
"authorization",
"token",
"key",
"secret",
"vertex_credentials",
"credentials",
"password",
"passwd",
)
def _mask_value(v: object) -> object:
if isinstance(v, dict):
@ -5117,12 +5124,7 @@ def get_masked_values(
return v[: unmasked_length // 2] + "*" * (len(v) - unmasked_length) + v[-unmasked_length // 2 :]
return {
k: (
v
if ignore_sensitive_values
or not any(sensitive_keyword in k.lower() for sensitive_keyword in sensitive_keywords)
else _mask_value(v)
)
k: (v if ignore_sensitive_values or not has_sensitive_keyword(k) else _mask_value(v))
for k, v in sensitive_object.items()
}

40
litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py
View file

@ -31,6 +31,7 @@ from ..constants import (
)
from ..errors import AnthropicContextManagementError
from ..placeholders import build_cleared_tool_result_content
from .compact import count_system_tokens, history_after_client_compaction
if TYPE_CHECKING:
from litellm.types.llms.openai import ChatCompletionToolParam
@ -203,9 +204,10 @@ def _trigger_met(
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
system: str | Sequence[Mapping[str, object]] | None,
counted_tool_uses: int,
) -> tuple[bool, int | None]:
"""Return (trigger_met, input_tokens if counted for reuse)."""
"""Return (trigger_met, message and tool tokens for reuse); like the API, input_tokens adds the system prompt."""
trigger_type: Final = trigger.get("type", "input_tokens")
threshold = trigger.get("value")
@ -216,14 +218,17 @@ def _trigger_met(
if not isinstance(threshold, int):
threshold = DEFAULT_INPUT_TOKENS_TRIGGER
current_tokens: Final = litellm.token_counter(
message_tokens: Final = litellm.token_counter(
model=model,
messages=messages,
tools=cast("list[ChatCompletionToolParam] | None", tools),
)
if message_tokens > threshold:
return True, message_tokens
current_tokens: Final = message_tokens + count_system_tokens(model, system)
verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens)
verbose_logger.debug("context_management polyfill: threshold: %s", threshold)
return current_tokens > threshold, current_tokens
return current_tokens > threshold, message_tokens
def _resolve_keep_count(keep: dict[str, object]) -> int:
@ -237,7 +242,7 @@ def _resolve_keep_count(keep: dict[str, object]) -> int:
def _last_completed_tool_use_id(
messages: list[dict[str, object]],
messages: Sequence[Mapping[str, object]],
) -> str | None:
"""Latest completed tool_result id; never cleared."""
last_id: str | None = None
@ -253,7 +258,7 @@ def _last_completed_tool_use_id(
def _clear_tool_results(
messages: list[dict[str, object]], ids_to_clear: set[str]
messages: list[dict[str, object]], ids_to_clear: Set[str]
) -> tuple[list[dict[str, object]], int]:
"""Clear matching tool_result content; return (messages, cleared_count)."""
cleared = 0
@ -311,7 +316,11 @@ def apply_clear_tool_uses_20250919(
system: str | list[dict[str, object]] | None,
edit_spec: ClearToolUsesEditSpec,
) -> tuple[list[dict[str, object]], AppliedEdit | None]:
"""Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None)."""
"""Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None).
Like the API, it reads only the history after a client compaction block and counts the block's summary as system.
"""
kept, counted_system = history_after_client_compaction(messages, system)
fields: Final = _ClearToolUsesFields.model_validate(edit_spec)
clear_at_least: Final = 0 if fields.clear_at_least is None else fields.clear_at_least.value
exclude_tools: Final = frozenset(fields.exclude_tools or ())
@ -326,7 +335,7 @@ def apply_clear_tool_uses_20250919(
"value": DEFAULT_KEEP_TOOL_USES,
}
blocks: Final = _all_blocks(messages)
blocks: Final = _all_blocks(kept)
results: Final = _blocks_by_id(blocks, "tool_result", "tool_use_id")
uses: Final = _blocks_by_id(blocks, "tool_use", "id")
tiny: Final = _tiny_result_ids(results)
@ -335,15 +344,15 @@ def apply_clear_tool_uses_20250919(
for tool_use_id, name in _tool_uses_in_order(blocks)
if name not in exclude_tools and (tool_use_id not in tiny or _clears_input(name, input_tools))
)
met, counted_tokens = _trigger_met(trigger, model, messages, tools, len(clearable))
if not met:
return messages, None
keep_count: Final = _resolve_keep_count(keep)
if len(clearable) <= keep_count:
return messages, None
last_completed_id: Final = _last_completed_tool_use_id(messages)
met, counted_tokens = _trigger_met(trigger, model, kept, tools, counted_system, len(clearable))
if not met:
return messages, None
last_completed_id: Final = _last_completed_tool_use_id(kept)
candidates: Final = tuple(
(tool_use_id, name)
for tool_use_id, name in clearable[: len(clearable) - keep_count]
@ -358,7 +367,7 @@ def apply_clear_tool_uses_20250919(
_select_in_batches(candidate_ids, tokens, clear_at_least) if clear_at_least > 0 else candidate_ids
)
edited, cleared_count = _clear_tool_results(messages, set(ids_to_clear))
edited, cleared_count = _clear_tool_results(messages, ids_to_clear)
verbose_logger.debug("context_management polyfill: edited: %s", edited)
if cleared_count == 0:
return messages, None
@ -370,11 +379,12 @@ def apply_clear_tool_uses_20250919(
counted_tokens
if counted_tokens is not None
else litellm.token_counter(
model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools)
model=model, messages=kept, tools=cast("list[ChatCompletionToolParam] | None", tools)
)
)
kept_after, _ = history_after_client_compaction(final_messages, system)
tokens_after: Final = litellm.token_counter(
model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools)
model=model, messages=kept_after, tools=cast("list[ChatCompletionToolParam] | None", tools)
)
cleared_input_tokens: Final = max(tokens_before - tokens_after, 0)

27
litellm/llms/anthropic/messages/context_management/editors/compact.py
View file

@ -189,6 +189,21 @@ def _strip_compaction_blocks(
return cleaned
def history_after_client_compaction(
messages: Sequence[dict[str, object]],
system: str | Sequence[Mapping[str, object]] | None,
) -> tuple[list[dict[str, object]], str | Sequence[Mapping[str, object]] | None]:
"""Return the history the API keeps after the latest compaction block, without it, and the system with its summary."""
sliced, block = _slice_around_compaction_block(messages)
if block is None:
return list(messages), system
tail: Final = _strip_compaction_blocks(sliced)
summary: Final = block.get("content")
if not isinstance(summary, str) or not summary:
return tail, system
return tail, _augment_system_with_summary(_system_to_text(system) or None, summary)
def _augment_system_with_summary(
system: str | list[dict[str, object]] | None,
summary_text: str,
@ -320,14 +335,16 @@ def _count_effective_tokens(
content: Final = compaction_block.get("content") or ""
if content:
total += litellm.token_counter(model=model, text=content)
return total + count_system_tokens(model, system)
def count_system_tokens(model: str, system: str | Sequence[Mapping[str, object]] | None) -> int:
system_text: Final = _system_to_text(system)
if system_text:
total += litellm.token_counter(model=model, text=system_text)
return total
return litellm.token_counter(model=model, text=system_text) if system_text else 0
def _system_to_text(
system: str | list[dict[str, object]] | None,
system: str | Sequence[Mapping[str, object]] | None,
) -> str:
"""Flatten an Anthropic-style ``system`` value into a single string for
token counting. Returns ``""`` when ``system`` carries no text."""
@ -338,7 +355,7 @@ def _system_to_text(
return "\n".join(
text
for block in system
if isinstance(block, dict)
if isinstance(block, Mapping)
and block.get("type") == "text"
and isinstance(text := block.get("text"), str)
and text

2
litellm/llms/openai/responses/transformation.py
View file

@ -1218,7 +1218,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
def compact_dropped_params(self, litellm_params: GenericLiteLLMParams) -> frozenset[str]:
if self.custom_llm_provider not in _PROVIDERS_WITH_STRICT_COMPACT_PARAMS:
return frozenset()
if _drops_params(litellm_params):
if drops_params(litellm_params):
return _COMPACT_TRACKING_PARAMS | _COMPACT_REJECTED_PARAMS
return _COMPACT_TRACKING_PARAMS

18
litellm/proxy/_types.py
View file

@ -3816,6 +3816,16 @@ class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase):
end_time: datetime | None = None
def proxy_exception_status_code(message: str, code: str) -> str:
# Litellm router.py returns "No healthy deployment available" when there are no deployments available
# Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487
if "No healthy deployment available" in message or "No deployments available" in message:
return "429"
if RouterErrors.no_deployments_with_tag_routing.value in message:
return "401"
return code
class ProxyException(Exception):
# NOTE: DO NOT MODIFY THIS
# This is used to map exactly to OPENAI Exceptions
@ -3844,13 +3854,7 @@ class ProxyException(Exception):
headers[k] = str(v)
self.headers = headers or {}
self.provider_specific_fields = provider_specific_fields
# rules for proxyExceptions
# Litellm router.py returns "No healthy deployment available" when there are no deployments available
# Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487
if "No healthy deployment available" in self.message or "No deployments available" in self.message:
self.code = "429"
elif RouterErrors.no_deployments_with_tag_routing.value in self.message:
self.code = "401"
self.code = proxy_exception_status_code(self.message, self.code)
def to_dict(self) -> dict:
"""Converts the ProxyException instance to a dictionary."""

112
litellm/proxy/auth/auth_checks.py
View file

@ -46,6 +46,11 @@ from litellm.constants import (
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.tool_fields import DropPolicy
from litellm.llms.base_llm.managed_resources.utils import (
extract_provider_resource_id_from_unified_id,
is_base64_encoded_unified_id,
)
from litellm.models.project import LiteLLM_ProjectTable
from litellm.proxy._types import (
MODEL_LIST_ROUTES,
@ -165,6 +170,7 @@ from litellm.types.proxy.auth.auth_checks import UserNotFoundError
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
from litellm.utils import get_utc_datetime
from litellm.vector_stores.vector_store_registry import listed_vector_store_ids, vector_store_ids_in_tool
from .auth_checks_organization import (
add_team_org_context_to_request_body,
@ -1514,6 +1520,7 @@ async def common_checks(
valid_token=valid_token,
user_object=user_object,
deny_by_default=_vector_store_deny_by_default(_typed_request_body(general_settings)),
llm_router=llm_router,
)
# 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
@ -7985,6 +7992,57 @@ def _strict_requested_vector_store_ids(request_body: Mapping[str, object]) -> tu
)
@dataclass(frozen=True, slots=True)
class _UnreadableVectorStoreIds:
pass
def _top_level_vector_store_ids(
vector_store_ids: object, drop_policies: Callable[[], tuple[DropPolicy, ...]]
) -> tuple[str, ...] | _UnreadableVectorStoreIds:
listed: Final = listed_vector_store_ids(vector_store_ids)
if vector_store_ids is None or vector_store_ids == [*listed]:
return listed
if any(policy.drops("vector_store_ids") for policy in drop_policies()):
return listed
if isinstance(vector_store_ids, str):
return (vector_store_ids,)
return _UnreadableVectorStoreIds()
def _legacy_requested_vector_store_ids(
request_body: Mapping[str, object], drop_policies: Callable[[], tuple[DropPolicy, ...]]
) -> tuple[str, ...] | _UnreadableVectorStoreIds:
"""
Reads tools the way VectorStoreRegistry does and leaves a malformed tool to the proxy's tool
validation. A top-level vector_store_ids string names one store, since a query parameter always
arrives as a string. Any other value that is not a list of strings names no readable stores unless a
drop policy the request can resolve to drops it, and then only the string entries of a list name
stores, as in registry retrieval.
"""
top_level: Final = _top_level_vector_store_ids(request_body.get("vector_store_ids"), drop_policies)
if isinstance(top_level, _UnreadableVectorStoreIds):
return top_level
return tuple(
chain(
top_level,
chain.from_iterable(vector_store_ids_in_tool(tool) for tool in _tools(request_body.get("tools"))),
)
)
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None]]) -> bool:
return any(grant is not None and bool(grant.vector_stores) for _, grant in grants)
def _request_drop_policies(
request_body: Mapping[str, object], llm_router: Router | None, valid_token: UserAPIKeyAuth | None
) -> tuple[DropPolicy, ...]:
from litellm.proxy.common_request_processing import counted_drop_policies
return counted_drop_policies(request_body, llm_router, valid_token.team_id if valid_token is not None else None)
def _require_vector_store_grant(
object_type: GrantLayer,
vector_store_ids_to_run: Sequence[str],
@ -8032,6 +8090,7 @@ async def vector_store_access_check(
*,
user_object: LiteLLM_UserTable | None = None,
deny_by_default: bool = False,
llm_router: Router | None = None,
):
"""
Checks whether the caller may use every vector store the request names.
@ -8040,11 +8099,19 @@ async def vector_store_access_check(
`retrieval_config.vector_store_id`. Grants come from each identity's
`object_permission.vector_stores`.
With `deny_by_default=False` (legacy), a key or team only restricts access when its list is
nonempty, and stores are only read from the request when a vector store registry is loaded.
Both modes read the same request fields, with or without a vector store registry, and answer 400
when the top-level `vector_store_ids` is not a list of strings, which legacy mode does only when the
key or team has an allowlist and the value is not a string, which it checks as one store id. A
LiteLLM-managed store id is compared by the provider store id it decodes to, on both sides. Legacy mode skips that value instead when the request, `litellm_settings`, the router defaults or a deployment of the requested
model drops it (`drop_params`, or `additional_drop_params` naming it), the same settings the proxy's
tool validation reads, and reads tools the way the registry does, where strict mode answers 400 on a
malformed one.
With `deny_by_default=True` (`general_settings.vector_store_deny_by_default`), stores are read
even without a registry, and a missing record, `null` or `[]` grants nothing:
With `deny_by_default=False` (legacy), a key or team only restricts access when its list is
nonempty.
With `deny_by_default=True` (`general_settings.vector_store_deny_by_default`), a missing record,
`null` or `[]` grants nothing:
- virtual key: the key must grant every store, and so must its team when it has one, even if
the team failed to load
@ -8056,8 +8123,8 @@ async def vector_store_access_check(
rescue nor restrict a key or team request.
Raises ProxyException (401, `{key,team,user}_vector_store_access_denied`) on the first identity
that does not grant a requested store, and with the flag on, ProxyException (400,
`invalid_request_error`) when a `vector_store_ids` field is not a list of strings.
that does not grant a requested store, and ProxyException (400, `invalid_request_error`) on a
malformed `vector_store_ids` the request does not drop and, in legacy mode, an allowlist applies to.
"""
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -8068,21 +8135,23 @@ async def vector_store_access_check(
verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check")
return True
registry_ids: Final = (
_strict_requested_vector_store_ids(_typed_request_body(request_body))
typed_request_body: Final = _typed_request_body(request_body)
requested_ids: Final = (
_strict_requested_vector_store_ids(typed_request_body)
if deny_by_default
else (
litellm.vector_store_registry.get_vector_store_ids_to_run(
non_default_params=request_body, tools=request_body.get("tools", None)
)
if litellm.vector_store_registry is not None
else None
else _legacy_requested_vector_store_ids(
typed_request_body, lambda: _request_drop_policies(typed_request_body, llm_router, valid_token)
)
or ()
)
rag_vector_store_id: Final = _get_rag_query_vector_store_id(_typed_request_body(request_body))
if isinstance(requested_ids, _UnreadableVectorStoreIds):
if _restricts_vector_stores(
await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache)
):
raise _malformed_vector_store_ids()
return True
rag_vector_store_id: Final = _get_rag_query_vector_store_id(typed_request_body)
rag_ids: Final = (rag_vector_store_id,) if rag_vector_store_id is not None else ()
vector_store_ids_to_run: Final = tuple(dict.fromkeys((*registry_ids, *rag_ids)))
vector_store_ids_to_run: Final = tuple(dict.fromkeys((*requested_ids, *rag_ids)))
if not vector_store_ids_to_run:
return True
@ -8109,6 +8178,12 @@ async def vector_store_access_check(
return True
def _provider_vector_store_id(vector_store_id: str) -> str:
if not is_base64_encoded_unified_id(vector_store_id):
return vector_store_id
return extract_provider_resource_id_from_unified_id(vector_store_id) or vector_store_id
def _can_object_call_vector_stores(
object_type: Literal["key", "team", "org", "user"],
vector_store_ids_to_run: Sequence[str],
@ -8127,8 +8202,9 @@ def _can_object_call_vector_stores(
if len(object_permissions.vector_stores) == 0:
return True
allowed_stores: Final = frozenset(map(_provider_vector_store_id, object_permissions.vector_stores))
for vector_store_id in vector_store_ids_to_run:
if vector_store_id not in object_permissions.vector_stores:
if _provider_vector_store_id(vector_store_id) not in allowed_stores:
raise ProxyException(
message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}",
type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type),

4
litellm/proxy/common_request_processing.py
View file

@ -725,7 +725,7 @@ def _router_defaults(llm_router: Router | None) -> Mapping[str, object]:
return {name: value for name, value in _REQUEST_FIELDS.validate_python(configured).items() if value is not None}
def _counted_drop_policies(
def counted_drop_policies(
data: Mapping[str, object], llm_router: Router | None, team_id: str | None
) -> tuple[DropPolicy, ...]:
requested: Final = {**_router_defaults(llm_router), **data}
@ -741,7 +741,7 @@ def _counted_drop_policies(
def resolve_counted_tool_fields(
data: Mapping[str, object], llm_router: Router | None, team_id: str | None
) -> Mapping[str, object] | ProxyException:
resolved: Final = resolve_tool_fields(data, lambda: _counted_drop_policies(data, llm_router, team_id))
resolved: Final = resolve_tool_fields(data, lambda: counted_drop_policies(data, llm_router, team_id))
if isinstance(resolved, MalformedField):
return ProxyException(
message=resolved.message,

96
litellm/proxy/common_utils/openai_error_payload.py
View file

@ -2,14 +2,17 @@
``type`` is a required string and ``param`` is nullable, neither of which the literal
string ``"None"`` satisfies."""
import sys
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from fastapi import status
from fastapi import HTTPException, status
from litellm.constants import STRINGIFIED_NONE
from litellm.proxy._types import ProxyException
from litellm.proxy._types import ProxyErrorTypes, ProxyException, proxy_exception_status_code
from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain
LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
@ -21,6 +24,19 @@ _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
}
)
_AUTH_FAILURE_STATUS_CODES: Final = frozenset({status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN})
_FAILOVER_SQLSTATES: Final = frozenset({"25006", "57P01", "57P02", "57P03"})
_FAILOVER_REFUSALS: Final = (
"is not currently accepting connections",
"the database system is starting up",
"the database system is shutting down",
"the database system is in recovery mode",
"the database system is not yet accepting connections",
"terminating connection due to administrator command",
)
def attribute_of(value: object, name: str, default: object = None) -> object:
return getattr(value, name, default)
@ -61,6 +77,82 @@ def openai_error_param(exc: object) -> str | None:
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
def proxy_exception_for(exc: Exception, default_status_code: int) -> ProxyException:
match exc:
case ProxyException():
return exc
case _ if _is_database_outage(exc):
return ProxyException(
message=PrismaDBExceptionHandler.database_unavailable_message(exc),
type=ProxyErrorTypes.no_db_connection.value,
param=None,
code=status.HTTP_503_SERVICE_UNAVAILABLE,
)
case HTTPException(status_code=status_code, detail=detail):
return _openai_shaped_proxy_exception(exc, str(detail), status_code)
case _:
return _openai_shaped_proxy_exception(
exc, str(exc) or type(exc).__name__, error_status_code(exc, default_status_code)
)
def _is_database_outage(exc: Exception) -> bool:
return any(map(_is_database_outage_link, exception_chain(exc)))
def _is_database_outage_link(exc: BaseException) -> bool:
match exc:
case DBLookupDeadlineExceeded():
return True
case Exception() if _prisma_is_loaded():
return _is_prisma_outage(exc)
case _:
return False
def _prisma_is_loaded() -> bool:
return sys.modules.get("prisma") is not None
def _is_prisma_outage(exc: Exception) -> bool:
if PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc):
return True
if type(exc).__module__.partition(".")[0] != "prisma":
return False
return PrismaDBExceptionHandler.is_database_service_unavailable_error(exc) or _is_failover_refusal(exc)
def _is_failover_refusal(exc: Exception) -> bool:
if PrismaDBExceptionHandler.is_read_only_transaction_error(exc):
return True
if PrismaDBExceptionHandler.postgres_sqlstate(exc) in _FAILOVER_SQLSTATES:
return True
message: Final = str(exc).lower()
return any(refusal in message for refusal in _FAILOVER_REFUSALS)
def _openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException:
answered_status: Final = int(proxy_exception_status_code(message, str(status_code)))
return ProxyException(
message=message,
type=_catch_all_error_type(exc, answered_status),
param=openai_error_param(exc),
code=answered_status,
)
def _catch_all_error_type(exc: Exception, status_code: int) -> str:
if _rejects_the_caller(exc, status_code):
return ProxyErrorTypes.auth_error.value
return openai_error_type(exc, status_code)
def _rejects_the_caller(exc: Exception, status_code: int) -> bool:
raised_as_http_answer: Final = isinstance(exc, HTTPException)
status_from_upstream: Final = error_status_code(exc, 0) == status_code
return status_code in _AUTH_FAILURE_STATUS_CODES and (raised_as_http_answer or not status_from_upstream)
def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers
if litellm_call_id is None:
return None

84
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -5,6 +5,7 @@ CRUD endpoints for storing reusable credentials.
import time
import unicodedata
from collections.abc import Mapping
from types import MappingProxyType
from typing import (
Annotated,
Final,
@ -23,7 +24,7 @@ from litellm.constants import (
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import get_masked_values
from litellm.litellm_core_utils.litellm_logging import get_masked_values, has_sensitive_keyword
from litellm.llms.anthropic.wif import (
ExportedJwks,
NotAnInternalIssuerCredential,
@ -83,6 +84,11 @@ _CREDENTIAL_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_GITHUB_COPILOT_PROVIDER: Final = "github_copilot"
_DISPLAY_NAME_MAX_LENGTH: Final = 255
_DISPLAY_NAME_FORBIDDEN_CATEGORIES: Final = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"})
_MASK_UNMASKED_LENGTH: Final = 4
_MASK_ASTERISKS: Final = 4
_SHORT_SECRET_MASK: Final = "*****"
_MASKED_READ_BACK_DEPTH: Final = 20
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
def _normalized_github_copilot_values(
@ -395,7 +401,14 @@ async def get_credentials(
"""
try:
masked_credentials: Final = [
_credential_view(credential, get_masked_values(credential.credential_values))
_credential_view(
credential,
get_masked_values(
credential.credential_values,
unmasked_length=_MASK_UNMASKED_LENGTH,
number_of_asterisks=_MASK_ASTERISKS,
),
)
for credential in litellm.credential_list
]
return {"success": True, "credentials": masked_credentials}
@ -423,7 +436,11 @@ async def get_credential_by_name(
if credential.credential_name == credential_name:
return _credential_view(
credential,
get_masked_values(credential.credential_values, unmasked_length=4, number_of_asterisks=4),
get_masked_values(
credential.credential_values,
unmasked_length=_MASK_UNMASKED_LENGTH,
number_of_asterisks=_MASK_ASTERISKS,
),
)
raise HTTPException(
status_code=404,
@ -521,8 +538,8 @@ async def get_credential_by_model(
raise HTTPException(status_code=404, detail="Model not found")
masked_credential_values: Final = get_masked_values(
credential_values,
unmasked_length=4,
number_of_asterisks=4,
unmasked_length=_MASK_UNMASKED_LENGTH,
number_of_asterisks=_MASK_ASTERISKS,
)
credential: Final = CredentialItem(
credential_name=f"{model.model_name}-credential-{model_id}",
@ -908,6 +925,48 @@ def update_db_credential(
return merged_credential
def _is_masked_read_back(value: object) -> bool:
if not isinstance(value, str):
return False
if value == _SHORT_SECRET_MASK:
return True
shown: Final = _MASK_UNMASKED_LENGTH // 2
return len(value) == _MASK_UNMASKED_LENGTH + _MASK_ASTERISKS and value[shown:-shown] == "*" * _MASK_ASTERISKS
def _as_fields(value: object) -> Mapping[str, object] | None:
try:
return _CREDENTIAL_DICT_ADAPTER.validate_python(value, strict=True)
except ValidationError:
return None
def _kept_value(key: str, value: object, stored: Mapping[str, object], depth: int) -> tuple[str, object] | None:
if not has_sensitive_keyword(key):
return key, value
fields: Final = _as_fields(value)
if fields is not None and depth < _MASKED_READ_BACK_DEPTH:
return key, _without_masked_read_backs(fields, _as_fields(stored.get(key)) or _NO_FIELDS, depth + 1)
if not _is_masked_read_back(value):
return key, value
if key in stored:
return key, stored[key]
return None
def _without_masked_read_backs(
requested: Mapping[str, object], stored: Mapping[str, object], depth: int = 0
) -> Mapping[str, object]:
kept: Final = (_kept_value(key, value, stored, depth) for key, value in requested.items())
return _CREDENTIAL_DICT_ADAPTER.validate_python(
MappingProxyType(dict(entry for entry in kept if entry is not None))
)
def _stored_objects(stored_values: Mapping[str, object]) -> Mapping[str, object]:
return MappingProxyType({key: value for key, value in stored_values.items() if _as_fields(value) is not None})
@router.patch(
"/credentials/{credential_name:path}",
dependencies=[Depends(user_api_key_auth)],
@ -936,12 +995,12 @@ async def update_credential(
)
requested_display_name: Final = _normalized_display_name(credential.display_name)
_reject_overlapping_credential_values(credential)
incoming_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
requested_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
_resolve_deployment_credentials(llm_router, credential.model_id)
if credential.model_id
else credential.credential_values or {}
)
_reject_non_admin_wif_fields(_incoming_wif_fields(incoming_values, credential), user_api_key_dict)
_reject_non_admin_wif_fields(_incoming_wif_fields(requested_values, credential), user_api_key_dict)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -952,19 +1011,20 @@ async def update_credential(
if db_credential is None:
raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.")
_reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict)
decrypted_stored_credential: Final = decrypted_credential(db_credential)
decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
decrypted_stored_credential.credential_values
)
incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values))
patch: Final = CredentialItem(
credential_name=credential_name,
display_name=(
requested_display_name if "display_name" in credential.model_fields_set else db_credential.display_name
),
credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info),
credential_values=incoming_values,
credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values),
credential_values_to_delete=credential.credential_values_to_delete,
)
decrypted_stored_credential: Final = decrypted_credential(db_credential)
decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
decrypted_stored_credential.credential_values
)
stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
)

4
litellm/proxy/db/exception_handler.py
View file

@ -36,7 +36,7 @@ _LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
def _exception_chain(e: BaseException) -> Iterator[BaseException]:
def exception_chain(e: BaseException) -> Iterator[BaseException]:
current = e # rebind-ok: advances one link per iteration of the bounded walk
for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH):
yield current
@ -49,7 +49,7 @@ def _exception_chain(e: BaseException) -> Iterator[BaseException]:
def _database_service_unavailable_errors(e: BaseException) -> tuple[Exception, ...]:
return tuple(
link
for link in _exception_chain(e)
for link in exception_chain(e)
if isinstance(link, Exception) and PrismaDBExceptionHandler.is_database_service_unavailable_error(link)
)

73
litellm/proxy/google_endpoints/endpoints.py
View file

@ -1,11 +1,11 @@
from collections.abc import Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final
from typing import Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, ORJSONResponse
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
from litellm.proxy._types import *
@ -27,14 +27,67 @@ from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
def _google_repeated_field(value: JsonValue) -> JsonValue:
match value:
case None:
return []
case dict():
return [value]
case list():
return [item for item in value if item is not None]
case _:
return value
def _google_content(value: JsonValue) -> JsonValue:
match value:
case {"parts": parts}:
return {**value, "parts": _google_repeated_field(parts)}
case _:
return value
def _google_contents(value: JsonValue) -> JsonValue:
match _google_repeated_field(value):
case list() as contents:
return [_google_content(content) for content in contents]
case other:
return other
_GOOGLE_PART_TEXT_AND_METADATA_FIELDS: Final = frozenset(
{"text", "thought", "thoughtSignature", "thought_signature", "partMetadata", "part_metadata"}
)
def _google_part_has_non_text_content(part: Mapping[str, JsonValue]) -> bool:
return any(field is not None for key, field in part.items() if key not in _GOOGLE_PART_TEXT_AND_METADATA_FIELDS)
def _google_system_part(value: JsonValue) -> JsonValue:
match value:
case dict() if value.get("text") is None and _google_part_has_non_text_content(value):
return {key: field for key, field in value.items() if key != "text"}
case dict():
return {"text": None, **value}
case _:
return value
_GOOGLE_REPEATED_FIELD: Final = BeforeValidator(_google_repeated_field)
class _GoogleCountTokensTextPart(BaseModel):
model_config = ConfigDict(frozen=True)
text: str
text: str | None
class _GoogleCountTokensSystemInstruction(BaseModel):
model_config = ConfigDict(frozen=True)
parts: tuple[_GoogleCountTokensTextPart, ...]
parts: Annotated[
tuple[Annotated[_GoogleCountTokensTextPart, BeforeValidator(_google_system_part)], ...],
_GOOGLE_REPEATED_FIELD,
] = ()
class _GoogleCountTokensFunctionDeclaration(BaseModel):
@ -47,16 +100,16 @@ class _GoogleCountTokensFunctionDeclaration(BaseModel):
class _GoogleCountTokensTool(BaseModel):
model_config = ConfigDict(frozen=True, populate_by_name=True)
function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field(
function_declarations: Annotated[tuple[_GoogleCountTokensFunctionDeclaration, ...], _GOOGLE_REPEATED_FIELD] = Field(
default=(), alias="functionDeclarations"
)
class _GoogleCountTokensRequest(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True)
contents: tuple[Mapping[str, JsonValue], ...] = ()
contents: Annotated[tuple[Mapping[str, JsonValue], ...], BeforeValidator(_google_contents)] = ()
system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction")
tools: tuple[_GoogleCountTokensTool, ...] = ()
tools: Annotated[tuple[_GoogleCountTokensTool, ...], _GOOGLE_REPEATED_FIELD] = ()
router: Final = APIRouter(
@ -112,7 +165,11 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq
GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", [])
)
system: Final = (
[{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts]
[
{"type": "text", "text": part.text}
for part in validated_data.system_instruction.parts
if part.text is not None
]
if validated_data.system_instruction is not None
else None
)

20
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -32,7 +32,6 @@ from litellm.proxy._types import (
EnterpriseLicenseData,
Litellm_EntityType,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
SpecialModelNames,
UserAPIKeyAuth,
@ -47,6 +46,7 @@ from litellm.proxy.auth.auth_utils import (
from litellm.proxy.auth.model_checks import get_key_models
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils import http_parsing_utils
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.health_check_latest import (
@ -683,21 +683,9 @@ async def health_services_endpoint(
except Exception as e:
verbose_proxy_logger.error("litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e
def convert_health_check_to_dict(check) -> dict:

19
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -46,6 +46,7 @@ from litellm.proxy.auth.password_policy import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.common_utils.user_api_key_cache import (
AUTH_OBJECTS_TARGET,
object_permission_cache_key,
@ -1729,21 +1730,9 @@ async def user_update(
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_update(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
async def bulk_update_processed_users(

55
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -75,6 +75,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
coordination_redis_cache,
publish_config_change,
)
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache
@ -4087,21 +4088,9 @@ async def update_key_fn(
# update based on remaining passed in values
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_key_fn(): Exception occured - %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
@router.post(
@ -7244,21 +7233,9 @@ async def list_keys(
except Exception as e:
verbose_proxy_logger.exception("Error in list_keys: %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"error({e})"),
type=ProxyErrorTypes.internal_server_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.internal_server_error,
param=getattr(e, "param", "None"),
code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR) from e
async def _apply_non_admin_alias_scope(
@ -7402,21 +7379,9 @@ async def key_aliases(
except Exception as e:
verbose_proxy_logger.exception("Error in key_aliases: %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"error({e})"),
type=ProxyErrorTypes.internal_server_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.internal_server_error,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e
def _validate_sort_params(

69
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -77,6 +77,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, is_team_admin
@ -2845,21 +2846,9 @@ async def delete_model(
except Exception as e:
verbose_proxy_logger.exception("Failed to delete model. Due to error - %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
async def delete_team_model_alias(
@ -2899,6 +2888,13 @@ async def delete_team_model_alias(
return removed_model_aliases
def _model_not_added_to_db() -> HTTPException:
return HTTPException(
status_code=500,
detail={"error": "Failed to add model to db. Check your server logs for more details."},
)
@router.post(
"/model/new",
description="Allows adding new models to the model list in the config.yaml",
@ -3080,6 +3076,8 @@ async def add_new_model(
if isinstance(e, HTTPException):
raise
verbose_proxy_logger.exception("Exception in add_new_model: %s", e)
if model_response is None:
raise _model_not_added_to_db() from e
else:
raise HTTPException(
@ -3088,10 +3086,7 @@ async def add_new_model(
)
if model_response is None:
raise HTTPException(
status_code=500,
detail={"error": "Failed to add model to db. Check your server logs for more details."},
)
raise _model_not_added_to_db()
## CREATE AUDIT LOG ##
asyncio.create_task(
@ -3121,21 +3116,9 @@ async def add_new_model(
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.add_new_model(): Exception occured - %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
#### MODEL MANAGEMENT ####
@ -3364,21 +3347,9 @@ async def update_model(
return None if model_response is None else _model_write_response(model_response, member_write)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_model(): Exception occured - %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
@router.post(

19
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_checks import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.budget_management_endpoints import (
@ -1396,21 +1397,9 @@ async def organization_member_add(
)
except Exception as e:
verbose_proxy_logger.exception("Error adding member to organization: %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e
async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> LiteLLM_UserTable:

19
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -132,6 +132,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
)
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
@ -5526,21 +5527,9 @@ async def team_info(
e,
traceback.format_exc(),
)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
@router.get(

59
litellm/proxy/proxy_server.py
View file

@ -524,6 +524,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
LITELLM_CALL_ID_HEADER,
headers_with_litellm_call_id,
litellm_call_id_headers,
proxy_exception_for,
with_litellm_call_id,
)
from litellm.proxy.common_utils.periodic_reload_schedule import (
@ -18536,25 +18537,13 @@ async def async_queue_request(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
elif isinstance(e, litellm.BadRequestError):
if isinstance(e, litellm.BadRequestError):
raise ProxyException(
message=e.message, type="invalid_request_error", param=getattr(e, "param", None), code=e.status_code
)
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
) from e
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
@app.get("/fallback/login", tags=["experimental"], include_in_schema=False)
@ -19833,21 +19822,9 @@ async def update_config(
except Exception as e:
verbose_proxy_logger.error("litellm.proxy.proxy_server.update_config(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
### CONFIG GENERAL SETTINGS
@ -21030,21 +21007,9 @@ async def get_config(
}
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.get_config(): Exception occured - %s", e)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({e})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
if isinstance(e, ProxyException):
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e
@router.get(

16
litellm/vector_stores/vector_store_registry.py
View file

@ -36,7 +36,7 @@ else:
_DB_ROW: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True))
_TOOL: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True))
_STORE_ID_ITEMS: Final = TypeAdapter(list[object], config=ConfigDict(strict=True))
_STORE_ID_ITEMS: Final = TypeAdapter(list[object] | tuple[object, ...], config=ConfigDict(strict=True))
def _named_store_ids(store_ids: object) -> tuple[str, ...]:
@ -49,7 +49,11 @@ def _named_store_ids(store_ids: object) -> tuple[str, ...]:
return tuple(store_id for store_id in items if isinstance(store_id, str))
def _tool_vector_store_ids(tool: object) -> tuple[str, ...]:
def listed_vector_store_ids(store_ids: object) -> tuple[str, ...]:
return _named_store_ids(store_ids) if isinstance(store_ids, (list, tuple)) else ()
def vector_store_ids_in_tool(tool: object) -> tuple[str, ...]:
try:
return _named_store_ids(_TOOL.validate_python(tool).get("vector_store_ids"))
except ValidationError:
@ -195,7 +199,7 @@ class VectorStoreRegistry:
tools_to_remove: Final[list[int]] = []
for i, tool in enumerate(tools):
tool_vector_store_ids = _tool_vector_store_ids(tool)
tool_vector_store_ids = vector_store_ids_in_tool(tool)
if not tool_vector_store_ids:
continue
@ -305,7 +309,7 @@ class VectorStoreRegistry:
List of vector stores with tool parameters merged into litellm_params
"""
# Pop vector_store_ids from params
vector_store_ids: Final[list[str]] = non_default_params.pop("vector_store_ids", None) or []
vector_store_ids: Final = [*listed_vector_store_ids(non_default_params.pop("vector_store_ids", None))]
# Extract params from tools and collect IDs
params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids)
@ -353,7 +357,7 @@ class VectorStoreRegistry:
List of vector stores with tool parameters merged into litellm_params
"""
# Pop vector_store_ids from params
vector_store_ids: Final[list[str]] = non_default_params.pop("vector_store_ids", None) or []
vector_store_ids: Final = [*listed_vector_store_ids(non_default_params.pop("vector_store_ids", None))]
# Extract params from tools and collect IDs
params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids)
@ -423,7 +427,7 @@ class VectorStoreRegistry:
validated_tools: Final = _ITEMS.validate_python(tools)
except ValidationError:
return vector_store_ids
return [*vector_store_ids, *chain.from_iterable(_tool_vector_store_ids(tool) for tool in validated_tools)]
return [*vector_store_ids, *chain.from_iterable(vector_store_ids_in_tool(tool) for tool in validated_tools)]
def load_vector_stores_from_config(self, vector_stores_config: list[dict]):
"""

26
tests/benchmarks/test_json_fragment_accumulator_benchmarks.py Normal file
View file

@ -0,0 +1,26 @@
from collections.abc import Iterator
from typing import Final
import pytest
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
CONCATENATED_VALUES: Final = 20_000
def _popped_values(accumulator: JSONFragmentAccumulator) -> Iterator[object]:
while True:
found, value = accumulator.pop_next_value()
if not found:
return
yield value
@pytest.mark.benchmark
def test_draining_many_concatenated_values() -> None:
accumulator: Final = JSONFragmentAccumulator()
accumulator.append('{"a": 1}' * CONCATENATED_VALUES)
drained: Final = sum(1 for _ in _popped_values(accumulator))
assert drained == CONCATENATED_VALUES

2
tests/integration/authorization/test_team_admin_gate.py
View file

@ -746,7 +746,7 @@ def test_internal_user_model_creation_prohibition_refuses_team_admin_but_not_pro
"error": {
"message": "Model creation is disabled for internal users by disable_model_add_for_internal_users.",
"type": "auth_error",
"param": "None",
"param": None,
"code": "403",
}
},

4
tests/integration/authorization/test_user_self_update_guard.py
View file

@ -38,7 +38,7 @@ def _escalation_error(field: str) -> dict[str, JsonValue]:
"error": {
"message": f"{{'error': \"Non-admin users cannot modify '{field}' on their own record. Contact your proxy admin.\"}}",
"type": "auth_error",
"param": "None",
"param": None,
"code": "403",
}
}
@ -49,7 +49,7 @@ def _role_change_error() -> dict[str, JsonValue]:
"error": {
"message": "Only proxy admins can modify user roles.",
"type": "auth_error",
"param": "None",
"param": None,
"code": "403",
}
}

4
tests/integration/management/test_team_delete_large_membership.py
View file

@ -255,8 +255,8 @@ def _team_not_found_body(team_id: str) -> dict[str, JsonValue]:
return {
"error": {
"message": f"{{'message': 'Team not found, passed team id: {team_id}.'}}",
"type": "auth_error",
"param": "None",
"type": "invalid_request_error",
"param": None,
"code": "404",
}
}

12
tests/unit/google_genai/test_google_genai_adapter.py
View file

@ -158,6 +158,18 @@ def test_native_function_response_preserves_json_when_media_is_malformed(invalid
assert messages == [{"role": "tool", "tool_call_id": "call_lookup", "content": '{"ok": true}'}]
@pytest.mark.parametrize("role", ["user", "model"])
def test_one_object_content_parts_translate_like_their_list_form(role: str) -> None:
part: Final = {"text": "Reply with OK"}
def messages(parts: object) -> object:
return GoogleGenAIAdapter().translate_generate_content_to_completion(
model="anthropic/claude-haiku-5-5", contents=[{"role": role, "parts": parts}]
)["messages"]
assert messages(part) == messages([part]) != []
@pytest.mark.parametrize(
("args", "expected_arguments"),
[(None, "null"), ('{"query":"x"}', '"{\\"query\\":\\"x\\"}"')],

71
tests/unit/litellm_core_utils/test_json_fragment_accumulator.py
View file

@ -1,5 +1,8 @@
import json
import time
import tracemalloc
from collections.abc import Iterator
from typing import Final
from unittest.mock import patch
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
@ -175,40 +178,48 @@ def test_accumulation_of_many_fragments_is_not_quadratic():
assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?"
def test_draining_many_concatenated_values_is_not_quadratic():
"""
Regression guard: peeling N JSON values already sitting in one buffer,
one pop_next_value() call per value with no new fragments in between,
must be O(n) total. Re-copying the shrinking remainder on every pop
(slicing a new string instead of advancing a cursor) makes total drain
time scale with the square of the buffer size.
def _popped_values(accumulator: JSONFragmentAccumulator) -> Iterator[object]:
while True:
found, value = accumulator.pop_next_value()
if not found:
return
yield value
Uses a doubling ratio rather than an absolute ms budget so it isn't
flaky on a slower or busier CI runner: doubling the input should
roughly double an O(n) drain's time but roughly quadruple an O(n^2)
drain's time, and that ratio holds regardless of machine speed.
"""
def drain_time_ms(n: int) -> float:
accumulator = JSONFragmentAccumulator()
accumulator.append('{"a": 1}' * n)
start = time.perf_counter()
drained = 0
while True:
found, _ = accumulator.pop_next_value()
if not found:
break
drained += 1
assert drained == n
return (time.perf_counter() - start) * 1000
def _streamed(fragments: tuple[str, ...]) -> JSONFragmentAccumulator:
accumulator: Final = JSONFragmentAccumulator()
for fragment in fragments:
accumulator.append(fragment)
return accumulator
small_ms = drain_time_ms(40_000)
large_ms = drain_time_ms(80_000)
ratio = large_ms / max(small_ms, 0.001)
assert ratio < 3.0, (
f"doubling drained values scaled time by {ratio:.2f}x ({small_ms:.1f} ms -> {large_ms:.1f} ms); "
"expected roughly 2x for O(n); O(n^2) regression?"
def _drain_peak_bytes(fragments: tuple[str, ...]) -> int:
accumulator: Final = _streamed(fragments)
first_found, _ = accumulator.pop_next_value()
assert first_found is True
already_tracing: Final = tracemalloc.is_tracing()
if not already_tracing:
tracemalloc.start()
tracemalloc.reset_peak()
baseline, _ = tracemalloc.get_traced_memory()
try:
drained: Final = sum(1 for _ in _popped_values(accumulator))
_, peak = tracemalloc.get_traced_memory()
finally:
if not already_tracing:
tracemalloc.stop()
assert drained == len(fragments) - 1
return peak - baseline
def test_draining_many_concatenated_values_never_copies_the_remainder():
value: Final = '{"text": "' + "x" * 20 + '"}'
fragments: Final = tuple(value if index % 2 else f" {value}" for index in range(40_000))
payload_bytes: Final = sum(len(fragment) for fragment in fragments)
peak_bytes: Final = _drain_peak_bytes(fragments)
assert peak_bytes < payload_bytes // 4, (
f"draining {len(fragments)} values allocated a peak of {peak_bytes} bytes over a {payload_bytes} byte payload, "
"copying the remainder on every pop?"
)

111
tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py
View file

@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm
from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled
from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled, prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import (
AnthropicContextManagementError,
PolyfillResult,
@ -690,6 +690,115 @@ def test_tool_uses_trigger_does_not_count_excluded_tools(trigger_value: int, tri
assert (applied is not None) == triggers
@pytest.mark.parametrize(
"system",
["Reference notes. " * 400, [{"type": "text", "text": "Reference notes. " * 400}]],
ids=["string", "blocks"],
)
def test_input_tokens_trigger_counts_the_system_prompt(system: str | list[dict[str, object]]) -> None:
"""Like the API, the trigger compares the whole prompt: a system prompt can push it over."""
messages = _make_history(n_pairs=3)
history_tokens: Final = litellm.token_counter(model=MODEL, messages=messages)
edit_spec: Final = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": history_tokens + 100},
"keep": {"type": "tool_uses", "value": 1},
}
_, without_system = apply_clear_tool_uses_20250919(
model=MODEL, messages=deepcopy(messages), tools=None, system=None, edit_spec=edit_spec
)
_, with_system = apply_clear_tool_uses_20250919(
model=MODEL, messages=deepcopy(messages), tools=None, system=system, edit_spec=edit_spec
)
_, met_without_system = apply_clear_tool_uses_20250919(
model=MODEL,
messages=deepcopy(messages),
tools=None,
system=None,
edit_spec={**edit_spec, "trigger": {"type": "input_tokens", "value": 0}},
)
assert without_system is None
assert with_system is not None
assert met_without_system is not None
assert with_system["cleared_tool_uses"] == 2
assert with_system["cleared_input_tokens"] == met_without_system["cleared_input_tokens"] > 0
@pytest.mark.parametrize("compact_position", [None, "after", "before"], ids=["alone", "before_compact", "after_compact"])
async def test_input_tokens_trigger_counts_a_client_compaction_summary(compact_position: str | None) -> None:
"""A client compaction summary counts toward the trigger and calls beside it stay clearable, in any edit order."""
tail: Final = _make_history(n_pairs=3)
edit_spec: Final = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": litellm.token_counter(model=MODEL, messages=tail) + 100},
"keep": {"type": "tool_uses", "value": 1},
}
compaction: Final = {"type": "compaction", "content": "Earlier work. " * 400}
compacted_history: Final = [tail[0], {**tail[1], "content": [compaction, *tail[1]["content"]]}, *tail[2:]]
compact_spec: Final = {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}}
edits: Final = {None: [edit_spec], "after": [edit_spec, compact_spec], "before": [compact_spec, edit_spec]}[
compact_position
]
async def prepare(messages: list[dict[str, object]]) -> PolyfillResult | None:
return await prepare_context_managed_request(
model=MODEL,
messages=messages,
tools=None,
system=None,
context_management_spec={"edits": edits},
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
def cleared(result: PolyfillResult | None) -> list[object]:
applied: Final = result.applied_edits if result is not None else []
return [edit.get("cleared_tool_uses") for edit in applied if edit.get("type") == "clear_tool_uses_20250919"]
without_summary: Final = await prepare(deepcopy(tail))
with_summary: Final = await prepare(deepcopy(compacted_history))
assert cleared(without_summary) == []
assert cleared(with_summary) == [2]
async def test_compact_after_clearing_replaces_the_client_compaction_summary(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy.proxy_server import general_settings
monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "summary-model")
summary: Final = MagicMock()
summary.choices = [MagicMock(message=MagicMock(content="<summary>Fresh summary</summary>"))]
summary.usage = MagicMock(prompt_tokens=50, completion_tokens=100)
router: Final = MagicMock()
router.acompletion = AsyncMock(return_value=summary)
tail: Final = _make_history(n_pairs=3)
result: Final = await prepare_context_managed_request(
model=MODEL,
messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Old summary. " * 20_000}]}, *tail],
tools=None,
system="You are a helpful assistant.",
context_management_spec={
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 1_000},
"keep": {"type": "tool_uses", "value": 1},
},
{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}},
]
},
litellm_metadata=None,
additional_drop_params=None,
llm_router=router,
)
assert result is not None
assert isinstance(result.system, str)
assert "Fresh summary" in result.system
assert "Old summary" not in result.system
assert result.system.endswith("You are a helpful assistant.")
@pytest.mark.parametrize(("clear_tool_inputs", "triggers"), [(False, False), (True, True)])
def test_tool_uses_trigger_counts_tiny_results_only_when_their_inputs_are_cleared(
clear_tool_inputs: bool, triggers: bool

1
tests/unit/llms/anthropic/messages/context_management/test_dispatcher.py
View file

@ -151,6 +151,7 @@ async def test_sync_editor_counts_tokens_off_the_event_loop():
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 10_000_000},
"keep": {"type": "tool_uses", "value": 0},
}
]
},

497
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -101,6 +101,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
tag_registry_cache_key,
)
from litellm.utils import get_utc_datetime
from litellm.llms.base_llm.managed_resources.utils import encode_unified_id, generate_unified_id_string
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
from tests._master_key import MASTER_KEY
@ -1551,37 +1552,24 @@ async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team,
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_client,vector_store_registry,expected_result",
[
(None, MagicMock(), True), # No prisma client
(MagicMock(), None, True), # No vector store registry
(MagicMock(), MagicMock(), True), # No vector stores to run
],
)
async def test_vector_store_access_check_early_returns(prisma_client, vector_store_registry, expected_result):
@pytest.mark.parametrize("prisma_client", [None, MagicMock()], ids=["no-prisma-client", "no-vector-stores-to-run"])
async def test_vector_store_access_check_early_returns(prisma_client):
"""Test vector_store_access_check returns True for early exit conditions"""
request_body = {"messages": [{"role": "user", "content": "test"}]}
if vector_store_registry:
vector_store_registry.get_vector_store_ids_to_run.return_value = None
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.vector_store_registry", vector_store_registry),
):
with patch("litellm.proxy.proxy_server.prisma_client", prisma_client):
result = await vector_store_access_check(
request_body=request_body,
team_object=None,
valid_token=None,
)
assert result == expected_result
assert result is True
@pytest.mark.asyncio
async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_requested():
"""Registry returns [] (not None) for plain requests; no object permission DB lookup should happen."""
"""A plain request names no vector store, so no object permission DB lookup should happen."""
valid_token = UserAPIKeyAuth(token="test-token", object_permission_id="perm-123")
team_object = MagicMock()
team_object.object_permission_id = "team-permission"
@ -1590,16 +1578,8 @@ async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_r
find_unique = AsyncMock()
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = find_unique
mock_vector_store_registry = MagicMock()
mock_vector_store_registry.get_vector_store_ids_to_run.return_value = []
with (
patch( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
),
patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists
"litellm.vector_store_registry", mock_vector_store_registry
),
with patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
):
result = await vector_store_access_check(
request_body={"messages": [{"role": "user", "content": "test"}]},
@ -1668,102 +1648,6 @@ def test_can_object_call_vector_stores_scenarios(object_permissions, vector_stor
assert result is True
@pytest.mark.asyncio
async def test_vector_store_access_check_with_permissions():
"""Test vector_store_access_check with actual permission checking"""
request_body = {"tools": [{"type": "function", "function": {"name": "test"}}]}
# Test with valid token that has access
valid_token = UserAPIKeyAuth(
token="test-token",
object_permission_id="perm-123",
models=["gpt-4"],
max_budget=100.0,
)
mock_prisma_client = MagicMock()
mock_permissions = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-123", vector_stores=["store-1", "store-2"]
)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=mock_permissions)
mock_vector_store_registry = MagicMock()
mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"]
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.vector_store_registry", mock_vector_store_registry),
):
result = await vector_store_access_check(
request_body=request_body,
team_object=None,
valid_token=valid_token,
)
assert result is True
# Test with denied access
mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-3"]
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.vector_store_registry", mock_vector_store_registry),
):
with pytest.raises(ProxyException) as exc_info:
await vector_store_access_check(
request_body=request_body,
team_object=None,
valid_token=valid_token,
)
assert exc_info.value.type == ProxyErrorTypes.key_vector_store_access_denied
@pytest.mark.asyncio
async def test_vector_store_access_check_with_team_permissions():
"""Ensure teams restricted to specific vector stores cannot access others."""
request_body = {}
valid_token = UserAPIKeyAuth(token="team-test-token", object_permission_id=None)
team_object = LiteLLM_TeamTable(team_id="team-1", object_permission_id="team-permission")
mock_prisma_client = MagicMock()
team_permissions = LiteLLM_ObjectPermissionTable(
object_permission_id="team-permission", vector_stores=["team-store-allowed"]
)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=team_permissions)
mock_vector_store_registry = MagicMock()
mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-allowed"]
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.vector_store_registry", mock_vector_store_registry),
):
result = await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
assert result is True
mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-denied"]
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.vector_store_registry", mock_vector_store_registry),
):
with pytest.raises(ProxyException) as exc_info:
await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
assert exc_info.value.type == ProxyErrorTypes.team_vector_store_access_denied
@pytest.mark.asyncio
@pytest.mark.parametrize(
"requested_vector_store_id,expected_error_type",
@ -1823,6 +1707,188 @@ async def test_vector_store_access_check_enforces_team_allowlist_for_rag_query(
assert exc_info.value.type == expected_error_type
def _file_search_tool_request(vector_store_id: str) -> dict[str, object]:
return {
"model": "gpt-5.4-mini",
"input": "what is in this store?",
"tools": [{"type": "file_search", "vector_store_ids": [vector_store_id]}],
}
def _file_search_tool_string_id_request(vector_store_id: str) -> dict[str, object]:
return {
"model": "gpt-5.4-mini",
"input": "what is in this store?",
"tools": [{"type": "file_search", "vector_store_ids": vector_store_id}],
}
def _vector_store_ids_request(vector_store_id: str) -> dict[str, object]:
return {
"model": "gpt-5.4-mini",
"messages": [{"role": "user", "content": "what is in this store?"}],
"vector_store_ids": [vector_store_id],
}
def _drop_params_request(vector_store_ids: object) -> dict[str, object]:
return {
"model": "gpt-5.4-mini",
"messages": [{"role": "user", "content": "what is in this store?"}],
"drop_params": True,
"vector_store_ids": vector_store_ids,
}
def _allowlist(object_permission_id: str, vector_stores: list[str] | None) -> LiteLLM_ObjectPermissionTable | None:
if vector_stores is None:
return None
return LiteLLM_ObjectPermissionTable(object_permission_id=object_permission_id, vector_stores=vector_stores)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key_vector_stores", "team_vector_stores", "request_body", "denied_by"),
[
pytest.param(
None,
["vs_allowed"],
_file_search_tool_request("vs_other_team"),
ProxyErrorTypes.team_vector_store_access_denied,
id="team-file-search-tool-other-store",
),
pytest.param(None, ["vs_allowed"], _file_search_tool_request("vs_allowed"), None, id="team-file-search-tool"),
pytest.param(
["vs_allowed"],
None,
_vector_store_ids_request("vs_other_team"),
ProxyErrorTypes.key_vector_store_access_denied,
id="key-vector-store-ids-other-store",
),
pytest.param(["vs_allowed"], None, _vector_store_ids_request("vs_allowed"), None, id="key-vector-store-ids"),
pytest.param(
None,
["vs_allowed"],
_file_search_tool_string_id_request("vs_other_team"),
ProxyErrorTypes.team_vector_store_access_denied,
id="team-file-search-tool-string-id-other-store",
),
pytest.param(
None,
None,
_file_search_tool_string_id_request("vs_other_team"),
None,
id="unrestricted-file-search-tool-string-id",
),
pytest.param(
None,
["vs_allowed"],
_drop_params_request(["vs_other_team", {"id": "vs_allowed"}]),
ProxyErrorTypes.team_vector_store_access_denied,
id="team-drop-params-list-with-an-object-other-store",
),
pytest.param(
None,
["vs_allowed"],
{
**_file_search_tool_request("vs_other_team"),
"tools": [{"type": "file_search", "vector_store_ids": ["vs_other_team", 1]}],
},
ProxyErrorTypes.team_vector_store_access_denied,
id="team-file-search-tool-with-an-int-other-store",
),
pytest.param(
None,
["vs_allowed"],
{
**_vector_store_ids_request("vs_other_team"),
"additional_drop_params": ["vector_store_ids"],
"vector_store_ids": {"vs_other_team": True},
},
None,
id="team-additional-drop-params-object-names-no-store",
),
pytest.param(
None,
["vs_allowed"],
_drop_params_request({"vs_other_team": True}),
None,
id="team-drop-params-object-names-no-store",
),
],
)
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_vector_store_access_check_enforces_allowlists_on_the_stores_the_request_names(
key_vector_stores: list[str] | None,
team_vector_stores: list[str] | None,
request_body: dict[str, object],
denied_by: ProxyErrorTypes | None,
vector_store_registry: VectorStoreRegistry | None,
) -> None:
valid_token: Final = UserAPIKeyAuth(
token="team-a-key",
team_id="team-a",
object_permission_id="key-permission",
object_permission=_allowlist("key-permission", key_vector_stores),
)
team_object: Final = LiteLLM_TeamTable(
team_id="team-a",
object_permission_id="team-permission",
object_permission=_allowlist("team-permission", team_vector_stores),
)
with (
patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists
"litellm.vector_store_registry", vector_store_registry
),
):
await _assert_rag_query_outcome(
vector_store_access_check(request_body=request_body, team_object=team_object, valid_token=valid_token),
denied_by,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_litellm_settings_drop_params_skips_a_malformed_vector_store_ids(
vector_store_registry: VectorStoreRegistry | None,
) -> None:
with patch("litellm.drop_params", True): # test-quality-ok: litellm_settings.drop_params is this module global
assert await _common_checks_for_rag_query(
{},
_virtual_key(object_permission_id="key-permission"),
SimpleNamespace(vector_stores=["KBSTOREA"]),
request_body={
"model": "gpt-4o-mini",
"input": "what is in this KB?",
"vector_store_ids": {"KBSTOREB": True},
},
vector_store_registry=vector_store_registry,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"deployment_params",
[{"drop_params": True}, {"additional_drop_params": ["vector_store_ids"]}],
ids=["drop-params", "additional-drop-params"],
)
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_a_deployment_that_drops_vector_store_ids_skips_a_malformed_one(
deployment_params: Mapping[str, object], vector_store_registry: VectorStoreRegistry | None
) -> None:
assert await _common_checks_for_rag_query(
{},
_virtual_key(object_permission_id="key-permission"),
SimpleNamespace(vector_stores=["KBSTOREA"]),
request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", "vector_store_ids": ["KBSTOREA", 1]},
vector_store_registry=vector_store_registry,
llm_router=_gpt_4o_mini_router(deployment_params),
)
_KEY_DENIED: Final = ProxyErrorTypes.key_vector_store_access_denied
_TEAM_DENIED: Final = ProxyErrorTypes.team_vector_store_access_denied
_USER_DENIED: Final = ProxyErrorTypes.user_vector_store_access_denied
@ -1866,6 +1932,7 @@ async def _common_checks_for_rag_query(
request_body: Mapping[str, object] | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
vector_store_registry: VectorStoreRegistry | None = None,
llm_router: "Router | None" = None,
) -> bool:
permissions: Final = {
"key-permission": _permission_row("key-permission", key_permission),
@ -1906,13 +1973,24 @@ async def _common_checks_for_rag_query(
global_proxy_spend=None,
general_settings=dict(general_settings),
route="/v1/rag/query",
llm_router=None,
llm_router=llm_router,
proxy_logging_obj=MagicMock(),
valid_token=valid_token,
request=MagicMock(spec=Request),
)
def _gpt_4o_mini_router(litellm_params: Mapping[str, object]) -> "Router":
return litellm.Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", **litellm_params},
}
]
)
async def _assert_rag_query_outcome(request: Awaitable[bool], denied_by: ProxyErrorTypes | None) -> None:
if denied_by is None:
assert await request is True
@ -2343,6 +2421,175 @@ async def test_deny_by_default_rejects_malformed_vector_store_ids_as_bad_request
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"request_body",
[
{"vector_store_ids": {"KBSTOREB": True}},
{"vector_store_ids": ["KBSTOREA", {"id": "KBSTOREB"}]},
{"vector_store_ids": 7},
],
ids=["object-top-level-ids", "object-top-level-id", "int-top-level-ids"],
)
@pytest.mark.parametrize(
"general_settings", [{"vector_store_deny_by_default": True}, {}], ids=["deny-by-default", "legacy"]
)
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_malformed_vector_store_ids_are_rejected_as_bad_request(
request_body: Mapping[str, object],
general_settings: Mapping[str, object],
vector_store_registry: VectorStoreRegistry | None,
):
with pytest.raises(ProxyException) as exc_info:
await _common_checks_for_rag_query(
general_settings,
_virtual_key(object_permission_id="key-permission"),
SimpleNamespace(vector_stores=["KBSTOREA"]),
request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", **request_body},
vector_store_registry=vector_store_registry,
llm_router=_gpt_4o_mini_router({}),
)
assert (exc_info.value.type, exc_info.value.param, exc_info.value.code) == (
"invalid_request_error",
"vector_store_ids",
"400",
)
_MALFORMED_TOP_LEVEL_IDS: Final = pytest.mark.parametrize(
"vector_store_ids",
[{"KBSTOREB": True}, ["KBSTOREA", {"id": "KBSTOREB"}], 7],
ids=["object-top-level-ids", "object-top-level-id", "int-top-level-ids"],
)
async def _team_key_query_with_top_level_ids(
key_vector_stores: list[str] | None,
team_vector_stores: list[str] | None,
vector_store_ids: object,
vector_store_registry: VectorStoreRegistry | None,
) -> bool:
return await _common_checks_for_rag_query(
{},
_virtual_key(object_permission_id=None if key_vector_stores is None else "key-permission", team_id="team-1"),
None if key_vector_stores is None else SimpleNamespace(vector_stores=key_vector_stores),
LiteLLM_TeamTable(
team_id="team-1", object_permission_id=None if team_vector_stores is None else "team-permission"
),
None if team_vector_stores is None else SimpleNamespace(vector_stores=team_vector_stores),
request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", "vector_store_ids": vector_store_ids},
vector_store_registry=vector_store_registry,
llm_router=_gpt_4o_mini_router({}),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"vector_store_ids",
["KBSTOREB", {"KBSTOREB": True}, ["KBSTOREA", {"id": "KBSTOREB"}]],
ids=["string-top-level-ids", "object-top-level-ids", "object-top-level-id"],
)
@pytest.mark.parametrize(
("key_vector_stores", "team_vector_stores"),
[(None, None), ([], [])],
ids=["no-permission-rows", "empty-allowlists"],
)
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_legacy_mode_passes_malformed_vector_store_ids_through_without_an_allowlist(
vector_store_ids: object,
key_vector_stores: list[str] | None,
team_vector_stores: list[str] | None,
vector_store_registry: VectorStoreRegistry | None,
):
assert (
await _team_key_query_with_top_level_ids(
key_vector_stores, team_vector_stores, vector_store_ids, vector_store_registry
)
is True
)
@pytest.mark.asyncio
@_MALFORMED_TOP_LEVEL_IDS
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_legacy_mode_rejects_malformed_vector_store_ids_under_a_team_allowlist(
vector_store_ids: object, vector_store_registry: VectorStoreRegistry | None
):
with pytest.raises(ProxyException) as exc_info:
await _team_key_query_with_top_level_ids(None, ["KBSTOREA"], vector_store_ids, vector_store_registry)
assert (exc_info.value.type, exc_info.value.param, exc_info.value.code) == (
"invalid_request_error",
"vector_store_ids",
"400",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("vector_store_ids", "denied_by"), [("KBSTOREA", None), ("KBSTOREB", _TEAM_DENIED)])
@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"])
async def test_legacy_mode_checks_a_top_level_vector_store_ids_string_as_one_store(
vector_store_ids: str, denied_by: ProxyErrorTypes | None, vector_store_registry: VectorStoreRegistry | None
):
await _assert_rag_query_outcome(
_team_key_query_with_top_level_ids(None, ["KBSTOREA"], vector_store_ids, vector_store_registry), denied_by
)
def _managed_vector_store_id(unified_uuid: str, provider_vector_store_id: str) -> str:
return encode_unified_id(
generate_unified_id_string(
resource_type="vector_store",
unified_uuid=unified_uuid,
target_model_names=["gpt-4o-mini"],
provider_resource_id=provider_vector_store_id,
model_id="deployment-1",
)
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("team_vector_stores", "requested_id", "denied_by"),
[
(["KBSTOREA"], _managed_vector_store_id("uuid-a", "KBSTOREA"), None),
([_managed_vector_store_id("uuid-a", "KBSTOREA")], "KBSTOREA", None),
([_managed_vector_store_id("uuid-a", "KBSTOREA")], _managed_vector_store_id("uuid-a2", "KBSTOREA"), None),
(["KBSTOREA"], _managed_vector_store_id("uuid-b", "KBSTOREB"), _TEAM_DENIED),
(["uuid-a"], _managed_vector_store_id("uuid-a", "KBSTOREB"), _TEAM_DENIED),
([_managed_vector_store_id("uuid-a", "KBSTOREA")], "KBSTOREB", _TEAM_DENIED),
(["KBSTOREA"], "litellm_proxy:vector_store;unified_id,uuid-a;resource_id,KBSTOREA", _TEAM_DENIED),
],
ids=[
"provider-id-allows-its-managed-id",
"managed-id-allows-its-provider-id",
"managed-ids-for-one-provider-store",
"managed-id-for-another-store",
"managed-uuid-cannot-vouch-for-another-store",
"managed-id-for-another-provider-store",
"undecoded-managed-id-is-its-own-store",
],
)
@pytest.mark.parametrize("deny_by_default", [True, False], ids=["strict", "default"])
async def test_allowlist_compares_a_managed_vector_store_id_by_the_provider_store_it_names(
team_vector_stores: list[str], requested_id: str, denied_by: ProxyErrorTypes | None, deny_by_default: bool
):
await _assert_rag_query_outcome(
_common_checks_for_rag_query(
{"vector_store_deny_by_default": deny_by_default},
_virtual_key(object_permission_id="key-permission", team_id="team-1"),
SimpleNamespace(vector_stores=[requested_id]),
LiteLLM_TeamTable(team_id="team-1", object_permission_id="team-permission"),
SimpleNamespace(vector_stores=team_vector_stores),
request_body={
"model": "gpt-4o-mini",
"input": "what is in this KB?",
"tools": [{"type": "file_search", "vector_store_ids": [requested_id]}],
},
),
denied_by,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("request_body", "denied_by"),

294
tests/unit/proxy/common_utils/test_openai_error_payload.py
View file

@ -1,16 +1,26 @@
import json
import sys
from typing import Final
import httpx
import prisma.errors
import pytest
from fastapi import HTTPException
from prisma._async_http import Response
from prisma.engine.utils import handle_response_errors
import litellm
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
proxy_exception_for,
with_litellm_call_id,
)
from litellm.types.router import RouterErrors
@pytest.mark.parametrize(
@ -203,3 +213,287 @@ def test_a_proxy_error_keeps_the_call_id_it_was_raised_with():
assert with_litellm_call_id(raised_with_id, "second").headers == {"x-litellm-call-id": "first"}
assert with_litellm_call_id(ProxyException(message="nope", type="None", param=None, code=400), None).headers == {}
@pytest.mark.parametrize(
"raised_status, expected_type",
[
(400, "invalid_request_error"),
(404, "invalid_request_error"),
(401, "auth_error"),
(403, "auth_error"),
(422, "invalid_request_error"),
(500, "internal_server_error"),
],
)
def test_a_catch_all_handler_keeps_the_status_its_route_answered_with(raised_status: int, expected_type: str) -> None:
wrapped: Final = proxy_exception_for(
HTTPException(status_code=raised_status, detail={"message": "boom"}),
default_status_code=400,
)
assert (wrapped.code, wrapped.type) == (str(raised_status), expected_type)
@pytest.mark.parametrize(
"default_status_code, expected_type", [(400, "invalid_request_error"), (500, "internal_server_error")]
)
def test_an_exception_carrying_no_status_falls_back_to_the_routes_default(default_status_code: int, expected_type: str) -> None:
wrapped: Final = proxy_exception_for(ValueError("boom"), default_status_code=default_status_code)
assert (wrapped.code, wrapped.type) == (str(default_status_code), expected_type)
def test_a_proxy_exception_passes_through_with_the_type_it_already_named() -> None:
raised: Final = ProxyException(
message="Budget has been exceeded",
type=ProxyErrorTypes.budget_exceeded.value,
param="max_budget",
code=400,
)
assert proxy_exception_for(raised, default_status_code=500) is raised
def test_the_message_is_what_the_route_raised_without_an_invented_prefix() -> None:
assert proxy_exception_for(ValueError("boom"), default_status_code=400).message == "boom"
def test_an_http_exceptions_detail_survives_the_wrapping() -> None:
wrapped: Final = proxy_exception_for(
HTTPException(status_code=400, detail={"error": "Invalid budget_duration 'not-a-duration'."}),
default_status_code=400,
)
assert "Invalid budget_duration" in wrapped.message
def test_a_carried_type_survives_the_wrapping() -> None:
class _Carrier(Exception):
type: str = "context_window_exceeded"
status_code: int = 400
assert proxy_exception_for(_Carrier("boom"), default_status_code=500).type == "context_window_exceeded"
@pytest.mark.parametrize("exc", [HTTPException(status_code=404, detail="nope"), ValueError("boom")])
def test_the_wrapped_error_serializes_as_an_openai_error_object(exc: Exception) -> None:
body: Final = json.loads(json.dumps({"error": proxy_exception_for(exc, default_status_code=400).to_dict()}))
assert body["error"]["param"] is None
assert body["error"]["type"] == "invalid_request_error"
@pytest.mark.parametrize(
"provider_failure, expected",
[
(
litellm.AuthenticationError(message="invalid x-api-key", llm_provider="anthropic", model="claude-haiku-5-5"),
("401", "authentication_error"),
),
(
litellm.PermissionDeniedError(
message="model not allowed",
llm_provider="anthropic",
model="claude-haiku-5-5",
response=httpx.Response(403, request=httpx.Request("POST", "https://api.anthropic.com/v1/messages")),
),
("403", "permission_error"),
),
],
)
def test_a_provider_credential_failure_is_labelled_the_way_chat_completions_labels_it(
provider_failure: Exception, expected: tuple[str, str]
) -> None:
wrapped: Final = proxy_exception_for(provider_failure, default_status_code=400)
assert (wrapped.code, wrapped.type) == expected
@pytest.mark.parametrize("status_code", [401, 403])
def test_a_rejection_the_route_raised_keeps_the_auth_label(status_code: int) -> None:
wrapped: Final = proxy_exception_for(HTTPException(status_code=status_code, detail="no"), default_status_code=400)
assert (wrapped.code, wrapped.type) == (str(status_code), "auth_error")
def test_an_exception_without_a_message_is_named_by_its_type() -> None:
wrapped: Final = proxy_exception_for(TimeoutError(), default_status_code=500)
assert wrapped.message == "TimeoutError"
def test_a_database_outage_answers_503_so_callers_retry_instead_of_fixing_their_request() -> None:
wrapped: Final = proxy_exception_for(
prisma.errors.PrismaError("Can't reach database server at `db:5432`"), default_status_code=400
)
assert (wrapped.code, wrapped.type) == ("503", "no_db_connection")
@pytest.mark.parametrize(
"postgres_refusal",
[
'FATAL: database "litellm" is not currently accepting connections',
"FATAL: the database system is starting up",
"FATAL: the database system is shutting down",
"FATAL: the database system is in recovery mode",
"FATAL: the database system is not yet accepting connections",
"FATAL: terminating connection due to administrator command",
"ERROR: cannot execute UPDATE in a read-only transaction",
],
)
def test_a_database_refusing_work_during_a_failover_answers_503(postgres_refusal: str) -> None:
refused: Final = prisma.errors.DataError(
{"user_facing_error": {"message": f"Error in connector: Error querying the database: {postgres_refusal}"}}
)
wrapped: Final = proxy_exception_for(refused, default_status_code=400)
assert (wrapped.code, wrapped.type) == ("503", "no_db_connection")
@pytest.mark.parametrize(
"localized_refusal",
[
{
"error_code": "P2010",
"message": "Raw query failed. Message: `FEHLER: kann UPDATE nicht in einer Nur-Lese-Transaktion ausführen`",
"meta": {"code": "25006"},
},
{
"message": 'PostgresError { code: "57P01", message: "Verbindung wird abgebrochen", severity: "FATAL" }',
},
],
)
def test_a_failover_refusal_in_a_localized_postgres_answers_503(localized_refusal: dict[str, object]) -> None:
refused: Final = prisma.errors.DataError({"user_facing_error": localized_refusal})
wrapped: Final = proxy_exception_for(refused, default_status_code=400)
assert (wrapped.code, wrapped.type) == ("503", "no_db_connection")
def test_a_failover_phrase_outside_the_database_is_not_reported_as_a_database_outage() -> None:
callback_down: Final = ValueError("FATAL: the database system is starting up")
wrapped: Final = proxy_exception_for(callback_down, default_status_code=400)
assert (wrapped.code, wrapped.type) == ("400", "invalid_request_error")
def test_a_database_dropped_mid_query_answers_503_though_prisma_fails_to_name_the_error() -> None:
engine_reply_with_null_meta: Final = [{"user_facing_error": {"error_code": "P1001", "message": "", "meta": None}}]
with pytest.raises(AttributeError) as raised:
handle_response_errors(Response(httpx.Response(500)), engine_reply_with_null_meta)
wrapped: Final = proxy_exception_for(raised.value, default_status_code=500)
assert (wrapped.code, wrapped.type) == ("503", "no_db_connection")
def test_a_network_failure_outside_the_database_is_not_reported_as_a_database_outage() -> None:
callback_unreachable: Final = httpx.ConnectError("[Errno 111] Connection refused")
wrapped: Final = proxy_exception_for(callback_unreachable, default_status_code=500)
assert (wrapped.code, wrapped.type) == ("500", "internal_server_error")
def _add_member_failing_with(db_error: Exception) -> None:
try:
raise db_error
except Exception as e:
raise ValueError(f"Error adding member to organization: {e}")
@pytest.mark.parametrize(
"underlying_error, expected",
[
(prisma.errors.PrismaError("Can't reach database server at `db:5432`"), ("503", "no_db_connection")),
(httpx.ConnectError("[Errno 111] Connection refused"), ("500", "internal_server_error")),
],
)
def test_a_route_that_rewraps_its_error_keeps_the_database_outage_it_wrapped(
underlying_error: Exception, expected: tuple[str, str]
) -> None:
with pytest.raises(ValueError, match="Error adding member") as raised:
_add_member_failing_with(underlying_error)
wrapped: Final = proxy_exception_for(raised.value, default_status_code=500)
assert (wrapped.code, wrapped.type) == expected
def test_a_proxy_without_prisma_installed_still_answers_with_an_error_object(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(sys.modules, "prisma", None)
wrapped: Final = proxy_exception_for(httpx.ConnectError("[Errno 111] Connection refused"), default_status_code=500)
assert (wrapped.code, wrapped.type) == ("500", "internal_server_error")
def _look_up_team_failing_with(lookup_error: Exception) -> None:
try:
raise lookup_error
except Exception:
raise HTTPException(status_code=404, detail="Team doesn't exist in db")
@pytest.mark.parametrize(
"lookup_error, expected",
[
(DBLookupDeadlineExceeded("team", 10.0), ("503", "no_db_connection")),
(prisma.errors.PrismaError("Can't reach database server at `db:5432`"), ("503", "no_db_connection")),
(KeyError("team_id"), ("404", "invalid_request_error")),
],
)
def test_a_not_found_answer_raised_while_the_database_was_down_reports_the_outage(
lookup_error: Exception, expected: tuple[str, str]
) -> None:
with pytest.raises(HTTPException) as raised:
_look_up_team_failing_with(lookup_error)
wrapped: Final = proxy_exception_for(raised.value, default_status_code=400)
assert (wrapped.code, wrapped.type) == expected
def test_a_lookup_that_missed_its_deadline_answers_503() -> None:
wrapped: Final = proxy_exception_for(DBLookupDeadlineExceeded("team", 10.0), default_status_code=400)
assert (wrapped.code, wrapped.type) == ("503", "no_db_connection")
def test_a_database_outage_answer_does_not_reveal_where_the_database_lives() -> None:
wrapped: Final = proxy_exception_for(
prisma.errors.PrismaError("Can't reach database server at `prod-db.internal:5432`"), default_status_code=400
)
assert wrapped.code == "503"
assert "prod-db.internal" not in wrapped.message
def test_an_upstream_failure_answers_with_its_own_status_not_the_routes_default() -> None:
upstream_overloaded: Final = litellm.ServiceUnavailableError(
message="upstream overloaded", llm_provider="openai", model="gpt-5.4-mini"
)
wrapped: Final = proxy_exception_for(upstream_overloaded, default_status_code=400)
assert (wrapped.code, wrapped.type) == ("503", "internal_server_error")
@pytest.mark.parametrize(
"router_rejection, expected",
[
(f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model=gpt-5.4-mini", ("401", "auth_error")),
(f"{RouterErrors.no_deployments_available.value}. Passed model=gpt-5.4-mini", ("429", "rate_limit_error")),
],
)
def test_a_router_rejection_is_labelled_by_the_status_it_answers_with(
router_rejection: str, expected: tuple[str, str]
) -> None:
wrapped: Final = proxy_exception_for(ValueError(router_rejection), default_status_code=400)
assert (wrapped.code, wrapped.type) == expected

265
tests/unit/proxy/credential_endpoints/test_endpoints.py
View file

@ -2,8 +2,9 @@
import json
from contextlib import contextmanager
from functools import reduce
from types import SimpleNamespace
from typing import Final
from typing import Final, Protocol
from unittest.mock import AsyncMock, MagicMock, call, patch
import httpx
@ -12,14 +13,16 @@ from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from fastapi import HTTPException, Request
from fastapi.testclient import TestClient
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import get_masked_values
from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key
from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router
from litellm.proxy.proxy_server import app
from litellm.types.utils import CredentialItem
@ -460,6 +463,264 @@ def test_update_credential_never_syncs_a_null_into_the_in_memory_credential(rest
assert synced.credential_values["api_key"] == "sk-rotated"
_WRITTEN_JSON: Final = TypeAdapter(dict[str, object])
class _InstallCredentialStore(Protocol):
def __call__(
self,
*,
connected: bool = True,
in_memory: tuple[object, ...] = (),
llm_router: object | None = None,
**repository_calls: AsyncMock,
) -> None: ...
def _encrypted_row(served: CredentialItem) -> CredentialItem:
return CredentialItem(
credential_name=served.credential_name,
credential_values={key: encrypt_value_helper(value) for key, value in served.credential_values.items()},
credential_info=dict(served.credential_info),
)
def _written_values(update_by_name: AsyncMock) -> dict[str, object]:
written: Final = _WRITTEN_JSON.validate_json(update_by_name.await_args.kwargs["data"]["credential_values"])
return {
key: decrypt_value_helper(value, key) if isinstance(value, str) else value for key, value in written.items()
}
def _written_info(update_by_name: AsyncMock) -> dict[str, object]:
return _WRITTEN_JSON.validate_json(update_by_name.await_args.kwargs["data"]["credential_info"])
def _serve_credential(
credential_store: _InstallCredentialStore,
served: CredentialItem,
in_memory: tuple[CredentialItem, ...] | None = None,
) -> AsyncMock:
update_by_name = AsyncMock(return_value=None)
credential_store(
in_memory=(served,) if in_memory is None else in_memory,
find_by_name=AsyncMock(side_effect=lambda _name: _encrypted_row(served)),
update_by_name=update_by_name,
)
return update_by_name
def _service_account(
private_key: str = "pk-1234567890", client_email: str = "svc@example.com"
) -> dict[str, dict[str, str]]:
return {"vertex_credentials": {"private_key": private_key, "client_email": client_email}}
def test_update_credential_keeps_the_stored_secret_when_the_patch_echoes_its_masked_read_back(
credential_store: _InstallCredentialStore,
):
served = CredentialItem(
credential_name="existing",
credential_values={"api_key": "sk-real-secret-1234"},
credential_info={"custom_llm_provider": "openai", "description": "original"},
)
update_by_name = _serve_credential(credential_store, served)
response = _patch_credential(
"existing",
{
"credential_values": {"api_key": "sk****34"},
"credential_info": {"custom_llm_provider": "openai", "description": "edited"},
},
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234"}
assert _written_info(update_by_name)["description"] == "edited"
assert CredentialAccessor.get_credential_values("existing") == {"api_key": "sk-real-secret-1234"}
assert litellm.credential_list[0].credential_info["description"] == "edited"
def test_update_credential_keeps_the_secret_when_the_patch_echoes_what_get_by_name_served(
credential_store: _InstallCredentialStore,
):
served = CredentialItem(
credential_name="existing",
credential_values={"api_key": "sk-real-secret-1234", "api_base": "https://api.example.com"},
credential_info={"custom_llm_provider": "openai", "description": "original"},
)
update_by_name = _serve_credential(credential_store, served)
read_back = _call_as("GET", "/credentials/by_name/existing").json()
assert read_back["credential_values"]["api_key"] != "sk-real-secret-1234"
response = _patch_credential(
"existing", {**read_back, "credential_info": {**read_back["credential_info"], "description": "edited"}}
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234", "api_base": "https://api.example.com"}
assert _written_info(update_by_name)["description"] == "edited"
def test_update_credential_keeps_a_short_secret_whose_read_back_is_all_asterisks(
credential_store: _InstallCredentialStore,
):
served = CredentialItem(credential_name="existing", credential_values={"api_key": "abc"}, credential_info={})
update_by_name = _serve_credential(credential_store, served)
read_back = _call_as("GET", "/credentials/by_name/existing").json()
assert read_back["credential_values"]["api_key"] != "abc"
response = _patch_credential(
"existing", {"credential_values": read_back["credential_values"], "credential_info": {}}
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {"api_key": "abc"}
assert CredentialAccessor.get_credential_values("existing") == {"api_key": "abc"}
@pytest.mark.parametrize(
"rotated_values",
(
{"api_key": "sk-rotated-key-5678"},
{"api_key": "new**secret"},
_service_account(private_key="pk**rotated"),
),
)
def test_update_credential_still_rotates_a_secret_sent_in_full(
credential_store: _InstallCredentialStore, rotated_values: dict[str, object]
):
served = CredentialItem(
credential_name="existing",
credential_values={"api_key": "sk-old-key-000011", **_service_account()},
credential_info={},
)
update_by_name = _serve_credential(credential_store, served)
response = _patch_credential("existing", {"credential_values": rotated_values, "credential_info": {}})
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {**served.credential_values, **rotated_values}
assert CredentialAccessor.get_credential_values("existing") == {**served.credential_values, **rotated_values}
def test_update_credential_saves_a_mask_shaped_value_on_a_key_get_never_masks(
credential_store: _InstallCredentialStore,
):
served = CredentialItem(
credential_name="existing",
credential_values={
"api_key": "sk-real-secret-1234",
"api_base": "https://api.example.com",
"extra_headers": {"x-team": "platform"},
},
credential_info={},
)
update_by_name = _serve_credential(credential_store, served)
sent: Final = {"api_base": "ab****cd", "organization": "*****", "extra_headers": {"x-team": "*****"}}
assert get_masked_values(sent, unmasked_length=4, number_of_asterisks=4) == sent
response = _patch_credential("existing", {"credential_values": sent, "credential_info": {}})
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234", **sent}
@pytest.mark.parametrize("client_email", ("svc@example.com", "new@example.com"))
def test_update_credential_keeps_a_masked_secret_inside_a_stored_object(
credential_store: _InstallCredentialStore, client_email: str
):
served = CredentialItem(credential_name="existing", credential_values=_service_account(), credential_info={})
update_by_name = _serve_credential(credential_store, served)
response = _patch_credential(
"existing",
{
"credential_values": _service_account(private_key="pk****90", client_email=client_email),
"credential_info": {},
},
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == _service_account(client_email=client_email)
assert CredentialAccessor.get_credential_values("existing") == _service_account(client_email=client_email)
def test_update_credential_drops_a_masked_leaf_the_stored_object_no_longer_has(
credential_store: _InstallCredentialStore,
):
served = CredentialItem(credential_name="existing", credential_values=_service_account(), credential_info={})
update_by_name = _serve_credential(credential_store, served)
response = _patch_credential(
"existing",
{
"credential_values": {
"vertex_credentials": {
"private_key": "pk****90",
"refresh_token": "rt****89",
"client_email": "svc@example.com",
}
},
"credential_info": {},
},
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == _service_account()
@pytest.mark.parametrize(
"in_memory",
(
pytest.param(
(
CredentialItem(
credential_name="existing",
credential_values={"api_key": "sk-old-key-000011"},
credential_info={"description": "original"},
),
),
id="replica_not_resynced",
),
pytest.param((), id="replica_never_loaded_it"),
),
)
def test_update_credential_keeps_the_table_secret_whatever_this_replica_last_rendered(
credential_store: _InstallCredentialStore, in_memory: tuple[CredentialItem, ...]
):
served = CredentialItem(
credential_name="existing",
credential_values={"api_key": "sk-rotated-key-000099"},
credential_info={"description": "original"},
)
update_by_name = _serve_credential(credential_store, served, in_memory=in_memory)
response = _patch_credential(
"existing",
{"credential_values": {"api_key": "sk****11"}, "credential_info": {"description": "edited"}},
)
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == {"api_key": "sk-rotated-key-000099"}
assert _written_info(update_by_name)["description"] == "edited"
@pytest.mark.parametrize("levels", (20, 21))
def test_update_credential_keeps_a_secret_exactly_as_deep_as_the_masker_masks(
credential_store: _InstallCredentialStore, levels: int
):
nested = reduce(lambda inner, _level: {"credentials": inner}, range(levels), {"api_key": "leaf-secret-1234"})
served = CredentialItem(credential_name="existing", credential_values=nested, credential_info={})
update_by_name = _serve_credential(credential_store, served)
read_back = get_masked_values(served.credential_values, unmasked_length=4, number_of_asterisks=4)
response = _patch_credential("existing", {"credential_values": read_back, "credential_info": {}})
assert response.status_code == 200, response.text
assert _written_values(update_by_name) == nested
def _generate_es256_pem() -> str:
key = ec.generate_private_key(ec.SECP256R1())
return key.private_bytes(

90
tests/unit/proxy/google_endpoints/test_endpoints.py
View file

@ -6,7 +6,7 @@ from typing import Final
import pytest
from dotenv import load_dotenv
from pydantic import JsonValue
from pydantic import JsonValue, ValidationError
from starlette.requests import Request
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
@ -105,3 +105,91 @@ async def test_proxy_gemini_to_openai_like_model_token_counting():
)
assert response.get("totalTokens") > 0
_DECLARATION: Final = {"name": "lookup", "description": "Find a record"}
_SYSTEM: Final = {"systemInstruction": {"parts": [{"text": "Follow these rules"}]}}
_NO_SYSTEM_PARTS: Final = {"systemInstruction": {"parts": []}}
_TOOLS: Final = {"tools": [{"functionDeclarations": [_DECLARATION]}]}
_TURN: Final = {"role": "user", "parts": [{"text": "Count this request"}]}
_CONTENTS: Final = {"contents": [_TURN]}
@pytest.mark.parametrize(
("spelled", "canonical"),
(
({"systemInstruction": {"parts": {"text": "Follow these rules"}}}, _SYSTEM),
({"system_instruction": {"parts": {"text": "Follow these rules"}}}, _SYSTEM),
({"systemInstruction": {"parts": [{"text": "Follow these rules"}, None]}}, _SYSTEM),
({"systemInstruction": {"parts": None}}, _NO_SYSTEM_PARTS),
({"systemInstruction": {"parts": [{"text": None}, {"text": "Follow these rules"}]}}, _SYSTEM),
({"systemInstruction": {"parts": {"text": None}}}, _NO_SYSTEM_PARTS),
({"systemInstruction": {"parts": {"text": None, "inlineData": None}}}, _NO_SYSTEM_PARTS),
({"systemInstruction": {"parts": {"text": None, "thought": False}}}, _NO_SYSTEM_PARTS),
({"systemInstruction": {"parts": [{"thoughtSignature": "c2ln"}, {"text": "Follow these rules"}]}}, _SYSTEM),
({"systemInstruction": {"parts": [{}]}}, _NO_SYSTEM_PARTS),
({"systemInstruction": {}}, _NO_SYSTEM_PARTS),
({"tools": None}, {}),
({"tools": [None]}, {}),
({"tools": [{"functionDeclarations": None}]}, {}),
({"tools": {"functionDeclarations": [_DECLARATION]}}, _TOOLS),
({"tools": [{"functionDeclarations": _DECLARATION}]}, _TOOLS),
({"tools": {"functionDeclarations": _DECLARATION}}, _TOOLS),
({"contents": _TURN}, _CONTENTS),
({"contents": [None, _TURN]}, _CONTENTS),
({"contents": None}, {"contents": []}),
({"contents": [{"role": "user", "parts": {"text": "Count this request"}}]}, _CONTENTS),
({"contents": [{"role": "user", "parts": [None, {"text": "Count this request"}]}]}, _CONTENTS),
({"contents": [{"role": "user", "parts": None}]}, {"contents": [{"role": "user", "parts": []}]}),
),
)
def test_google_count_tokens_counts_every_google_spelling_of_a_repeated_field_like_its_list_form(
spelled: dict[str, JsonValue], canonical: dict[str, JsonValue]
) -> None:
contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}]
assert _google_count_tokens_request("counted", {"contents": contents, **spelled}) == _google_count_tokens_request(
"counted", {"contents": contents, **canonical}
)
@pytest.mark.parametrize(
"malformed",
(
{"tools": "lookup"},
{"tools": [{"functionDeclarations": "lookup"}]},
{"systemInstruction": {"parts": 7}},
{"contents": 5},
{"contents": ["Count this request"]},
),
)
def test_google_count_tokens_rejects_a_scalar_where_google_expects_an_object_or_list(
malformed: dict[str, JsonValue],
) -> None:
with pytest.raises(ValidationError):
_google_count_tokens_request("counted", {"contents": [_TURN], **malformed})
@pytest.mark.parametrize(
"parts",
(
{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}},
[{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}],
{"text": None, "inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}},
[{"text": "Follow these rules"}, {"text": None, "fileData": {"fileUri": "gs://bucket/rules.pdf"}}],
),
)
def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None:
with pytest.raises(ValidationError):
_google_count_tokens_request(
"counted",
{"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}},
)
def test_google_count_tokens_keeps_counting_system_text_that_sits_beside_media() -> None:
text_with_media: Final = {"text": "Follow these rules", "inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}
assert _google_count_tokens_request(
"counted", {**_CONTENTS, "systemInstruction": {"parts": [text_with_media]}}
) == _google_count_tokens_request("counted", {**_CONTENTS, **_SYSTEM})

9
tests/unit/proxy/health_endpoints/test_health_endpoints.py
View file

@ -5480,6 +5480,15 @@ async def test_health_services_endpoint_langfuse_missing_keys_errors(monkeypatch
assert "connection refused by lf.internal.example" in str(raised.value.message)
@pytest.mark.asyncio
async def test_health_services_reports_an_unknown_service_as_a_request_error_not_an_auth_error() -> None:
with pytest.raises(ProxyException) as exc_info:
await health_services_endpoint(service="totally_unknown_service_xyz")
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize("admin", [False, True])
async def test_unsaved_chatgpt_account_probe_requires_proxy_admin(

15
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -5211,3 +5211,18 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke
await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller)
revoke_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None:
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
with pytest.raises(ProxyException) as exc_info:
await user_update(
data=UpdateUserRequest(user_id="target-user", budget_duration="not-a-duration"),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Invalid budget_duration" in exc_info.value.message
assert "Authentication Error" not in exc_info.value.message

72
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -6538,7 +6538,7 @@ async def test_list_keys_with_invalid_status():
)
# Verify ProxyException properties
assert exc_info.value.code == "400"
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Invalid status value" in str(exc_info.value.message)
assert "deleted" in str(exc_info.value.message)
@ -23377,3 +23377,73 @@ def test_key_info_database_not_connected_sanitized_500(
assert "Database not connected" in caplog.text
finally:
proxy_server.app.dependency_overrides.pop(user_api_key_auth, None)
@pytest.mark.asyncio
async def test_update_key_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
from fastapi import Request
from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn
monkeypatch.setattr( # test-quality-ok: the endpoint reads proxy_server.prisma_client itself, no parameter to inject
"litellm.proxy.proxy_server.prisma_client", MagicMock()
)
with pytest.raises(ProxyException) as exc_info:
await update_key_fn(
request=MagicMock(spec=Request),
data=UpdateKeyRequest(key="sk-whatever", budget_duration="not-a-duration"),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
async def test_update_key_passes_its_own_rejection_through_without_making_it_its_own_cause(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from fastapi import Request
from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn
prisma_client: Final = MagicMock()
prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr( # test-quality-ok: the endpoint reads proxy_server.prisma_client itself, no parameter to inject
"litellm.proxy.proxy_server.prisma_client", prisma_client
)
with pytest.raises(ProxyException) as exc_info:
await update_key_fn(
request=MagicMock(spec=Request),
data=UpdateKeyRequest(key="sk-missing", max_budget=5),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("404", "not_found_error")
assert exc_info.value.__cause__ is None
@pytest.mark.asyncio
async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypatch: pytest.MonkeyPatch) -> None:
import prisma.errors
from litellm.proxy.management_endpoints.key_management_endpoints import key_aliases
unreachable_db: Final = MagicMock()
unreachable_db.db.query_raw = AsyncMock(
side_effect=prisma.errors.PrismaError("Can't reach database server at `db:5432`")
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", unreachable_db)
with pytest.raises(ProxyException) as exc_info:
await key_aliases(
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
page=1,
size=50,
search=None,
team_id=None,
)
assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection")

75
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -311,7 +311,7 @@ async def test_model_creation_fails_closed_when_authoritative_policy_is_unavaila
Deployment(model_name="cannot-create", litellm_params={"model": "openai/test-model"}),
UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert error.value.code == "503"
assert (error.value.code, error.value.type) == ("503", "internal_server_error")
prisma.db.litellm_proxymodeltable.create.assert_not_called()
prisma.db.litellm_teamtable.find_unique.assert_not_called()
@ -10913,3 +10913,76 @@ class TestNonAdminCannotSetPerUserOauthOnModel:
user_api_key_dict=admin,
)
assert result is created_row
@pytest.mark.asyncio
async def test_delete_model_reports_an_unknown_model_id_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy._types import ModelInfoDelete
from litellm.proxy.management_endpoints.model_management_endpoints import delete_model
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
with pytest.raises(ProxyException) as exc_info:
await delete_model(
model_info=ModelInfoDelete(id="no-such-model"),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
async def test_update_model_reports_an_unknown_model_id_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", None)
with pytest.raises(ProxyException) as exc_info:
await update_model(
model_params=updateDeployment(model_info=ModelInfo(id="no-such-model")),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert exc_info.value.message == "model not found"
@pytest.mark.asyncio
async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import prisma.errors
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
unreachable_db: Final = MagicMock()
unreachable_db.db.query_raw = AsyncMock(return_value=[])
unreachable_db.db.litellm_proxymodeltable.create = AsyncMock(
side_effect=prisma.errors.PrismaError("Can't reach database server at `db:5432`")
)
monkeypatch.setattr(proxy_server, "prisma_client", unreachable_db)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "premium_user", True)
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-outage-salt")
with pytest.raises(ProxyException) as exc_info:
await add_new_model(
model_params=Deployment(
model_name="outage-model",
litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="fake-key"),
model_info={"id": "outage-model-1"},
),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection")

31
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1617,6 +1617,37 @@ async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu
prisma_client.db.litellm_organizationtable.delete.assert_not_called()
@pytest.mark.asyncio
async def test_organization_member_add_reports_an_unknown_org_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
from fastapi import Request
from litellm.proxy._types import (
LitellmUserRoles,
OrganizationMemberAddRequest,
OrgMember,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
with pytest.raises(ProxyException) as exc_info:
await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id="no-such-org",
member=OrgMember(role="internal_user", user_id="someone"),
),
http_request=MagicMock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("404", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize("budget_duration", ["30d", None])
async def test_new_organization_schedules_budget_reset_from_duration(monkeypatch, budget_duration):

45
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -19372,3 +19372,48 @@ async def test_default_budget_cache_eviction_does_not_fail_after_redis_error() -
)
user_api_key_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:budget_456")
@pytest.mark.asyncio
async def test_team_info_reports_a_missing_team_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
from fastapi import Request
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints import team_endpoints
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
with pytest.raises(ProxyException) as exc_info:
await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="no-such-team",
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("404", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
async def test_team_info_keeps_the_auth_label_for_a_caller_outside_the_team(monkeypatch: pytest.MonkeyPatch) -> None:
from fastapi import Request
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints import team_endpoints
other_team: Final = MagicMock()
other_team.model_dump.return_value = {"team_id": "team-a", "members_with_roles": []}
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=other_team)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
with pytest.raises(ProxyException) as exc_info:
await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-a",
user_api_key_dict=UserAPIKeyAuth(user_id="outsider", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert (exc_info.value.code, exc_info.value.type) == ("403", "auth_error")

99
tests/unit/proxy/test_proxy_server.py
View file

@ -3,6 +3,7 @@ import sys
import traceback
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Final, Literal, NoReturn, cast
from unittest import mock
@ -3817,6 +3818,48 @@ def test_live_session_preserves_pretty_printed_start_frame(
assert routed_data["live_session_start"] == expected_start
@pytest.mark.asyncio
async def test_update_config_reports_an_unsupported_router_setting_as_a_request_error_not_an_auth_error() -> None:
from fastapi import Request
from litellm.proxy._types import ConfigYAML, ProxyException
from litellm.proxy.proxy_server import update_config
body: Final = {"router_settings": {"not_a_router_setting": True}}
request: Final = MagicMock(spec=Request)
request.json = AsyncMock(return_value=body)
with pytest.raises(ProxyException) as exc_info:
await update_config(
config_info=ConfigYAML(**body),
request=request,
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Authentication Error" not in exc_info.value.message
@pytest.mark.asyncio
async def test_get_config_callbacks_reports_an_unreadable_config_as_a_server_error_not_an_auth_error(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException
from litellm.proxy.proxy_server import get_config
empty_config: Final = tmp_path / "config.yaml"
empty_config.write_text("")
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(proxy_server, "user_config_file_path", str(empty_config))
with pytest.raises(ProxyException) as exc_info:
await get_config(user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN))
assert (exc_info.value.code, exc_info.value.type) == ("500", "internal_server_error")
assert exc_info.value.message == "Config cannot be None or Empty."
@pytest.mark.parametrize("scope", [None, "expand"])
@pytest.mark.parametrize(
"general_settings, wildcard_query, wildcard_listed",
@ -3861,6 +3904,62 @@ async def test_model_list_return_wildcard_routes_defaults_to_general_setting(
assert ("openai/*" in listed) is wildcard_listed, listed
@pytest.mark.asyncio
async def test_queue_request_labels_an_upstream_failure_a_server_error_not_an_auth_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from starlette.requests import Request
from starlette.responses import Response
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException
from litellm.proxy.proxy_server import async_queue_request
body: Final = json.dumps(
{"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}], "priority": 0}
).encode()
async def receive() -> dict[str, object]:
return {"type": "http.request", "body": body, "more_body": False}
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": "/queue/chat/completions",
"headers": [(b"content-type", b"application/json")],
"query_string": b"",
"scheme": "http",
"server": ("testserver", 80),
"root_path": "",
},
receive,
)
router: Final = litellm.Router(
model_list=[
{
"model_name": "gpt-5.4-mini",
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_key": "sk-test",
"mock_response": "litellm.InternalServerError",
},
}
],
num_retries=0,
)
monkeypatch.setattr(proxy_server, "llm_router", router)
with pytest.raises(ProxyException) as exc_info:
await async_queue_request(
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert (exc_info.value.code, exc_info.value.type) == ("500", "internal_server_error")
_DEBUG_LOG_SENTINEL: Final = "sentinel-request-body-7f3a"
_LONG_PROMPT: Final = f"{_DEBUG_LOG_SENTINEL} " + "Keep edits minimal and explain them briefly. " * 20
_TOOL_PARAMETERS: Final = {"type": "object", "properties": {"cmd": {"type": "string"}}}

34
tests/unit/vector_stores/test_vector_store_registry.py
View file

@ -519,6 +519,40 @@ def test_vector_store_tools_pop_exactly_the_store_ids_the_registry_reports_to_ru
assert ran == checked
@pytest.mark.asyncio
async def test_vector_stores_to_run_ignore_a_vector_store_ids_object() -> None:
now: Final = datetime.now(timezone.utc)
registry: Final = VectorStoreRegistry(
[LiteLLM_ManagedVectorStore(vector_store_id="vs_registered", created_at=now, updated_at=now)]
)
ran: Final = (
registry.pop_vector_stores_to_run(non_default_params={"vector_store_ids": {"vs_registered": True}}),
await registry.pop_vector_stores_to_run_with_db_fallback(
non_default_params={"vector_store_ids": {"vs_registered": True}}
),
)
assert ran == ([], [])
@pytest.mark.asyncio
async def test_vector_stores_to_run_read_a_tuple_of_vector_store_ids() -> None:
now: Final = datetime.now(timezone.utc)
registry: Final = VectorStoreRegistry(
[LiteLLM_ManagedVectorStore(vector_store_id="vs_registered", created_at=now, updated_at=now)]
)
ran: Final = (
registry.pop_vector_stores_to_run(non_default_params={"vector_store_ids": ("vs_registered",)}),
await registry.pop_vector_stores_to_run_with_db_fallback(
non_default_params={"vector_store_ids": ("vs_registered",)}
),
)
assert [[store["vector_store_id"] for store in stores] for stores in ran] == [["vs_registered"], ["vs_registered"]]
class TestRegistriesDoNotShareDefaultState:
def test_adding_to_one_registry_leaves_the_next_one_empty(self):
VectorStoreRegistry().add_vector_store_to_registry(