diff --git a/private_gpt/celery/callback.py b/private_gpt/celery/callback.py index c1cf1410..a6a4c2b4 100644 --- a/private_gpt/celery/callback.py +++ b/private_gpt/celery/callback.py @@ -18,6 +18,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _callback_task_name(task: Any) -> str: + name: str = getattr(task, "callback_task_name", task.name) + return name + + def _publish_callback( exchange: str, routing_key: str, @@ -47,18 +52,20 @@ def run_callback( final = False if state == states.SUCCESS: + task_name = _callback_task_name(task) async_response = AsyncResponse( data=result, - type=f"pgpt.{task.name}.done", + type=f"pgpt.{task_name}.done", error=None, callback_properties=callback.properties, ) routing_key = callback_amqp.routing_key_done or async_response.type final = True elif state == custom_states.PROGRESS: + task_name = _callback_task_name(task) async_response = AsyncResponse( data=result.model_dump(), - type=f"pgpt.{task.name}.progress", + type=f"pgpt.{task_name}.progress", error=None, callback_properties=callback.properties, ) @@ -77,7 +84,7 @@ def run_callback( async_response = AsyncResponse( data=None, - type=f"pgpt.{task.name}.error", + type=f"pgpt.{_callback_task_name(task)}.error", error=error_result.dict(), callback_properties=callback.properties, ) diff --git a/private_gpt/celery/dispatch.py b/private_gpt/celery/dispatch.py new file mode 100644 index 00000000..e25d2548 --- /dev/null +++ b/private_gpt/celery/dispatch.py @@ -0,0 +1,23 @@ +from typing import Any + + +def dispatch_task( + *, + task_name: str, + queue: str, + args: tuple[Any, ...] | list[Any] | None = None, + kwargs: dict[str, Any] | None = None, + task_id: str | 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 + + return celery_app.send_task( + task_name, + args=args, + kwargs=kwargs, + **options, + ) diff --git a/private_gpt/celery/task_helper.py b/private_gpt/celery/task_helper.py index 72274019..41865bc0 100644 --- a/private_gpt/celery/task_helper.py +++ b/private_gpt/celery/task_helper.py @@ -69,7 +69,11 @@ class IngestionTaskHelper: def is_ingestion_cancel_task_scheduled( celery_app: Any, collection: str, artifact: str ) -> bool: - for task in find_tasks(celery_app, task_name="delete_ingested_task"): + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + + for task in find_tasks(celery_app, task_name=DELETE_INGESTED_TASK_NAME): if not task.args or not isinstance( task.args[0], DeleteIngestedDocumentAsyncBody ): @@ -86,9 +90,12 @@ class IngestionTaskHelper: @staticmethod def revoke_ingestion_task(celery_app: Any, collection: str, artifact: str) -> bool: + from private_gpt.celery.tasks.ingestion.extraction_tasks import ( + VECTOR_INDEX_TASK_NAME, + ) from private_gpt.server.ingest.ingest_router import IngestAsyncBody - for task in find_tasks(celery_app, task_name="vector_index_task"): + for task in find_tasks(celery_app, task_name=VECTOR_INDEX_TASK_NAME): if not task.args or not isinstance(task.args[0], IngestAsyncBody): continue @@ -97,7 +104,6 @@ class IngestionTaskHelper: task_body.ingest_body.collection == collection and task_body.ingest_body.artifact == artifact ): - revoke_task(celery_app, task.task_id) return True @@ -105,7 +111,11 @@ class IngestionTaskHelper: @staticmethod def revoke_deletion_task(celery_app: Any, collection: str, artifact: str) -> bool: - for task in find_tasks(celery_app, task_name="delete_ingested_task"): + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + + for task in find_tasks(celery_app, task_name=DELETE_INGESTED_TASK_NAME): if not task.args or not isinstance( task.args[0], DeleteIngestedDocumentAsyncBody ): diff --git a/private_gpt/celery/tasks/chat/chat_run_task.py b/private_gpt/celery/tasks/chat/chat_run_task.py index 46c020d0..034b0b36 100644 --- a/private_gpt/celery/tasks/chat/chat_run_task.py +++ b/private_gpt/celery/tasks/chat/chat_run_task.py @@ -7,7 +7,6 @@ API becomes a pure Redis -> SSE proxy and its event loop never contends with CPU-bound work (LLM calls, tools, semantic search, retrieval, tokenization). """ import logging -from typing import Any from private_gpt.celery.base import StatefulBackgroundTask from private_gpt.celery.celery import celery_app @@ -26,19 +25,17 @@ logger = logging.getLogger(__name__) base=StatefulBackgroundTask, ) async def chat_run_task( - body: dict[str, Any], + body: ChatBody, correlation_id: str, stream_type: str, - metadata: dict[str, Any], + metadata: dict[str, object], ) -> None: """Run a chat completion loop and push events to the stream. Args: - body: JSON-safe :class:`ChatBody` data received by the API. It is re-mapped - to a :class:`ChatRequest` inside the worker so that tool - implementations and output schemas are built from the worker's - warm dependency injector (no closures or dynamic classes cross - the process boundary). + body: :class:`ChatBody` received by the API. It is re-mapped to a + :class:`ChatRequest` inside the worker so that tool implementations + and output schemas are built from the worker's warm dependency injector. correlation_id: The Redis stream correlation ID already created by the API via ``StreamService.create_stream``. stream_type: The stream type (e.g. ``"chat_completion"``). @@ -57,8 +54,7 @@ async def chat_run_task( event_handler = StreamingEventHandler() try: - chat_body = ChatBody.model_validate(body) - request = await chat_request_mapper.create_request_from_body(chat_body) + request = await chat_request_mapper.create_request_from_body(body) completion_gen = await chat_service.stream_chat(request) except Exception as e: error_event = event_handler.error_event(correlation_id, e) diff --git a/private_gpt/celery/tasks/ingestion/delete_tasks.py b/private_gpt/celery/tasks/ingestion/delete_tasks.py index 5b6f595d..c7fbeed8 100644 --- a/private_gpt/celery/tasks/ingestion/delete_tasks.py +++ b/private_gpt/celery/tasks/ingestion/delete_tasks.py @@ -16,9 +16,12 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG if settings().server.debug_mode else logging.INFO) +DELETE_INGESTED_TASK_NAME = "private_gpt.ingestion.delete" +DELETE_INGESTED_CALLBACK_TASK_NAME = "delete_ingested_task" + @celery_app.task( - name="delete_ingested_task", + name=DELETE_INGESTED_TASK_NAME, base=StatelessBackgroundTask, # Retry on ValueError and IndexNotReadyException. # ValueError is thrown when the index is not initialized @@ -51,3 +54,6 @@ def delete_ingested_task(body: "DeleteIngestedDocumentAsyncBody") -> None: # If the task was not revoked, we follow the normal # flow and raise the exception. raise + + +delete_ingested_task.callback_task_name = DELETE_INGESTED_CALLBACK_TASK_NAME # type: ignore[attr-defined] diff --git a/private_gpt/celery/tasks/ingestion/extraction_tasks.py b/private_gpt/celery/tasks/ingestion/extraction_tasks.py index 919cde5d..d7781c74 100644 --- a/private_gpt/celery/tasks/ingestion/extraction_tasks.py +++ b/private_gpt/celery/tasks/ingestion/extraction_tasks.py @@ -18,6 +18,9 @@ from private_gpt.settings.settings import settings logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG if settings().server.debug_mode else logging.INFO) +VECTOR_INDEX_TASK_NAME = "private_gpt.ingestion.vector_index" +VECTOR_INDEX_CALLBACK_TASK_NAME = "vector_index_task" + T = TypeVar("T") @@ -45,7 +48,7 @@ def cleanup_temporal_files(func: Callable[..., T]) -> Callable[..., T]: @celery_app.task( - name="vector_index_task", + name=VECTOR_INDEX_TASK_NAME, base=StatelessBackgroundTask, autoretry_for=AUTORETRY_EXCEPTIONS, ) @@ -135,6 +138,9 @@ def vector_index_task(body: IngestAsyncBody) -> Any: ) +vector_index_task.callback_task_name = VECTOR_INDEX_CALLBACK_TASK_NAME # type: ignore[attr-defined] + + def ensure_to_remove_temporal_files(body: IngestAsyncBody) -> None: """Remove temporal files from S3 if the input was a URI. diff --git a/private_gpt/celery/tasks/tools/tool_run_task.py b/private_gpt/celery/tasks/tools/tool_run_task.py index cd36b1a8..0abf3030 100644 --- a/private_gpt/celery/tasks/tools/tool_run_task.py +++ b/private_gpt/celery/tasks/tools/tool_run_task.py @@ -3,17 +3,19 @@ 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.components.tools.remote_execution import ( - ToolExecutionRequest, ToolExecutionResponse, execute_tool_request, ) from private_gpt.events.models import TextBlock +if TYPE_CHECKING: + from private_gpt.components.tools.remote_execution import ToolExecutionRequest + logger = logging.getLogger(__name__) @@ -21,9 +23,7 @@ logger = logging.getLogger(__name__) name="private_gpt.tools.run", base=StatefulBackgroundTask, ) -async def tool_run_task(**payload: Any) -> dict[str, Any]: - request = ToolExecutionRequest.model_validate(payload) - +async def tool_run_task(request: ToolExecutionRequest) -> dict[str, Any]: try: response = await execute_tool_request(request) return response.model_dump(mode="json") diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py index 8a33e200..9cf3b353 100644 --- a/private_gpt/components/tools/tool_scheduler.py +++ b/private_gpt/components/tools/tool_scheduler.py @@ -7,12 +7,15 @@ from typing import TYPE_CHECKING from injector import inject, singleton +from private_gpt.celery.dispatch import dispatch_task from private_gpt.components.tools.remote_execution import ( ToolExecutionResponse, execute_tool_request, ) from private_gpt.settings.settings import Settings +TOOL_TASK_NAME = "private_gpt.tools.run" + if TYPE_CHECKING: from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( ChatLoopState, @@ -57,10 +60,7 @@ class CeleryToolScheduler(BaseToolScheduler): @inject def __init__(self, settings: Settings) -> None: - from private_gpt.celery.celery import celery_app - - self._celery_app = celery_app - self._tools_queue = settings.scheduler.tools.celery_queue + self._settings = settings async def execute( self, @@ -70,10 +70,10 @@ class CeleryToolScheduler(BaseToolScheduler): ) -> ToolExecutionResponse: del state_ctx, interceptors - result = self._celery_app.send_task( - "private_gpt.tools.run", - kwargs=request.model_dump(mode="json"), - queue=self._tools_queue, + result = dispatch_task( + task_name=TOOL_TASK_NAME, + args=(request,), + queue=self._settings.scheduler.tools.celery_queue, ) while not result.ready(): diff --git a/private_gpt/server/chat_async/chat_async_service.py b/private_gpt/server/chat_async/chat_async_service.py index 050f80e9..03268a93 100644 --- a/private_gpt/server/chat_async/chat_async_service.py +++ b/private_gpt/server/chat_async/chat_async_service.py @@ -3,6 +3,7 @@ from typing import Any from injector import inject, singleton +from private_gpt.celery.dispatch import dispatch_task from private_gpt.components.streaming.providers.models import ( StreamMetadata, StreamStatus, @@ -15,6 +16,8 @@ from private_gpt.server.chat.chat_models import ChatBody from private_gpt.server.chat.chat_request_mapper import ChatRequestMapper from private_gpt.settings.settings import settings as _settings +CHAT_TASK_NAME = "private_gpt.chat.run" + @singleton class ChatAsyncService: @@ -93,8 +96,6 @@ class ChatAsyncService: using its own warm injector so tool implementations and output schemas are built fresh in-process. """ - from private_gpt.celery.celery import celery_app - stream_service = self.stream_manager.stream_service correlation_id = await stream_service.create_stream( @@ -109,13 +110,12 @@ class ChatAsyncService: "message_count": len(body.messages), } - body_data = body.model_dump(mode="json") try: # Use the correlation_id as the Celery task id so revocation is trivial. - celery_app.send_task( - "private_gpt.chat.run", - args=[body_data, correlation_id, "chat_completion", metadata], - queue="chat", + dispatch_task( + task_name=CHAT_TASK_NAME, + args=[body, correlation_id, "chat_completion", metadata], + queue=_settings().scheduler.chat.celery_queue, task_id=correlation_id, ) except Exception as e: diff --git a/private_gpt/server/ingest/ingest_router.py b/private_gpt/server/ingest/ingest_router.py index b724c7a8..f42bd2f2 100644 --- a/private_gpt/server/ingest/ingest_router.py +++ b/private_gpt/server/ingest/ingest_router.py @@ -17,7 +17,14 @@ from private_gpt.settings.settings import settings try: from private_gpt.celery.celery import celery_app + from private_gpt.celery.dispatch import dispatch_task from private_gpt.celery.model import Task, TaskStatus + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + from private_gpt.celery.tasks.ingestion.extraction_tasks import ( + VECTOR_INDEX_TASK_NAME, + ) if TYPE_CHECKING: from celery.result import AsyncResult @@ -666,9 +673,10 @@ if _CELERY_AVAILABLE: body.ingest_body.input = UriArtifact(value=s3_url) # Enqueue ingestion tasks (index population) - async_result = celery_app.send_task( - "vector_index_task", + async_result = dispatch_task( + task_name=VECTOR_INDEX_TASK_NAME, args=(body,), + queue=config.scheduler.ingestion.celery_queue, ) return Task(task_id=async_result.task_id) @@ -966,9 +974,10 @@ if _CELERY_AVAILABLE: return Task(task_id="revoked") # Enqueue deletion task - async_result: AsyncResult[None] = celery_app.send_task( - "delete_ingested_task", + async_result: AsyncResult[None] = dispatch_task( + task_name=DELETE_INGESTED_TASK_NAME, args=(body,), + queue=settings().scheduler.ingestion.celery_queue, ) return Task(task_id=async_result.task_id) @@ -1033,7 +1042,7 @@ if _CELERY_AVAILABLE: description="Unique identifier of the deletion task to check status for", examples=["123e4567-e89b-12d3-a456-426614174000"], ), - ] + ], ) -> TaskStatus[None]: """Retrieves the current status of an asynchronous deletion task. diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py index 7de1491d..6b3817fc 100644 --- a/private_gpt/settings/settings.py +++ b/private_gpt/settings/settings.py @@ -473,6 +473,10 @@ class SchedulerSettings(BaseModel): class SchedulerConfig(BaseModel): + ingestion: SchedulerSettings = Field( + default_factory=lambda: SchedulerSettings(celery_queue="ingestion"), + description="Ingestion worker scheduler configuration.", + ) chat: SchedulerSettings = Field( default_factory=lambda: SchedulerSettings(celery_queue="chat"), description="Chat worker scheduler configuration.", diff --git a/settings.yaml b/settings.yaml index 886a351f..59d873e7 100644 --- a/settings.yaml +++ b/settings.yaml @@ -290,6 +290,8 @@ semaphore: mode: ${PGPT_SEMAPHORE_MODE:memory} scheduler: + ingestion: + celery_queue: ${PGPT_INGESTION_CELERY_QUEUE:ingestion} chat: mode: ${PGPT_CHAT_SCHEDULER_MODE:local} celery_queue: ${PGPT_CHAT_CELERY_QUEUE:chat} diff --git a/tests/server/chat/test_chat_async_worker.py b/tests/server/chat/test_chat_async_worker.py index 6a533d9e..f56fc3cb 100644 --- a/tests/server/chat/test_chat_async_worker.py +++ b/tests/server/chat/test_chat_async_worker.py @@ -59,7 +59,7 @@ async def test_worker_handoff_sends_json_safe_body( assert sent["name"] == "private_gpt.chat.run" assert sent["kwargs"]["queue"] == "chat" assert sent["kwargs"]["task_id"] == "msg-1" - assert isinstance(sent["args"][0], dict) + assert sent["args"][0] == _chat_body() assert sent["args"][1] == "msg-1" diff --git a/tests/server/ingest/test_ingestion_async.py b/tests/server/ingest/test_ingestion_async.py index 1cc31d2d..3c6b67bb 100644 --- a/tests/server/ingest/test_ingestion_async.py +++ b/tests/server/ingest/test_ingestion_async.py @@ -79,7 +79,7 @@ def test_delete_during_ingestion(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], } ] @@ -146,7 +146,7 @@ def test_delete_different_artifact_ingesting(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [different_ingestion_body], } ] @@ -195,17 +195,17 @@ def test_delete_with_multiple_ingestion_tasks(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [different_artifact], }, { "id": "task2", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [same_artifact_diff_collection], }, { "id": "task3", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], }, {"id": "task4", "name": "different_task", "args": [{}]}, @@ -234,7 +234,7 @@ def test_delete_terminates_pending_tasks(mock_setup, test_bodies): pending_task = { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], "status": "PENDING", } @@ -270,7 +270,7 @@ def test_delete_scheduled_when_ingestion_will_run(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "delete_ingested_task", + "name": "private_gpt.ingestion.delete", "args": [delete_body], } ]