phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-12 12:06:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: update VectorIORouter to get provider before calling openai_* methods

Browse source 
VectorDBsRoutingTable was removed in a165b8b5, so VectorIORouter needs to get the provider directly using routing_table.get_provider_impl() before calling provider methods, consistent with how insert_chunks() already works.
This commit is contained in:
Ashwin Bharambe 2025-10-11 17:37:10 -07:00
parent 3568ccdc81
commit 58fcaa445e
2 changed files with 44 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

42
llama_stack/core/routers/vector_io.py
View file

@ -214,7 +214,8 @@ class VectorIORouter(VectorIO):
vector_store_id: str, vector_store_id: str,
) -> VectorStoreObject: ) -> VectorStoreObject:
logger.debug(f"VectorIORouter.openai_retrieve_vector_store: {vector_store_id}") 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( async def openai_update_vector_store(
self, self,
@ -224,7 +225,8 @@ class VectorIORouter(VectorIO):
metadata: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None,
) -> VectorStoreObject: ) -> VectorStoreObject:
logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}") 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, vector_store_id=vector_store_id,
name=name, name=name,
expires_after=expires_after, expires_after=expires_after,
@ -236,7 +238,8 @@ class VectorIORouter(VectorIO):
vector_store_id: str, vector_store_id: str,
) -> VectorStoreDeleteResponse: ) -> VectorStoreDeleteResponse:
logger.debug(f"VectorIORouter.openai_delete_vector_store: {vector_store_id}") 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( async def openai_search_vector_store(
self, self,
@ -249,7 +252,8 @@ class VectorIORouter(VectorIO):
search_mode: str | None = "vector", search_mode: str | None = "vector",
) -> VectorStoreSearchResponsePage: ) -> VectorStoreSearchResponsePage:
logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}") 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, vector_store_id=vector_store_id,
query=query, query=query,
filters=filters, filters=filters,
@ -267,7 +271,8 @@ class VectorIORouter(VectorIO):
chunking_strategy: VectorStoreChunkingStrategy | None = None, chunking_strategy: VectorStoreChunkingStrategy | None = None,
) -> VectorStoreFileObject: ) -> VectorStoreFileObject:
logger.debug(f"VectorIORouter.openai_attach_file_to_vector_store: {vector_store_id}, {file_id}") 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, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
attributes=attributes, attributes=attributes,
@ -284,7 +289,8 @@ class VectorIORouter(VectorIO):
filter: VectorStoreFileStatus | None = None, filter: VectorStoreFileStatus | None = None,
) -> list[VectorStoreFileObject]: ) -> list[VectorStoreFileObject]:
logger.debug(f"VectorIORouter.openai_list_files_in_vector_store: {vector_store_id}") 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, vector_store_id=vector_store_id,
limit=limit, limit=limit,
order=order, order=order,
@ -299,7 +305,8 @@ class VectorIORouter(VectorIO):
file_id: str, file_id: str,
) -> VectorStoreFileObject: ) -> VectorStoreFileObject:
logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file: {vector_store_id}, {file_id}") 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, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
) )
@ -310,7 +317,8 @@ class VectorIORouter(VectorIO):
file_id: str, file_id: str,
) -> VectorStoreFileContentsResponse: ) -> VectorStoreFileContentsResponse:
logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}") 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, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
) )
@ -322,7 +330,8 @@ class VectorIORouter(VectorIO):
attributes: dict[str, Any], attributes: dict[str, Any],
) -> VectorStoreFileObject: ) -> VectorStoreFileObject:
logger.debug(f"VectorIORouter.openai_update_vector_store_file: {vector_store_id}, {file_id}") 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, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
attributes=attributes, attributes=attributes,
@ -334,7 +343,8 @@ class VectorIORouter(VectorIO):
file_id: str, file_id: str,
) -> VectorStoreFileDeleteResponse: ) -> VectorStoreFileDeleteResponse:
logger.debug(f"VectorIORouter.openai_delete_vector_store_file: {vector_store_id}, {file_id}") 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, vector_store_id=vector_store_id,
file_id=file_id, file_id=file_id,
) )
@ -370,7 +380,8 @@ class VectorIORouter(VectorIO):
logger.debug( logger.debug(
f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files" 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, vector_store_id=vector_store_id,
params=params, params=params,
) )
@ -381,7 +392,8 @@ class VectorIORouter(VectorIO):
vector_store_id: str, vector_store_id: str,
) -> VectorStoreFileBatchObject: ) -> VectorStoreFileBatchObject:
logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_batch: {batch_id}, {vector_store_id}") 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, batch_id=batch_id,
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
) )
@ -397,7 +409,8 @@ class VectorIORouter(VectorIO):
order: str | None = "desc", order: str | None = "desc",
) -> VectorStoreFilesListInBatchResponse: ) -> VectorStoreFilesListInBatchResponse:
logger.debug(f"VectorIORouter.openai_list_files_in_vector_store_file_batch: {batch_id}, {vector_store_id}") 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, batch_id=batch_id,
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
after=after, after=after,
@ -413,7 +426,8 @@ class VectorIORouter(VectorIO):
vector_store_id: str, vector_store_id: str,
) -> VectorStoreFileBatchObject: ) -> VectorStoreFileBatchObject:
logger.debug(f"VectorIORouter.openai_cancel_vector_store_file_batch: {batch_id}, {vector_store_id}") 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, batch_id=batch_id,
vector_store_id=vector_store_id, vector_store_id=vector_store_id,
) )

43
tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py
View file

@ -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_dbs import VectorDB
from llama_stack.apis.vector_io import ( from llama_stack.apis.vector_io import (
Chunk, Chunk,
OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
QueryChunksResponse, QueryChunksResponse,
VectorStoreChunkingStrategyAuto, VectorStoreChunkingStrategyAuto,
VectorStoreFileObject, VectorStoreFileObject,
@ -326,8 +327,7 @@ async def test_create_vector_store_file_batch(vector_io_adapter):
vector_io_adapter._process_file_batch_async = AsyncMock() vector_io_adapter._process_file_batch_async = AsyncMock()
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
assert batch.vector_store_id == store_id 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 # Create batch first
created_batch = await vector_io_adapter.openai_create_vector_store_file_batch( created_batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Retrieve batch # Retrieve batch
@ -388,8 +387,7 @@ async def test_cancel_vector_store_file_batch(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Cancel batch # Cancel batch
@ -434,8 +432,7 @@ async def test_list_files_in_vector_store_file_batch(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# List files # List files
@ -454,8 +451,7 @@ async def test_file_batch_validation_errors(vector_io_adapter):
# Test nonexistent vector store # Test nonexistent vector store
with pytest.raises(VectorStoreNotFoundError): with pytest.raises(VectorStoreNotFoundError):
await vector_io_adapter.openai_create_vector_store_file_batch( await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id="nonexistent", vector_store_id="nonexistent", params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"])
file_ids=["file_1"],
) )
# Setup store for remaining tests # 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 # Test wrong vector store for batch
vector_io_adapter.openai_attach_file_to_vector_store = AsyncMock() vector_io_adapter.openai_attach_file_to_vector_store = AsyncMock()
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=["file_1"])
file_ids=["file_1"],
) )
# Create wrong_store so it exists but the batch doesn't belong to it # 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 # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Test pagination with limit # Test pagination with limit
@ -593,8 +587,7 @@ async def test_file_batch_status_filtering(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Test filtering by completed status # Test filtering by completed status
@ -636,8 +629,7 @@ async def test_cancel_completed_batch_fails(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Manually update status to completed # Manually update status to completed
@ -671,8 +663,7 @@ async def test_file_batch_persistence_across_restarts(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
batch_id = batch.id batch_id = batch.id
@ -727,8 +718,7 @@ async def test_cancelled_batch_persists_in_storage(vector_io_adapter):
# Create batch # Create batch
batch = await vector_io_adapter.openai_create_vector_store_file_batch( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
batch_id = batch.id batch_id = batch.id
@ -775,10 +765,10 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter):
# Create multiple batches # Create multiple batches
batch1 = await vector_io_adapter.openai_create_vector_store_file_batch( 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( 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) # 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 # Create a third batch that stays in progress
batch3 = await vector_io_adapter.openai_create_vector_store_file_batch( 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 # 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 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( batch = await vector_io_adapter.openai_create_vector_store_file_batch(
vector_store_id=store_id, vector_store_id=store_id, params=OpenAICreateVectorStoreFileBatchRequestWithExtraBody(file_ids=file_ids)
file_ids=file_ids,
) )
# Give time for the semaphore logic to start processing files # Give time for the semaphore logic to start processing files