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
This commit is contained in:
yuneng-jiang 2026-10-17 12:51:47 -07:00 • committed by GitHub
parent 4143e0869e
commit 4846cfabc7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 782 additions and 6 deletions

View file

@ -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

View file

@ -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())

View file

@ -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"

View file

@ -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
)

View file

@ -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

View file

@ -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}
)

View file

View file

@ -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

View file

@ -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

View file

@ -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.

View file

@ -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]