more fixes

This commit is contained in:
Ashwin Bharambe 2025-10-11 22:05:20 -07:00
parent bf59d26362
commit e5a1cdf554
9 changed files with 43 additions and 146 deletions

View file

@ -491,11 +491,13 @@ class OpenAICreateVectorStoreRequestWithExtraBody(BaseModel, extra="allow"):
class OpenAICreateVectorStoreFileBatchRequestWithExtraBody(BaseModel, extra="allow"):
"""Request to create a vector store file batch with extra_body support.
:param vector_store_id: The ID of the vector store to create the file batch for
:param file_ids: A list of File IDs that the vector store should use
:param attributes: (Optional) Key-value attributes to store with the files
:param chunking_strategy: (Optional) The chunking strategy used to chunk the file(s). Defaults to auto
"""
vector_store_id: str
file_ids: list[str]
attributes: dict[str, Any] | None = None
chunking_strategy: VectorStoreChunkingStrategy | None = None
@ -847,13 +849,11 @@ class VectorIO(Protocol):
)
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
"""Create a vector store file batch.
Generate an OpenAI-compatible vector store file batch for the given vector store.
:param vector_store_id: The ID of the vector store to create the file batch for.
:returns: A VectorStoreFileBatchObject representing the created file batch.
"""
...

View file

@ -135,7 +135,7 @@ class VectorIORouter(VectorIO):
logger.debug(f"VectorIORouter.openai_create_vector_store: name={params.name}, provider_id={provider_id}")
# If no embedding model is provided, use the first available one
# TODO: this branch will soon be deleted so you _must_ provide the embedding_model when
# TODO: this branch will soon be deleted so you _must_ provide the embedding_model when
# creating a vector store
if embedding_model is None:
embedding_model_info = await self._get_first_embedding_model()
@ -383,17 +383,13 @@ class VectorIORouter(VectorIO):
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
logger.debug(
f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files"
)
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,
f"VectorIORouter.openai_create_vector_store_file_batch: {params.vector_store_id}, {len(params.file_ids)} files"
)
provider = await self.routing_table.get_provider_impl(params.vector_store_id)
return await provider.openai_create_vector_store_file_batch(params)
async def openai_retrieve_vector_store_file_batch(
self,

View file

@ -978,10 +978,10 @@ class OpenAIVectorStoreMixin(ABC):
async def openai_create_vector_store_file_batch(
self,
vector_store_id: str,
params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)],
) -> VectorStoreFileBatchObject:
"""Create a vector store file batch."""
vector_store_id = params.vector_store_id
if vector_store_id not in self.openai_vector_stores:
raise VectorStoreNotFoundError(vector_store_id)