diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index b779df1b7..41d32b10c 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -214,7 +214,8 @@ class VectorIORouter(VectorIO): vector_store_id: str, ) -> VectorStoreObject: logger.debug(f"VectorIORouter.openai_retrieve_vector_store: {vector_store_id}") - return await self.routing_table.openai_retrieve_vector_store(vector_store_id) + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_retrieve_vector_store(vector_store_id) async def openai_update_vector_store( self, @@ -224,7 +225,8 @@ class VectorIORouter(VectorIO): metadata: dict[str, Any] | None = None, ) -> VectorStoreObject: logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}") - return await self.routing_table.openai_update_vector_store( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_update_vector_store( vector_store_id=vector_store_id, name=name, expires_after=expires_after, @@ -236,7 +238,8 @@ class VectorIORouter(VectorIO): vector_store_id: str, ) -> VectorStoreDeleteResponse: logger.debug(f"VectorIORouter.openai_delete_vector_store: {vector_store_id}") - return await self.routing_table.openai_delete_vector_store(vector_store_id) + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_delete_vector_store(vector_store_id) async def openai_search_vector_store( self, @@ -249,7 +252,8 @@ class VectorIORouter(VectorIO): search_mode: str | None = "vector", ) -> VectorStoreSearchResponsePage: logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}") - return await self.routing_table.openai_search_vector_store( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_search_vector_store( vector_store_id=vector_store_id, query=query, filters=filters, @@ -267,7 +271,8 @@ class VectorIORouter(VectorIO): chunking_strategy: VectorStoreChunkingStrategy | None = None, ) -> VectorStoreFileObject: logger.debug(f"VectorIORouter.openai_attach_file_to_vector_store: {vector_store_id}, {file_id}") - return await self.routing_table.openai_attach_file_to_vector_store( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_attach_file_to_vector_store( vector_store_id=vector_store_id, file_id=file_id, attributes=attributes, @@ -284,7 +289,8 @@ class VectorIORouter(VectorIO): filter: VectorStoreFileStatus | None = None, ) -> list[VectorStoreFileObject]: logger.debug(f"VectorIORouter.openai_list_files_in_vector_store: {vector_store_id}") - return await self.routing_table.openai_list_files_in_vector_store( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_list_files_in_vector_store( vector_store_id=vector_store_id, limit=limit, order=order, @@ -299,7 +305,8 @@ class VectorIORouter(VectorIO): file_id: str, ) -> VectorStoreFileObject: logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file: {vector_store_id}, {file_id}") - return await self.routing_table.openai_retrieve_vector_store_file( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_retrieve_vector_store_file( vector_store_id=vector_store_id, file_id=file_id, ) @@ -310,7 +317,8 @@ class VectorIORouter(VectorIO): file_id: str, ) -> VectorStoreFileContentsResponse: logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}") - return await self.routing_table.openai_retrieve_vector_store_file_contents( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_retrieve_vector_store_file_contents( vector_store_id=vector_store_id, file_id=file_id, ) @@ -322,7 +330,8 @@ class VectorIORouter(VectorIO): attributes: dict[str, Any], ) -> VectorStoreFileObject: logger.debug(f"VectorIORouter.openai_update_vector_store_file: {vector_store_id}, {file_id}") - return await self.routing_table.openai_update_vector_store_file( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_update_vector_store_file( vector_store_id=vector_store_id, file_id=file_id, attributes=attributes, @@ -334,7 +343,8 @@ class VectorIORouter(VectorIO): file_id: str, ) -> VectorStoreFileDeleteResponse: logger.debug(f"VectorIORouter.openai_delete_vector_store_file: {vector_store_id}, {file_id}") - return await self.routing_table.openai_delete_vector_store_file( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_delete_vector_store_file( vector_store_id=vector_store_id, file_id=file_id, ) @@ -370,7 +380,8 @@ class VectorIORouter(VectorIO): logger.debug( f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files" ) - return await self.routing_table.openai_create_vector_store_file_batch( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_create_vector_store_file_batch( vector_store_id=vector_store_id, params=params, ) @@ -381,7 +392,8 @@ class VectorIORouter(VectorIO): vector_store_id: str, ) -> VectorStoreFileBatchObject: logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_batch: {batch_id}, {vector_store_id}") - return await self.routing_table.openai_retrieve_vector_store_file_batch( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_retrieve_vector_store_file_batch( batch_id=batch_id, vector_store_id=vector_store_id, ) @@ -397,7 +409,8 @@ class VectorIORouter(VectorIO): order: str | None = "desc", ) -> VectorStoreFilesListInBatchResponse: logger.debug(f"VectorIORouter.openai_list_files_in_vector_store_file_batch: {batch_id}, {vector_store_id}") - return await self.routing_table.openai_list_files_in_vector_store_file_batch( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_list_files_in_vector_store_file_batch( batch_id=batch_id, vector_store_id=vector_store_id, after=after, @@ -413,7 +426,8 @@ class VectorIORouter(VectorIO): vector_store_id: str, ) -> VectorStoreFileBatchObject: logger.debug(f"VectorIORouter.openai_cancel_vector_store_file_batch: {batch_id}, {vector_store_id}") - return await self.routing_table.openai_cancel_vector_store_file_batch( + provider = await self.routing_table.get_provider_impl(vector_store_id) + return await provider.openai_cancel_vector_store_file_batch( batch_id=batch_id, vector_store_id=vector_store_id, ) diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index ed0934224..a5c491f53 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -15,6 +15,7 @@ from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_io import ( Chunk, + OpenAICreateVectorStoreFileBatchRequestWithExtraBody, QueryChunksResponse, VectorStoreChunkingStrategyAuto, VectorStoreFileObject, @@ -326,8 +327,7 @@ async def test_create_vector_store_file_batch(vector_io_adapter): vector_io_adapter._process_file_batch_async = AsyncMock() batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) assert batch.vector_store_id == store_id @@ -354,8 +354,7 @@ async def test_retrieve_vector_store_file_batch(vector_io_adapter): # Create batch first created_batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Retrieve batch @@ -388,8 +387,7 @@ async def test_cancel_vector_store_file_batch(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Cancel batch @@ -434,8 +432,7 @@ async def test_list_files_in_vector_store_file_batch(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # List files @@ -454,8 +451,7 @@ async def test_file_batch_validation_errors(vector_io_adapter): # Test nonexistent vector store with pytest.raises(VectorStoreNotFoundError): await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id="nonexistent", - file_ids=["file_1"], + vector_store_id="nonexistent", params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"]) ) # Setup store for remaining tests @@ -472,8 +468,7 @@ async def test_file_batch_validation_errors(vector_io_adapter): # Test wrong vector store for batch vector_io_adapter.openai_attach_file_to_vector_store = AsyncMock() batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=["file_1"], + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"]) ) # Create wrong_store so it exists but the batch doesn't belong to it @@ -520,8 +515,7 @@ async def test_file_batch_pagination(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Test pagination with limit @@ -593,8 +587,7 @@ async def test_file_batch_status_filtering(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Test filtering by completed status @@ -636,8 +629,7 @@ async def test_cancel_completed_batch_fails(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Manually update status to completed @@ -671,8 +663,7 @@ async def test_file_batch_persistence_across_restarts(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) batch_id = batch.id @@ -727,8 +718,7 @@ async def test_cancelled_batch_persists_in_storage(vector_io_adapter): # Create batch batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) batch_id = batch.id @@ -775,10 +765,10 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter): # Create multiple batches batch1 = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, file_ids=["file_1"] + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"]) ) batch2 = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, file_ids=["file_2"] + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_2"]) ) # Complete one batch (should persist with completed status) @@ -791,7 +781,7 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter): # Create a third batch that stays in progress batch3 = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, file_ids=["file_3"] + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_3"]) ) # Simulate restart - clear memory and reload from persistence @@ -952,8 +942,7 @@ async def test_max_concurrent_files_per_batch(vector_io_adapter): file_ids = [f"file_{i}" for i in range(8)] # 8 files, but limit should be 5 batch = await vector_io_adapter.openai_create_vector_store_file_batch( - vector_store_id=store_id, - file_ids=file_ids, + vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids) ) # Give time for the semaphore logic to start processing files