Skip to content

Commit 523ca4b

Browse files
joeldierkesjoeldierkes
authored
Merge pull request #245 from mixedbread-ai/next-update
chore: add the stores helper functions
2 parents da4d3b6 + 153269e commit 523ca4b

1 file changed

Lines changed: 220 additions & 2 deletions

File tree

‎src/mixedbread/resources/stores/files.py‎

Lines changed: 220 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22

33
from __future__ import annotations
44

5-
from typing import List, Union, Iterable, Optional
5+
import functools
6+
from typing import Any, List, Union, Iterable, Optional
67

78
import httpx
89

9-
from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
10+
from ...lib import polling
11+
from ..._types import Body, FileTypes, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
1012
from ..._utils import maybe_transform, async_maybe_transform
1113
from ..._compat import cached_property
1214
from ..._resource import SyncAPIResource, AsyncAPIResource
@@ -340,6 +342,112 @@ def search(
340342
cast_to=FileSearchResponse,
341343
)
342344

345+
def poll(
346+
self,
347+
file_id: str,
348+
*,
349+
store_identifier: str,
350+
poll_interval_ms: int | NotGiven = not_given,
351+
poll_timeout_ms: float | NotGiven = not_given,
352+
**kwargs: Any,
353+
) -> StoreFile:
354+
"""
355+
Poll for a file's status until it reaches a terminal state.
356+
Args:
357+
file_id: The ID of the file to poll
358+
store_identifier: The ID of the store
359+
poll_interval_ms: The interval between polls in milliseconds
360+
poll_timeout_ms: The maximum time to poll for in milliseconds
361+
Returns:
362+
The file object once it reaches a terminal state
363+
"""
364+
polling_interval_ms = poll_interval_ms or 500
365+
polling_timeout_ms = poll_timeout_ms or None
366+
return polling.poll(
367+
fn=functools.partial(self.retrieve, file_id, store_identifier=store_identifier, **kwargs),
368+
condition=lambda res: res.status == "completed" or res.status == "failed" or res.status == "cancelled",
369+
interval_seconds=polling_interval_ms / 1000,
370+
timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None,
371+
)
372+
373+
def create_and_poll(
374+
self,
375+
file_id: str,
376+
*,
377+
store_identifier: str,
378+
metadata: Optional[object] | NotGiven = not_given,
379+
experimental: file_create_params.Experimental | Omit = omit,
380+
poll_interval_ms: int | NotGiven = not_given,
381+
poll_timeout_ms: float | NotGiven = not_given,
382+
**kwargs: Any,
383+
) -> StoreFile:
384+
"""
385+
Attach a file to the given store and wait for it to be processed.
386+
Args:
387+
file_id: The ID of the file to poll
388+
store_identifier: The ID of the store
389+
metadata: The metadata to attach to the file
390+
poll_interval_ms: The interval between polls in milliseconds
391+
poll_timeout_ms: The maximum time to poll for in milliseconds
392+
Returns:
393+
The file object once it reaches a terminal state
394+
"""
395+
self.create(
396+
store_identifier=store_identifier, file_id=file_id, metadata=metadata, experimental=experimental, **kwargs
397+
)
398+
return self.poll(
399+
file_id,
400+
store_identifier=store_identifier,
401+
poll_interval_ms=poll_interval_ms,
402+
poll_timeout_ms=poll_timeout_ms,
403+
**kwargs,
404+
)
405+
406+
def upload(
407+
self,
408+
*,
409+
store_identifier: str,
410+
file: FileTypes,
411+
metadata: Optional[object] | Omit = omit,
412+
experimental: file_create_params.Experimental | Omit = omit,
413+
**kwargs: Any,
414+
) -> StoreFile:
415+
"""Upload a file to the `files` API and then attach it to the given store.
416+
Note the file will be asynchronously processed (you can use the alternative
417+
polling helper method to wait for processing to complete).
418+
"""
419+
file_obj = self._client.files.create(file=file, **kwargs)
420+
return self.create(
421+
store_identifier=store_identifier,
422+
file_id=file_obj.id,
423+
metadata=metadata,
424+
experimental=experimental,
425+
**kwargs,
426+
)
427+
428+
def upload_and_poll(
429+
self,
430+
*,
431+
store_identifier: str,
432+
file: FileTypes,
433+
metadata: Optional[object] | NotGiven = not_given,
434+
experimental: file_create_params.Experimental | Omit = omit,
435+
poll_interval_ms: int | NotGiven = not_given,
436+
poll_timeout_ms: float | NotGiven = not_given,
437+
**kwargs: Any,
438+
) -> StoreFile:
439+
"""Add a file to a store and poll until processing is complete."""
440+
file_obj = self._client.files.create(file=file, **kwargs)
441+
return self.create_and_poll(
442+
store_identifier=store_identifier,
443+
file_id=file_obj.id,
444+
metadata=metadata,
445+
experimental=experimental,
446+
poll_interval_ms=poll_interval_ms,
447+
poll_timeout_ms=poll_timeout_ms,
448+
**kwargs,
449+
)
450+
343451

344452
class AsyncFilesResource(AsyncAPIResource):
345453
@cached_property
@@ -656,6 +764,116 @@ async def search(
656764
cast_to=FileSearchResponse,
657765
)
658766

767+
async def poll(
768+
self,
769+
file_id: str,
770+
*,
771+
store_identifier: str,
772+
poll_interval_ms: int | NotGiven = not_given,
773+
poll_timeout_ms: float | NotGiven = not_given,
774+
**kwargs: Any,
775+
) -> StoreFile:
776+
"""
777+
Poll for a file's status until it reaches a terminal state.
778+
Args:
779+
file_id: The ID of the file to poll
780+
store_identifier: The ID of the store
781+
poll_interval_ms: The interval between polls in milliseconds
782+
poll_timeout_ms: The maximum time to poll for in milliseconds
783+
Returns:
784+
The file object once it reaches a terminal state
785+
"""
786+
polling_interval_ms = poll_interval_ms or 500
787+
polling_timeout_ms = poll_timeout_ms or None
788+
return await polling.poll_async(
789+
fn=functools.partial(self.retrieve, file_id, store_identifier=store_identifier, **kwargs),
790+
condition=lambda res: res.status == "completed" or res.status == "failed" or res.status == "cancelled",
791+
interval_seconds=polling_interval_ms / 1000,
792+
timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None,
793+
)
794+
795+
async def create_and_poll(
796+
self,
797+
file_id: str,
798+
*,
799+
store_identifier: str,
800+
metadata: Optional[object] | NotGiven = not_given,
801+
experimental: file_create_params.Experimental | Omit = omit,
802+
poll_interval_ms: int | NotGiven = not_given,
803+
poll_timeout_ms: float | NotGiven = not_given,
804+
**kwargs: Any,
805+
) -> StoreFile:
806+
"""
807+
Attach a file to the given vector store and wait for it to be processed.
808+
Args:
809+
file_id: The ID of the file to poll
810+
store_identifier: The ID of the store
811+
metadata: The metadata to attach to the file
812+
poll_interval_ms: The interval between polls in milliseconds
813+
poll_timeout_ms: The maximum time to poll for in milliseconds
814+
Returns:
815+
The file object once it reaches a terminal state
816+
"""
817+
await self.create(
818+
store_identifier=store_identifier,
819+
file_id=file_id,
820+
metadata=metadata,
821+
experimental=experimental,
822+
**kwargs,
823+
)
824+
return await self.poll(
825+
file_id,
826+
store_identifier=store_identifier,
827+
poll_interval_ms=poll_interval_ms,
828+
poll_timeout_ms=poll_timeout_ms,
829+
**kwargs,
830+
)
831+
832+
async def upload(
833+
self,
834+
*,
835+
store_identifier: str,
836+
file: FileTypes,
837+
metadata: Optional[object] | NotGiven = not_given,
838+
experimental: file_create_params.Experimental | Omit = omit,
839+
**kwargs: Any,
840+
) -> StoreFile:
841+
"""Upload a file to the `files` API and then attach it to the given vector store.
842+
Note the file will be asynchronously processed (you can use the alternative
843+
polling helper method to wait for processing to complete).
844+
"""
845+
file_obj = await self._client.files.create(file=file, **kwargs)
846+
return await self.create(
847+
store_identifier=store_identifier,
848+
file_id=file_obj.id,
849+
metadata=metadata,
850+
experimental=experimental,
851+
**kwargs,
852+
)
853+
854+
async def upload_and_poll(
855+
self,
856+
*,
857+
store_identifier: str,
858+
file: FileTypes,
859+
metadata: Optional[object] | NotGiven = not_given,
860+
experimental: file_create_params.Experimental | Omit = omit,
861+
poll_interval_ms: int | NotGiven = not_given,
862+
poll_timeout_ms: float | NotGiven = not_given,
863+
**kwargs: Any,
864+
) -> StoreFile:
865+
"""Add a file to a store and poll until processing is complete."""
866+
file_obj = await self._client.files.create(file=file, **kwargs)
867+
return await self.create_and_poll(
868+
store_identifier=store_identifier,
869+
file_id=file_obj.id,
870+
metadata=metadata,
871+
experimental=experimental,
872+
poll_interval_ms=poll_interval_ms,
873+
poll_timeout_ms=poll_timeout_ms,
874+
**kwargs,
875+
)
876+
659877

660878
class FilesResourceWithRawResponse:
661879
def __init__(self, files: FilesResource) -> None:

0 commit comments

Comments
 (0)