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]