From 9385af1b1805cd67101077fb87c0e9b188647e1b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:32:11 -0700 Subject: [PATCH 1/4] test(proxy): cover the mapped Azure AI Search index name on the azure_ai passthrough --- .../test_llm_pass_through_endpoints.py | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 0186a33dbc4..16abeb3dda4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5470,6 +5470,7 @@ class TestAzureProxyRouteCrossIndexAuthorization: async def test_authorizes_the_targeted_index(self): index_object = MagicMock() index_object.litellm_params.vector_store_name = "my-store" + index_object.litellm_params.vector_store_index = "my-real-index" vector_store = {"litellm_params": {"api_base": "https://svc.search.windows.net"}} with ( @@ -5556,6 +5557,119 @@ class TestAzureProxyRouteCrossIndexAuthorization: assert mock_handler.await_args.kwargs["custom_llm_provider"] == litellm.LlmProviders.AZURE +@pytest.fixture +def azure_ai_search_virtual_index_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[UserAPIKeyAuth], TestClient]]: + from litellm.proxy import proxy_server + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[])) + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + vector_store_indexes=[ + LiteLLM_ManagedVectorStoreIndex( + id="idx-1", + index_name="virtual-idx", + litellm_params=IndexCreateLiteLLMParams(vector_store_index="real-idx", vector_store_name="search"), + ) + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[ + { + "vector_store_id": "can-be-anything", + "vector_store_name": "search", + "custom_llm_provider": "azure_ai", + "litellm_params": { + "vector_store_name": "search", + "api_base": "https://svc.search.windows.net", + "api_key": "search-admin-key", + }, + } + ] + ), + ) + + def client_as(caller: UserAPIKeyAuth) -> TestClient: + monkeypatch.setitem(proxy_server.app.dependency_overrides, user_api_key_auth, lambda: caller) + return TestClient(proxy_server.app) + + yield client_as + + +DEVELOPER_KEY_ON_VIRTUAL_INDEX: Final = UserAPIKeyAuth( + api_key="sk-developer", + metadata={"allowed_vector_store_indexes": [{"index_name": "virtual-idx", "index_permissions": ["read", "write"]}]}, +) + + +@pytest.mark.parametrize( + "method, path, body, upstream_path", + [ + ("POST", "indexes/virtual-idx/docs/search", {"search": "refund"}, "/indexes/real-idx/docs/search"), + ( + "POST", + "indexes/virtual-idx/docs/index", + {"value": [{"@search.action": "upload", "id": "1", "content": "refund policy"}]}, + "/indexes/real-idx/docs/index", + ), + ("GET", "indexes/virtual-idx/docs/$count", None, "/indexes/real-idx/docs/$count"), + ("GET", "indexes/virtual-idx", None, "/indexes/real-idx"), + ], +) +def test_azure_ai_search_virtual_index_reaches_azure_as_the_mapped_index( + azure_ai_search_virtual_index_client: Callable[[UserAPIKeyAuth], TestClient], + method: str, + path: str, + body: dict[str, object] | None, + upstream_path: str, +) -> None: + client: Final = azure_ai_search_virtual_index_client(DEVELOPER_KEY_ON_VIRTUAL_INDEX) + with respx.mock(assert_all_called=True) as upstream: + azure: Final = upstream.route(host="svc.search.windows.net").mock( + return_value=httpx.Response(200, json={"value": []}) + ) + response: Final = client.request(method, f"/azure_ai/{path}?api-version=2024-07-01", json=body) + + assert response.status_code == 200 + sent: Final = azure.calls.last.request + assert (sent.method, sent.url.path, sent.url.params["api-version"]) == (method, upstream_path, "2024-07-01") + assert sent.headers["api-key"] == "search-admin-key" + assert json.loads(sent.content or b"null") == body + + +@pytest.mark.parametrize( + "sent_name, upstream_name", + [("virtual-idx", "real-idx"), ("some-other-idx", "some-other-idx")], +) +def test_azure_ai_search_index_definition_put_names_the_mapped_index( + azure_ai_search_virtual_index_client: Callable[[UserAPIKeyAuth], TestClient], sent_name: str, upstream_name: str +) -> None: + client: Final = azure_ai_search_virtual_index_client( + UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + ) + fields: Final = [{"name": "id", "type": "Edm.String", "key": True}] + with respx.mock(assert_all_called=True) as upstream: + azure: Final = upstream.route(host="svc.search.windows.net").mock( + return_value=httpx.Response(201, json={"name": upstream_name}) + ) + response: Final = client.put( + "/azure_ai/indexes/virtual-idx?api-version=2024-07-01", json={"name": sent_name, "fields": fields} + ) + + assert response.status_code == 201 + sent: Final = azure.calls.last.request + assert (sent.method, sent.url.path) == ("PUT", "/indexes/real-idx") + assert json.loads(sent.content) == {"name": upstream_name, "fields": fields} + + class TestAzureProxyRouteServiceLevelIndexCreate: """``POST /indexes`` carries no index name, so the managed-index branch cannot claim it and it would otherwise reach the generic Azure passthrough on the From 3fecab24fb5501eb0d7f99db459a46c3c6ecd204 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:40:20 -0700 Subject: [PATCH 2/4] fix(proxy): send the mapped Azure AI Search index name on the azure_ai passthrough --- .../llm_passthrough_endpoints.py | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..11fb2f5a102 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2160,6 +2160,26 @@ def is_azure_ai_search_service_level_index_create(method: str, endpoint: str) -> return path == "indexes" or path.endswith("/indexes") +def azure_ai_search_endpoint_on_index(endpoint: str, index_name: str) -> str: + segments: Final = endpoint.split("/") + index_position: Final = segments.index("indexes") + 1 + return "/".join((*segments[:index_position], index_name, *segments[index_position + 1 :])) + + +_AZURE_AI_SEARCH_INDEX_DEFINITION: Final = TypeAdapter(Mapping[str, object]) + + +async def name_azure_ai_search_index_definition( + request: Request, endpoint: str, virtual_index: str, vector_store_index: str +) -> None: + if endpoint.strip("/") != f"indexes/{virtual_index}": + return + index_definition: Final = _AZURE_AI_SEARCH_INDEX_DEFINITION.validate_python(await read_request_body(request)) + if index_definition.get("name") != virtual_index: + return + safe_set_request_parsed_body(request, {**index_definition, "name": vector_store_index}) + + async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> AsyncGenerator[bytes, None]: try: async for chunk in upstream: @@ -2468,8 +2488,16 @@ async def azure_proxy_route( base_target_url = _optional_str(litellm_params.get("api_base")) if base_target_url is None: raise Exception(f"API base not found for {part}") - return await BaseOpenAIPassThroughHandler.base_openai_pass_through_handler( + await name_azure_ai_search_index_definition( + request=request, endpoint=endpoint, + virtual_index=part, + vector_store_index=index_object.litellm_params.vector_store_index, + ) + return await BaseOpenAIPassThroughHandler.base_openai_pass_through_handler( + endpoint=azure_ai_search_endpoint_on_index( + endpoint, index_object.litellm_params.vector_store_index + ), request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, From 78e5a0a738cc910d02590e9a994d0d8dab559483 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:43:14 -0700 Subject: [PATCH 3/4] test(proxy): pin the mapped Azure AI Search index on index delete and document-key paths --- .../test_llm_pass_through_endpoints.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 16abeb3dda4..cbebb615bf3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5622,6 +5622,13 @@ DEVELOPER_KEY_ON_VIRTUAL_INDEX: Final = UserAPIKeyAuth( ), ("GET", "indexes/virtual-idx/docs/$count", None, "/indexes/real-idx/docs/$count"), ("GET", "indexes/virtual-idx", None, "/indexes/real-idx"), + ("GET", "indexes/virtual-idx/docs/virtual-idx", None, "/indexes/real-idx/docs/virtual-idx"), + ( + "POST", + "indexes/virtual-idx/docs/search", + {"search": "*", "name": "virtual-idx"}, + "/indexes/real-idx/docs/search", + ), ], ) def test_azure_ai_search_virtual_index_reaches_azure_as_the_mapped_index( @@ -5670,6 +5677,21 @@ def test_azure_ai_search_index_definition_put_names_the_mapped_index( assert json.loads(sent.content) == {"name": upstream_name, "fields": fields} + +def test_azure_ai_search_index_delete_on_the_virtual_index_deletes_the_mapped_index( + azure_ai_search_virtual_index_client: Callable[[UserAPIKeyAuth], TestClient], +) -> None: + client: Final = azure_ai_search_virtual_index_client( + UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + ) + with respx.mock(assert_all_called=True) as upstream: + azure: Final = upstream.route(host="svc.search.windows.net").mock(return_value=httpx.Response(204)) + response: Final = client.delete("/azure_ai/indexes/virtual-idx?api-version=2024-07-01") + + assert response.status_code == 204 + sent: Final = azure.calls.last.request + assert (sent.method, sent.url.path, sent.headers["api-key"]) == ("DELETE", "/indexes/real-idx", "search-admin-key") + class TestAzureProxyRouteServiceLevelIndexCreate: """``POST /indexes`` carries no index name, so the managed-index branch cannot claim it and it would otherwise reach the generic Azure passthrough on the From 2b53ff4ac1f0cf75beb1dcab22e819d2176883f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:02:22 -0700 Subject: [PATCH 4/4] style(tests): format the Azure AI Search passthrough tests --- .../test_llm_pass_through_endpoints.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index cbebb615bf3..9ac83e8e702 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2142,7 +2142,6 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) - @pytest.mark.asyncio async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( self, monkeypatch: pytest.MonkeyPatch @@ -2183,6 +2182,7 @@ class TestBedrockPassthroughAwsErrors: await self._answer(error) assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") + @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), @@ -5558,7 +5558,9 @@ class TestAzureProxyRouteCrossIndexAuthorization: @pytest.fixture -def azure_ai_search_virtual_index_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[UserAPIKeyAuth], TestClient]]: +def azure_ai_search_virtual_index_client( + monkeypatch: pytest.MonkeyPatch, +) -> Iterator[Callable[[UserAPIKeyAuth], TestClient]]: from litellm.proxy import proxy_server from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry @@ -5677,7 +5679,6 @@ def test_azure_ai_search_index_definition_put_names_the_mapped_index( assert json.loads(sent.content) == {"name": upstream_name, "fields": fields} - def test_azure_ai_search_index_delete_on_the_virtual_index_deletes_the_mapped_index( azure_ai_search_virtual_index_client: Callable[[UserAPIKeyAuth], TestClient], ) -> None: @@ -5690,7 +5691,12 @@ def test_azure_ai_search_index_delete_on_the_virtual_index_deletes_the_mapped_in assert response.status_code == 204 sent: Final = azure.calls.last.request - assert (sent.method, sent.url.path, sent.headers["api-key"]) == ("DELETE", "/indexes/real-idx", "search-admin-key") + assert (sent.method, sent.url.path, sent.headers["api-key"]) == ( + "DELETE", + "/indexes/real-idx", + "search-admin-key", + ) + class TestAzureProxyRouteServiceLevelIndexCreate: """``POST /indexes`` carries no index name, so the managed-index branch cannot