diff --git a/private_gpt/celery/base.py b/private_gpt/celery/base.py index 4f927e2b..0faa30dd 100644 --- a/private_gpt/celery/base.py +++ b/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) diff --git a/private_gpt/celery/tasks/ingestion/delete_tasks.py b/private_gpt/celery/tasks/ingestion/delete_tasks.py index c7fbeed8..ceae61c1 100644 --- a/private_gpt/celery/tasks/ingestion/delete_tasks.py +++ b/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] diff --git a/private_gpt/celery/tasks/ingestion/extraction_tasks.py b/private_gpt/celery/tasks/ingestion/extraction_tasks.py index 786d7834..a31105a1 100644 --- a/private_gpt/celery/tasks/ingestion/extraction_tasks.py +++ b/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. diff --git a/tests/celery/test_tasks.py b/tests/celery/test_tasks.py index 972d17eb..991effff 100644 --- a/tests/celery/test_tasks.py +++ b/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)