Start using Dependency injection to allow injecting the config mo…

service/config.template.yaml

@@ -20,3 +20,8 @@ database:
 # Path to prompts config. Path is set as the relative path from this config.
 files:
   prompts_path: "prompts.yaml"
+
+# This specifies how many DMPs can be processed at the same time.
+# Each DMP process can create further parallel lines of execution depending on the tenants setting.
+# This limit is server-wide, that means it has the same counter for all tenants.
+max_parallel_executions: 2

service/src/ai_document_plugin_service/ai/common/config.py

@@ -49,6 +49,7 @@ class Config:
     section_id: SystemAndUserPrompt
     dmp_generation: SystemPrompt
     dmp_polishing: SystemAndUserPrompt
+    max_parallel_executions: int
 
 
 @dataclass(frozen=True)
@@ -217,4 +218,5 @@ def load_config(config_path: str | None = None) -> Config:
             system_message=_get(prompts, 'dmp_polishing', 'system_message'),
             user_message=_get(prompts, 'dmp_polishing', 'user_message'),
         ),
+        max_parallel_executions=int(_get(config, 'max_parallel_executions')),
     )

service/src/ai_document_plugin_service/api/routes.py

@@ -1,11 +1,8 @@
-from typing import Annotated
 from uuid import uuid4
 
 import fastapi
 
-from ai_document_plugin_service.ai.common.config import Config, LLMConfig
-from ai_document_plugin_service.ai.persistence.database import Database
-from ai_document_plugin_service.api.auth import AuthenticatedUser, verify_authenticated
+from ai_document_plugin_service.api.auth import verify_authenticated
 from ai_document_plugin_service.api.jwt import extract_identity_from_token
 from ai_document_plugin_service.api.types import (
     PipelineRunRequest,
@@ -18,34 +15,24 @@
     TemplateListItem,
     _model_from_fields,
 )
-from ai_document_plugin_service.service import pipeline_service as pipeline
+from ai_document_plugin_service.di import AuthenticatedDI, ConfigDI, DatabaseDI, PipelineServiceDI
 
 public_router = fastapi.APIRouter()
 protected_router = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_authenticated)])
 
 
-def _load_app_config(request: fastapi.Request) -> Config:
-    return request.app.state.config
-
-
-def _load_database(request: fastapi.Request) -> Database:
-    return request.app.state.database
-
-
 @public_router.get('/health')
 def health_check() -> dict[str, str]:
     return {'status': 'healthy'}
 
 
 @protected_router.get('/templates')
-async def list_templates(request: fastapi.Request) -> list[TemplateListItem]:
-    database = _load_database(request)
+async def list_templates(database: DatabaseDI) -> list[TemplateListItem]:
     return [_model_from_fields(TemplateListItem, **item) for item in await database.list_templates()]
 
 
 @protected_router.get('/templates/{template_uuid}')
-async def get_template(template_uuid: str, request: fastapi.Request) -> TemplateDetail:
-    database = _load_database(request)
+async def get_template(template_uuid: str, database: DatabaseDI) -> TemplateDetail:
     template = await database.get_template(template_uuid)
 
     if template is None:
@@ -60,7 +47,7 @@ async def get_template(template_uuid: str, request: fastapi.Request) -> Template
 
 
 @protected_router.post('/templates', status_code=201)
-async def create_template(payload: TemplateCreateRequest, request: fastapi.Request) -> TemplateDetail:
+async def create_template(payload: TemplateCreateRequest, database: DatabaseDI) -> TemplateDetail:
     trimmed_title = payload.title.strip()
     if not trimmed_title:
         raise fastapi.HTTPException(status_code=400, detail='Template title is required')
@@ -72,7 +59,6 @@ async def create_template(payload: TemplateCreateRequest, request: fastapi.Reque
             detail='Template JSON must contain a top-level "sections" array.',
         )
 
-    database = _load_database(request)
     template_uuid = str(uuid4())
 
     try:
@@ -95,11 +81,11 @@ async def create_template(payload: TemplateCreateRequest, request: fastapi.Reque
 @protected_router.post('/pipelines/run')
 async def start_pipeline(
     payload: PipelineRunRequest,
-    request: fastapi.Request,
-    auth: Annotated[AuthenticatedUser, fastapi.Depends(verify_authenticated)],
+    auth: AuthenticatedDI,
+    config: ConfigDI,
+    database: DatabaseDI,
+    pipeline: PipelineServiceDI,
 ) -> PipelineRunResponse:
-    config = _load_app_config(request)
-    database = _load_database(request)
     template = await database.get_template(payload.template_uuid)
 
     if template is None:
@@ -113,19 +99,11 @@ async def start_pipeline(
     run_id = str(uuid4())
     pipeline.enqueue_pipeline_job(
         run_id,
-        payload.questionnaire_uuid,
-        payload.template_uuid,
+        payload,
         template['title'],
         user_uuid,
         tenant_uuid,
-        auth.token,
-        auth.api_url,
-        LLMConfig(
-            model=payload.llm_model,
-            api_key=payload.llm_api_key,
-            api_url=payload.llm_api_url,
-            parallel_workers=payload.llm_max_workers,
-        ),
+        auth,
         config,
     )
     return _model_from_fields(
@@ -141,9 +119,7 @@ async def start_pipeline(
 
 
 @protected_router.get('/pipelines/status/{run_id}')
-def get_pipeline_status(
-    run_id: str,
-) -> PipelineStatusResponse:
+def get_pipeline_status(run_id: str, pipeline: PipelineServiceDI) -> PipelineStatusResponse:
     status = pipeline.get_pipeline_status(run_id)
     if status is None:
         raise fastapi.HTTPException(status_code=404, detail='Pipeline run not found')
@@ -152,39 +128,6 @@ def get_pipeline_status(
 
 @protected_router.post('/pipelines/status/{run_id}/save')
 async def save_pipeline_result(
-    run_id: str,
-    payload: PipelineSaveRequest,
-    request: fastapi.Request,
+    run_id: str, save_request: PipelineSaveRequest, pipeline: PipelineServiceDI
 ) -> PipelineStatusResponse:
-    status = pipeline.get_pipeline_status(run_id)
-    if status is None:
-        raise fastapi.HTTPException(status_code=404, detail='Pipeline run not found')
-
-    if status.knowledge_model_uuid is None:
-        raise fastapi.HTTPException(status_code=500, detail='Missing knowledge_model_uuid')
-
-    database = _load_database(request)
-
-    await database.update_result(
-        template_uuid=status.template_uuid,
-        knowledge_model_uuid=status.knowledge_model_uuid,
-        user_uuid=status.user_uuid,
-        tenant_uuid=status.tenant_uuid,
-        markdown=payload.result_markdown,
-    )
-
-    updated_status = pipeline.build_pipeline_status(
-        run_id=status.run_id,
-        status=status.status,
-        questionnaire_uuid=status.questionnaire_uuid,
-        knowledge_model_uuid=status.knowledge_model_uuid,
-        user_uuid=status.user_uuid,
-        tenant_uuid=status.tenant_uuid,
-        template_uuid=status.template_uuid,
-        template_title=status.template_title,
-        error=status.error,
-        result_format='markdown',
-        result_markdown=payload.result_markdown,
-    )
-    pipeline.set_pipeline_status(run_id, updated_status)
-    return updated_status
+    return await pipeline.update_pipeline_result(run_id, save_request)

service/src/ai_document_plugin_service/app.py

@@ -6,6 +6,7 @@
 from ai_document_plugin_service.ai.persistence.database import PostgresDB
 from ai_document_plugin_service.ai.persistence.migrations import run_startup_migrations
 from ai_document_plugin_service.api.routes import protected_router, public_router
+from ai_document_plugin_service.di import setup_app_state
 
 
 def create_app(*, run_migrations: bool = True) -> fastapi.FastAPI:
@@ -17,8 +18,7 @@ def create_app(*, run_migrations: bool = True) -> fastapi.FastAPI:
         run_startup_migrations(config, config_path)
 
     app = fastapi.FastAPI(title='Plugin Service', version='1.0.0')
-    app.state.config = config
-    app.state.config_path = config_path
+    setup_app_state(app, config)
     app.state.database = PostgresDB(config.database)
 
     app.add_middleware(

service/src/ai_document_plugin_service/di.py

@@ -0,0 +1,40 @@
+from typing import Annotated
+
+import fastapi
+
+from ai_document_plugin_service.ai.common import Config
+from ai_document_plugin_service.ai.persistence.database import Database, PostgresDB
+from ai_document_plugin_service.api.auth import AuthenticatedUser, verify_authenticated
+from ai_document_plugin_service.service.pipeline_queue_manager import PipelineQueueManager
+from ai_document_plugin_service.service.pipeline_service import PipelineService
+
+
+def setup_app_state(app: fastapi.FastAPI, config: Config) -> None:
+    app.state.config = config
+    app.state.database = PostgresDB(config.database)
+    app.state.pipeline_queue_manager = PipelineQueueManager(config.max_parallel_executions)
+    app.state.pipeline_service = PipelineService(app.state.pipeline_queue_manager, app.state.database)
+
+
+AuthenticatedDI = Annotated[AuthenticatedUser, fastapi.Depends(verify_authenticated)]
+
+
+def _get_pipeline_service(request: fastapi.Request) -> PipelineService:
+    return request.app.state.pipeline_service
+
+
+PipelineServiceDI = Annotated[PipelineService, fastapi.Depends(_get_pipeline_service)]
+
+
+def _get_app_config(request: fastapi.Request) -> Config:
+    return request.app.state.config
+
+
+ConfigDI = Annotated[Config, fastapi.Depends(_get_app_config)]
+
+
+def _get_database(request: fastapi.Request) -> Database:
+    return request.app.state.database
+
+
+DatabaseDI = Annotated[Database, fastapi.Depends(_get_database)]

service/src/ai_document_plugin_service/service/pipeline_queue_manager.py

@@ -7,8 +7,6 @@
 
 logger = logging.getLogger(__name__)
 
-MAX_CONCURRENT_PIPELINE_JOBS = 2
-
 JobFactory = Callable[[], Coroutine[Any, Any, None]]
 
 
@@ -23,10 +21,11 @@ def format_queue_progress(jobs_ahead: int) -> str:
 class PipelineQueueManager:
     """FIFO pipeline job queue running coroutines on a dedicated event loop.
 
-    Jobs are coroutines scheduled onto a single background event loop and gated by an
+    Jobs are coroutines scheduled onto a single background event loop and gated by a
+    semaphore so at most ``max_concurrent_jobs`` run at once.
     """
 
-    def __init__(self, max_concurrent_jobs: int = MAX_CONCURRENT_PIPELINE_JOBS) -> None:
+    def __init__(self, max_concurrent_jobs: int) -> None:
         self._max_concurrent_jobs = max_concurrent_jobs
         self._order: list[str] = []
         self._order_lock = threading.Lock()
@@ -58,10 +57,8 @@ def progress_message(self, run_id: str) -> str | None:
 
     def remove(self, run_id: str) -> None:
         with self._order_lock:
-            try:
+            if run_id in self._order:
                 self._order.remove(run_id)
-            except ValueError:
-                return
 
     def _jobs_waiting_ahead(self, run_id: str) -> int | None:
         with self._order_lock:
@@ -85,6 +82,3 @@ def _log_job_failure(future: Future[None]) -> None:
         error = future.exception()
         if error is not None:
             logger.error('Pipeline job crashed without handling its error', exc_info=error)
-
-
-pipeline_queue_manager = PipelineQueueManager()