fix: add logs in celery runner

This commit is contained in:
Javier Martinez
2026-07-16 15:20:25 +02:00
parent 05b983ef63
commit 469bb2b103
3 changed files with 54 additions and 28 deletions

View File

@@ -26,7 +26,12 @@ logger = logging.getLogger(__name__)
ignore_result=True,
)
async def tool_run_task(*, request_data: dict[str, Any]) -> None:
request = ToolExecutionRequest.model_validate(request_data)
try:
request = ToolExecutionRequest.model_validate(request_data)
except Exception:
logger.exception("Invalid tool execution request")
raise
try:
response = await execute_tool_request(
request,

View File

@@ -89,9 +89,13 @@ class LocalToolScheduler(BaseToolScheduler):
state_ctx: ChatState | None = None,
interceptors: list[ToolExecutionInterceptor] | None = None,
) -> ToolExecutionResponse:
return await execute_tool_request(
request, state_ctx=state_ctx, interceptors=interceptors
)
try:
return await execute_tool_request(
request, state_ctx=state_ctx, interceptors=interceptors
)
except Exception:
logger.exception("Local tool '%s' execution failed", request.tool_name)
raise
async def cancel(
self,

View File

@@ -1,11 +1,48 @@
import logging
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock
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 CeleryToolScheduler
from private_gpt.components.tools.tool_scheduler import (
CeleryToolScheduler,
LocalToolScheduler,
)
def tool_request() -> ToolExecutionRequest:
return ToolExecutionRequest.model_validate(
{
"tool_id": "tool-1",
"tool_name": "bash",
"tool_kwargs": {},
"tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
"context": {"correlation_id": "msg-1", "messages": []},
}
)
@pytest.mark.anyio
async def test_local_tool_scheduler_logs_execution_failure(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
execute_tool_request = AsyncMock(side_effect=RuntimeError("boom"))
monkeypatch.setattr(
"private_gpt.components.tools.tool_scheduler.execute_tool_request",
execute_tool_request,
)
with caplog.at_level(
logging.ERROR,
logger="private_gpt.components.tools.tool_scheduler",
), pytest.raises(RuntimeError, match="boom"):
await LocalToolScheduler().execute(tool_request())
assert "Local tool 'bash' execution failed" in caplog.text
assert "RuntimeError: boom" in caplog.text
@pytest.mark.anyio
@@ -16,18 +53,8 @@ async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None:
),
)
request = ToolExecutionRequest.model_validate(
{
"tool_id": "tool-1",
"tool_name": "bash",
"tool_kwargs": {},
"tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
"context": {"correlation_id": "msg-1", "messages": []},
}
)
with pytest.raises(NotImplementedError):
await scheduler.execute(request)
await scheduler.execute(tool_request())
@pytest.mark.anyio
@@ -43,17 +70,7 @@ async def test_celery_tool_scheduler_cancel_revokes_task(
),
)
request = ToolExecutionRequest.model_validate(
{
"tool_id": "tool-1",
"tool_name": "bash",
"tool_kwargs": {},
"tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
"context": {"correlation_id": "msg-1", "messages": []},
}
)
cancelled = await scheduler.cancel(request, task_id="task-abc")
cancelled = await scheduler.cancel(tool_request(), task_id="task-abc")
assert cancelled is True
celery_app.control.revoke.assert_called_once_with("task-abc", terminate=True)