fix: remove tool result (#2327)

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

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)

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))

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)

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