feat(api): Add vector store file batches api #3642
Conversation
meta-cla
bot
added
the
CLA Signed
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
from
ac1aa70 to
0f459be
Compare
slekkala1
requested review from
ashwinb,
yanxi0830,
hardikjshah,
raghotham,
ehhuang,
terrytangyuan,
leseb,
bbrowning,
reluctantfuturist and
mattf
as code owners
| asyncio.create_task(self._process_file_batch_async(batch_id, batch_info)) | ||
|
|
||
| # Run cleanup if needed (throttled to once every 1 day) | ||
| asyncio.create_task(self._cleanup_expired_file_batches_if_needed()) |
There was a problem hiding this comment.
We'd only want to spawn one such task, right? Now every create request adds one of these tasks.
There was a problem hiding this comment.
Moved the throttling check before create_task
| VectorStoreChunkingStrategyStatic(**chunking_strategy) | ||
| if chunking_strategy.get("type") == "static" | ||
| else VectorStoreChunkingStrategyAuto(**chunking_strategy) |
There was a problem hiding this comment.
I think pydantic has a way of handling thing for you. TypeAdapter or something?
There was a problem hiding this comment.
yes TypeAdapter works, thanks for recommendation.
| await self.openai_attach_file_to_vector_store( | ||
| vector_store_id=vector_store_id, | ||
| file_id=file_id, | ||
| attributes=attributes, | ||
| chunking_strategy=chunking_strategy_obj, | ||
| ) |
There was a problem hiding this comment.
for... await... would process each task in sequence. Do we want them in parallel instead?
There was a problem hiding this comment.
add concurrency for attaching files
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
from
0f459be to
d0bfdc7
Compare
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
from
d0bfdc7 to
42ebf66
Compare
| self.openai_file_batches: dict[str, dict[str, Any]] = {} | ||
| self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} | ||
| self._last_file_batch_cleanup_time = 0 |
There was a problem hiding this comment.
I know this has been the pattern, but should we just put this in the mixin's init and have subclasses call it?
There was a problem hiding this comment.
yes I agree, this list is growing
| # If cancellation fails because batch is already completed, that's acceptable | ||
| if "Cannot cancel" in str(e) or "already completed" in str(e): | ||
| pytest.skip(f"Batch completed too quickly to cancel: {e}") |
There was a problem hiding this comment.
this doesn't inspire confidence for this test.
There was a problem hiding this comment.
improved by increasing the number of file batches and removing the try catch.
There was a problem hiding this comment.
cancel test is finishing pretty fast even with 1000 files in a batch, hard cancel it before it completes.
| ) | ||
| assert search_response is not None | ||
|
|
||
|
|
There was a problem hiding this comment.
Do we have a test that actually tries to retrieve some context that verifies that the files are actually saved to the vector store?
There was a problem hiding this comment.
so this was tested via cc-vec integration, where it retrieved files from vector store after attaching via these apis.
Let me add one integration test for this.
There was a problem hiding this comment.
added integration test
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
2 times, most recently
from
c9be8c1 to
510ace2
Compare
|
Cancel test doesnt work as the batches complete too fast even with 1000 files, so expected failure in CI |
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
from
8da19a9 to
310dacd
Compare
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
2 times, most recently
from
16742e2 to
569ede7
Compare
slekkala1
force-pushed
the
add-vs-file-batches-and-tests
branch
from
130eb85 to
6533340
Compare
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