mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-16 17:00:12 +00:00
fix: add logs in celery runner
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user