From 4846cfabc7f4b8cf6a83a960b7d1123eebef4315 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 12:51:47 -0700 Subject: [PATCH] fix(vector-stores): keep proxy context out of provider vector store metadata (#47021) * fix(vector-stores): keep proxy context out of provider vector store metadata The proxy merged its request context (user_api_key_* fields, endpoint, user_agent, requester_ip_address, litellm_api_version) and the router's deployment info into the vector store request's metadata. For vector stores that field is also OpenAI's own metadata, which is capped at 16 keys, so a store created with a team key was full on creation and every later metadata update returned 400. Vector store routes now keep that context in litellm_metadata, like batches, files and responses, and the router does the same for avector_store_create. Provider-side metadata holds only what the caller sent. Adds db_scripts/strip_vector_store_internal_metadata.py, an opt-in script that removes the proxy-written keys from stores created before this change. It is a dry run unless --apply is passed. * fix(vector-stores): page the metadata cleanup script without recursion and tolerate non-string values Listing every store recursed once per page, which overflows the interpreter stack on large projects. Pages are now walked iteratively. Metadata values are no longer required to be strings, and a malformed response exits 2 with a fixed message instead of a traceback that could echo metadata values. * fix(vector-stores): read proxy agent id, tags and spend_logs_metadata from litellm_metadata Vector store routes now keep proxy context in litellm_metadata, but the v3 limiter only read agent_id from metadata, and logging let provider metadata tags or spend_logs_metadata replace the proxy's values. A string tags value also crashed get_request_tags. The limiter now prefers the proxy-resolved agent_id in litellm_metadata. Logging combines list tags from both buckets (metadata first, deduplicated) and merges dict spend_logs_metadata with litellm_metadata winning, ignoring values of the wrong type. * fix(vector-stores): narrow the logging fallback to mistyped caller values The tag union and spend_logs_metadata merge changed logged tags and spend_logs_metadata on every route that has both metadata buckets, deduped tags quadratically, and stopped background interaction polling from keeping string tags as sent. Caller values now win exactly as before, and the proxy's litellm_metadata tags or spend_logs_metadata only replace a caller value of the wrong type. The limiter reads agent ids through a pydantic model instead of a cast and ignores non-string ids. * fix(vector-stores): log the proxy's full tag list on litellm_metadata routes On routes that keep proxy context in litellm_metadata, the proxy builds litellm_metadata.tags from the caller's tags plus key, team and header tags, but get_request_tags preferred the caller's metadata.tags. A caller sending metadata.tags on a vector store call dropped the key's tags from request_tags, which tag spend tracking reads. get_request_tags now prefers a non-empty litellm_metadata tag list, and falls back to metadata.tags as before. * fix(vector-stores): strip LiteLLM-written keys at the provider boundary instead of moving metadata buckets Moving vector store proxy context into litellm_metadata broke shared readers (agent limits, request tags, merged logging metadata), and fixing them widened the change to every litellm_metadata route. This reverts the bucket move, the router change and the reader fixes, so proxy context stays in metadata exactly as on main. Vector store create and update now drop the keys LiteLLM itself writes (user_api_key*, requester and deployment keys, agent ids) from the metadata sent to the provider, in the shared optional-param helper both paths use. The cleanup script reuses the same key list. * fix(vector-stores): pass non-string metadata keys through the provider metadata filter * fix(vector-stores): strip LiteLLM-written metadata keys only on proxy and router calls * fix(vector-stores): send the provider exactly the metadata the caller sent through the proxy * fix(vector-stores): snapshot caller metadata from the raw request body and keep it out of logged request bodies * fix(vector-stores): copy the parsed-body metadata when the raw request body cannot be reparsed --- .github/workflows/test-unit.yml | 1 + .../strip_vector_store_internal_metadata.py | 201 ++++++++++++++++++ litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 3 +- .../proxy/vector_store_endpoints/endpoints.py | 9 + litellm/vector_stores/utils.py | 91 +++++++- tests/unit/db_scripts/__init__.py | 0 ...st_strip_vector_store_internal_metadata.py | 198 +++++++++++++++++ .../test_vector_store_endpoints.py | 157 ++++++++++++++ .../unit/test_router/test_router_endpoints.py | 37 ++++ tests/unit/vector_stores/test_main.py | 90 +++++++- 11 files changed, 782 insertions(+), 6 deletions(-) create mode 100644 db_scripts/strip_vector_store_internal_metadata.py create mode 100644 tests/unit/db_scripts/__init__.py create mode 100644 tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1b486d78f04..cb1b595f3f0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -658,6 +658,7 @@ jobs: - shard: proxy-extras test-path: >- tests/unit/litellm_proxy_extras + tests/unit/db_scripts python-version: "3.12" workers: 2 reruns: 0 diff --git a/db_scripts/strip_vector_store_internal_metadata.py b/db_scripts/strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..34ffa85dcd1 --- /dev/null +++ b/db_scripts/strip_vector_store_internal_metadata.py @@ -0,0 +1,201 @@ +import argparse +import os +import sys +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from itertools import accumulate, repeat, takewhile +from typing import Final, Literal, TextIO + +import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError + +from litellm.vector_stores.utils import is_litellm_written_metadata_key + +DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +LIST_PAGE_SIZE: Final = "100" + + +class VectorStore(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + metadata: dict[str, JsonValue] | None = None + + +class VectorStorePage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[VectorStore, ...] + has_more: bool = False + + +class OpenAIErrorBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str = "" + + +class OpenAIErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + error: OpenAIErrorBody + + +@dataclass(frozen=True, slots=True) +class Options: + api_key: str + api_base: str + vector_store_ids: tuple[str, ...] + extra_keys: frozenset[str] + apply: bool + + +@dataclass(frozen=True, slots=True) +class StoreResult: + store_id: str + matched_keys: tuple[str, ...] + status: Literal["dry_run", "updated", "failed"] + error: str = "" + + +def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool: + return is_litellm_written_metadata_key(key) or key in extra_keys + + +def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]: + return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys))) + + +def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage: + params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after} + return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content) + + +def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None: + if page is None or not page.has_more or not page.data: + return None + return fetch_page(client, after=page.data[-1].id) + + +def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]: + pages: Final = accumulate( + repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None) + ) + for page in takewhile(lambda page: page is not None, pages): + yield from page.data + + +def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]: + for store_id in store_ids: + yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content) + + +def describe_http_error(response: httpx.Response) -> str: + try: + message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message + except ValidationError: + return f"HTTP {response.status_code}" + return f"HTTP {response.status_code}: {message}" + + +def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult: + try: + response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)}) + except httpx.TransportError as error: + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__) + if response.is_success: + return StoreResult(store_id=store_id, matched_keys=keys, status="updated") + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response)) + + +def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None: + keys: Final = matching_keys(store, options.extra_keys) + if not keys: + return None + if not options.apply: + return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run") + return delete_keys(client, store.id, keys) + + +def format_result(result: StoreResult) -> str: + key_list: Final = ", ".join(result.matched_keys) + count: Final = len(result.matched_keys) + match result.status: + case "dry_run": + return f"{result.store_id}: would delete {count} keys: {key_list}" + case "updated": + return f"{result.store_id}: deleted {count} keys: {key_list}" + case "failed": + return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}" + + +def clean_and_report( + client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO +) -> Iterator[StoreResult | None]: + for store in stores: + result: Final = clean_store(client, store, options) + if result is not None: + print(format_result(result), file=out) + yield result + + +def run(options: Options, client: httpx.Client, out: TextIO) -> int: + stores: Final = ( + retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client) + ) + scanned: Final = tuple(clean_and_report(client, stores, options, out)) + results: Final = tuple(result for result in scanned if result is not None) + updated: Final = sum(1 for result in results if result.status == "updated") + failed: Final = sum(1 for result in results if result.status == "failed") + mode: Final = "applied" if options.apply else "dry run, pass --apply to delete" + print( + f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})", + file=out, + ) + return 1 if failed else 0 + + +def parse_options(argv: Sequence[str] | None) -> Options: + parser: Final = argparse.ArgumentParser( + description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default." + ) + parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY") + parser.add_argument("--api-base", default=DEFAULT_API_BASE) + parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store") + parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete") + parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them") + args: Final = parser.parse_args(argv) + if not args.api_key: + parser.error("--api-key or OPENAI_API_KEY is required") + return Options( + api_key=args.api_key, + api_base=args.api_base.rstrip("/"), + vector_store_ids=tuple(args.vector_store_id), + extra_keys=frozenset(args.extra_key), + apply=args.apply, + ) + + +def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int: + options: Final = parse_options(argv) + with httpx.Client( + base_url=options.api_base, + headers={"Authorization": f"Bearer {options.api_key}"}, + timeout=60.0, + transport=transport, + ) as client: + try: + return run(options, client, sys.stdout) + except httpx.HTTPStatusError as error: + sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n") + return 2 + except httpx.TransportError as error: + sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n") + return 2 + except ValidationError: + sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/litellm/constants.py b/litellm/constants.py index 77f2961d6a1..94ff02a8c26 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1631,6 +1631,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3ad5caade1b..29fe636f5e4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -41,6 +41,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, + VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) from litellm.integrations.custom_guardrail import without_server_streaming_classification @@ -2160,7 +2161,7 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) + frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 22721dc26bc..1b54dddcf4f 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -9,12 +9,14 @@ from typing import ( from fastapi import APIRouter, Depends, HTTPException, Request, Response +from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body from litellm.proxy.utils import jsonify_object from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, @@ -23,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse +from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() @@ -213,6 +216,9 @@ async def vector_store_create( ) data: Final = await read_request_body(request=request) + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) # Check for target_model_names parameter target_model_names: Final = data.pop("target_model_names", None) @@ -448,6 +454,9 @@ async def vector_store_update( data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line if "vector_store_id" not in data: data["vector_store_id"] = vector_store_id + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index a2b5a6732a8..eb61491853f 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -1,11 +1,92 @@ +import copy +from collections.abc import Container, Mapping from typing import Final, cast, get_type_hints +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.types.vector_stores import ( VectorStoreCreateOptionalRequestParams, VectorStoreSearchOptionalRequestParams, ) +LITELLM_WRITTEN_METADATA_KEYS: Final = frozenset( + { + "requester_ip_address", + "user_agent", + "endpoint", + "litellm_api_version", + "deployment", + "deployment_model_name", + "api_base", + "agent_id", + "actor_agent_id", + "target_agent_id", + "billing_agent_id", + "agent_execution_mode", + "verified_human_user_id", + "litellm_caller_session_id", + } +) + + +def is_litellm_written_metadata_key(key: str) -> bool: + return key == "user_api_key" or key.startswith("user_api_key_") or key in LITELLM_WRITTEN_METADATA_KEYS + + +def _is_callers_key(key: object) -> bool: + return not isinstance(key, str) or not is_litellm_written_metadata_key(key) + + +def _sent_by_proxy_or_router(params: Mapping[str, object], metadata: Container[object]) -> bool: + return params.get("proxy_server_request") is not None or ROUTING_REQUEST_TAGS_METADATA_KEY in metadata + + +def _callers_metadata(metadata: Mapping[object, object]) -> Mapping[object, object]: + return {key: value for key, value in metadata.items() if _is_callers_key(key)} + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + +class VectorStoreWriteBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: object = None + + +def _json_object(raw: str | bytes) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_json(raw) + except ValidationError: + return None + + +def _decoded_metadata(metadata: object) -> object: + if not isinstance(metadata, str): + return metadata + decoded: Final = _json_object(metadata) + return metadata if decoded is None else decoded + + +def callers_metadata_snapshot(raw_json_body: bytes | None, parsed_body: VectorStoreWriteBody) -> object: + raw_body: Final = None if raw_json_body is None else _json_object(raw_json_body) + metadata: Final = copy.deepcopy(parsed_body.metadata) if raw_body is None else raw_body.get("metadata") + return _decoded_metadata(metadata) + + +def _metadata_for_provider(params: Mapping[str, object]) -> object: + if VECTOR_STORE_CALLER_METADATA_KWARG in params: + return params[VECTOR_STORE_CALLER_METADATA_KWARG] + metadata: Final = params.get("metadata") + if not isinstance(metadata, Mapping): + return metadata + request_metadata: Final = cast(Mapping[object, object], metadata) # cast-ok: request metadata is a JSON object + if not _sent_by_proxy_or_router(params, request_metadata): + return request_metadata + return _callers_metadata(request_metadata) + class VectorStoreRequestUtils: """Helper utils for constructing Vector Store search requests""" @@ -48,7 +129,11 @@ class VectorStoreRequestUtils: Returns: VectorStoreCreateOptionalRequestParams instance with only the valid parameters """ - valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() + valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - {"metadata"} filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} - - return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + metadata: Final = _metadata_for_provider(params) + if metadata is None: + return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + return cast( # cast-ok: same params as above plus the metadata the provider should store + VectorStoreCreateOptionalRequestParams, {**filtered_params, "metadata": metadata} + ) diff --git a/tests/unit/db_scripts/__init__.py b/tests/unit/db_scripts/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..05b055f75a1 --- /dev/null +++ b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py @@ -0,0 +1,198 @@ +import importlib.util +import json +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest + +_SCRIPT: Final = Path(__file__).resolve().parents[3] / "db_scripts" / "strip_vector_store_internal_metadata.py" +_spec: Final = importlib.util.spec_from_file_location("strip_vector_store_internal_metadata", _SCRIPT) +strip_script: Final = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = strip_script +_spec.loader.exec_module(strip_script) + +POLLUTED: Final = { + "id": "vs_1", + "metadata": { + "user_api_key": "hash-value-1", + "user_api_key_hash": "hash-value-1", + "user_api_key_user_email": "member@example.com", + "endpoint": "http://proxy/v1/vector_stores", + "deployment": "openai/gpt-5", + "team": "llmproxy", + "trace_id": "trace-123", + }, +} +CLEAN: Final = {"id": "vs_2", "metadata": {"team": "llmproxy"}} +NO_METADATA: Final = {"id": "vs_3", "metadata": None} + + +def _fake_openai( + stores: tuple[Mapping[str, object], ...], page_size: int = 100, post_status: int = 200 +) -> tuple[httpx.MockTransport, list[httpx.Request]]: + seen: Final[list[httpx.Request]] = [] + by_id: Final = {store["id"]: store for store in stores} + + def handle(request: httpx.Request) -> httpx.Response: + seen.append(request) + path: Final = request.url.path.removeprefix("/v1") + if request.method == "POST": + if post_status != 200: + return httpx.Response(post_status, json={"error": {"message": "upstream rejected the update"}}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + if path == "/vector_stores": + after: Final = request.url.params.get("after") + start: Final = 0 if after is None else [store["id"] for store in stores].index(after) + 1 + page: Final = stores[start : start + page_size] + return httpx.Response(200, json={"data": list(page), "has_more": start + page_size < len(stores)}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + + return httpx.MockTransport(handle), seen + + +def _run(args: list[str], transport: httpx.MockTransport) -> int: + return strip_script.main(["--api-key", "sk-test", *args], transport=transport) + + +def _posts(seen: list[httpx.Request]) -> list[tuple[str, dict[str, object]]]: + return [(request.url.path, json.loads(request.content)) for request in seen if request.method == "POST"] + + +def test_dry_run_lists_matching_key_names_without_values_and_sends_no_update(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert ( + "vs_1: would delete 5 keys: deployment, endpoint, user_api_key, user_api_key_hash, user_api_key_user_email" + in out + ) + assert "scanned=2 matched=1 updated=0 failed=0" in out + assert "hash-value-1" not in out + assert "member@example.com" not in out + assert _posts(seen) == [] + assert {request.headers["authorization"] for request in seen} == {"Bearer sk-test"} + + +def test_apply_nulls_exactly_the_proxy_written_keys_and_leaves_caller_keys_alone(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + exit_code: Final = _run(["--apply"], transport) + + assert exit_code == 0 + assert _posts(seen) == [ + ( + "/v1/vector_stores/vs_1", + { + "metadata": { + "deployment": None, + "endpoint": None, + "user_api_key": None, + "user_api_key_hash": None, + "user_api_key_user_email": None, + } + }, + ) + ] + assert "scanned=3 matched=1 updated=1 failed=0" in capsys.readouterr().out + + +def test_extra_key_is_deleted_alongside_the_built_in_keys(): + transport, seen = _fake_openai((POLLUTED,)) + + _run(["--apply", "--extra-key", "trace_id"], transport) + + assert _posts(seen)[0][1]["metadata"]["trace_id"] is None + assert "team" not in _posts(seen)[0][1]["metadata"] + + +def test_listing_follows_has_more_with_the_last_store_id(capsys): + transport, seen = _fake_openai((CLEAN, POLLUTED, NO_METADATA), page_size=2) + + _run([], transport) + + assert [request.url.params.get("after") for request in seen] == [None, "vs_1"] + assert "scanned=3 matched=1" in capsys.readouterr().out + + +def test_listing_walks_more_pages_than_the_interpreter_recursion_limit(capsys): + stores: Final = tuple({"id": f"vs_{index}", "metadata": None} for index in range(sys.getrecursionlimit() + 50)) + transport, _ = _fake_openai(stores, page_size=1) + + assert _run([], transport) == 0 + assert f"scanned={len(stores)} matched=0" in capsys.readouterr().out + + +def test_non_string_metadata_values_are_scanned_without_printing_them(capsys): + transport, _ = _fake_openai(({"id": "vs_9", "metadata": {"user_api_key": "secret-value", "count": 7}},)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert "vs_9: would delete 1 keys: user_api_key" in out + assert "secret-value" not in out + + +def test_an_unexpected_response_shape_exits_non_zero_without_echoing_metadata(capsys): + def malformed(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"metadata": {"user_api_key": "secret-value"}}], "has_more": False}) + + exit_code: Final = _run([], httpx.MockTransport(malformed)) + + captured: Final = capsys.readouterr() + assert exit_code == 2 + assert "unexpected vector store response shape" in captured.err + assert "secret-value" not in captured.out + captured.err + + +def test_vector_store_id_retrieves_only_the_named_stores_without_listing(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + _run(["--vector-store-id", "vs_2", "--vector-store-id", "vs_1"], transport) + + assert [request.url.path for request in seen] == ["/v1/vector_stores/vs_2", "/v1/vector_stores/vs_1"] + assert "scanned=2 matched=1" in capsys.readouterr().out + + +def test_a_failed_update_is_reported_and_exits_non_zero(capsys): + transport, _ = _fake_openai((POLLUTED,), post_status=400) + + exit_code: Final = _run(["--apply"], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 1 + assert "vs_1: failed to delete 5 keys (HTTP 400: upstream rejected the update)" in out + assert "updated=0 failed=1" in out + + +def test_a_listing_error_exits_non_zero_without_updating_anything(): + seen: Final[list[httpx.Request]] = [] + + def reject(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(401, json={"error": {"message": "bad key"}}) + + assert _run(["--apply"], httpx.MockTransport(reject)) == 2 + assert [request.method for request in seen] == ["GET"] + + +def test_an_unreachable_api_exits_non_zero(): + def unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + assert _run([], httpx.MockTransport(unreachable)) == 2 + + +def test_an_api_key_is_required(monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(SystemExit) as exit_info: + strip_script.main([]) + + assert exit_info.value.code == 2 diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 2f70c1f5188..48642a9b50d 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException, Request import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -3256,3 +3257,159 @@ class TestConfigOwnedVectorStores: assert response["status"] == "success", response prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID}) assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None + + +CALLERS_METADATA = {"team": "llmproxy", "endpoint": "ingest", "trace_id": "callers-trace"} + + +class _PreCallBodies(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_bodies = [] + + def log_pre_api_call(self, model, messages, kwargs): + self.logged_bodies.append(kwargs["litellm_params"]["proxy_server_request"]["body"]) + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize("model", [None, "gpt-5"]) +@pytest.mark.parametrize( + ("metadata_field", "expected_metadata"), + [ + ({"metadata": CALLERS_METADATA}, CALLERS_METADATA), + ({"metadata": {}}, {}), + ({"metadata": json.dumps(CALLERS_METADATA)}, CALLERS_METADATA), + ({"metadata": None}, None), + ({}, None), + ], +) +def test_proxy_vector_store_writes_send_the_provider_exactly_the_callers_metadata( + monkeypatch, operation, model, metadata_field, expected_metadata +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ), + ) + pre_call_bodies = _PreCallBodies() + monkeypatch.setattr(litellm, "callbacks", [pre_call_bodies]) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + routing = {} if model is None else {"model": model} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post( + path, + json={"name": "x", **routing, **metadata_field}, + headers={ + "Authorization": "Bearer sk-master-test", + "x-litellm-tags": "header-tag", + "x-litellm-trace-id": "header-trace", + "x-litellm-session-id": "header-session", + }, + ) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == expected_metadata + assert len(pre_call_bodies.logged_bodies) == 1 + assert "name" in pre_call_bodies.logged_bodies[0] + assert VECTOR_STORE_CALLER_METADATA_KWARG not in pre_call_bodies.logged_bodies[0] + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "request_body", + [ + {"data": {"name": "x", "metadata": json.dumps(CALLERS_METADATA)}}, + { + "content": json.dumps({"name": "x", "metadata": CALLERS_METADATA})[:-1] + + ', "nested": ' + + "[" * 300 + + "]" * 300 + + "}", + "headers": {"Content-Type": "application/json"}, + }, + ], + ids=["form", "json-too-deep-to-reparse"], +) +def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_body_cannot_be_reparsed( + monkeypatch, operation, request_body +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[])) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + headers = { + "Authorization": "Bearer sk-master-test", + "x-litellm-trace-id": "header-trace", + **request_body.get("headers", {}), + } + body_kwargs = {key: value for key, value in request_body.items() if key != "headers"} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post(path, headers=headers, **body_kwargs) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == CALLERS_METADATA diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 74a1bfefeee..815b32b33f7 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -1,10 +1,14 @@ import asyncio +import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import litellm import pytest from litellm import Router +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @pytest.fixture @@ -573,6 +577,39 @@ async def test_init_vector_store_api_endpoints(): mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai") +@pytest.mark.asyncio +async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider(): + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "vs_echo", + "object": "vector_store", + "created_at": 0, + "name": body["name"], + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + router: Final = Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ) + + response: Final = await router.avector_store_create( + model="gpt-5", + name="x", + metadata={"team": "llmproxy"}, + client=AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)), + ) + + assert response["id"] == "vs_echo" + assert response["metadata"] == {"team": "llmproxy"} + + def test_apply_default_settings(): """ Test the apply_default_settings method. diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 1c968126c42..25a378270b1 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -8,16 +8,19 @@ model_dump() it (the #19550 serialization trap). """ import json +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.vector_stores.main import search +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.vector_stores.main import acreate, aupdate, search MOCK_SEARCH_RESPONSE = { "object": "vector_store.search_results.page", @@ -114,3 +117,86 @@ def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): posted = json.loads(client.post.call_args.kwargs["data"]) assert posted["userContext"] == {"userId": "alice@example.com"} assert posted["retrievalQuery"] == {"text": "q"} + + +PROXY_STAMPED_METADATA: Final = { + "team": "llmproxy", + "user_api_key": "hashed-key", + "user_api_key_hash": "hashed-key", + "user_api_key_team_id": "team-1", + "requester_ip_address": "127.0.0.1", + "user_agent": "curl/8.7.1", + "endpoint": "http://localhost:4000/v1/vector_stores", + "litellm_api_version": "1.0.0", + "deployment": "openai/gpt-5", + "deployment_model_name": "gpt-5", + "api_base": "https://api.openai.com/v1", + "agent_id": "agent-1", + "actor_agent_id": "agent-2", + "target_agent_id": "agent-3", + "billing_agent_id": "agent-4", + "agent_execution_mode": "autonomous", + "verified_human_user_id": "user-1", + "litellm_caller_session_id": "session-1", + "user_api_keyboard": "qwerty", + 1: "a", +} + + +def _copy(metadata: Mapping[object, object] | None) -> dict[object, object] | None: + return None if metadata is None else dict(metadata) + + +async def _metadata_sent_to_provider( + operation: str, metadata: Mapping[object, object] | None, **kwargs: object +) -> list[object]: + sent_metadata: Final[list[object]] = [] + + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + sent_metadata.append(body.get("metadata")) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)) + if operation == "create": + await acreate(name="x", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + else: + await aupdate(vector_store_id="vs_1", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + return sent_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +async def test_proxy_vector_store_writes_send_only_the_callers_metadata_to_the_provider(operation: str): + sent_metadata: Final = await _metadata_sent_to_provider( + operation, + PROXY_STAMPED_METADATA, + proxy_server_request={"url": "http://localhost:4000/v1/vector_stores", "method": "POST", "body": {}}, + ) + + assert sent_metadata == [{"team": "llmproxy", "user_api_keyboard": "qwerty", "1": "a"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "callers_metadata", [{"team": "llmproxy", "endpoint": "ingest", "deployment": "prod", "agent_id": "bot-7"}, None] +) +async def test_direct_sdk_vector_store_writes_send_the_callers_metadata_unchanged( + operation: str, callers_metadata: Mapping[object, object] | None +): + sent_metadata: Final = await _metadata_sent_to_provider(operation, callers_metadata) + + assert sent_metadata == [callers_metadata]