From c526601b1d8dc4e31edccc434a95d4531e6b2faa Mon Sep 17 00:00:00 2001
From: Javier Martinez <javiermartinezalvarez98@gmail.com>
Date: Wed, 8 Jul 2026 12:37:55 +0200
Subject: [PATCH] fix: mypy

---
 private_gpt/celery/base.py                      | 7 +++----
 private_gpt/celery/tasks/tools/tool_run_task.py | 9 +++++++--
 private_gpt/components/tools/tool_scheduler.py  | 6 ++++--
 tests/server/chat/test_chat_async_worker.py     | 4 +++-
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/private_gpt/celery/base.py b/private_gpt/celery/base.py
index 6d21cdc0..4643c9a4 100644
--- a/private_gpt/celery/base.py
+++ b/private_gpt/celery/base.py
@@ -387,9 +387,9 @@ class StatefulBackgroundTask(_BackgroundTask):
 
         if not loop.is_closed():
             with contextlib.suppress(Exception):
-                asyncio.run_coroutine_threadsafe(
-                    cls._shutdown_async(), loop
-                ).result(timeout=10)
+                asyncio.run_coroutine_threadsafe(cls._shutdown_async(), loop).result(
+                    timeout=10
+                )
             with contextlib.suppress(RuntimeError):
                 loop.call_soon_threadsafe(loop.stop)
 
@@ -408,4 +408,3 @@ class StatefulBackgroundTask(_BackgroundTask):
             if self.rollback_fn:
                 self.rollback_fn(*args, **kwargs)
             raise e
-
diff --git a/private_gpt/celery/tasks/tools/tool_run_task.py b/private_gpt/celery/tasks/tools/tool_run_task.py
index 01f68db2..18cdc757 100644
--- a/private_gpt/celery/tasks/tools/tool_run_task.py
+++ b/private_gpt/celery/tasks/tools/tool_run_task.py
@@ -3,13 +3,18 @@
 When ``tool_scheduler.mode`` is ``"celery"``, the chat worker dispatches tool
 calls to this task on the ``tools`` queue instead of running them in-process.
 """
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from private_gpt.celery.base import StatefulBackgroundTask
 from private_gpt.celery.celery import celery_app
 from private_gpt.di import get_global_injector
 
+if TYPE_CHECKING:
+    from private_gpt.server.tools.tool_service import ToolService
+
 logger = logging.getLogger(__name__)
 
 
@@ -61,7 +66,7 @@ async def tool_run_task(
 
 
 async def _execute_by_name(
-    tool_service: "ToolService",
+    tool_service: ToolService,
     tool_name: str,
     kwargs: dict[str, Any],
 ) -> Any:
diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py
index 236105ea..3c8009c7 100644
--- a/private_gpt/components/tools/tool_scheduler.py
+++ b/private_gpt/components/tools/tool_scheduler.py
@@ -6,11 +6,13 @@ import logging
 import time
 import uuid
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from injector import inject, singleton
 
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
 from private_gpt.components.tools.tool_names import (
     BASH_TOOL_NAME,
     DATABASE_QUERY_TOOL_NAME,
diff --git a/tests/server/chat/test_chat_async_worker.py b/tests/server/chat/test_chat_async_worker.py
index 8ae8481b..95ba10d3 100644
--- a/tests/server/chat/test_chat_async_worker.py
+++ b/tests/server/chat/test_chat_async_worker.py
@@ -35,7 +35,9 @@ def _chat_body() -> ChatBody:
 
 
 @pytest.mark.anyio
-async def test_worker_handoff_sends_json_safe_body(monkeypatch: pytest.MonkeyPatch) -> None:
+async def test_worker_handoff_sends_json_safe_body(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     stream_service = InMemoryStreamService()
     sent: dict[str, Any] = {}