diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a35a8881a28..bff2742e2b6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2609,7 +2609,7 @@ async def azure_proxy_route( if vector_store is None: raise Exception(f"Vector store not found for {vector_store_name}") await assert_user_can_call_vector_store_passthrough( - vector_store, user_api_key_dict, provider_config, request.method, endpoint + vector_store, user_api_key_dict, provider_config, request.method, endpoint, store_segment="indexes" ) litellm_params = vector_store.get("litellm_params") or {} auth_credentials = provider_config.get_auth_credentials(litellm_params=litellm_params) @@ -3145,6 +3145,7 @@ async def _base_vertex_proxy_route( else ProviderConfigManager.get_provider_vector_stores_config(LlmProviders.VERTEX_AI, "search_api"), request.method, endpoint, + store_segment="dataStores", ) vertex_project: str | None = get_vertex_project_id_from_url(endpoint) diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index eeadab99737..cabd3144751 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,6 +1,7 @@ import json import re -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence +from itertools import accumulate from types import MappingProxyType from typing import Final, Literal @@ -432,7 +433,20 @@ async def assert_user_can_read_vector_store_id( return vector_store -def _is_provider_read_request(endpoints: VectorStoreIndexEndpoints, method: str, endpoint: str) -> bool: +def _store_marker_count(segments: Sequence[str], marker: str) -> int: + """How many stores the path names: a segment ending in the marker, as an unanchored extractor matches it, + followed by an id, which may itself contain the marker.""" + + def opens_a_store(after_marker: bool, segment_and_next: tuple[str, str | None]) -> bool: + segment, following = segment_and_next + return segment.endswith(marker) and following is not None and not after_marker + + return sum(accumulate(zip(segments, (*segments[1:], None)), opens_a_store, initial=False)) + + +def _is_provider_read_request( + endpoints: VectorStoreIndexEndpoints, method: str, endpoint: str, store_segment: str | None +) -> bool: """A substring match could hide a write, so a read needs a plain path the forwarder keeps as is, ending with a read entry. A GET only has to contain one.""" path: Final = endpoint if endpoint.startswith("/") else f"/{endpoint}" @@ -440,6 +454,8 @@ def _is_provider_read_request(endpoints: VectorStoreIndexEndpoints, method: str, return False if httpx.URL(endpoint).path != endpoint: return False + if store_segment is not None and _store_marker_count(path.split("/"), store_segment) > 1: + return False if any(write_method == method and write_path in path for write_method, write_path in endpoints["write"]): return False return any( @@ -454,10 +470,12 @@ async def assert_user_can_call_vector_store_passthrough( provider_config: BaseVectorStoreConfig | None, method: str, endpoint: str, + store_segment: str | None = None, ) -> None: - """Pass-through routes also carry writes, so groups count only on a read the provider's table names.""" + """Pass-through routes also carry writes, so groups count only on a read the provider's table names. + A path naming a second store through `store_segment` is never a read, since the gate checked only the first.""" if provider_config is not None and _is_provider_read_request( - provider_config.get_vector_store_endpoints_by_type(), method, endpoint + provider_config.get_vector_store_endpoints_by_type(), method, endpoint, store_segment ): await assert_user_can_read_vector_store(vector_store, user_api_key_dict) return 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 709bfc7cb7e..a8597a84c8d 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 @@ -9984,6 +9984,7 @@ _GROUP_READS: Final = ( pytest.param("POST", "/azure_ai/indexes/{index}/docs/suggest", {"search": "hi"}, id="azure-suggest"), pytest.param("GET", "/azure_ai/indexes/{index}", None, id="azure-get"), pytest.param("GET", "/azure_ai/indexes/{index}/docs/$count", None, id="azure-count"), + pytest.param("GET", "/azure_ai/indexes/{index}/docs/indexes", None, id="azure-document-key-named-indexes"), pytest.param( "POST", _GROUP_DISCOVERY + "/{store}/servingConfigs/default_config:search", {"query": "hi"}, id="vertex-search" ), @@ -9994,6 +9995,18 @@ _GROUP_WRITES: Final = ( pytest.param("POST", "/milvus/v2/vectordb/entities/insert", {"collectionName": "{index}"}, id="milvus-insert"), pytest.param("POST", "/azure_ai/indexes/{index}/docs/index", {"value": []}, id="azure-index"), pytest.param("POST", _GROUP_DISCOVERY + "/{store}/branches/0/documents:import", {}, id="vertex-import"), + pytest.param( + "POST", _GROUP_DISCOVERY + "/{store}/dataStores/store-b/servingConfigs/c:search", {}, id="vertex-second-store" + ), + pytest.param("POST", "/azure_ai/indexes/{index}/indexes/idx-b/docs/search", {}, id="azure-second-index"), + pytest.param( + "POST", + "/vertex_ai/discovery/v1/projects/p/locations/global/collections/c" + "/xdataStores/{store}/dataStores/store-b/c:search", + {}, + id="vertex-store-hidden-in-a-longer-segment", + ), + pytest.param("GET", "/azure_ai/indexes/{index}/indexes/idx-b", None, id="azure-second-index-get"), pytest.param("DELETE", _GROUP_DISCOVERY + "/{store}", None, id="vertex-delete"), pytest.param("POST", _GROUP_DISCOVERY + "/{store}/x:search/%2e%2e/%2e%2e/documents:purge", {}, id="vertex-traversal"), pytest.param("POST", _GROUP_DISCOVERY + "/{store}/documents:purge%3Fx:search", {}, id="vertex-encoded-query"), @@ -10108,6 +10121,38 @@ class TestVectorStorePassthroughGroupReads: assert await group_routes.send_as(key, method, path, body, "idx-a", "store-a") == 403 assert group_routes.upstream.call_count == 0 + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("path", "body"), + [ + pytest.param("/azure_ai/indexes/indexes/docs/search", {"search": "hi"}, id="azure-index-named-indexes"), + pytest.param( + _GROUP_DISCOVERY + "/dataStores/servingConfigs/c:search", + {"query": "hi"}, + id="vertex-store-named-datastores", + ), + ], + ) + async def test_a_store_named_like_the_path_marker_is_still_one_store( + self, group_routes: _GroupRoutes, monkeypatch: pytest.MonkeyPatch, path: str, body: Mapping[str, object] + ) -> None: + await group_routes.proxy.add_group("g1", ("store-a", "dataStores")) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[_group_test_store("store-a", "milvus"), _group_test_store("dataStores", "vertex_ai")] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry(vector_store_indexes=[_group_test_index("indexes", "store-a")]), + ) + + assert await group_routes.send_as(_group_test_key(indexes=("indexes",)), "POST", path, body, "", "") == 200 + assert group_routes.upstream.call_count == 1 + @pytest.mark.asyncio async def test_group_grant_does_not_replace_the_index_permission(self, group_routes: _GroupRoutes) -> None: status: Final = await group_routes.send_as(