fix: celery callbacks

private_gpt/celery/base.py

@@ -15,7 +15,12 @@ from celery.utils.log import get_task_logger
 
 from private_gpt.celery.callback import task_after_return
 from private_gpt.celery.config import celery_settings
-from private_gpt.di import clean_global_injector, get_global_injector
+from private_gpt.di import (
+    clean_global_injector,
+    create_application_injector,
+    get_global_injector,
+    set_global_injector,
+)
 
 logger = get_task_logger(__name__)
 logger.setLevel("DEBUG")
@@ -257,8 +262,7 @@ class StatelessBackgroundTask(_BackgroundTask):
         if run_method is None:
             raise NotImplementedError("Subclass must implement 'run' method")
 
-        # Inject injector
-        get_global_injector(allow_to_generate_new_injectors=True)
+        set_global_injector(create_application_injector())
 
         if asyncio.iscoroutinefunction(run_method):
             result = await run_method(*args, **kwargs)

private_gpt/celery/tasks/ingestion/delete_tasks.py

@@ -23,6 +23,7 @@ DELETE_INGESTED_CALLBACK_TASK_NAME = "delete_ingested_task"
 @celery_app.task(
     name=DELETE_INGESTED_TASK_NAME,
     base=StatelessBackgroundTask,
+    callback_task_name=DELETE_INGESTED_CALLBACK_TASK_NAME,
     # Retry on ValueError and IndexNotReadyException.
     # ValueError is thrown when the index is not initialized
     #   and we cannot guarantee that the index will not be ready.
@@ -54,6 +55,3 @@ 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]

private_gpt/celery/tasks/ingestion/extraction_tasks.py

@@ -50,6 +50,7 @@ def cleanup_temporal_files(func: Callable[..., T]) -> Callable[..., T]:
 @celery_app.task(
     name=VECTOR_INDEX_TASK_NAME,
     base=StatelessBackgroundTask,
+    callback_task_name=VECTOR_INDEX_CALLBACK_TASK_NAME,
     autoretry_for=AUTORETRY_EXCEPTIONS,
 )
 @cleanup_temporal_files
@@ -137,9 +138,6 @@ 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.
 

tests/celery/test_tasks.py

@@ -82,6 +82,46 @@ def test_success_task_posts_to_success_broker_queue(injector: MockInjector):
     )
 
 
+def test_success_task_posts_to_callback_task_name_queue(injector: MockInjector):
+    broker_mock = Mock(BrokerComponent)
+    injector.bind_mock(BrokerComponent, broker_mock)
+
+    set_global_injector(injector.test_injector)
+
+    @celery_app.task(
+        name="renamed_callback_task",
+        callback_task_name="legacy_callback_task",
+        after_return=task_after_return,
+    )
+    def renamed_task(input_with_callback: CallbackInput) -> CallbackResponse:
+        return CallbackResponse(
+            result=input_with_callback.x * input_with_callback.y,
+            label="test",
+        )
+
+    celery_app.send_task(
+        "renamed_callback_task",
+        args=(
+            CallbackInput(
+                x=2,
+                y=3,
+                callback=Callback(amqp=AMQP(exchange="main")),
+            ),
+        ),
+    )
+
+    expected_response = AsyncResponse(
+        data=CallbackResponse(result=6, label="test"),
+        type="pgpt.legacy_callback_task.done",
+    )
+
+    broker_mock.publish.assert_called_once_with(
+        exchange="main",
+        routing_key="pgpt.legacy_callback_task.done",
+        body=bytes(expected_response.model_dump_json(), "utf-8"),
+    )
+
+
 def test_failing_task_posts_to_error_handler_queue(injector: MockInjector):
     broker_mock = Mock(BrokerComponent)
     injector.bind_mock(BrokerComponent, broker_mock)