Skip to content

chore: add the stores helper functions #245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
joeldierkes merged 1 commit into from
Oct 1, 2025
Merged

chore: add the stores helper functions #245

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
222 changes: 220 additions & 2 deletions src/mixedbread/resources/stores/files.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@

from __future__ import annotations

from typing import List, Union, Iterable, Optional
import functools
from typing import Any, List, Union, Iterable, Optional

import httpx

from ..._types import Body, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
from ...lib import polling
from ..._types import Body, FileTypes, Omit, Query, Headers, NotGiven, SequenceNotStr, omit, not_given
from ..._utils import maybe_transform, async_maybe_transform
from ..._compat import cached_property
from ..._resource import SyncAPIResource, AsyncAPIResource
Expand Down Expand Up @@ -340,6 +342,112 @@ def search(
cast_to=FileSearchResponse,
)

def poll(
self,
file_id: str,
*,
store_identifier: str,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""
Poll for a file's status until it reaches a terminal state.
Args:
file_id: The ID of the file to poll
store_identifier: The ID of the store
poll_interval_ms: The interval between polls in milliseconds
poll_timeout_ms: The maximum time to poll for in milliseconds
Returns:
The file object once it reaches a terminal state
"""
polling_interval_ms = poll_interval_ms or 500
polling_timeout_ms = poll_timeout_ms or None
return polling.poll(
fn=functools.partial(self.retrieve, file_id, store_identifier=store_identifier, **kwargs),
condition=lambda res: res.status == "completed" or res.status == "failed" or res.status == "cancelled",
interval_seconds=polling_interval_ms / 1000,
timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None,

@cursor cursor Bot Oct 1, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Poll Parameters Misinterpret Zero Values

The poll_interval_ms and poll_timeout_ms parameters in the poll methods (sync and async) incorrectly treat explicit 0 values as falsy. This causes 0 to be overridden by the default 500ms interval or None (no timeout), preventing users from setting an immediate interval or timeout and potentially leading to unexpected infinite polling.

Additional Locations (1)

Fix in Cursor Fix in Web

)

def create_and_poll(
self,
file_id: str,
*,
store_identifier: str,
metadata: Optional[object] | NotGiven = not_given,
experimental: file_create_params.Experimental | Omit = omit,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""
Attach a file to the given store and wait for it to be processed.
Args:
file_id: The ID of the file to poll
store_identifier: The ID of the store
metadata: The metadata to attach to the file
poll_interval_ms: The interval between polls in milliseconds
poll_timeout_ms: The maximum time to poll for in milliseconds
Returns:
The file object once it reaches a terminal state
"""
self.create(
store_identifier=store_identifier, file_id=file_id, metadata=metadata, experimental=experimental, **kwargs

@cursor cursor Bot Oct 1, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Metadata Parameter Type Mismatch

The new create_and_poll, upload, and upload_and_poll methods (and their async versions) define the metadata parameter using NotGiven. However, they pass this parameter to the underlying create method, which expects Omit. This type annotation inconsistency may lead to type checking errors or unexpected runtime behavior.

Additional Locations (4)

Fix in Cursor Fix in Web

)
return self.poll(
file_id,
store_identifier=store_identifier,
poll_interval_ms=poll_interval_ms,
poll_timeout_ms=poll_timeout_ms,
**kwargs,
)

def upload(
self,
*,
store_identifier: str,
file: FileTypes,
metadata: Optional[object] | Omit = omit,
experimental: file_create_params.Experimental | Omit = omit,
**kwargs: Any,
) -> StoreFile:
"""Upload a file to the `files` API and then attach it to the given store.
Note the file will be asynchronously processed (you can use the alternative
polling helper method to wait for processing to complete).
"""
file_obj = self._client.files.create(file=file, **kwargs)
return self.create(
store_identifier=store_identifier,
file_id=file_obj.id,
metadata=metadata,
experimental=experimental,
**kwargs,
)

def upload_and_poll(
self,
*,
store_identifier: str,
file: FileTypes,
metadata: Optional[object] | NotGiven = not_given,
experimental: file_create_params.Experimental | Omit = omit,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""Add a file to a store and poll until processing is complete."""
file_obj = self._client.files.create(file=file, **kwargs)
return self.create_and_poll(
store_identifier=store_identifier,
file_id=file_obj.id,
metadata=metadata,
experimental=experimental,
poll_interval_ms=poll_interval_ms,
poll_timeout_ms=poll_timeout_ms,
**kwargs,
)


class AsyncFilesResource(AsyncAPIResource):
@cached_property
Expand Down Expand Up @@ -656,6 +764,116 @@ async def search(
cast_to=FileSearchResponse,
)

async def poll(
self,
file_id: str,
*,
store_identifier: str,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""
Poll for a file's status until it reaches a terminal state.
Args:
file_id: The ID of the file to poll
store_identifier: The ID of the store
poll_interval_ms: The interval between polls in milliseconds
poll_timeout_ms: The maximum time to poll for in milliseconds
Returns:
The file object once it reaches a terminal state
"""
polling_interval_ms = poll_interval_ms or 500
polling_timeout_ms = poll_timeout_ms or None
return await polling.poll_async(
fn=functools.partial(self.retrieve, file_id, store_identifier=store_identifier, **kwargs),
condition=lambda res: res.status == "completed" or res.status == "failed" or res.status == "cancelled",
interval_seconds=polling_interval_ms / 1000,
timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None,
)

async def create_and_poll(
self,
file_id: str,
*,
store_identifier: str,
metadata: Optional[object] | NotGiven = not_given,
experimental: file_create_params.Experimental | Omit = omit,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""
Attach a file to the given vector store and wait for it to be processed.
Args:
file_id: The ID of the file to poll
store_identifier: The ID of the store
metadata: The metadata to attach to the file
poll_interval_ms: The interval between polls in milliseconds
poll_timeout_ms: The maximum time to poll for in milliseconds
Returns:
The file object once it reaches a terminal state
"""
await self.create(
store_identifier=store_identifier,
file_id=file_id,
metadata=metadata,
experimental=experimental,
**kwargs,
)
return await self.poll(
file_id,
store_identifier=store_identifier,
poll_interval_ms=poll_interval_ms,
poll_timeout_ms=poll_timeout_ms,
**kwargs,
)

async def upload(
self,
*,
store_identifier: str,
file: FileTypes,
metadata: Optional[object] | NotGiven = not_given,

@cursor cursor Bot Oct 1, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Sync/Async upload Methods Have Inconsistent Metadata Handling

The metadata parameter in the upload methods has an inconsistent type annotation and default value between the sync and async versions. One uses Optional[object] | Omit = omit while the other uses Optional[object] | NotGiven = not_given. This also creates a type mismatch when upload passes metadata to the create method, which expects object | Omit.

Fix in Cursor Fix in Web

experimental: file_create_params.Experimental | Omit = omit,
**kwargs: Any,
) -> StoreFile:
"""Upload a file to the `files` API and then attach it to the given vector store.
Note the file will be asynchronously processed (you can use the alternative
polling helper method to wait for processing to complete).
"""
file_obj = await self._client.files.create(file=file, **kwargs)
return await self.create(
store_identifier=store_identifier,
file_id=file_obj.id,
metadata=metadata,
experimental=experimental,
**kwargs,
)

async def upload_and_poll(
self,
*,
store_identifier: str,
file: FileTypes,
metadata: Optional[object] | NotGiven = not_given,
experimental: file_create_params.Experimental | Omit = omit,
poll_interval_ms: int | NotGiven = not_given,
poll_timeout_ms: float | NotGiven = not_given,
**kwargs: Any,
) -> StoreFile:
"""Add a file to a store and poll until processing is complete."""
file_obj = await self._client.files.create(file=file, **kwargs)
return await self.create_and_poll(
store_identifier=store_identifier,
file_id=file_obj.id,
metadata=metadata,
experimental=experimental,
poll_interval_ms=poll_interval_ms,
poll_timeout_ms=poll_timeout_ms,
**kwargs,
)


class FilesResourceWithRawResponse:
def __init__(self, files: FilesResource) -> None:
Expand Down
Loading