Skip to content

Commit bba9957

Browse files
authored
feat(api): Add vector store file batches api (#3642)
# What does this PR do? Add Open AI Compatible vector store file batches api. This functionality is needed to attach many files to a vector store as a batch. #3533 API Stubs have been merged #3615 Adds persistence for file batches as discussed in diff #3544 (Used claude code for generation and reviewed by me) ## Test Plan 1. Unit tests pass 2. Also verified the cc-vec integration with LLamaStackClient works with the file batches api. https://github.com/raghotham/cc-vec 2. Integration tests pass
1 parent 597d405 commit bba9957

File tree

37 files changed

+10322
-53
lines changed

37 files changed

+10322
-53
lines changed

llama_stack/core/routing_tables/vector_dbs.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,3 +245,65 @@ async def openai_delete_vector_store_file(
245245
vector_store_id=vector_store_id,
246246
file_id=file_id,
247247
)
248+
249+
async def openai_create_vector_store_file_batch(
250+
self,
251+
vector_store_id: str,
252+
file_ids: list[str],
253+
attributes: dict[str, Any] | None = None,
254+
chunking_strategy: Any | None = None,
255+
):
256+
await self.assert_action_allowed("update", "vector_db", vector_store_id)
257+
provider = await self.get_provider_impl(vector_store_id)
258+
return await provider.openai_create_vector_store_file_batch(
259+
vector_store_id=vector_store_id,
260+
file_ids=file_ids,
261+
attributes=attributes,
262+
chunking_strategy=chunking_strategy,
263+
)
264+
265+
async def openai_retrieve_vector_store_file_batch(
266+
self,
267+
batch_id: str,
268+
vector_store_id: str,
269+
):
270+
await self.assert_action_allowed("read", "vector_db", vector_store_id)
271+
provider = await self.get_provider_impl(vector_store_id)
272+
return await provider.openai_retrieve_vector_store_file_batch(
273+
batch_id=batch_id,
274+
vector_store_id=vector_store_id,
275+
)
276+
277+
async def openai_list_files_in_vector_store_file_batch(
278+
self,
279+
batch_id: str,
280+
vector_store_id: str,
281+
after: str | None = None,
282+
before: str | None = None,
283+
filter: str | None = None,
284+
limit: int | None = 20,
285+
order: str | None = "desc",
286+
):
287+
await self.assert_action_allowed("read", "vector_db", vector_store_id)
288+
provider = await self.get_provider_impl(vector_store_id)
289+
return await provider.openai_list_files_in_vector_store_file_batch(
290+
batch_id=batch_id,
291+
vector_store_id=vector_store_id,
292+
after=after,
293+
before=before,
294+
filter=filter,
295+
limit=limit,
296+
order=order,
297+
)
298+
299+
async def openai_cancel_vector_store_file_batch(
300+
self,
301+
batch_id: str,
302+
vector_store_id: str,
303+
):
304+
await self.assert_action_allowed("update", "vector_db", vector_store_id)
305+
provider = await self.get_provider_impl(vector_store_id)
306+
return await provider.openai_cancel_vector_store_file_batch(
307+
batch_id=batch_id,
308+
vector_store_id=vector_store_id,
309+
)

llama_stack/providers/inline/vector_io/faiss/faiss.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,12 +200,10 @@ async def query_hybrid(
200200

201201
class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate):
202202
def __init__(self, config: FaissVectorIOConfig, inference_api: Inference, files_api: Files | None) -> None:
203+
super().__init__(files_api=files_api, kvstore=None)
203204
self.config = config
204205
self.inference_api = inference_api
205-
self.files_api = files_api
206206
self.cache: dict[str, VectorDBWithIndex] = {}
207-
self.kvstore: KVStore | None = None
208-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
209207

210208
async def initialize(self) -> None:
211209
self.kvstore = await kvstore_impl(self.config.kvstore)

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -410,12 +410,10 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
410410
"""
411411

412412
def __init__(self, config, inference_api: Inference, files_api: Files | None) -> None:
413+
super().__init__(files_api=files_api, kvstore=None)
413414
self.config = config
414415
self.inference_api = inference_api
415-
self.files_api = files_api
416416
self.cache: dict[str, VectorDBWithIndex] = {}
417-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
418-
self.kvstore: KVStore | None = None
419417

420418
async def initialize(self) -> None:
421419
self.kvstore = await kvstore_impl(self.config.kvstore)

llama_stack/providers/remote/vector_io/chroma/chroma.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,14 +140,13 @@ def __init__(
140140
inference_api: Api.inference,
141141
files_api: Files | None,
142142
) -> None:
143+
super().__init__(files_api=files_api, kvstore=None)
143144
log.info(f"Initializing ChromaVectorIOAdapter with url: {config}")
144145
self.config = config
145146
self.inference_api = inference_api
146147
self.client = None
147148
self.cache = {}
148-
self.kvstore: KVStore | None = None
149149
self.vector_db_store = None
150-
self.files_api = files_api
151150

152151
async def initialize(self) -> None:
153152
self.kvstore = await kvstore_impl(self.config.kvstore)

llama_stack/providers/remote/vector_io/milvus/milvus.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -309,14 +309,12 @@ def __init__(
309309
inference_api: Inference,
310310
files_api: Files | None,
311311
) -> None:
312+
super().__init__(files_api=files_api, kvstore=None)
312313
self.config = config
313314
self.cache = {}
314315
self.client = None
315316
self.inference_api = inference_api
316-
self.files_api = files_api
317-
self.kvstore: KVStore | None = None
318317
self.vector_db_store = None
319-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
320318
self.metadata_collection_name = "openai_vector_stores_metadata"
321319

322320
async def initialize(self) -> None:

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -345,14 +345,12 @@ def __init__(
345345
inference_api: Api.inference,
346346
files_api: Files | None = None,
347347
) -> None:
348+
super().__init__(files_api=files_api, kvstore=None)
348349
self.config = config
349350
self.inference_api = inference_api
350351
self.conn = None
351352
self.cache = {}
352-
self.files_api = files_api
353-
self.kvstore: KVStore | None = None
354353
self.vector_db_store = None
355-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
356354
self.metadata_collection_name = "openai_vector_stores_metadata"
357355

358356
async def initialize(self) -> None:

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
from llama_stack.log import get_logger
2828
from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
2929
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
30-
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
30+
from llama_stack.providers.utils.kvstore import kvstore_impl
3131
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
3232
from llama_stack.providers.utils.memory.vector_store import (
3333
ChunkForDeletion,
@@ -162,14 +162,12 @@ def __init__(
162162
inference_api: Api.inference,
163163
files_api: Files | None = None,
164164
) -> None:
165+
super().__init__(files_api=files_api, kvstore=None)
165166
self.config = config
166167
self.client: AsyncQdrantClient = None
167168
self.cache = {}
168169
self.inference_api = inference_api
169-
self.files_api = files_api
170170
self.vector_db_store = None
171-
self.kvstore: KVStore | None = None
172-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
173171
self._qdrant_lock = asyncio.Lock()
174172

175173
async def initialize(self) -> None:

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -284,14 +284,12 @@ def __init__(
284284
inference_api: Api.inference,
285285
files_api: Files | None,
286286
) -> None:
287+
super().__init__(files_api=files_api, kvstore=None)
287288
self.config = config
288289
self.inference_api = inference_api
289290
self.client_cache = {}
290291
self.cache = {}
291-
self.files_api = files_api
292-
self.kvstore: KVStore | None = None
293292
self.vector_db_store = None
294-
self.openai_vector_stores: dict[str, dict[str, Any]] = {}
295293
self.metadata_collection_name = "openai_vector_stores_metadata"
296294

297295
def _get_client(self) -> weaviate.WeaviateClient:

0 commit comments

Comments
 (0)