From a5c6d241b27f9b66ab18d91b07d23f76b978befd Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Wed, 5 Aug 2026 09:03:26 +0200 Subject: [PATCH] fix: remove tool result (#2327) --- private_gpt/celery/result.py | 31 +++++++++++++ .../ingestion/ingestion_scheduler.py | 34 +++++--------- .../components/tools/tool_scheduler.py | 6 ++- tests/celery/test_result.py | 46 +++++++++++++++++++ tests/components/tools/test_tool_scheduler.py | 3 +- 5 files changed, 94 insertions(+), 26 deletions(-) create mode 100644 private_gpt/celery/result.py create mode 100644 tests/celery/test_result.py diff --git a/private_gpt/celery/result.py b/private_gpt/celery/result.py new file mode 100644 index 00000000..9c775870 --- /dev/null +++ b/private_gpt/celery/result.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +from celery.exceptions import TimeoutError as CeleryTimeoutError + +if TYPE_CHECKING: + from celery.result import AsyncResult + + +def wait_for_celery_result( + result: AsyncResult[Any], + timeout: float | None = None, + poll_interval: float = 0.1, +) -> Any: + deadline = time.monotonic() + timeout if timeout is not None else None + + while not result.ready(): + if deadline is None: + time.sleep(poll_interval) + continue + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise CeleryTimeoutError(f"Task {result.id} timed out") + time.sleep(min(poll_interval, remaining)) + + if result.failed(): + raise result.result + return result.result diff --git a/private_gpt/components/ingestion/ingestion_scheduler.py b/private_gpt/components/ingestion/ingestion_scheduler.py index 20a6d3e5..d262d882 100644 --- a/private_gpt/components/ingestion/ingestion_scheduler.py +++ b/private_gpt/components/ingestion/ingestion_scheduler.py @@ -4,17 +4,17 @@ import asyncio import logging from abc import ABC, abstractmethod from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from injector import Injector, inject, singleton +from private_gpt.celery.result import wait_for_celery_result from private_gpt.components.ingest.utils import get_extension, get_file_name from private_gpt.components.storage.s3_helper import S3Helper from private_gpt.server.ingest.ingest_service import IngestService from private_gpt.settings.settings import Settings, settings if TYPE_CHECKING: - from celery.result import AsyncResult from llama_index.core.schema import BaseNode from private_gpt.server.ingest.ingest_router import ( @@ -467,7 +467,6 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): def ingest(self, ingest_body: IngestBody) -> IngestResponse: """Parse then store synchronously, blocking until both complete.""" - import time import uuid from private_gpt.server.ingest.ingest_router import ( @@ -507,8 +506,6 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): async_body.ingest_body.input = UriArtifact(value=s3_url) # 1. Parse on worker, wait — fast step. - from celery.result import AsyncResult - from private_gpt.celery.dispatch import dispatch_task from private_gpt.celery.tasks.ingestion.extraction_tasks import PARSE_TASK_NAME @@ -517,19 +514,14 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): args=(async_body,), queue=config.scheduler.ingestion.celery_queue, ) - while not parse_result.ready(): - time.sleep(0.1) - if parse_result.failed(): - raise parse_result.result + parse_result_value = wait_for_celery_result(parse_result) # 2. parse_task returns the store_vectors task_id; poll it. - assert isinstance(parse_result.result, str) - store_result: AsyncResult[Any] = AsyncResult(parse_result.result) - while not store_result.ready(): - time.sleep(0.1) - if store_result.failed(): - raise store_result.result - return IngestResponse.model_validate(store_result.result) + assert isinstance(parse_result_value, str) + from celery.result import AsyncResult + + store_result = AsyncResult(parse_result_value) + return IngestResponse.model_validate(wait_for_celery_result(store_result)) async def ingest_for_request(self, ingest_body: IngestBody) -> IngestResponse: @@ -581,7 +573,6 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): def bytes_to_text(self, raw: bytes, ext: str) -> str: """Dispatch parse_task in parse-only mode on the worker, return text.""" import base64 - import time import uuid from private_gpt.celery.dispatch import dispatch_task @@ -604,12 +595,9 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): kwargs={"dispatch_store": False}, queue=config.scheduler.ingestion.celery_queue, ) - while not result.ready(): - time.sleep(0.1) - if result.failed(): - raise result.result - assert isinstance(result.result, str) - return result.result + result_value = wait_for_celery_result(result) + assert isinstance(result_value, str) + return result_value register_ingestion_scheduler("local", LocalIngestionScheduler) diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py index 5868e1af..11f60594 100644 --- a/private_gpt/components/tools/tool_scheduler.py +++ b/private_gpt/components/tools/tool_scheduler.py @@ -11,6 +11,7 @@ from celery.exceptions import TimeoutError as CeleryTimeoutError from injector import Injector, inject, singleton from private_gpt.celery.dispatch import dispatch_task +from private_gpt.celery.result import wait_for_celery_result from private_gpt.components.tools.remote_execution import ( execute_tool_request, invoke_execution_hook, @@ -173,8 +174,9 @@ class CeleryToolScheduler(BaseToolScheduler): ) try: response_data = await to_thread( - result.get, - timeout=self._settings.scheduler.tools.callback_timeout_seconds, + wait_for_celery_result, + result, + self._settings.scheduler.tools.callback_timeout_seconds, ) except (CancelledError, CeleryTimeoutError): await self.cancel_task(str(result.id)) diff --git a/tests/celery/test_result.py b/tests/celery/test_result.py new file mode 100644 index 00000000..4b2dadc3 --- /dev/null +++ b/tests/celery/test_result.py @@ -0,0 +1,46 @@ +from unittest.mock import MagicMock + +import pytest +from celery.exceptions import TimeoutError as CeleryTimeoutError + +from private_gpt.celery.result import wait_for_celery_result + + +def test_wait_for_celery_result_polls_until_ready( + monkeypatch: pytest.MonkeyPatch, +) -> None: + result = MagicMock(id="task-1") + result.ready.side_effect = [False, True] + result.failed.return_value = False + result.result = {"result": "worker result"} + sleep = MagicMock() + monkeypatch.setattr("private_gpt.celery.result.time.sleep", sleep) + + response = wait_for_celery_result(result, timeout=42) + + assert response == {"result": "worker result"} + assert result.ready.call_count == 2 + result.failed.assert_called_once_with() + sleep.assert_called_once_with(0.1) + + +def test_wait_for_celery_result_raises_worker_exception() -> None: + result = MagicMock(id="task-1") + result.ready.return_value = True + result.failed.return_value = True + result.result = ValueError("worker failed") + + with pytest.raises(ValueError, match="worker failed"): + wait_for_celery_result(result) + + +def test_wait_for_celery_result_enforces_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + result = MagicMock(id="task-1") + result.ready.return_value = False + monotonic = MagicMock(side_effect=[0.0, 1.0]) + monkeypatch.setattr("private_gpt.celery.result.time.monotonic", monotonic) + + with pytest.raises(CeleryTimeoutError, match="task-1"): + wait_for_celery_result(result, timeout=0.5) diff --git a/tests/components/tools/test_tool_scheduler.py b/tests/components/tools/test_tool_scheduler.py index d37c52e4..e24830e6 100644 --- a/tests/components/tools/test_tool_scheduler.py +++ b/tests/components/tools/test_tool_scheduler.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from private_gpt.celery.result import wait_for_celery_result from private_gpt.components.chat.models.chat_config_models import ToolSpec from private_gpt.components.tools.remote_execution import ToolExecutionRequest from private_gpt.components.tools.tool_scheduler import ( @@ -96,7 +97,7 @@ async def test_celery_tool_scheduler_execute_dispatches_and_waits( queue="tools", ignore_result=False, ) - to_thread.assert_awaited_once_with(async_result.get, timeout=42) + to_thread.assert_awaited_once_with(wait_for_celery_result, async_result, 42) @pytest.mark.anyio