mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge remote-tracking branch 'origin/main' into litellm_mcp_auto_execute_parallel_slot
This commit is contained in:
mateo-berri
commit
a297a7d7ad
13 changed files with 940 additions and 216 deletions
7
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
7
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -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>)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -5052,6 +5052,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,
|
||||
|
|
@ -5070,16 +5087,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):
|
||||
|
|
@ -5106,12 +5113,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()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -724,7 +724,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}
|
||||
|
|
@ -740,7 +740,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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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?"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue