github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-17 20:03:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: unify all logic

Browse Source
This commit is contained in:
Javier Martinez
2026-07-08 18:17:12 +02:00
parent ad3d1e0655
commit 67bbd86a03
14 changed files with 115 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
private_gpt/celery/callback.py
View File

@@ -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,
)

23
private_gpt/celery/dispatch.py Normal file
View File

@@ -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,
)

18
private_gpt/celery/task_helper.py
View File

@@ -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
):

16
private_gpt/celery/tasks/chat/chat_run_task.py
View File

@@ -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)

8
private_gpt/celery/tasks/ingestion/delete_tasks.py
View File

@@ -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]

8
private_gpt/celery/tasks/ingestion/extraction_tasks.py
View File

@@ -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.

10
private_gpt/celery/tasks/tools/tool_run_task.py
View File

@@ -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")

16
private_gpt/components/tools/tool_scheduler.py
View File

@@ -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():

14
private_gpt/server/chat_async/chat_async_service.py
View File

@@ -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:

19
private_gpt/server/ingest/ingest_router.py
View File

@@ -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.

4
private_gpt/settings/settings.py
View File

@@ -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.",

2
settings.yaml
View File

@@ -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}

2
tests/server/chat/test_chat_async_worker.py
View File

@@ -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"

14
tests/server/ingest/test_ingestion_async.py
View File

@@ -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],
}
]