diff --git a/private_gpt/celery/dispatch.py b/private_gpt/celery/dispatch.py index e25d2548..14e3ad92 100644 --- a/private_gpt/celery/dispatch.py +++ b/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, diff --git a/private_gpt/celery/tasks/tools/tool_run_task.py b/private_gpt/celery/tasks/tools/tool_run_task.py index a377279d..e295e252 100644 --- a/private_gpt/celery/tasks/tools/tool_run_task.py +++ b/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( diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py index e34f8326..f66ca004 100644 --- a/private_gpt/components/tools/tool_scheduler.py +++ b/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, diff --git a/private_gpt/server/tools/tool_router.py b/private_gpt/server/tools/tool_router.py index a1c43cff..f7d80879 100644 --- a/private_gpt/server/tools/tool_router.py +++ b/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) diff --git a/private_gpt/server/tools/tool_service.py b/private_gpt/server/tools/tool_service.py index 5d5f3482..786114fe 100644 --- a/private_gpt/server/tools/tool_service.py +++ b/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, diff --git a/tests/arq/tasks/chat/test_callback.py b/tests/arq/tasks/chat/test_callback.py new file mode 100644 index 00000000..7a2b6a6c --- /dev/null +++ b/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"), + ) diff --git a/tests/arq/tasks/chat/test_resume.py b/tests/arq/tasks/chat/test_resume.py new file mode 100644 index 00000000..ad29ca26 --- /dev/null +++ b/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" diff --git a/tests/celery/tasks/tools/test_tool_run_task.py b/tests/celery/tasks/tools/test_tool_run_task.py new file mode 100644 index 00000000..3522a4a9 --- /dev/null +++ b/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) diff --git a/tests/components/tools/test_tool_scheduler.py b/tests/components/tools/test_tool_scheduler.py index 0b492fe1..a84d3e2b 100644 --- a/tests/components/tools/test_tool_scheduler.py +++ b/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 diff --git a/tests/server/chat/test_chat_routes.py b/tests/server/chat/test_chat_routes.py index d73c5606..2b33ba34 100644 --- a/tests/server/chat/test_chat_routes.py +++ b/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