fix: tools endpoint

private_gpt/celery/dispatch.py

@@ -8,12 +8,15 @@ def dispatch_task(
     args: tuple[Any, ...] | list[Any] | None = None,
     kwargs: dict[str, Any] | None = None,
     task_id: str | None = None,
+    ignore_result: bool | None = None,
 ) -> Any:
     from private_gpt.celery.celery import celery_app
 
     options: dict[str, Any] = {"queue": queue}
     if task_id is not None:
         options["task_id"] = task_id
+    if ignore_result is not None:
+        options["ignore_result"] = ignore_result
 
     return celery_app.send_task(
         task_name,

private_gpt/celery/tasks/tools/tool_run_task.py

@@ -25,7 +25,7 @@ logger = logging.getLogger(__name__)
     base=StatefulBackgroundTask,
     ignore_result=True,
 )
-async def tool_run_task(*, request_data: dict[str, Any]) -> None:
+async def tool_run_task(*, request_data: dict[str, Any]) -> dict[str, Any]:
     try:
         request = ToolExecutionRequest.model_validate(request_data)
     except Exception:
@@ -37,7 +37,6 @@ async def tool_run_task(*, request_data: dict[str, Any]) -> None:
             request,
             interceptors=resolve_tool_execution_interceptors(request.interceptor_paths),
         )
-        await _notify_completion(request, response)
     except Exception as exc:
         logger.exception("Tool '%s' execution failed", request.tool_name)
         response = ToolExecutionResponse(
@@ -47,7 +46,9 @@ async def tool_run_task(*, request_data: dict[str, Any]) -> None:
             is_error=True,
             tool_message=request_error_message(request, str(exc)),
         )
-        await _notify_completion(request, response)
+
+    await _notify_completion(request, response)
+    return response.model_dump(mode="json")
 
 
 async def _notify_completion(
@@ -57,11 +58,8 @@ async def _notify_completion(
     correlation_id = request.context.get("correlation_id")
     if not correlation_id or not request.tool_id:
         return
-    try:
-        scheduler = get_global_injector(True).get(ToolSchedulerFactory).get()
-        await scheduler.complete(request, response)
-    except Exception:
-        logger.exception("Tool completion notify failed for %s", request.tool_id)
+    scheduler = get_global_injector(True).get(ToolSchedulerFactory).get()
+    await scheduler.complete(request, response)
 
 
 def request_error_message(

private_gpt/components/tools/tool_scheduler.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
+from asyncio import to_thread
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
@@ -27,6 +28,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+
 ToolSchedulerProvider = (
     type["BaseToolScheduler"] | Callable[[Injector], "BaseToolScheduler"]
 )
@@ -131,10 +133,23 @@ class CeleryToolScheduler(BaseToolScheduler):
         state_ctx: ChatState | None = None,
         interceptors: list[ToolExecutionInterceptor] | None = None,
     ) -> ToolExecutionResponse:
-        del request, state_ctx, interceptors
-        raise NotImplementedError(
-            "CeleryToolScheduler only supports async_execute() in resumable chat mode."
+        del state_ctx
+        request = request.model_copy(
+            update={"interceptor_paths": tool_execution_interceptor_paths(interceptors)}
         )
+        result = dispatch_task(
+            task_name=TOOL_TASK_NAME,
+            kwargs={"request_data": request.model_dump(mode="json")},
+            queue=self._settings.scheduler.tools.celery_queue,
+            ignore_result=False,
+        )
+        response_data = await to_thread(
+            result.get,
+            timeout=self._settings.scheduler.tools.callback_timeout_seconds,
+        )
+        from private_gpt.components.tools.remote_execution import ToolExecutionResponse
+
+        return ToolExecutionResponse.model_validate(response_data)
 
     async def cancel(
         self,

private_gpt/server/tools/tool_router.py

@@ -1,9 +1,15 @@
 from typing import Literal
+from uuid import uuid4
 
 from fastapi import APIRouter, Depends, Request
 from pydantic import BaseModel, Field
 
 from private_gpt.chat.extensions.context_filter import ContextFilter
+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 (
+    ToolSchedulerFactory,
+)
 from private_gpt.events.models import ResultContentBlockType
 from private_gpt.server.tools.tool_service import ToolService
 from private_gpt.server.utils.artifact_input import ArtifactType, SqlDatabaseArtifact
@@ -251,12 +257,12 @@ async def semantic_search(
     - Return ContentBlocks as MCP does
        - It returns Zylon custom content blocks, be careful if you use directly with MCP
     """
-    service: ToolService = request.state.injector.get(ToolService)
-    result = await service.semantic_search_tool(
-        body.query,
+    service = request.state.injector.get(ToolService)
+    tool = await service.build_semantic_search_tool(
         context_filter=body.context_filter,
         generate_citations=body.format == "citations",
     )
+    result = await _execute_tool(request, tool, {"query": body.query})
     return ToolResponse(content=result.content, is_error=result.is_error)
 
 
@@ -359,11 +365,11 @@ async def tabular_data_analysis(
     - Return ContentBlocks as MCP does
        - It returns Zylon custom content blocks, be careful if you use directly with MCP
     """
-    service: ToolService = request.state.injector.get(ToolService)
-    result = await service.tabular_data_analysis_tool(
-        body.query,
+    service = request.state.injector.get(ToolService)
+    tool = await service.build_tabular_data_analysis_tool(
         context_filter=body.context_filter,
     )
+    result = await _execute_tool(request, tool, {"query": body.query})
     return ToolResponse(content=result.content, is_error=result.is_error)
 
 
@@ -466,13 +472,15 @@ async def database_query(
     - Return ContentBlocks as MCP does
        - It returns Zylon custom content blocks, be careful if you use directly with MCP
     """
-    service: ToolService = request.state.injector.get(ToolService)
-    result = await service.database_query_tool(
-        body.query,
+    service = request.state.injector.get(ToolService)
+    tool = await service.build_database_query_tool(
         sql_artifacts=[
-            ctx for ctx in body.artifacts if isinstance(ctx, SqlDatabaseArtifact)
+            artifact
+            for artifact in body.artifacts
+            if isinstance(artifact, SqlDatabaseArtifact)
         ],
     )
+    result = await _execute_tool(request, tool, {"query": body.query})
     return ToolResponse(content=result.content, is_error=result.is_error)
 
 
@@ -550,8 +558,12 @@ async def web_fetch(
     - Return ContentBlocks as MCP does
        - It returns Zylon custom content blocks, be careful if you use directly with MCP
     """
-    service: ToolService = request.state.injector.get(ToolService)
-    result = await service.web_fetch_tool(body.url)
+    service = request.state.injector.get(ToolService)
+    result = await _execute_tool(
+        request,
+        service.build_web_fetch_tool(),
+        {"url": body.url},
+    )
     return ToolResponse(content=result.content, is_error=result.is_error)
 
 
@@ -651,6 +663,28 @@ async def web_search(
     Returns search results in a structured format with both source objects
     and text representations for easy consumption.
     """
-    service: ToolService = request.state.injector.get(ToolService)
-    result = await service.web_search_tool(body.query)
+    service = request.state.injector.get(ToolService)
+    result = await _execute_tool(
+        request,
+        await service.build_web_search_tool(),
+        {"query": body.query},
+    )
     return ToolResponse(content=result.content, is_error=result.is_error)
+
+
+async def _execute_tool(
+    request: Request,
+    tool: ToolSpec,
+    tool_kwargs: dict[str, object],
+) -> ToolResponse:
+    tool_name = tool.name or tool.get_original_tool_name()
+    scheduler = request.state.injector.get(ToolSchedulerFactory).get()
+    response = await scheduler.execute(
+        ToolExecutionRequest(
+            tool_id=f"api-{uuid4().hex}",
+            tool_name=tool_name,
+            tool_kwargs=tool_kwargs,
+            tool_spec=tool,
+        )
+    )
+    return ToolResponse(content=response.result_content, is_error=response.is_error)

private_gpt/server/tools/tool_service.py

@@ -5,6 +5,7 @@ from llama_index.core.base.llms.types import ChatMessage
 from pydantic import BaseModel
 
 from private_gpt.chat.extensions.context_filter import ContextFilter
+from private_gpt.components.chat.models.chat_config_models import ToolSpec
 from private_gpt.components.tools.tool_factories import (
     DatabaseQueryToolBuilderFactory,
     SemanticSearchToolBuilderFactory,
@@ -51,6 +52,40 @@ class ToolService:
         self._web_fetch_tool_builder_factory = web_fetch_tool_builder_factory
         self._web_search_tool_builder_factory = web_search_tool_builder_factory
 
+    async def build_semantic_search_tool(
+        self,
+        context_filter: ContextFilter,
+        generate_citations: bool = False,
+    ) -> ToolSpec:
+        token_limit = self.settings.chat.maximum_context_length or None
+        return await self._semantic_search_tool_builder_factory.create().build_tool(
+            context_filter=context_filter,
+            generate_citations=generate_citations,
+            token_limit=token_limit,
+        )
+
+    async def build_tabular_data_analysis_tool(
+        self,
+        context_filter: ContextFilter,
+    ) -> ToolSpec:
+        return await self._tabular_data_tool_builder_factory.create().build_tool(
+            context_filter=context_filter,
+        )
+
+    async def build_database_query_tool(
+        self,
+        sql_artifacts: list[SqlDatabaseArtifact],
+    ) -> ToolSpec:
+        return await self._database_query_tool_builder_factory.create().build_tool(
+            sql_artifacts=sql_artifacts,
+        )
+
+    def build_web_fetch_tool(self) -> ToolSpec:
+        return self._web_fetch_tool_builder_factory.create().build_tool()
+
+    async def build_web_search_tool(self) -> ToolSpec:
+        return await self._web_search_tool_builder_factory.create().build_tool()
+
     async def semantic_search_tool(
         self,
         query: str,

tests/arq/tasks/chat/test_callback.py

@@ -0,0 +1,69 @@
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from private_gpt.arq.tasks.chat.callback import resume_chat_callback
+from private_gpt.components.engines.chat.execution_scheduler import (
+    ChatExecutionSchedulerFactory,
+)
+from private_gpt.components.engines.chat.models.execution_hooks import ExecutionHooks
+from private_gpt.components.tools.remote_execution import (
+    ToolExecutionRequest,
+    ToolExecutionResponse,
+)
+from private_gpt.events.models import TextBlock
+
+
+@pytest.mark.anyio
+async def test_resume_chat_callback_sends_error_tool_result_to_scheduler(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    scheduler = MagicMock()
+    scheduler.callback = AsyncMock()
+    scheduler_factory = MagicMock()
+    scheduler_factory.get.return_value = scheduler
+    injector = MagicMock()
+    injector.get.return_value = scheduler_factory
+    monkeypatch.setattr(
+        "private_gpt.di.get_global_injector",
+        MagicMock(return_value=injector),
+    )
+    request = ToolExecutionRequest.model_validate(
+        {
+            "tool_id": "semantic-search-1",
+            "tool_name": "semantic_search",
+            "tool_kwargs": {},
+            "tool_spec": {
+                "name": "semantic_search",
+                "runtime": "server",
+                "input_schema": {},
+            },
+            "context": {"correlation_id": "chat-1"},
+            "hooks": ExecutionHooks(),
+        }
+    )
+    response = ToolExecutionResponse(
+        tool_name="semantic_search",
+        tool_id="semantic-search-1",
+        result_content=[TextBlock(text="query: Field required")],
+        is_error=True,
+        tool_message={
+            "role": "tool",
+            "content": "query: Field required",
+            "additional_kwargs": {
+                "tool_call_id": "semantic-search-1",
+                "tool_call_name": "semantic_search",
+                "tool_call_args": {},
+                "raw_output": "query: Field required",
+            },
+        },
+    )
+
+    await resume_chat_callback(request=request, response=response)
+
+    injector.get.assert_called_once_with(ChatExecutionSchedulerFactory)
+    scheduler.callback.assert_awaited_once_with(
+        execution_id="chat-1",
+        tool_id="semantic-search-1",
+        result=response.model_dump(mode="json"),
+    )

tests/arq/tasks/chat/test_resume.py

@@ -0,0 +1,35 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from private_gpt.arq.tasks.chat.resume import enqueue_tool_resume_job
+
+
+@pytest.mark.anyio
+async def test_enqueue_tool_resume_job_passes_error_result_as_arq_argument(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    enqueue_job = AsyncMock()
+    monkeypatch.setattr("private_gpt.arq.tasks.chat.resume.enqueue_job", enqueue_job)
+    result = {
+        "tool_name": "semantic_search",
+        "tool_id": "semantic-search-1",
+        "result_content": [{"type": "text", "text": "query: Field required"}],
+        "is_error": True,
+        "tool_message": {
+            "role": "tool",
+            "content": "query: Field required",
+            "additional_kwargs": {"tool_call_id": "semantic-search-1"},
+        },
+    }
+
+    await enqueue_tool_resume_job(
+        correlation_id="chat-1",
+        tool_id="semantic-search-1",
+        result=result,
+    )
+
+    enqueue_job.assert_awaited_once()
+    call = enqueue_job.await_args.kwargs
+    assert call["args"] == ("chat-1", "semantic-search-1", result)
+    assert call["correlation_id"] == "chat-1"

tests/celery/tasks/tools/test_tool_run_task.py

@@ -0,0 +1,58 @@
+import importlib
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from private_gpt.celery.tasks.tools.tool_run_task import _notify_completion
+from private_gpt.components.engines.chat.models.execution_hooks import ExecutionHooks
+from private_gpt.components.tools.remote_execution import (
+    ToolExecutionRequest,
+    ToolExecutionResponse,
+)
+from private_gpt.events.models import TextBlock
+
+
+@pytest.mark.anyio
+async def test_notify_completion_propagates_callback_failure(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    scheduler = MagicMock()
+    scheduler.complete = AsyncMock(side_effect=RuntimeError("ARQ enqueue failed"))
+    scheduler_factory = MagicMock()
+    scheduler_factory.get.return_value = scheduler
+    injector = MagicMock()
+    injector.get.return_value = scheduler_factory
+    task_module = importlib.import_module(
+        "private_gpt.celery.tasks.tools.tool_run_task"
+    )
+    monkeypatch.setattr(
+        task_module, "get_global_injector", MagicMock(return_value=injector)
+    )
+    request = ToolExecutionRequest.model_validate(
+        {
+            "tool_id": "semantic-search-1",
+            "tool_name": "semantic_search",
+            "tool_kwargs": {},
+            "tool_spec": {
+                "name": "semantic_search",
+                "runtime": "server",
+                "input_schema": {},
+            },
+            "context": {"correlation_id": "chat-1"},
+            "hooks": ExecutionHooks(),
+        }
+    )
+    response = ToolExecutionResponse(
+        tool_name="semantic_search",
+        tool_id="semantic-search-1",
+        result_content=[TextBlock(text="query: Field required")],
+        is_error=True,
+        tool_message={
+            "role": "tool",
+            "content": "query: Field required",
+            "additional_kwargs": {"tool_call_id": "semantic-search-1"},
+        },
+    )
+
+    with pytest.raises(RuntimeError, match="ARQ enqueue failed"):
+        await _notify_completion(request, response)

tests/components/tools/test_tool_scheduler.py

@@ -7,9 +7,11 @@ import pytest
 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 (
+    TOOL_TASK_NAME,
     CeleryToolScheduler,
     LocalToolScheduler,
 )
+from private_gpt.events.models import TextBlock
 
 
 def tool_request() -> ToolExecutionRequest:
@@ -49,15 +51,51 @@ async def test_local_tool_scheduler_logs_execution_failure(
 
 
 @pytest.mark.anyio
-async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None:
+async def test_celery_tool_scheduler_execute_dispatches_and_waits(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    async_result = MagicMock()
+    dispatch_task = MagicMock(return_value=async_result)
+    response_data = {
+        "tool_name": "bash",
+        "tool_id": "tool-1",
+        "result_content": [TextBlock(text="worker result").model_dump(mode="json")],
+        "is_error": False,
+        "tool_message": {
+            "role": "tool",
+            "content": "worker result",
+            "additional_kwargs": {"tool_call_id": "tool-1"},
+        },
+    }
+    to_thread = AsyncMock(return_value=response_data)
+    monkeypatch.setattr(
+        "private_gpt.components.tools.tool_scheduler.dispatch_task", dispatch_task
+    )
+    monkeypatch.setattr(
+        "private_gpt.components.tools.tool_scheduler.to_thread", to_thread
+    )
     scheduler = CeleryToolScheduler(
         settings=SimpleNamespace(
-            scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools"))
+            scheduler=SimpleNamespace(
+                tools=SimpleNamespace(
+                    celery_queue="tools",
+                    callback_timeout_seconds=42,
+                )
+            )
         ),
     )
+    request = tool_request()
 
-    with pytest.raises(NotImplementedError):
-        await scheduler.execute(tool_request())
+    response = await scheduler.execute(request)
+
+    assert response.result_content == [TextBlock(text="worker result")]
+    dispatch_task.assert_called_once_with(
+        task_name=TOOL_TASK_NAME,
+        kwargs={"request_data": request.model_dump(mode="json")},
+        queue="tools",
+        ignore_result=False,
+    )
+    to_thread.assert_awaited_once_with(async_result.get, timeout=42)
 
 
 @pytest.mark.anyio

tests/server/chat/test_chat_routes.py

@@ -274,6 +274,47 @@ async def test_chat_with_non_existent_artifact_context_fails_even_in_lazy_mode(
     assert tool_result.is_error
 
 
+@pytest.mark.anyio
+async def test_chat_semantic_search_without_query_returns_error_and_continues(
+    async_test_client: AsyncClient, injector: MockInjector
+) -> None:
+    await mock_llm(
+        injector,
+        deltas=[
+            [
+                ToolSelection(
+                    tool_id=SEMANTIC_SEARCH_TOOL_NAME,
+                    tool_name=SEMANTIC_SEARCH_TOOL_NAME,
+                    tool_kwargs={},
+                )
+            ],
+            ["I could not search because the query was missing."],
+        ],
+    )
+
+    body = ChatBody(
+        messages=[MessageInput(content="Lorem ipsum", role="user")],
+        stream=False,
+        tools=tools(use_context=True),
+        tool_choice=tool_choice(use_context=True, validation_mode="lazy"),
+        tool_context=tool_context(use_context=True),
+    )
+
+    result = await async_test_client.post("/v1/messages", json=body.model_dump())
+
+    assert result.status_code == 200
+    completion = Message.model_validate(result.json())
+    tool_result = next(
+        block for block in completion.content if isinstance(block, ToolResultBlock)
+    )
+    assert tool_result.is_error
+    assert any(
+        isinstance(block, TextBlock)
+        and block.text == "I could not search because the query was missing."
+        for block in completion.content
+    )
+
+
 @pytest.mark.anyio
 async def test_chat_with_metadata_context(
     async_test_client: AsyncClient, injector: MockInjector