diff --git a/src/openai/lib/_vector_stores.py b/src/openai/lib/_vector_stores.py index 7357a4908e..1451219305 100644 --- a/src/openai/lib/_vector_stores.py +++ b/src/openai/lib/_vector_stores.py @@ -1,9 +1,11 @@ from __future__ import annotations +import math +import time from typing import TYPE_CHECKING, Mapping from typing_extensions import assert_never -from .._types import Omit +from .._types import Omit, omit from .._utils import is_given from ..types.vector_stores.vector_store_file import VectorStoreFile from ..types.vector_stores.vector_store_file_batch import VectorStoreFileBatch @@ -21,19 +23,37 @@ def _get_poll_interval_ms(headers: Mapping[str, str]) -> int: return 1000 +def validate_max_wait_seconds(max_wait_seconds: float | Omit) -> None: + if is_given(max_wait_seconds) and (not math.isfinite(max_wait_seconds) or max_wait_seconds < 0): + raise ValueError("Expected a finite, non-negative value for `max_wait_seconds`") + + +def _remaining_poll_time(deadline: float, file_id: str) -> float: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for vector store file {file_id!r} to finish processing") + return remaining + + def poll_vector_store_file( resource: Files, file_id: str, *, vector_store_id: str, poll_interval_ms: int | Omit, + max_wait_seconds: float | Omit = omit, ) -> VectorStoreFile: """Poll a vector-store file using the caller's resource hooks.""" + validate_max_wait_seconds(max_wait_seconds) + deadline = time.monotonic() + max_wait_seconds if is_given(max_wait_seconds) else None headers: dict[str, str] = {"X-Stainless-Poll-Helper": "true"} if is_given(poll_interval_ms): headers["X-Stainless-Custom-Poll-Interval"] = str(poll_interval_ms) while True: + if deadline is not None: + _remaining_poll_time(deadline, file_id) + response = resource.with_raw_response.retrieve( file_id, vector_store_id=vector_store_id, @@ -45,7 +65,10 @@ def poll_vector_store_file( if not is_given(poll_interval_ms): poll_interval_ms = _get_poll_interval_ms(response.headers) - resource._sleep(poll_interval_ms / 1000) + sleep_seconds = poll_interval_ms / 1000 + if deadline is not None: + sleep_seconds = min(sleep_seconds, _remaining_poll_time(deadline, file_id)) + resource._sleep(sleep_seconds) elif file.status == "cancelled" or file.status == "completed" or file.status == "failed": return file else: @@ -61,13 +84,19 @@ async def async_poll_vector_store_file( *, vector_store_id: str, poll_interval_ms: int | Omit, + max_wait_seconds: float | Omit = omit, ) -> VectorStoreFile: """Poll a vector-store file using the caller's async resource hooks.""" + validate_max_wait_seconds(max_wait_seconds) + deadline = time.monotonic() + max_wait_seconds if is_given(max_wait_seconds) else None headers: dict[str, str] = {"X-Stainless-Poll-Helper": "true"} if is_given(poll_interval_ms): headers["X-Stainless-Custom-Poll-Interval"] = str(poll_interval_ms) while True: + if deadline is not None: + _remaining_poll_time(deadline, file_id) + response = await resource.with_raw_response.retrieve( file_id, vector_store_id=vector_store_id, @@ -79,7 +108,10 @@ async def async_poll_vector_store_file( if not is_given(poll_interval_ms): poll_interval_ms = _get_poll_interval_ms(response.headers) - await resource._sleep(poll_interval_ms / 1000) + sleep_seconds = poll_interval_ms / 1000 + if deadline is not None: + sleep_seconds = min(sleep_seconds, _remaining_poll_time(deadline, file_id)) + await resource._sleep(sleep_seconds) elif file.status == "cancelled" or file.status == "completed" or file.status == "failed": return file else: diff --git a/src/openai/resources/vector_stores/files.py b/src/openai/resources/vector_stores/files.py index 9ce285f4ba..4f18f96f39 100644 --- a/src/openai/resources/vector_stores/files.py +++ b/src/openai/resources/vector_stores/files.py @@ -18,6 +18,7 @@ from ..._base_client import AsyncPaginator, make_request_options from ...lib._vector_stores import ( poll_vector_store_file as _poll_vector_store_file, + validate_max_wait_seconds as _validate_max_wait_seconds, async_poll_vector_store_file as _async_poll_vector_store_file, ) from ...types.vector_stores import file_list_params, file_create_params, file_update_params @@ -335,6 +336,7 @@ def create_and_poll( vector_store_id: str, attributes: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, chunking_strategy: FileChunkingStrategyParam | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. @@ -343,7 +345,11 @@ def create_and_poll( extra_body: Body | None = None, timeout: float | httpx2.Timeout | None | NotGiven = not_given, ) -> VectorStoreFile: - """Attach a file to the given vector store and wait for it to be processed.""" + """Attach a file to the given vector store and wait for it to be processed. + + `max_wait_seconds` limits polling after attachment; see `poll` for details. + """ + _validate_max_wait_seconds(max_wait_seconds) self.create( vector_store_id=vector_store_id, file_id=file_id, @@ -359,6 +365,7 @@ def create_and_poll( file_id, vector_store_id=vector_store_id, poll_interval_ms=poll_interval_ms, + max_wait_seconds=max_wait_seconds, ) def poll( @@ -367,9 +374,16 @@ def poll( *, vector_store_id: str, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, ) -> VectorStoreFile: """Wait for the vector store file to finish processing. + `max_wait_seconds` must be finite and non-negative. When provided, it limits + polling and raises `TimeoutError` on expiry; zero prevents any polling requests. + Omit it to wait indefinitely. In-flight requests and retries use the client's + HTTP timeout settings and may finish after the polling deadline. Timing out + does not cancel processing on the server. + Note: this will return even if the file failed to process, you need to check file.last_error and file.status to handle these cases """ @@ -378,6 +392,7 @@ def poll( file_id, vector_store_id=vector_store_id, poll_interval_ms=poll_interval_ms, + max_wait_seconds=max_wait_seconds, ) def upload( @@ -402,15 +417,21 @@ def upload_and_poll( file: FileTypes, attributes: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, chunking_strategy: FileChunkingStrategyParam | Omit = omit, ) -> VectorStoreFile: - """Add a file to a vector store and poll until processing is complete.""" + """Add a file to a vector store and poll until processing is complete. + + `max_wait_seconds` limits polling after upload and attachment; see `poll` for details. + """ + _validate_max_wait_seconds(max_wait_seconds) 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, + max_wait_seconds=max_wait_seconds, attributes=attributes, ) @@ -767,6 +788,7 @@ async def create_and_poll( vector_store_id: str, attributes: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, chunking_strategy: FileChunkingStrategyParam | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. @@ -775,7 +797,11 @@ async def create_and_poll( extra_body: Body | None = None, timeout: float | httpx2.Timeout | None | NotGiven = not_given, ) -> VectorStoreFile: - """Attach a file to the given vector store and wait for it to be processed.""" + """Attach a file to the given vector store and wait for it to be processed. + + `max_wait_seconds` limits polling after attachment; see `poll` for details. + """ + _validate_max_wait_seconds(max_wait_seconds) await self.create( vector_store_id=vector_store_id, file_id=file_id, @@ -791,6 +817,7 @@ async def create_and_poll( file_id, vector_store_id=vector_store_id, poll_interval_ms=poll_interval_ms, + max_wait_seconds=max_wait_seconds, ) async def poll( @@ -799,9 +826,16 @@ async def poll( *, vector_store_id: str, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, ) -> VectorStoreFile: """Wait for the vector store file to finish processing. + `max_wait_seconds` must be finite and non-negative. When provided, it limits + polling and raises `TimeoutError` on expiry; zero prevents any polling requests. + Omit it to wait indefinitely. In-flight requests and retries use the client's + HTTP timeout settings and may finish after the polling deadline. Timing out + does not cancel processing on the server. + Note: this will return even if the file failed to process, you need to check file.last_error and file.status to handle these cases """ @@ -810,6 +844,7 @@ async def poll( file_id, vector_store_id=vector_store_id, poll_interval_ms=poll_interval_ms, + max_wait_seconds=max_wait_seconds, ) async def upload( @@ -836,14 +871,20 @@ async def upload_and_poll( file: FileTypes, attributes: Optional[Dict[str, Union[str, float, bool]]] | Omit = omit, poll_interval_ms: int | Omit = omit, + max_wait_seconds: float | Omit = omit, chunking_strategy: FileChunkingStrategyParam | Omit = omit, ) -> VectorStoreFile: - """Add a file to a vector store and poll until processing is complete.""" + """Add a file to a vector store and poll until processing is complete. + + `max_wait_seconds` limits polling after upload and attachment; see `poll` for details. + """ + _validate_max_wait_seconds(max_wait_seconds) file_obj = await self._client.files.create(file=file, purpose="assistants") return await self.create_and_poll( vector_store_id=vector_store_id, file_id=file_obj.id, poll_interval_ms=poll_interval_ms, + max_wait_seconds=max_wait_seconds, chunking_strategy=chunking_strategy, attributes=attributes, ) diff --git a/tests/lib/test_vector_store_file_poll_timeout.py b/tests/lib/test_vector_store_file_poll_timeout.py new file mode 100644 index 0000000000..0278bf9949 --- /dev/null +++ b/tests/lib/test_vector_store_file_poll_timeout.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +from typing import Any, cast +from unittest import mock + +import pytest + +from openai import OpenAI, AsyncOpenAI +from openai._types import omit +from openai._models import construct_type_unchecked +from openai.resources.vector_stores.files import Files, AsyncFiles +from openai.types.vector_stores.vector_store_file import VectorStoreFile + +FILE_ID = "file-synthetic" +VECTOR_STORE_ID = "vs-synthetic" + + +class Clock: + now: float = 0 + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> Clock: + clock = Clock() + monkeypatch.setattr("openai.lib._vector_stores.time", clock) + return clock + + +@pytest.fixture(params=[False, True], ids=["sync", "async"]) +def resource(request: pytest.FixtureRequest, client: OpenAI, async_client: AsyncOpenAI) -> Files | AsyncFiles: + return (async_client if request.param else client).vector_stores.files + + +def response(status: str = "in_progress") -> mock.Mock: + file = construct_type_unchecked(type_=VectorStoreFile, value={"id": FILE_ID, "status": status}) + return mock.Mock(headers={}, parse=mock.Mock(return_value=file)) + + +async def call(resource: Files | AsyncFiles, method: str = "poll", **kwargs: Any) -> VectorStoreFile: + params = {"vector_store_id": VECTOR_STORE_ID, **kwargs} + if method == "upload_and_poll": + params["file"] = ("probe.md", b"synthetic content") + else: + params["file_id"] = FILE_ID + result = getattr(resource, method)(**params) + if isinstance(resource, AsyncFiles): + result = await result + return cast(VectorStoreFile, result) + + +@pytest.mark.parametrize("interval", [omit, 250], ids=["server-hint", "explicit"]) +async def test_deadline_caps_sleep_and_prevents_next_request( + resource: Files | AsyncFiles, clock: Clock, interval: Any +) -> None: + pending = response() + pending.headers = {"openai-poll-after-ms": "10000"} + with ( + mock.patch.object(resource.with_raw_response, "retrieve", return_value=pending) as retrieve, + mock.patch.object(resource, "_sleep", side_effect=clock.sleep) as sleep, + ): + with pytest.raises(TimeoutError, match=FILE_ID): + await call(resource, max_wait_seconds=0.625, poll_interval_ms=interval) + assert clock.now == 0.625 + assert retrieve.call_count == (1 if interval is omit else 3) + assert [args.args[0] for args in sleep.call_args_list] == ([0.625] if interval is omit else [0.25, 0.25, 0.125]) + + +@pytest.mark.parametrize("request_seconds", [0.5, 3.0]) +async def test_request_time_counts_toward_deadline( + resource: Files | AsyncFiles, clock: Clock, request_seconds: float +) -> None: + def retrieve(*_args: Any, **_kwargs: Any) -> mock.Mock: + clock.now += request_seconds + return response() + + with ( + mock.patch.object(resource.with_raw_response, "retrieve", side_effect=retrieve) as get, + mock.patch.object(resource, "_sleep", side_effect=clock.sleep) as sleep, + ): + with pytest.raises(TimeoutError): + await call(resource, max_wait_seconds=2.5) + assert get.call_count == (2 if request_seconds == 0.5 else 1) + assert [args.args[0] for args in sleep.call_args_list] == ([1.0, 0.5] if request_seconds == 0.5 else []) + # Preserve the client's per-request timeout and retry policy. + assert all("timeout" not in args.kwargs for args in get.call_args_list) + + +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled"]) +async def test_terminal_result_is_preserved(resource: Files | AsyncFiles, clock: Clock, status: str) -> None: + terminal = response(status) + + def retrieve(*_args: Any, **_kwargs: Any) -> mock.Mock: + clock.now += 2 + return terminal + + with ( + mock.patch.object(resource.with_raw_response, "retrieve", side_effect=retrieve), + mock.patch.object(resource, "_sleep") as sleep, + ): + # A terminal response from an in-flight request is still useful after expiry. + assert await call(resource, max_wait_seconds=1) is terminal.parse.return_value + sleep.assert_not_called() + + +async def test_omitted_limit_preserves_unbounded_wait(resource: Files | AsyncFiles, clock: Clock) -> None: + terminal = response("completed") + with ( + mock.patch.object(resource.with_raw_response, "retrieve", side_effect=[response(), response(), terminal]), + mock.patch.object(resource, "_sleep", side_effect=clock.sleep), + ): + assert await call(resource, poll_interval_ms=3600000) is terminal.parse.return_value + assert clock.now == 7200 + + +@pytest.mark.parametrize("method", ["poll", "create_and_poll", "upload_and_poll"]) +@pytest.mark.parametrize("limit", [-1, float("nan"), float("inf"), -float("inf")]) +async def test_invalid_limit_fails_before_side_effects(resource: Files | AsyncFiles, method: str, limit: float) -> None: + with ( + mock.patch.object(resource.with_raw_response, "retrieve") as retrieve, + mock.patch.object(resource, "create") as attach, + mock.patch.object(resource._client.files, "create") as upload, + ): + with pytest.raises(ValueError, match="finite, non-negative"): + await call(resource, method, max_wait_seconds=limit) + retrieve.assert_not_called() + attach.assert_not_called() + upload.assert_not_called() + + +async def test_zero_limit_prevents_polling(resource: Files | AsyncFiles, clock: Clock) -> None: + with ( + mock.patch.object(resource.with_raw_response, "retrieve") as retrieve, + mock.patch.object(resource, "_sleep") as sleep, + ): + with pytest.raises(TimeoutError): + await call(resource, max_wait_seconds=0) + assert clock.now == 0 + retrieve.assert_not_called() + sleep.assert_not_called() + + +@pytest.mark.parametrize("method", ["create_and_poll", "upload_and_poll"]) +async def test_helpers_start_deadline_after_upload_and_attachment( + resource: Files | AsyncFiles, clock: Clock, method: str +) -> None: + def create(*_args: Any, **kwargs: Any) -> mock.Mock: + assert "max_wait_seconds" not in kwargs + clock.now += 10 + return mock.Mock(id=FILE_ID) + + with ( + mock.patch.object(resource._client.files, "create", side_effect=create) as upload, + mock.patch.object(resource, "create", side_effect=create) as attach, + mock.patch.object(resource.with_raw_response, "retrieve", return_value=response()) as retrieve, + mock.patch.object(resource, "_sleep", side_effect=clock.sleep), + ): + with pytest.raises(TimeoutError): + await call(resource, method, max_wait_seconds=0.5) + assert upload.call_count == (1 if method == "upload_and_poll" else 0) + attach.assert_called_once() + retrieve.assert_called_once() + assert clock.now == (20.5 if method == "upload_and_poll" else 10.5)