Commit cc4e1b00

pstern-sl <157847713+pstern-sl@users.noreply.github.com>
2024-06-08 05:36:40
feat(api): adding chunking_strategy to polling helpers (#1478)
1 parent b8a85a0
Changed files (2)
src
openai
resources
beta
src/openai/resources/beta/vector_stores/file_batches.py
@@ -174,11 +174,13 @@ class FileBatches(SyncAPIResource):
         *,
         file_ids: List[str],
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_batch_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFileBatch:
         """Create a vector store batch and poll until all files have been processed."""
         batch = self.create(
             vector_store_id=vector_store_id,
             file_ids=file_ids,
+            chunking_strategy=chunking_strategy,
         )
         # TODO: don't poll unless necessary??
         return self.poll(
@@ -306,6 +308,7 @@ class FileBatches(SyncAPIResource):
         max_concurrency: int = 5,
         file_ids: List[str] = [],
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_batch_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFileBatch:
         """Uploads the given files concurrently and then creates a vector store file batch.
 
@@ -343,6 +346,7 @@ class FileBatches(SyncAPIResource):
             vector_store_id=vector_store_id,
             file_ids=[*file_ids, *(f.id for f in results)],
             poll_interval_ms=poll_interval_ms,
+            chunking_strategy=chunking_strategy,
         )
         return batch
 
@@ -488,11 +492,13 @@ class AsyncFileBatches(AsyncAPIResource):
         *,
         file_ids: List[str],
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_batch_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFileBatch:
         """Create a vector store batch and poll until all files have been processed."""
         batch = await self.create(
             vector_store_id=vector_store_id,
             file_ids=file_ids,
+            chunking_strategy=chunking_strategy,
         )
         # TODO: don't poll unless necessary??
         return await self.poll(
@@ -620,6 +626,7 @@ class AsyncFileBatches(AsyncAPIResource):
         max_concurrency: int = 5,
         file_ids: List[str] = [],
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_batch_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFileBatch:
         """Uploads the given files concurrently and then creates a vector store file batch.
 
@@ -680,6 +687,7 @@ class AsyncFileBatches(AsyncAPIResource):
             vector_store_id=vector_store_id,
             file_ids=[*file_ids, *(f.id for f in uploaded_files)],
             poll_interval_ms=poll_interval_ms,
+            chunking_strategy=chunking_strategy,
         )
         return batch
 
src/openai/resources/beta/vector_stores/files.py
@@ -245,9 +245,10 @@ class Files(SyncAPIResource):
         *,
         vector_store_id: str,
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Attach a file to the given vector store and wait for it to be processed."""
-        self.create(vector_store_id=vector_store_id, file_id=file_id)
+        self.create(vector_store_id=vector_store_id, file_id=file_id, chunking_strategy=chunking_strategy)
 
         return self.poll(
             file_id,
@@ -301,6 +302,7 @@ class Files(SyncAPIResource):
         *,
         vector_store_id: str,
         file: FileTypes,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Upload a file to the `files` API and then attach it to the given vector store.
 
@@ -308,7 +310,7 @@ class Files(SyncAPIResource):
         polling helper method to wait for processing to complete).
         """
         file_obj = self._client.files.create(file=file, purpose="assistants")
-        return self.create(vector_store_id=vector_store_id, file_id=file_obj.id)
+        return self.create(vector_store_id=vector_store_id, file_id=file_obj.id, chunking_strategy=chunking_strategy)
 
     def upload_and_poll(
         self,
@@ -316,12 +318,14 @@ class Files(SyncAPIResource):
         vector_store_id: str,
         file: FileTypes,
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Add a file to a vector store and poll until processing is complete."""
         file_obj = self._client.files.create(file=file, purpose="assistants")
         return self.create_and_poll(
             vector_store_id=vector_store_id,
             file_id=file_obj.id,
+            chunking_strategy=chunking_strategy,
             poll_interval_ms=poll_interval_ms,
         )
 
@@ -542,9 +546,10 @@ class AsyncFiles(AsyncAPIResource):
         *,
         vector_store_id: str,
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Attach a file to the given vector store and wait for it to be processed."""
-        await self.create(vector_store_id=vector_store_id, file_id=file_id)
+        await self.create(vector_store_id=vector_store_id, file_id=file_id, chunking_strategy=chunking_strategy)
 
         return await self.poll(
             file_id,
@@ -598,6 +603,7 @@ class AsyncFiles(AsyncAPIResource):
         *,
         vector_store_id: str,
         file: FileTypes,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Upload a file to the `files` API and then attach it to the given vector store.
 
@@ -605,7 +611,7 @@ class AsyncFiles(AsyncAPIResource):
         polling helper method to wait for processing to complete).
         """
         file_obj = await self._client.files.create(file=file, purpose="assistants")
-        return await self.create(vector_store_id=vector_store_id, file_id=file_obj.id)
+        return await self.create(vector_store_id=vector_store_id, file_id=file_obj.id, chunking_strategy=chunking_strategy)
 
     async def upload_and_poll(
         self,
@@ -613,6 +619,7 @@ class AsyncFiles(AsyncAPIResource):
         vector_store_id: str,
         file: FileTypes,
         poll_interval_ms: int | NotGiven = NOT_GIVEN,
+        chunking_strategy: file_create_params.ChunkingStrategy | NotGiven = NOT_GIVEN,
     ) -> VectorStoreFile:
         """Add a file to a vector store and poll until processing is complete."""
         file_obj = await self._client.files.create(file=file, purpose="assistants")
@@ -620,6 +627,7 @@ class AsyncFiles(AsyncAPIResource):
             vector_store_id=vector_store_id,
             file_id=file_obj.id,
             poll_interval_ms=poll_interval_ms,
+            chunking_strategy=chunking_strategy
         )