From 8bb73c1878e94438d50b3bba3f674538ab7b81a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 13:21:02 +0200 Subject: [PATCH 1/5] 85: Start using Dependency injection to allow injecting the config more easily --- service/config.template.yaml | 5 + .../ai/common/config.py | 2 + .../ai_document_plugin_service/api/routes.py | 42 +- service/src/ai_document_plugin_service/app.py | 4 +- service/src/ai_document_plugin_service/di.py | 40 ++ .../service/pipeline_queue_manager.py | 7 +- .../service/pipeline_service.py | 470 +++++++++--------- 7 files changed, 301 insertions(+), 269 deletions(-) create mode 100644 service/src/ai_document_plugin_service/di.py diff --git a/service/config.template.yaml b/service/config.template.yaml index 9b24035..5b4036a 100644 --- a/service/config.template.yaml +++ b/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 \ No newline at end of file diff --git a/service/src/ai_document_plugin_service/ai/common/config.py b/service/src/ai_document_plugin_service/ai/common/config.py index 4be249f..43e0018 100644 --- a/service/src/ai_document_plugin_service/ai/common/config.py +++ b/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')) ) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 8f5ad78..df78e8f 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -1,11 +1,9 @@ -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.ai.common.config import LLMConfig +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 +16,25 @@ 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 +49,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 +61,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 +83,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: @@ -143,6 +131,7 @@ async def start_pipeline( @protected_router.get('/pipelines/status/{run_id}') def get_pipeline_status( run_id: str, + pipeline: PipelineServiceDI ) -> PipelineStatusResponse: status = pipeline.get_pipeline_status(run_id) if status is None: @@ -154,7 +143,8 @@ def get_pipeline_status( async def save_pipeline_result( run_id: str, payload: PipelineSaveRequest, - request: fastapi.Request, + database: DatabaseDI, + pipeline: PipelineServiceDI ) -> PipelineStatusResponse: status = pipeline.get_pipeline_status(run_id) if status is None: @@ -163,8 +153,6 @@ async def save_pipeline_result( 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, @@ -173,7 +161,7 @@ async def save_pipeline_result( markdown=payload.result_markdown, ) - updated_status = pipeline.build_pipeline_status( + updated_status = pipeline._build_pipeline_status( run_id=status.run_id, status=status.status, questionnaire_uuid=status.questionnaire_uuid, diff --git a/service/src/ai_document_plugin_service/app.py b/service/src/ai_document_plugin_service/app.py index 58ab20f..c111e47 100644 --- a/service/src/ai_document_plugin_service/app.py +++ b/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( diff --git a/service/src/ai_document_plugin_service/di.py b/service/src/ai_document_plugin_service/di.py new file mode 100644 index 0000000..54acbcc --- /dev/null +++ b/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) + + +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)] diff --git a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py b/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py index ea41890..71f3145 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py +++ b/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]] @@ -26,7 +24,7 @@ class PipelineQueueManager: Jobs are coroutines scheduled onto a single background event loop and gated by an """ - 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() @@ -85,6 +83,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() diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index a599aec..7ab1b8e 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -21,13 +21,10 @@ PipelineStatusResponse, _model_from_fields, ) -from ai_document_plugin_service.service.pipeline_queue_manager import pipeline_queue_manager +from ai_document_plugin_service.service.pipeline_queue_manager import PipelineQueueManager logger = logging.getLogger(__name__) -_pipeline_runs: dict[str, PipelineStatusResponse] = {} -_pipeline_runs_lock = threading.Lock() - AUTHORIZATION_ERROR_MESSAGE = 'Authorization error, invalid or expired token.' SERVER_ERROR_MESSAGE = 'The action could not be completed. Please try again later.' TEMPLATE_NOT_FOUND_MESSAGE = 'Template not found.' @@ -46,40 +43,20 @@ def _pipeline_error_from_exception(error: Exception) -> PipelineErrorResponse: ) -def set_pipeline_status(run_id: str, status: PipelineStatusResponse) -> None: - with _pipeline_runs_lock: - _pipeline_runs[run_id] = status - - -def get_pipeline_status(run_id: str) -> PipelineStatusResponse | None: - with _pipeline_runs_lock: - status = _pipeline_runs.get(run_id) - - if status is None or status.status != PipelineStatus.QUEUED: - return status - - progress_message = pipeline_queue_manager.progress_message(run_id) - if progress_message is None: - return status - - return status.model_copy(update={'progress_message': progress_message}) - - -def build_pipeline_status( - *, - run_id: str, - status: PipelineStatus, - questionnaire_uuid: str, - user_uuid: str, - tenant_uuid: str, - template_uuid: str, - template_title: str, - knowledge_model_uuid: str | None = None, - error: PipelineErrorResponse | None = None, - result_format: str | None = None, - result_markdown: str | None = None, - progress_message: str | None = None, -) -> PipelineStatusResponse: +def _build_pipeline_status(*, + run_id: str, + status: PipelineStatus, + questionnaire_uuid: str, + user_uuid: str, + tenant_uuid: str, + template_uuid: str, + template_title: str, + knowledge_model_uuid: str | None = None, + error: PipelineErrorResponse | None = None, + result_format: str | None = None, + result_markdown: str | None = None, + progress_message: str | None = None, + ) -> PipelineStatusResponse: return _model_from_fields( PipelineStatusResponse, run_id=run_id, @@ -98,223 +75,248 @@ def build_pipeline_status( ) -def enqueue_pipeline_job( - run_id: str, - questionnaire_uuid: str, - template_uuid: str, - template_title: str, - user_uuid: str, - tenant_uuid: str, - token: str, - api_url: str, - llm_config: LLMConfig, - config: Config, -) -> None: - """Queue a pipeline job; concurrency is limited by ``pipeline_queue_manager``.""" - set_pipeline_status( - run_id, - build_pipeline_status( - run_id=run_id, - status=PipelineStatus.QUEUED, - questionnaire_uuid=questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - ), - ) - - pipeline_queue_manager.enqueue( - run_id, - lambda: _run_pipeline_job( +class PipelineService: + def __init__(self, pipeline_queue_manager: PipelineQueueManager) -> None: + self.pipeline_queue_manager = pipeline_queue_manager + self._pipeline_runs: dict[str, PipelineStatusResponse] = {} + self._pipeline_runs_lock = threading.Lock() + + def set_pipeline_status(self, run_id: str, status: PipelineStatusResponse) -> None: + with self._pipeline_runs_lock: + self._pipeline_runs[run_id] = status + + def get_pipeline_status(self, run_id: str) -> PipelineStatusResponse | None: + with self._pipeline_runs_lock: + status = self._pipeline_runs.get(run_id) + + if status is None or status.status != PipelineStatus.QUEUED: + return status + + progress_message = self.pipeline_queue_manager.progress_message(run_id) + if progress_message is None: + return status + + return status.model_copy(update={'progress_message': progress_message}) + + def enqueue_pipeline_job( + self, + run_id: str, + questionnaire_uuid: str, + template_uuid: str, + template_title: str, + user_uuid: str, + tenant_uuid: str, + token: str, + api_url: str, + llm_config: LLMConfig, + config: Config, + ) -> None: + """Queue a pipeline job; concurrency is limited by ``pipeline_queue_manager``.""" + self.set_pipeline_status( run_id, - questionnaire_uuid, - template_uuid, - template_title, - user_uuid, - tenant_uuid, - token, - api_url, - llm_config, - config, - ), - ) - - -def _update_running_progress( - run_id: str, - *, - questionnaire_uuid: str, - user_uuid: str, - tenant_uuid: str, - template_uuid: str, - template_title: str, - progress_message: str, -) -> None: - with _pipeline_runs_lock: - current = _pipeline_runs.get(run_id) - if current is None: - return - - set_pipeline_status( - run_id, - build_pipeline_status( - run_id=run_id, - status=PipelineStatus.RUNNING, - questionnaire_uuid=questionnaire_uuid, - knowledge_model_uuid=current.knowledge_model_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - progress_message=progress_message, - ), - ) + _build_pipeline_status( + run_id=run_id, + status=PipelineStatus.QUEUED, + questionnaire_uuid=questionnaire_uuid, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + template_uuid=template_uuid, + template_title=template_title, + ), + ) + self.pipeline_queue_manager.enqueue( + run_id, + lambda: self._pipeline_entrypoint( + run_id, + questionnaire_uuid, + template_uuid, + template_title, + user_uuid, + tenant_uuid, + token, + api_url, + llm_config, + config, + ), + ) -async def _run_pipeline_job( - run_id: str, - questionnaire_uuid: str, - template_uuid: str, - template_title: str, - user_uuid: str, - tenant_uuid: str, - token: str, - dsw_api_url: str, - llm_config: LLMConfig, - config: Config, -) -> None: - database = PostgresDB(config.database) - saver = DBSaver(database) - llm_client = LLMClient(llm_config.model, llm_config.api_key, llm_config.api_url, llm_config.parallel_workers) - try: - template = await database.get_template(template_uuid) - if template is None: - _fail_template_not_found(questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, user_uuid) + def _update_running_progress( + self, + run_id: str, + *, + questionnaire_uuid: str, + user_uuid: str, + tenant_uuid: str, + template_uuid: str, + template_title: str, + progress_message: str, + ) -> None: + with self._pipeline_runs_lock: + current = self._pipeline_runs.get(run_id) + if current is None: return - await _start_pipeline( - config, - database, - dsw_api_url, - llm_client, - questionnaire_uuid, - run_id, - saver, - template, - template_title, - template_uuid, - tenant_uuid, - token, - user_uuid, - ) - except Exception as error: - set_pipeline_status( + self.set_pipeline_status( run_id, - build_pipeline_status( + _build_pipeline_status( run_id=run_id, - status=PipelineStatus.FAILED, + status=PipelineStatus.RUNNING, questionnaire_uuid=questionnaire_uuid, + knowledge_model_uuid=current.knowledge_model_uuid, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, template_uuid=template_uuid, template_title=template_title, + progress_message=progress_message, + ), + ) + + async def _pipeline_entrypoint( + self, + run_id: str, + questionnaire_uuid: str, + template_uuid: str, + template_title: str, + user_uuid: str, + tenant_uuid: str, + token: str, + dsw_api_url: str, + llm_config: LLMConfig, + config: Config, + ) -> None: + database = PostgresDB(config.database) + saver = DBSaver(database) + llm_client = LLMClient(llm_config.model, llm_config.api_key, llm_config.api_url, llm_config.parallel_workers) + try: + template = await database.get_template(template_uuid) + if template is None: + self._fail_template_not_found(questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, + user_uuid) + return + + await self._run_pipeline( + config, + database, + dsw_api_url, + llm_client, + questionnaire_uuid, + run_id, + saver, + template, + template_title, + template_uuid, + tenant_uuid, + token, + user_uuid, + ) + except Exception as error: + self.set_pipeline_status( + run_id, + _build_pipeline_status( + run_id=run_id, + status=PipelineStatus.FAILED, + questionnaire_uuid=questionnaire_uuid, + template_uuid=template_uuid, + template_title=template_title, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + error=_pipeline_error_from_exception(error), + ), + ) + logger.exception('Pipeline run failed') + finally: + await database.dispose() + + async def _run_pipeline( + self, + config: Config, + database: PostgresDB, + dsw_api_url: str, + llm_client: LLMClient, + questionnaire_uuid: str, + run_id: str, + saver: DBSaver, + template: dict[str, Any], + template_title: str, + template_uuid: str, + tenant_uuid: str, + token: str, + user_uuid: str, + ) -> None: + self.set_pipeline_status( + run_id, + _build_pipeline_status( + run_id=run_id, + status=PipelineStatus.RUNNING, + questionnaire_uuid=questionnaire_uuid, user_uuid=user_uuid, tenant_uuid=tenant_uuid, - error=_pipeline_error_from_exception(error), + template_uuid=template_uuid, + template_title=template_title, + progress_message='Starting pipeline...', ), ) - logger.exception('Pipeline run failed') - finally: - await database.dispose() + def on_progress(message: str) -> None: + self._update_running_progress( + run_id, + questionnaire_uuid=questionnaire_uuid, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + template_uuid=template_uuid, + template_title=template_title, + progress_message=message, + ) -async def _start_pipeline( - config: Config, - database: PostgresDB, - dsw_api_url: str, - llm_client: LLMClient, - questionnaire_uuid: str, - run_id: str, - saver: DBSaver, - template: dict[str, Any], - template_title: str, - template_uuid: str, - tenant_uuid: str, - token: str, - user_uuid: str, -) -> None: - set_pipeline_status( - run_id, - build_pipeline_status( - run_id=run_id, - status=PipelineStatus.RUNNING, + pipeline = build_pipeline(database=database, saver=saver, config=config, llm_client=llm_client) + knowledge_model_uuid, result = await run_pipeline( questionnaire_uuid=questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, template_uuid=template_uuid, - template_title=template_title, - progress_message='Starting pipeline...', - ), - ) - - def on_progress(message: str) -> None: - _update_running_progress( - run_id, - questionnaire_uuid=questionnaire_uuid, + template_title=template['title'], + template_data=template['content'], user_uuid=user_uuid, tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - progress_message=message, + pipeline=pipeline, + database=database, + on_progress=on_progress, + model_name=llm_client.get_model_name(), + dsw_client=DSWClient(token, dsw_api_url), ) - pipeline = build_pipeline(database=database, saver=saver, config=config, llm_client=llm_client) - knowledge_model_uuid, result = await run_pipeline( - questionnaire_uuid=questionnaire_uuid, - template_uuid=template_uuid, - template_title=template['title'], - template_data=template['content'], - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - pipeline=pipeline, - database=database, - on_progress=on_progress, - model_name=llm_client.get_model_name(), - dsw_client=DSWClient(token, dsw_api_url), - ) - - set_pipeline_status( - run_id, - build_pipeline_status( - run_id=run_id, - status=PipelineStatus.SUCCEEDED, - questionnaire_uuid=questionnaire_uuid, - knowledge_model_uuid=knowledge_model_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - result_format='markdown', - result_markdown=result, - ), - ) - + self.set_pipeline_status( + run_id, + _build_pipeline_status( + run_id=run_id, + status=PipelineStatus.SUCCEEDED, + questionnaire_uuid=questionnaire_uuid, + knowledge_model_uuid=knowledge_model_uuid, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + template_uuid=template_uuid, + template_title=template_title, + result_format='markdown', + result_markdown=result, + ), + ) -def _fail_template_not_found( - questionnaire_uuid: str, run_id: str, template_title: str, template_uuid: str, tenant_uuid: str, user_uuid: str -) -> None: - set_pipeline_status( - run_id, - build_pipeline_status( - run_id=run_id, - status=PipelineStatus.FAILED, - questionnaire_uuid=questionnaire_uuid, - template_uuid=template_uuid, - template_title=template_title, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - error=PipelineErrorResponse( - type=ErrorType.TEMPLATE_NOT_FOUND, - message=TEMPLATE_NOT_FOUND_MESSAGE, + def _fail_template_not_found( + self, questionnaire_uuid: str, run_id: str, template_title: str, template_uuid: str, tenant_uuid: str, + user_uuid: str + ) -> None: + self.set_pipeline_status( + run_id, + _build_pipeline_status( + run_id=run_id, + status=PipelineStatus.FAILED, + questionnaire_uuid=questionnaire_uuid, + template_uuid=template_uuid, + template_title=template_title, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + error=PipelineErrorResponse( + type=ErrorType.TEMPLATE_NOT_FOUND, + message=TEMPLATE_NOT_FOUND_MESSAGE, + ), ), - ), - ) + ) From 48347dc65f0f685dac2b1af7e9e23b6698178c9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 14:16:58 +0200 Subject: [PATCH 2/5] 85: Handle usage of private function --- .../ai_document_plugin_service/api/routes.py | 33 +----------- service/src/ai_document_plugin_service/di.py | 2 +- .../service/pipeline_service.py | 54 ++++++++++++++----- 3 files changed, 45 insertions(+), 44 deletions(-) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index df78e8f..45209ec 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -142,37 +142,8 @@ def get_pipeline_status( @protected_router.post('/pipelines/status/{run_id}/save') async def save_pipeline_result( run_id: str, - payload: PipelineSaveRequest, - database: DatabaseDI, + 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') - - 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) diff --git a/service/src/ai_document_plugin_service/di.py b/service/src/ai_document_plugin_service/di.py index 54acbcc..abd1ce6 100644 --- a/service/src/ai_document_plugin_service/di.py +++ b/service/src/ai_document_plugin_service/di.py @@ -13,7 +13,7 @@ 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.pipeline_service = PipelineService(app.state.pipeline_queue_manager, app.state.database) AuthenticatedDI = Annotated[AuthenticatedUser, fastapi.Depends(verify_authenticated)] diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index 7ab1b8e..e90f741 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from typing import Any +import fastapi from openai import AuthenticationError from ai_document_plugin_service.ai.common.config import ( @@ -12,14 +13,14 @@ from ai_document_plugin_service.ai.common.llm_client import LLMClient from ai_document_plugin_service.ai.knowledgemodel.dsw_client import DSWClient from ai_document_plugin_service.ai.persistence.assignment_saver_component import DBSaver -from ai_document_plugin_service.ai.persistence.database import PostgresDB +from ai_document_plugin_service.ai.persistence.database import Database from ai_document_plugin_service.ai.run_pipeline import build_pipeline, run_pipeline from ai_document_plugin_service.api.types import ( ErrorType, PipelineErrorResponse, PipelineStatus, PipelineStatusResponse, - _model_from_fields, + _model_from_fields, PipelineSaveRequest, ) from ai_document_plugin_service.service.pipeline_queue_manager import PipelineQueueManager @@ -76,8 +77,9 @@ def _build_pipeline_status(*, class PipelineService: - def __init__(self, pipeline_queue_manager: PipelineQueueManager) -> None: + def __init__(self, pipeline_queue_manager: PipelineQueueManager, database: Database) -> None: self.pipeline_queue_manager = pipeline_queue_manager + self.database = database self._pipeline_runs: dict[str, PipelineStatusResponse] = {} self._pipeline_runs_lock = threading.Lock() @@ -141,6 +143,39 @@ def enqueue_pipeline_job( ), ) + async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRequest)->PipelineStatusResponse: + pipeline_status = self.get_pipeline_status(run_id) + if pipeline_status is None: + raise fastapi.HTTPException(status_code=404, detail='Pipeline run not found') + + if pipeline_status.knowledge_model_uuid is None: + raise fastapi.HTTPException(status_code=500, detail='Missing knowledge_model_uuid') + + await self.database.update_result( + template_uuid=pipeline_status.template_uuid, + knowledge_model_uuid=pipeline_status.knowledge_model_uuid, + user_uuid=pipeline_status.user_uuid, + tenant_uuid=pipeline_status.tenant_uuid, + markdown=save_request.result_markdown, + ) + + updated_status = _build_pipeline_status( + run_id=pipeline_status.run_id, + status=pipeline_status.status, + questionnaire_uuid=pipeline_status.questionnaire_uuid, + knowledge_model_uuid=pipeline_status.knowledge_model_uuid, + user_uuid=pipeline_status.user_uuid, + tenant_uuid=pipeline_status.tenant_uuid, + template_uuid=pipeline_status.template_uuid, + template_title=pipeline_status.template_title, + error=pipeline_status.error, + result_format='markdown', + result_markdown=save_request.result_markdown, + ) + self.set_pipeline_status(run_id, updated_status) + return updated_status + + def _update_running_progress( self, run_id: str, @@ -185,11 +220,10 @@ async def _pipeline_entrypoint( llm_config: LLMConfig, config: Config, ) -> None: - database = PostgresDB(config.database) - saver = DBSaver(database) + saver = DBSaver(self.database) llm_client = LLMClient(llm_config.model, llm_config.api_key, llm_config.api_url, llm_config.parallel_workers) try: - template = await database.get_template(template_uuid) + template = await self.database.get_template(template_uuid) if template is None: self._fail_template_not_found(questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, user_uuid) @@ -197,7 +231,6 @@ async def _pipeline_entrypoint( await self._run_pipeline( config, - database, dsw_api_url, llm_client, questionnaire_uuid, @@ -225,13 +258,10 @@ async def _pipeline_entrypoint( ), ) logger.exception('Pipeline run failed') - finally: - await database.dispose() async def _run_pipeline( self, config: Config, - database: PostgresDB, dsw_api_url: str, llm_client: LLMClient, questionnaire_uuid: str, @@ -269,7 +299,7 @@ def on_progress(message: str) -> None: progress_message=message, ) - pipeline = build_pipeline(database=database, saver=saver, config=config, llm_client=llm_client) + pipeline = build_pipeline(database=self.database, saver=saver, config=config, llm_client=llm_client) knowledge_model_uuid, result = await run_pipeline( questionnaire_uuid=questionnaire_uuid, template_uuid=template_uuid, @@ -278,7 +308,7 @@ def on_progress(message: str) -> None: user_uuid=user_uuid, tenant_uuid=tenant_uuid, pipeline=pipeline, - database=database, + database=self.database, on_progress=on_progress, model_name=llm_client.get_model_name(), dsw_client=DSWClient(token, dsw_api_url), From 4afdde544a37d36b4935643a53f7782546ee6b4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 14:17:57 +0200 Subject: [PATCH 3/5] 85: Fix formatting --- .../ai/common/config.py | 2 +- .../ai_document_plugin_service/api/routes.py | 14 ++---- .../service/pipeline_service.py | 49 +++++++++++-------- 3 files changed, 33 insertions(+), 32 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/config.py b/service/src/ai_document_plugin_service/ai/common/config.py index 43e0018..5292eb1 100644 --- a/service/src/ai_document_plugin_service/ai/common/config.py +++ b/service/src/ai_document_plugin_service/ai/common/config.py @@ -218,5 +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')) + max_parallel_executions=int(_get(config, 'max_parallel_executions')), ) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 45209ec..a680939 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -16,8 +16,7 @@ TemplateListItem, _model_from_fields, ) -from ai_document_plugin_service.di import AuthenticatedDI, \ - ConfigDI, DatabaseDI, PipelineServiceDI +from ai_document_plugin_service.di import AuthenticatedDI, ConfigDI, DatabaseDI, PipelineServiceDI public_router = fastapi.APIRouter() protected_router = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_authenticated)]) @@ -86,7 +85,7 @@ async def start_pipeline( auth: AuthenticatedDI, config: ConfigDI, database: DatabaseDI, - pipeline: PipelineServiceDI + pipeline: PipelineServiceDI, ) -> PipelineRunResponse: template = await database.get_template(payload.template_uuid) @@ -129,10 +128,7 @@ async def start_pipeline( @protected_router.get('/pipelines/status/{run_id}') -def get_pipeline_status( - run_id: str, - pipeline: PipelineServiceDI -) -> 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') @@ -141,9 +137,7 @@ def get_pipeline_status( @protected_router.post('/pipelines/status/{run_id}/save') async def save_pipeline_result( - run_id: str, - save_request: PipelineSaveRequest, - pipeline: PipelineServiceDI + run_id: str, save_request: PipelineSaveRequest, pipeline: PipelineServiceDI ) -> PipelineStatusResponse: return await pipeline.update_pipeline_result(run_id, save_request) diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index e90f741..8c9493a 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -18,9 +18,10 @@ from ai_document_plugin_service.api.types import ( ErrorType, PipelineErrorResponse, + PipelineSaveRequest, PipelineStatus, PipelineStatusResponse, - _model_from_fields, PipelineSaveRequest, + _model_from_fields, ) from ai_document_plugin_service.service.pipeline_queue_manager import PipelineQueueManager @@ -44,20 +45,21 @@ def _pipeline_error_from_exception(error: Exception) -> PipelineErrorResponse: ) -def _build_pipeline_status(*, - run_id: str, - status: PipelineStatus, - questionnaire_uuid: str, - user_uuid: str, - tenant_uuid: str, - template_uuid: str, - template_title: str, - knowledge_model_uuid: str | None = None, - error: PipelineErrorResponse | None = None, - result_format: str | None = None, - result_markdown: str | None = None, - progress_message: str | None = None, - ) -> PipelineStatusResponse: +def _build_pipeline_status( + *, + run_id: str, + status: PipelineStatus, + questionnaire_uuid: str, + user_uuid: str, + tenant_uuid: str, + template_uuid: str, + template_title: str, + knowledge_model_uuid: str | None = None, + error: PipelineErrorResponse | None = None, + result_format: str | None = None, + result_markdown: str | None = None, + progress_message: str | None = None, +) -> PipelineStatusResponse: return _model_from_fields( PipelineStatusResponse, run_id=run_id, @@ -143,7 +145,7 @@ def enqueue_pipeline_job( ), ) - async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRequest)->PipelineStatusResponse: + async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRequest) -> PipelineStatusResponse: pipeline_status = self.get_pipeline_status(run_id) if pipeline_status is None: raise fastapi.HTTPException(status_code=404, detail='Pipeline run not found') @@ -175,7 +177,6 @@ async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRe self.set_pipeline_status(run_id, updated_status) return updated_status - def _update_running_progress( self, run_id: str, @@ -225,8 +226,9 @@ async def _pipeline_entrypoint( try: template = await self.database.get_template(template_uuid) if template is None: - self._fail_template_not_found(questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, - user_uuid) + self._fail_template_not_found( + questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, user_uuid + ) return await self._run_pipeline( @@ -331,8 +333,13 @@ def on_progress(message: str) -> None: ) def _fail_template_not_found( - self, questionnaire_uuid: str, run_id: str, template_title: str, template_uuid: str, tenant_uuid: str, - user_uuid: str + self, + questionnaire_uuid: str, + run_id: str, + template_title: str, + template_uuid: str, + tenant_uuid: str, + user_uuid: str, ) -> None: self.set_pipeline_status( run_id, From 91a0cc504f5828c85b366d6ca13310eda635152e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 14:53:18 +0200 Subject: [PATCH 4/5] 85: Refactor pipeline service - simplified --- .../ai_document_plugin_service/api/routes.py | 14 +- .../service/pipeline_queue_manager.py | 7 +- .../service/pipeline_service.py | 332 ++++++------------ 3 files changed, 105 insertions(+), 248 deletions(-) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index a680939..13054ba 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -2,7 +2,6 @@ import fastapi -from ai_document_plugin_service.ai.common.config import LLMConfig 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 ( @@ -100,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( @@ -139,5 +130,4 @@ def get_pipeline_status(run_id: str, pipeline: PipelineServiceDI) -> PipelineSta async def save_pipeline_result( run_id: str, save_request: PipelineSaveRequest, pipeline: PipelineServiceDI ) -> PipelineStatusResponse: - return await pipeline.update_pipeline_result(run_id, save_request) diff --git a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py b/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py index 71f3145..18d7e67 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py +++ b/service/src/ai_document_plugin_service/service/pipeline_queue_manager.py @@ -21,7 +21,8 @@ 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) -> None: @@ -56,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: diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index 8c9493a..849ff7b 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -1,7 +1,6 @@ import logging import threading from datetime import UTC, datetime -from typing import Any import fastapi from openai import AuthenticationError @@ -15,9 +14,11 @@ from ai_document_plugin_service.ai.persistence.assignment_saver_component import DBSaver from ai_document_plugin_service.ai.persistence.database import Database from ai_document_plugin_service.ai.run_pipeline import build_pipeline, run_pipeline +from ai_document_plugin_service.api.auth import AuthenticatedUser from ai_document_plugin_service.api.types import ( ErrorType, PipelineErrorResponse, + PipelineRunRequest, PipelineSaveRequest, PipelineStatus, PipelineStatusResponse, @@ -45,54 +46,44 @@ def _pipeline_error_from_exception(error: Exception) -> PipelineErrorResponse: ) -def _build_pipeline_status( - *, - run_id: str, - status: PipelineStatus, - questionnaire_uuid: str, - user_uuid: str, - tenant_uuid: str, - template_uuid: str, - template_title: str, - knowledge_model_uuid: str | None = None, - error: PipelineErrorResponse | None = None, - result_format: str | None = None, - result_markdown: str | None = None, - progress_message: str | None = None, -) -> PipelineStatusResponse: - return _model_from_fields( - PipelineStatusResponse, - run_id=run_id, - status=status, - questionnaire_uuid=questionnaire_uuid, - knowledge_model_uuid=knowledge_model_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - error=error, - result_format=result_format, - result_markdown=result_markdown, - progress_message=progress_message, - updated_at=datetime.now(tz=UTC).isoformat(), - ) +def _now() -> str: + return datetime.now(tz=UTC).isoformat() + + +class PipelineRunStore: + """Thread-safe in-memory store of pipeline run statuses.""" + + def __init__(self) -> None: + self._runs: dict[str, PipelineStatusResponse] = {} + self._lock = threading.Lock() + + def get(self, run_id: str) -> PipelineStatusResponse | None: + with self._lock: + return self._runs.get(run_id) + + def set(self, run_id: str, status: PipelineStatusResponse) -> None: + with self._lock: + self._runs[run_id] = status + + def update(self, run_id: str, **updates: object) -> PipelineStatusResponse | None: + """Store a copy of the current status with ``updates`` applied and a fresh ``updated_at``.""" + with self._lock: + current = self._runs.get(run_id) + if current is None: + return None + status = current.model_copy(update={**updates, 'updated_at': _now()}) + self._runs[run_id] = status + return status class PipelineService: def __init__(self, pipeline_queue_manager: PipelineQueueManager, database: Database) -> None: self.pipeline_queue_manager = pipeline_queue_manager self.database = database - self._pipeline_runs: dict[str, PipelineStatusResponse] = {} - self._pipeline_runs_lock = threading.Lock() - - def set_pipeline_status(self, run_id: str, status: PipelineStatusResponse) -> None: - with self._pipeline_runs_lock: - self._pipeline_runs[run_id] = status + self._runs = PipelineRunStore() def get_pipeline_status(self, run_id: str) -> PipelineStatusResponse | None: - with self._pipeline_runs_lock: - status = self._pipeline_runs.get(run_id) - + status = self._runs.get(run_id) if status is None or status.status != PipelineStatus.QUEUED: return status @@ -105,44 +96,36 @@ def get_pipeline_status(self, run_id: str) -> PipelineStatusResponse | None: def enqueue_pipeline_job( self, run_id: str, - questionnaire_uuid: str, - template_uuid: str, + payload: PipelineRunRequest, template_title: str, user_uuid: str, tenant_uuid: str, - token: str, - api_url: str, - llm_config: LLMConfig, + auth: AuthenticatedUser, config: Config, ) -> None: """Queue a pipeline job; concurrency is limited by ``pipeline_queue_manager``.""" - self.set_pipeline_status( - run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.QUEUED, - questionnaire_uuid=questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - ), + run = _model_from_fields( + PipelineStatusResponse, + run_id=run_id, + status=PipelineStatus.QUEUED, + questionnaire_uuid=payload.questionnaire_uuid, + user_uuid=user_uuid, + tenant_uuid=tenant_uuid, + template_uuid=payload.template_uuid, + template_title=template_title, + updated_at=_now(), ) + self._runs.set(run_id, run) + llm_config = LLMConfig( + model=payload.llm_model, + api_key=payload.llm_api_key, + api_url=payload.llm_api_url, + parallel_workers=payload.llm_max_workers, + ) self.pipeline_queue_manager.enqueue( run_id, - lambda: self._pipeline_entrypoint( - run_id, - questionnaire_uuid, - template_uuid, - template_title, - user_uuid, - tenant_uuid, - token, - api_url, - llm_config, - config, - ), + lambda: self._run_pipeline_job(run, auth.token, auth.api_url, llm_config, config), ) async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRequest) -> PipelineStatusResponse: @@ -161,199 +144,84 @@ async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRe markdown=save_request.result_markdown, ) - updated_status = _build_pipeline_status( - run_id=pipeline_status.run_id, - status=pipeline_status.status, - questionnaire_uuid=pipeline_status.questionnaire_uuid, - knowledge_model_uuid=pipeline_status.knowledge_model_uuid, - user_uuid=pipeline_status.user_uuid, - tenant_uuid=pipeline_status.tenant_uuid, - template_uuid=pipeline_status.template_uuid, - template_title=pipeline_status.template_title, - error=pipeline_status.error, + updated_status = self._runs.update( + run_id, result_format='markdown', result_markdown=save_request.result_markdown, + progress_message=None, ) - self.set_pipeline_status(run_id, updated_status) + if updated_status is None: + raise fastapi.HTTPException(status_code=404, detail='Pipeline run not found') return updated_status - def _update_running_progress( + async def _run_pipeline_job( self, - run_id: str, - *, - questionnaire_uuid: str, - user_uuid: str, - tenant_uuid: str, - template_uuid: str, - template_title: str, - progress_message: str, - ) -> None: - with self._pipeline_runs_lock: - current = self._pipeline_runs.get(run_id) - if current is None: - return - - self.set_pipeline_status( - run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.RUNNING, - questionnaire_uuid=questionnaire_uuid, - knowledge_model_uuid=current.knowledge_model_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - progress_message=progress_message, - ), - ) - - async def _pipeline_entrypoint( - self, - run_id: str, - questionnaire_uuid: str, - template_uuid: str, - template_title: str, - user_uuid: str, - tenant_uuid: str, + run: PipelineStatusResponse, token: str, dsw_api_url: str, llm_config: LLMConfig, config: Config, ) -> None: - saver = DBSaver(self.database) - llm_client = LLMClient(llm_config.model, llm_config.api_key, llm_config.api_url, llm_config.parallel_workers) try: - template = await self.database.get_template(template_uuid) - if template is None: - self._fail_template_not_found( - questionnaire_uuid, run_id, template_title, template_uuid, tenant_uuid, user_uuid - ) - return - - await self._run_pipeline( - config, - dsw_api_url, - llm_client, - questionnaire_uuid, - run_id, - saver, - template, - template_title, - template_uuid, - tenant_uuid, - token, - user_uuid, - ) + await self._run_pipeline(run, token, dsw_api_url, llm_config, config) except Exception as error: - self.set_pipeline_status( - run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.FAILED, - questionnaire_uuid=questionnaire_uuid, - template_uuid=template_uuid, - template_title=template_title, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - error=_pipeline_error_from_exception(error), - ), - ) logger.exception('Pipeline run failed') + self._runs.update( + run.run_id, + status=PipelineStatus.FAILED, + error=_pipeline_error_from_exception(error), + progress_message=None, + ) async def _run_pipeline( self, - config: Config, - dsw_api_url: str, - llm_client: LLMClient, - questionnaire_uuid: str, - run_id: str, - saver: DBSaver, - template: dict[str, Any], - template_title: str, - template_uuid: str, - tenant_uuid: str, + run: PipelineStatusResponse, token: str, - user_uuid: str, + dsw_api_url: str, + llm_config: LLMConfig, + config: Config, ) -> None: - self.set_pipeline_status( - run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.RUNNING, - questionnaire_uuid=questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - progress_message='Starting pipeline...', - ), - ) - - def on_progress(message: str) -> None: - self._update_running_progress( + run_id = run.run_id + template = await self.database.get_template(run.template_uuid) + if template is None: + self._runs.update( run_id, - questionnaire_uuid=questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - progress_message=message, + status=PipelineStatus.FAILED, + error=PipelineErrorResponse( + type=ErrorType.TEMPLATE_NOT_FOUND, + message=TEMPLATE_NOT_FOUND_MESSAGE, + ), ) + return - pipeline = build_pipeline(database=self.database, saver=saver, config=config, llm_client=llm_client) + self._runs.update(run_id, status=PipelineStatus.RUNNING, progress_message='Starting pipeline...') + + llm_client = LLMClient(llm_config.model, llm_config.api_key, llm_config.api_url, llm_config.parallel_workers) + pipeline = build_pipeline( + database=self.database, + saver=DBSaver(self.database), + config=config, + llm_client=llm_client, + ) knowledge_model_uuid, result = await run_pipeline( - questionnaire_uuid=questionnaire_uuid, - template_uuid=template_uuid, + questionnaire_uuid=run.questionnaire_uuid, + template_uuid=run.template_uuid, template_title=template['title'], template_data=template['content'], - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, + user_uuid=run.user_uuid, + tenant_uuid=run.tenant_uuid, pipeline=pipeline, database=self.database, - on_progress=on_progress, + on_progress=lambda message: self._runs.update(run_id, progress_message=message), model_name=llm_client.get_model_name(), dsw_client=DSWClient(token, dsw_api_url), ) - self.set_pipeline_status( + self._runs.update( run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.SUCCEEDED, - questionnaire_uuid=questionnaire_uuid, - knowledge_model_uuid=knowledge_model_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - template_uuid=template_uuid, - template_title=template_title, - result_format='markdown', - result_markdown=result, - ), - ) - - def _fail_template_not_found( - self, - questionnaire_uuid: str, - run_id: str, - template_title: str, - template_uuid: str, - tenant_uuid: str, - user_uuid: str, - ) -> None: - self.set_pipeline_status( - run_id, - _build_pipeline_status( - run_id=run_id, - status=PipelineStatus.FAILED, - questionnaire_uuid=questionnaire_uuid, - template_uuid=template_uuid, - template_title=template_title, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, - error=PipelineErrorResponse( - type=ErrorType.TEMPLATE_NOT_FOUND, - message=TEMPLATE_NOT_FOUND_MESSAGE, - ), - ), + status=PipelineStatus.SUCCEEDED, + knowledge_model_uuid=knowledge_model_uuid, + result_format='markdown', + result_markdown=result, + progress_message=None, ) From b17446a97afe79dc10a1be7f827f76f34d1d1a0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 14:59:49 +0200 Subject: [PATCH 5/5] 85: Fixed ty check --- .../ai_document_plugin_service/service/pipeline_service.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/service/src/ai_document_plugin_service/service/pipeline_service.py b/service/src/ai_document_plugin_service/service/pipeline_service.py index 849ff7b..8cf2ff8 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -203,6 +203,10 @@ async def _run_pipeline( config=config, llm_client=llm_client, ) + + def on_progress(message: str) -> None: + self._runs.update(run_id, progress_message=message) + knowledge_model_uuid, result = await run_pipeline( questionnaire_uuid=run.questionnaire_uuid, template_uuid=run.template_uuid, @@ -212,7 +216,7 @@ async def _run_pipeline( tenant_uuid=run.tenant_uuid, pipeline=pipeline, database=self.database, - on_progress=lambda message: self._runs.update(run_id, progress_message=message), + on_progress=on_progress, model_name=llm_client.get_model_name(), dsw_client=DSWClient(token, dsw_api_url), )