diff --git a/tests/integration/authorization/test_passthrough_response_id_issuance.py b/tests/integration/authorization/test_passthrough_response_id_issuance.py index 1f18f56c9e6..fb03aaf6ae4 100644 --- a/tests/integration/authorization/test_passthrough_response_id_issuance.py +++ b/tests/integration/authorization/test_passthrough_response_id_issuance.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import uuid from collections.abc import Iterator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -46,6 +47,10 @@ _BRING_OWNER_TABLE_BACK: Final = ( 'ALTER TABLE "LiteLLM_ManagedObjectTable_offline" RENAME TO "LiteLLM_ManagedObjectTable"' ) +_SEARCH_STORE: Final = "integration-search-store" +_SEARCH_KEY: Final = "synthetic-search-key" +_SEARCH_RESULTS: Final = Reply(body=json.dumps({"value": []}).encode()) + Route = Literal[ "retrieve", "input_items", @@ -155,6 +160,51 @@ def _owner_table_offline(database_url: str) -> Iterator[None]: admin.execute(_BRING_OWNER_TABLE_BACK) +def _search_settings(search: Wire) -> str: + return ( + _MANAGED + + "vector_store_registry:\n" + + f" - vector_store_name: {_SEARCH_STORE}\n" + + " litellm_params:\n" + + f" vector_store_id: {_SEARCH_STORE}\n" + + " custom_llm_provider: azure_ai\n" + + f" api_key: {_SEARCH_KEY}\n" + + f" api_base: {search.url}\n" + ) + + +@contextmanager +def _search_index(gateway: Gateway) -> Iterator[str]: + index: Final = f"integration-index-{uuid.uuid4().hex}" + gateway.post( + "/v1/indexes", + {"index_name": index, "litellm_params": {"vector_store_index": index, "vector_store_name": _SEARCH_STORE}}, + ) + try: + yield index + finally: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute('DELETE FROM "LiteLLM_ManagedVectorStoreIndexTable" WHERE index_name = %s', (index,)) + + +@contextmanager +def _search_rig(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, Wire, str, str]]: + with ( + wire_server(lambda _request: _SEARCH_RESULTS) as search, + wire_server(_respond) as openai_upstream, + _search_index(gateway) as index, + _proxy(gateway, directory, openai_upstream, settings=_search_settings(search)) as candidate, + candidate.scenario() as scenario, + ): + scenario.model(model=f"azure/{_MODEL}", api_base=openai_upstream.url, api_key=_UPSTREAM_KEY) + reader: Final = scenario.key( + user_id=scenario.user(), + metadata={"allowed_vector_store_indexes": [{"index_name": index, "index_permissions": ["read"]}]}, + ) + yield candidate, search, index, reader + assert _forwarded(openai_upstream) == () + + def _call(route: Route, vendor_id: str) -> tuple[str, str, bytes | None]: match route: case "retrieve": @@ -384,6 +434,31 @@ def test_the_model_list_relay_hands_out_and_checks_gateway_ids(gateway: Gateway, assert _forwarded(openai_upstream) == () +def test_an_azure_ai_search_query_reaches_the_index_unchanged(gateway: Gateway, tmp_path: Path) -> None: + query: Final = {"search": "mango", "top": 3} + with _search_rig(gateway, tmp_path) as (candidate, search, index, reader): + answered: Final = _send( + candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode() + ) + assert answered.status_code == 200, answered.text + assert answered.json() == {"value": []}, answered.text + received: Final = search.drain() + assert [(request.method, urlsplit(request.target).path) for request in received] == [ + ("POST", f"/indexes/{index}/docs/search") + ] + assert [_JSON_OBJECT.validate_json(request.body) for request in received] == [query] + assert [request.headers.get("api-key") for request in received] == [_SEARCH_KEY] + + +def test_an_azure_ai_search_query_carrying_a_vendor_response_id_is_refused(gateway: Gateway, tmp_path: Path) -> None: + query: Final = {"search": "mango", "previous_response_id": _vendor_id()} + with _search_rig(gateway, tmp_path) as (candidate, search, index, reader): + _assert_unmanaged_refusal( + _send(candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode()) + ) + assert _forwarded(search) == () + + def test_every_verdict_holds_through_a_burst_after_a_worker_is_replaced(gateway: Gateway, tmp_path: Path) -> None: vendor_ids: Final = tuple(_vendor_id() for _ in range(10)) config: Final = tmp_path / "passthrough-burst.yaml"