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 01/20] 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 02/20] 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 03/20] 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 04/20] 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 05/20] 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), ) From daa0d4c2c44f9d06ae2d4541988c588b9a3c6ef4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 16:19:21 +0200 Subject: [PATCH 06/20] 75: Removed global semaphore - moved do local semaphore for each llm client --- .../ai/common/llm_client.py | 57 ++++++++++++++++--- .../ai_document_plugin_service/api/auth.py | 10 +++- .../ai_document_plugin_service/api/routes.py | 12 +--- .../ai_document_plugin_service/api/types.py | 2 - .../service/pipeline_service.py | 52 +++++++++++------ 5 files changed, 95 insertions(+), 38 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index 4281b95..ec408fb 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -15,8 +15,6 @@ logger = logging.getLogger(__name__) -semaphore = DynamicSemaphore(1) - class MissingTokenUsageError(ValueError): """Raised when a model response has no usage token information.""" @@ -77,13 +75,56 @@ def add_usage(stats: 'AssignmentStats | None', response: object) -> None: class LLMClient: - def __init__(self, model: str, api_key: str, api_url: str, parallel_workers: int | None) -> None: + """ + LLM Client for calling llm server using the OpenAI API standard. + It is able to update its config on the run. + + There should always be at most one instance of llm client per tenant! + This is because LLM client handles throttling to avoid spamming the LLM API. + """ + + def __init__(self, tenant_uuid: str) -> None: + """ + Initializes the client with empty config. Call update_config before using it. + """ + self.tenant_uuid = tenant_uuid + self.model = None + self.max_workers = None + self.semaphore = DynamicSemaphore(1) + self.api_key = None + self.api_url = None + self.client = None + + def update_config(self, model: str, api_key: str, api_url: str, parallel_workers: int | None) -> None: + """ + Changes LLMClient config. Can be called even while this class is being used in parallel by asyncio elsewhere. + If all inputs are the same as they were, nothing updates. + :param model: updated model name (or the previous) + :param api_key: updated api_key + :param api_url: updated api_url + :param parallel_workers: updated parallel workers. If worker count is being reduced, it may take a while before + the llm client reaches the reduced state. This is because when going for example from 8 to 5 workers, + LLMClient does not kill any requests, instead, it stops queueing new requests until the 3 extra requests + finish running. + :return: + """ + if ( + self.model == model + and self.api_key == api_key + and self.api_url == api_url + and self.max_workers == parallel_workers + ): + # nothing has changed, we can return + return self.model = model - self.max_workers = parallel_workers or 1 - semaphore.set_limit(max(1, self.max_workers)) + self.api_key = api_key + self.api_url = api_url + self.max_workers = max(1, parallel_workers or 1) + self.semaphore.set_limit(self.max_workers) self.client = AsyncOpenAI(api_key=api_key, base_url=api_url, max_retries=0) logger.debug( - 'Initializing LLM client, setting semaphore limit to %s', + '[llm] tenant=%s: Updated LLM client config, setting semaphore limit to %s', + self.tenant_uuid, self.max_workers, ) @@ -101,13 +142,13 @@ async def completion( req_id = uuid.uuid4().hex[:8] wait_start = time.perf_counter() logger.debug('[llm] req=%s model=%s queueing', req_id, self.model) - async with semaphore: + async with self.semaphore: wait_s = time.perf_counter() - wait_start logger.debug( '[llm] req=%s acquired semaphore after %.3fs (limit=%s)', req_id, wait_s, - semaphore.limit, + self.semaphore.limit, ) call_start = time.perf_counter() result = await self.client.chat.completions.create(*args, model=self.model, **kwargs) diff --git a/service/src/ai_document_plugin_service/api/auth.py b/service/src/ai_document_plugin_service/api/auth.py index d328da3..0244eac 100644 --- a/service/src/ai_document_plugin_service/api/auth.py +++ b/service/src/ai_document_plugin_service/api/auth.py @@ -5,6 +5,7 @@ import httpx from ai_document_plugin_service.ai.common.config import Config, normalize_project_url +from ai_document_plugin_service.api.jwt import extract_identity_from_token DSW_API_URL_HEADER = 'X-Dsw-Api-Url' DSW_USER_VALIDATION_TIMEOUT_SECONDS = 10.0 @@ -15,6 +16,8 @@ class AuthenticatedUser: token: str api_url: str + user_uuid: str + tenant_uuid: str def is_allowed_project_url(api_url: str, allowed_project_urls: tuple[str, ...]) -> bool: @@ -67,4 +70,9 @@ def verify_authenticated( if not _validate_dsw_user(normalized_api_url, token): raise fastapi.HTTPException(status_code=401, detail='Unauthorized') - return AuthenticatedUser(token=token, api_url=normalized_api_url) + try: + user_uuid, tenant_uuid = extract_identity_from_token(token) + except ValueError as error: + raise fastapi.HTTPException(status_code=400, detail=str(error)) from error + + return AuthenticatedUser(token=token, api_url=normalized_api_url, user_uuid=user_uuid, tenant_uuid=tenant_uuid) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 13054ba..5485ae2 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -3,7 +3,6 @@ import fastapi 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, PipelineRunResponse, @@ -91,18 +90,11 @@ async def start_pipeline( if template is None: raise fastapi.HTTPException(status_code=404, detail='Template not found') - try: - user_uuid, tenant_uuid = extract_identity_from_token(auth.token) - except ValueError as error: - raise fastapi.HTTPException(status_code=400, detail=str(error)) from error - run_id = str(uuid4()) pipeline.enqueue_pipeline_job( run_id, payload, template['title'], - user_uuid, - tenant_uuid, auth, config, ) @@ -111,8 +103,8 @@ async def start_pipeline( status=PipelineStatus.ACCEPTED, run_id=run_id, questionnaire_uuid=payload.questionnaire_uuid, - user_uuid=user_uuid, - tenant_uuid=tenant_uuid, + user_uuid=auth.user_uuid, + tenant_uuid=auth.tenant_uuid, template_uuid=payload.template_uuid, template_title=template['title'], ) diff --git a/service/src/ai_document_plugin_service/api/types.py b/service/src/ai_document_plugin_service/api/types.py index 463fb04..70a8fff 100644 --- a/service/src/ai_document_plugin_service/api/types.py +++ b/service/src/ai_document_plugin_service/api/types.py @@ -75,8 +75,6 @@ class PipelineStatusResponse(ApiModel): status: PipelineStatus questionnaire_uuid: str = Field(alias='questionnaireUuid') knowledge_model_uuid: str | None = Field(default=None, alias='knowledgeModelUuid') - user_uuid: str = Field(alias='userUuid') - tenant_uuid: str = Field(alias='tenantUuid') template_uuid: str = Field(alias='templateUuid') template_title: str = Field(alias='templateTitle') error: PipelineErrorResponse | None = None 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 8cf2ff8..119b7d7 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -50,6 +50,27 @@ def _now() -> str: return datetime.now(tz=UTC).isoformat() +class LlmClientTenantStore: + """ + Manages LLM Clients for different tenants. Each tenant has its own LLM client with its own config and limits + """ + + def __init__(self) -> None: + # tenant id -> llm client + self._clients: dict[str, LLMClient] = {} + self._lock = threading.Lock() + + def get_llm_client(self, tenant_uuid: str) -> LLMClient: + """ + Returns LLM client, creates a new one if it currently doesn't exist + :param tenant_uuid: Tenant to get the LLM client for. + """ + with self._lock: + if tenant_uuid not in self._clients: + self._clients[tenant_uuid] = LLMClient(tenant_uuid) + return self._clients[tenant_uuid] + + class PipelineRunStore: """Thread-safe in-memory store of pipeline run statuses.""" @@ -81,6 +102,7 @@ def __init__(self, pipeline_queue_manager: PipelineQueueManager, database: Datab self.pipeline_queue_manager = pipeline_queue_manager self.database = database self._runs = PipelineRunStore() + self._llm_clients = LlmClientTenantStore() def get_pipeline_status(self, run_id: str) -> PipelineStatusResponse | None: status = self._runs.get(run_id) @@ -98,8 +120,6 @@ def enqueue_pipeline_job( run_id: str, payload: PipelineRunRequest, template_title: str, - user_uuid: str, - tenant_uuid: str, auth: AuthenticatedUser, config: Config, ) -> None: @@ -109,8 +129,6 @@ def enqueue_pipeline_job( 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(), @@ -125,10 +143,11 @@ def enqueue_pipeline_job( ) self.pipeline_queue_manager.enqueue( run_id, - lambda: self._run_pipeline_job(run, auth.token, auth.api_url, llm_config, config), + lambda: self._run_pipeline_job(run, auth, llm_config, config), ) - 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, + auth: AuthenticatedUser) -> 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') @@ -139,8 +158,8 @@ async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRe 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, + user_uuid=auth.user_uuid, + tenant_uuid=auth.tenant_uuid, markdown=save_request.result_markdown, ) @@ -157,13 +176,12 @@ async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRe async def _run_pipeline_job( self, run: PipelineStatusResponse, - token: str, - dsw_api_url: str, + auth: AuthenticatedUser, llm_config: LLMConfig, config: Config, ) -> None: try: - await self._run_pipeline(run, token, dsw_api_url, llm_config, config) + await self._run_pipeline(run, auth, llm_config, config) except Exception as error: logger.exception('Pipeline run failed') self._runs.update( @@ -176,8 +194,7 @@ async def _run_pipeline_job( async def _run_pipeline( self, run: PipelineStatusResponse, - token: str, - dsw_api_url: str, + auth: AuthenticatedUser, llm_config: LLMConfig, config: Config, ) -> None: @@ -196,7 +213,8 @@ async def _run_pipeline( 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) + llm_client = self._llm_clients.get_llm_client(auth.tenant_uuid) + llm_client.update_config(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), @@ -212,13 +230,13 @@ def on_progress(message: str) -> None: template_uuid=run.template_uuid, template_title=template['title'], template_data=template['content'], - user_uuid=run.user_uuid, - tenant_uuid=run.tenant_uuid, + user_uuid=auth.user_uuid, + tenant_uuid=auth.tenant_uuid, pipeline=pipeline, database=self.database, on_progress=on_progress, model_name=llm_client.get_model_name(), - dsw_client=DSWClient(token, dsw_api_url), + dsw_client=DSWClient(auth.token, auth.api_url), ) self._runs.update( From b53c996bf36272503cb53574bb531ce3a122ca19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 16:29:03 +0200 Subject: [PATCH 07/20] 75: Removed global semaphore - moved do local semaphore for each llm client --- .../ai/common/llm_client.py | 25 ++++++++++++++++++- .../ai_document_plugin_service/api/routes.py | 4 +-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index ec408fb..70b4916 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -93,7 +93,7 @@ def __init__(self, tenant_uuid: str) -> None: self.semaphore = DynamicSemaphore(1) self.api_key = None self.api_url = None - self.client = None + self.client: AsyncOpenAI | None = None def update_config(self, model: str, api_key: str, api_url: str, parallel_workers: int | None) -> None: """ @@ -129,9 +129,13 @@ def update_config(self, model: str, api_key: str, api_url: str, parallel_workers ) def get_max_workers(self) -> int: + if self.max_workers is None: + raise RuntimeError('max_workers is None for tenant %s. `get_max_workers` was accessed before calling update_config', self.tenant_uuid) return self.max_workers def get_model_name(self) -> str: + if self.model is None: + raise RuntimeError('model is None for tenant %s. `get_model_name` was accessed before calling update_config', self.tenant_uuid) return self.model async def completion( @@ -139,6 +143,25 @@ async def completion( *args: Any, # noqa: ANN401 **kwargs: Any, # noqa: ANN401 ) -> ChatCompletion: + if self.model is None: + raise ValueError( + "LLM 'model' for tenant % is None, did you call `update_config` before using the client?", + self.tenant_uuid) + if self.max_workers is None: + raise ValueError( + "LLM 'max_workers' for tenant % is None, did you call `update_config` before using the client?", + self.tenant_uuid) + if self.api_url is None: + raise ValueError( + "LLM .api_url' for tenant % is None, did you call `update_config` before using the client?", + self.tenant_uuid) + if self.api_key is None: + raise ValueError( + "LLM 'api_key' for tenant % is None, did you call `update_config` before using the client?", + self.tenant_uuid) + if self.client is None: + raise RuntimeError('LLM internal client is null but api_key and api_url is set for tenant %s.', + self.tenant_uuid) req_id = uuid.uuid4().hex[:8] wait_start = time.perf_counter() logger.debug('[llm] req=%s model=%s queueing', req_id, self.model) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 5485ae2..101b87f 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -120,6 +120,6 @@ def get_pipeline_status(run_id: str, pipeline: PipelineServiceDI) -> PipelineSta @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, auth: AuthenticatedDI ) -> PipelineStatusResponse: - return await pipeline.update_pipeline_result(run_id, save_request) + return await pipeline.update_pipeline_result(run_id, save_request, auth) From a202b96e5bc434412b6d6a4951e6ebb9575bc331 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 16:41:15 +0200 Subject: [PATCH 08/20] 75: Fix ruff errors --- .../ai/common/llm_client.py | 33 ++++++++++--------- .../service/pipeline_service.py | 5 +-- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index 70b4916..85aecdb 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -16,6 +16,13 @@ logger = logging.getLogger(__name__) +class InvalidLLMConfigError(ValueError): + def __init__(self, var_name: str, tenant: str) -> None: + super().__init__( + f"LLM '{var_name}' for tenant {tenant} is None, did you call `update_config` before using the client?" + ) + + class MissingTokenUsageError(ValueError): """Raised when a model response has no usage token information.""" @@ -130,12 +137,14 @@ def update_config(self, model: str, api_key: str, api_url: str, parallel_workers def get_max_workers(self) -> int: if self.max_workers is None: - raise RuntimeError('max_workers is None for tenant %s. `get_max_workers` was accessed before calling update_config', self.tenant_uuid) + msg = 'max_workers is None for tenant %s. `get_max_workers` was accessed before calling update_config' + raise RuntimeError(msg, self.tenant_uuid) return self.max_workers def get_model_name(self) -> str: if self.model is None: - raise RuntimeError('model is None for tenant %s. `get_model_name` was accessed before calling update_config', self.tenant_uuid) + msg = 'model is None for tenant %s. `get_model_name` was accessed before calling update_config' + raise RuntimeError(msg, self.tenant_uuid) return self.model async def completion( @@ -144,24 +153,16 @@ async def completion( **kwargs: Any, # noqa: ANN401 ) -> ChatCompletion: if self.model is None: - raise ValueError( - "LLM 'model' for tenant % is None, did you call `update_config` before using the client?", - self.tenant_uuid) + raise InvalidLLMConfigError('model', self.tenant_uuid) # noqa: EM101 if self.max_workers is None: - raise ValueError( - "LLM 'max_workers' for tenant % is None, did you call `update_config` before using the client?", - self.tenant_uuid) + raise InvalidLLMConfigError('max_workers', self.tenant_uuid) # noqa: EM101 if self.api_url is None: - raise ValueError( - "LLM .api_url' for tenant % is None, did you call `update_config` before using the client?", - self.tenant_uuid) + raise InvalidLLMConfigError('api_url', self.tenant_uuid) # noqa: EM101 if self.api_key is None: - raise ValueError( - "LLM 'api_key' for tenant % is None, did you call `update_config` before using the client?", - self.tenant_uuid) + raise InvalidLLMConfigError('api_key', self.tenant_uuid) # noqa: EM101 if self.client is None: - raise RuntimeError('LLM internal client is null but api_key and api_url is set for tenant %s.', - self.tenant_uuid) + msg = f'LLM internal client is null but api_key and api_url is set for tenant {self.tenant_uuid}.' + raise RuntimeError(msg) req_id = uuid.uuid4().hex[:8] wait_start = time.perf_counter() logger.debug('[llm] req=%s model=%s queueing', req_id, self.model) 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 119b7d7..1fc568c 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -146,8 +146,9 @@ def enqueue_pipeline_job( lambda: self._run_pipeline_job(run, auth, llm_config, config), ) - async def update_pipeline_result(self, run_id: str, save_request: PipelineSaveRequest, - auth: AuthenticatedUser) -> PipelineStatusResponse: + async def update_pipeline_result( + self, run_id: str, save_request: PipelineSaveRequest, auth: AuthenticatedUser + ) -> 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') From afbf9f4f2ffc671e9772b5b9015864292ad4e777 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 17:20:22 +0200 Subject: [PATCH 09/20] 75: Improved logging --- .../ai_document_plugin_service/ai/common/llm_client.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index 85aecdb..2ca13ea 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -165,11 +165,12 @@ async def completion( raise RuntimeError(msg) req_id = uuid.uuid4().hex[:8] wait_start = time.perf_counter() - logger.debug('[llm] req=%s model=%s queueing', req_id, self.model) + logger.debug('[llm] tenant=%s req=%s model=%s queueing', self.tenant_uuid, req_id, self.model) async with self.semaphore: wait_s = time.perf_counter() - wait_start logger.debug( - '[llm] req=%s acquired semaphore after %.3fs (limit=%s)', + '[llm] tenant=%s req=%s acquired semaphore after %.3fs (limit=%s)', + self.tenant_uuid, req_id, wait_s, self.semaphore.limit, @@ -177,7 +178,8 @@ async def completion( call_start = time.perf_counter() result = await self.client.chat.completions.create(*args, model=self.model, **kwargs) logger.debug( - '[llm] req=%s completed in %.3fs (releasing semaphore)', + '[llm] tenant=%s req=%s completed in %.3fs (releasing semaphore)', + self.tenant_uuid, req_id, time.perf_counter() - call_start, ) From fefce4eed44ab4560217895db4a1fa7ef7d1d926 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 3 Jul 2026 17:49:59 +0200 Subject: [PATCH 10/20] 73: Moved strings to UUIDs --- .../ai/common/llm_client.py | 4 +- .../ai/knowledgemodel/dsw_client.py | 4 +- .../ai/persistence/database.py | 49 ++++++++++--------- .../ai/run_pipeline.py | 17 ++++--- .../ai_document_plugin_service/api/auth.py | 5 +- .../src/ai_document_plugin_service/api/jwt.py | 5 +- .../ai_document_plugin_service/api/routes.py | 15 ++---- .../ai_document_plugin_service/api/types.py | 13 ++--- .../service/pipeline_service.py | 9 ++-- 9 files changed, 62 insertions(+), 59 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/common/llm_client.py b/service/src/ai_document_plugin_service/ai/common/llm_client.py index 2ca13ea..111e580 100644 --- a/service/src/ai_document_plugin_service/ai/common/llm_client.py +++ b/service/src/ai_document_plugin_service/ai/common/llm_client.py @@ -17,7 +17,7 @@ class InvalidLLMConfigError(ValueError): - def __init__(self, var_name: str, tenant: str) -> None: + def __init__(self, var_name: str, tenant: uuid.UUID) -> None: super().__init__( f"LLM '{var_name}' for tenant {tenant} is None, did you call `update_config` before using the client?" ) @@ -90,7 +90,7 @@ class LLMClient: This is because LLM client handles throttling to avoid spamming the LLM API. """ - def __init__(self, tenant_uuid: str) -> None: + def __init__(self, tenant_uuid: uuid.UUID) -> None: """ Initializes the client with empty config. Call update_config before using it. """ diff --git a/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py b/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py index 0b35303..57f9f3a 100644 --- a/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py +++ b/service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py @@ -1,3 +1,5 @@ +from uuid import UUID + import httpx @@ -6,7 +8,7 @@ def __init__(self, token: str, api_url: str) -> None: self.token = token self.api_url = api_url.rstrip('/') - async def get_questionnaire_detail(self, questionnaire_uuid: str) -> dict: + async def get_questionnaire_detail(self, questionnaire_uuid: str | UUID) -> dict: url = f'{self.api_url}/projects/{questionnaire_uuid}/questionnaire' headers: dict[str, str] = {} diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index d34e094..f6bee9a 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import Any +from uuid import UUID from sqlalchemy import Connection, inspect from sqlalchemy.dialects.postgresql import insert as postgresql_insert @@ -12,6 +13,7 @@ from ai_document_plugin_service.ai.common.config import DatabaseConfig from ai_document_plugin_service.ai.persistence.schema import create_persistence_schema +from ai_document_plugin_service.api.types import TemplateDetail logger = logging.getLogger(__name__) @@ -63,7 +65,7 @@ async def list_templates(self) -> list[dict[str, str]]: """List available templates from a database backend.""" @abstractmethod - async def get_template(self, template_uuid: str) -> dict[str, Any] | None: + async def get_template(self, template_uuid: UUID) -> TemplateDetail | None: """Get a template record from a database backend.""" @abstractmethod @@ -81,10 +83,10 @@ async def save_result( @abstractmethod async def save_stats( self, - template_uuid: str, - knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + template_uuid: UUID, + knowledge_model_uuid: UUID, + user_uuid: UUID, + tenant_uuid: UUID, stats: JsonValue, ) -> None: """Persist a stats result in a database backend.""" @@ -92,10 +94,10 @@ async def save_stats( @abstractmethod async def update_result( self, - template_uuid: str, - knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + template_uuid: UUID, + knowledge_model_uuid: UUID, + user_uuid: UUID, + tenant_uuid: UUID, markdown: str, ) -> None: """Persist a markdown result in a database backend.""" @@ -294,7 +296,7 @@ async def list_templates(self) -> list[dict[str, str]]: for row in rows ] - async def get_template(self, template_uuid: str) -> dict[str, Any] | None: + async def get_template(self, template_uuid: UUID) -> TemplateDetail | None: await self._ensure_schema() statement = self.template_table.select().where(self.template_table.c.uuid == template_uuid) @@ -309,12 +311,11 @@ async def get_template(self, template_uuid: str) -> dict[str, Any] | None: self.schema_name, ) return None - - return { - 'uuid': str(row.uuid), - 'title': row.title, - 'content': row.content, - } + return TemplateDetail( + uuid=row.uuid, + title=row.title, + content=row.content, + ) async def save_result( self, @@ -359,10 +360,10 @@ async def save_result( async def save_stats( self, - template_uuid: str, - knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + template_uuid: UUID, + knowledge_model_uuid: UUID, + user_uuid: UUID, + tenant_uuid: UUID, stats: JsonValue, ) -> None: await self._ensure_schema() @@ -394,10 +395,10 @@ async def save_stats( async def update_result( self, - template_uuid: str, - knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + template_uuid: UUID, + knowledge_model_uuid: UUID, + user_uuid: UUID, + tenant_uuid: UUID, markdown: str, ) -> None: await self._ensure_schema() diff --git a/service/src/ai_document_plugin_service/ai/run_pipeline.py b/service/src/ai_document_plugin_service/ai/run_pipeline.py index 652fa7d..f5ba565 100644 --- a/service/src/ai_document_plugin_service/ai/run_pipeline.py +++ b/service/src/ai_document_plugin_service/ai/run_pipeline.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from collections.abc import Mapping + from uuid import UUID from haystack.components.routers.conditional_router import Route @@ -110,12 +111,12 @@ def build_pipeline(database: Database, saver: DBSaver, config: Config, llm_clien async def run_pipeline( - questionnaire_uuid: str, - template_uuid: str, + questionnaire_uuid: UUID, + template_uuid: UUID, template_title: str, template_data: Mapping[str, object], - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, pipeline: AsyncPipeline, database: Database, dsw_client: DSWClient, @@ -198,10 +199,10 @@ async def run_pipeline( async def write_metrics( database: Database, - template_uuid: str, - knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + template_uuid: UUID, + knowledge_model_uuid: UUID, + user_uuid: UUID, + tenant_uuid: UUID, result: Mapping[str, object], model_name: str, t1: float, diff --git a/service/src/ai_document_plugin_service/api/auth.py b/service/src/ai_document_plugin_service/api/auth.py index 0244eac..3211001 100644 --- a/service/src/ai_document_plugin_service/api/auth.py +++ b/service/src/ai_document_plugin_service/api/auth.py @@ -1,5 +1,6 @@ from dataclasses import dataclass from typing import Annotated +from uuid import UUID import fastapi import httpx @@ -16,8 +17,8 @@ class AuthenticatedUser: token: str api_url: str - user_uuid: str - tenant_uuid: str + user_uuid: UUID + tenant_uuid: UUID def is_allowed_project_url(api_url: str, allowed_project_urls: tuple[str, ...]) -> bool: diff --git a/service/src/ai_document_plugin_service/api/jwt.py b/service/src/ai_document_plugin_service/api/jwt.py index 7c769ef..d3cd431 100644 --- a/service/src/ai_document_plugin_service/api/jwt.py +++ b/service/src/ai_document_plugin_service/api/jwt.py @@ -1,5 +1,6 @@ import base64 import json +import uuid from uuid import UUID JWT_PART_COUNT = 2 @@ -42,8 +43,8 @@ def _get_required_uuid_claim(payload: dict[str, object], *keys: str) -> str: raise ValueError(msg) -def extract_identity_from_token(token: str) -> tuple[str, str]: +def extract_identity_from_token(token: str) -> tuple[UUID, UUID]: payload = decode_jwt_payload(token) user_uuid = _get_required_uuid_claim(payload, 'user_uuid', 'userUuid') tenant_uuid = _get_required_uuid_claim(payload, 'tenant_uuid', 'tenantUuid') - return user_uuid, tenant_uuid + return uuid.UUID(user_uuid), uuid.UUID(tenant_uuid) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 101b87f..a5be9bf 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -1,4 +1,4 @@ -from uuid import uuid4 +from uuid import UUID, uuid4 import fastapi @@ -31,18 +31,13 @@ async def list_templates(database: DatabaseDI) -> list[TemplateListItem]: @protected_router.get('/templates/{template_uuid}') -async def get_template(template_uuid: str, database: DatabaseDI) -> TemplateDetail: +async def get_template(template_uuid: UUID, database: DatabaseDI) -> TemplateDetail: template = await database.get_template(template_uuid) if template is None: raise fastapi.HTTPException(status_code=404, detail='Template not found') - return _model_from_fields( - TemplateDetail, - uuid=template['uuid'], - title=template['title'], - content=template['content'], - ) + return template @protected_router.post('/templates', status_code=201) @@ -94,7 +89,7 @@ async def start_pipeline( pipeline.enqueue_pipeline_job( run_id, payload, - template['title'], + template.title, auth, config, ) @@ -106,7 +101,7 @@ async def start_pipeline( user_uuid=auth.user_uuid, tenant_uuid=auth.tenant_uuid, template_uuid=payload.template_uuid, - template_title=template['title'], + template_title=template.title, ) diff --git a/service/src/ai_document_plugin_service/api/types.py b/service/src/ai_document_plugin_service/api/types.py index 70a8fff..4f33b26 100644 --- a/service/src/ai_document_plugin_service/api/types.py +++ b/service/src/ai_document_plugin_service/api/types.py @@ -1,4 +1,5 @@ from enum import StrEnum +from uuid import UUID from pydantic import BaseModel, ConfigDict, Field @@ -34,7 +35,7 @@ class TemplateListItem(ApiModel): class TemplateDetail(ApiModel): - uuid: str + uuid: UUID title: str content: dict @@ -45,8 +46,8 @@ class TemplateCreateRequest(ApiModel): class PipelineRunRequest(ApiModel): - questionnaire_uuid: str = Field(alias='questionnaireUuid') - template_uuid: str = Field(alias='templateUuid') + questionnaire_uuid: UUID = Field(alias='questionnaireUuid') + template_uuid: UUID = Field(alias='templateUuid') llm_model: str = Field(alias='llmModel') llm_api_key: str = Field(alias='llmApiKey') llm_api_url: str = Field(alias='llmApiUrl') @@ -73,9 +74,9 @@ class PipelineErrorResponse(ApiModel): class PipelineStatusResponse(ApiModel): run_id: str = Field(alias='runId') status: PipelineStatus - questionnaire_uuid: str = Field(alias='questionnaireUuid') - knowledge_model_uuid: str | None = Field(default=None, alias='knowledgeModelUuid') - template_uuid: str = Field(alias='templateUuid') + questionnaire_uuid: UUID = Field(alias='questionnaireUuid') + knowledge_model_uuid: UUID | None = Field(default=None, alias='knowledgeModelUuid') + template_uuid: UUID = Field(alias='templateUuid') template_title: str = Field(alias='templateTitle') error: PipelineErrorResponse | None = None result_format: str | None = Field(default=None, alias='resultFormat') 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 1fc568c..daa6a0d 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -1,6 +1,7 @@ import logging import threading from datetime import UTC, datetime +from uuid import UUID import fastapi from openai import AuthenticationError @@ -57,10 +58,10 @@ class LlmClientTenantStore: def __init__(self) -> None: # tenant id -> llm client - self._clients: dict[str, LLMClient] = {} + self._clients: dict[UUID, LLMClient] = {} self._lock = threading.Lock() - def get_llm_client(self, tenant_uuid: str) -> LLMClient: + def get_llm_client(self, tenant_uuid: UUID) -> LLMClient: """ Returns LLM client, creates a new one if it currently doesn't exist :param tenant_uuid: Tenant to get the LLM client for. @@ -229,8 +230,8 @@ def on_progress(message: str) -> None: knowledge_model_uuid, result = await run_pipeline( questionnaire_uuid=run.questionnaire_uuid, template_uuid=run.template_uuid, - template_title=template['title'], - template_data=template['content'], + template_title=template.title, + template_data=template.content, user_uuid=auth.user_uuid, tenant_uuid=auth.tenant_uuid, pipeline=pipeline, From 27bf46da568f2158fbd37d8b8b6837bb772a1b04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Thu, 9 Jul 2026 23:03:41 +0200 Subject: [PATCH 11/20] Added db migration --- ...20260709_01_add_tenant_uuid_to_template.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py diff --git a/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py b/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py new file mode 100644 index 0000000..5d66292 --- /dev/null +++ b/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py @@ -0,0 +1,55 @@ +"""Add tenant_uuid to template table. + +Templates were not previously scoped per tenant, so there is no way to +backfill tenant_uuid for existing rows. This migration drops all existing +data in template (and, via cascade, assignment and result which reference +it) before adding the column as NOT NULL. + +Revision ID: 20260709_01 +Revises: 20260521_02 +Create Date: 2026-07-09 00:00:00 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import context, op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '20260709_01' +down_revision = '20260521_02' +branch_labels = None +depends_on = None + + +def _qualified_table_reference(schema: str | None, table: str) -> str: + if schema: + return f'{schema}.{table}' + return table + + +def upgrade() -> None: + schema = context.get_context().version_table_schema + # Deletes the current template table - + # I think it is ok trade of to making this backwards compatible with nullable tenant_uuid. + # Alternative is making the tenant_uuid nullable and showing null templates to everyone, + # but then there are issues with deleting... + op.execute( + sa.text( + f'TRUNCATE TABLE {_qualified_table_reference(schema, "result")}, ' + f'{_qualified_table_reference(schema, "assignment")}, ' + f'{_qualified_table_reference(schema, "template")} CASCADE', + ), + ) + + op.add_column( + 'template', + sa.Column('tenant_uuid', postgresql.UUID(as_uuid=True), nullable=False), + schema=schema, + ) + + +def downgrade() -> None: + schema = context.get_context().version_table_schema + op.drop_column('template', 'tenant_uuid', schema=schema) From 1c59629b1c4b44d75dd10ced65e4419cc1b7101d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Thu, 9 Jul 2026 23:28:51 +0200 Subject: [PATCH 12/20] Updated db --- .../persistence/assignment_saver_component.py | 10 +++++- .../ai/persistence/database.py | 31 +++++++++++++------ .../ai/persistence/schema.py | 3 +- .../ai/run_pipeline.py | 1 + .../ai_document_plugin_service/api/routes.py | 15 +++++---- ...20260709_01_add_tenant_uuid_to_template.py | 12 +++++++ .../service/pipeline_service.py | 2 +- 7 files changed, 56 insertions(+), 18 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index 738b105..c888f68 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING, TypedDict +from uuid import UUID from haystack import component @@ -47,6 +48,7 @@ async def run_async( template_uuid: str, template_title: str, template_data: JsonValue, + tenant_uuid: UUID, assignments: list[SectionAssignment], stats: AssignmentStats | None = None, ) -> AssignmentSaverComponentResult: @@ -63,6 +65,7 @@ async def run_async( template_uuid=template_uuid, template_title=template_title, template_data=template_data, + tenant_uuid=tenant_uuid, created_at=datetime.now(tz=UTC), ) @@ -81,6 +84,7 @@ def run( template_uuid: str, template_title: str, template_data: JsonValue, + tenant_uuid: UUID, assignments: list[SectionAssignment], stats: AssignmentStats | None = None, ) -> AssignmentSaverComponentResult: @@ -103,6 +107,7 @@ async def save( template_uuid: str, template_title: str, template_data: JsonValue, + tenant_uuid: UUID, created_at: datetime | None = None, ) -> None: """Persist assignments and their template.""" @@ -119,9 +124,10 @@ async def save( template_uuid: str, template_title: str, template_data: JsonValue, + tenant_uuid: UUID, created_at: datetime | None = None, ) -> None: - _ = (template_uuid, template_title, template_data) + _ = (template_uuid, template_title, template_data, tenant_uuid) output_name = self._build_filename( knowledge_model_uuid, knowledge_model_name, @@ -176,12 +182,14 @@ async def save( template_uuid: str, template_title: str, template_data: JsonValue, + tenant_uuid: UUID, created_at: datetime | None = None, ) -> None: await self.database.save_template( uuid=template_uuid, title=template_title, content=template_data, + tenant_uuid=tenant_uuid, ) await self.database.save_assignments( knowledge_model_uuid=knowledge_model_uuid, diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index f6bee9a..76eb5bc 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -27,6 +27,7 @@ async def create_template( uuid: str, title: str, content: JsonValue, + tenant_uuid: UUID, ) -> None: """Create a new template in a database backend.""" @@ -49,6 +50,7 @@ async def save_template( uuid: str, title: str, content: JsonValue, + tenant_uuid: UUID, ) -> None: """Persist a template in a database backend.""" @@ -61,11 +63,11 @@ async def get_assignments( """Get assignments from a database backend.""" @abstractmethod - async def list_templates(self) -> list[dict[str, str]]: + async def list_templates(self, tenant_uuid: UUID) -> list[dict[str, str]]: """List available templates from a database backend.""" @abstractmethod - async def get_template(self, template_uuid: UUID) -> TemplateDetail | None: + async def get_template(self, template_uuid: UUID, tenant_uuid: UUID) -> TemplateDetail | None: """Get a template record from a database backend.""" @abstractmethod @@ -74,7 +76,7 @@ async def save_result( template_uuid: str, knowledge_model_uuid: str, user_uuid: str, - tenant_uuid: str, + tenant_uuid: UUID, prepolished_markdown: str, markdown: str, ) -> None: @@ -198,12 +200,14 @@ async def create_template( uuid: str, title: str, content: JsonValue, + tenant_uuid: UUID, ) -> None: await self._ensure_schema() statement = postgresql_insert(self.template_table).values( uuid=uuid, title=title, content=content, + tenant_uuid=tenant_uuid, ) try: @@ -224,12 +228,14 @@ async def save_template( uuid: str, title: str, content: JsonValue, + tenant_uuid: UUID, ) -> None: await self._ensure_schema() statement = postgresql_insert(self.template_table).values( uuid=uuid, title=title, content=content, + tenant_uuid=tenant_uuid, ) upsert_statement = statement.on_conflict_do_update( index_elements=[self.template_table.c.uuid], @@ -280,9 +286,13 @@ async def get_assignments( return row.assignments - async def list_templates(self) -> list[dict[str, str]]: + async def list_templates(self, tenant_uuid: UUID) -> list[dict[str, str]]: await self._ensure_schema() - statement = self.template_table.select().order_by(self.template_table.c.title.asc()) + statement = ( + self.template_table.select() + .where(self.template_table.c.tenant_uuid == tenant_uuid) + .order_by(self.template_table.c.title.asc()) + ) async with self.engine.begin() as connection: result = await connection.execute(statement) @@ -296,9 +306,11 @@ async def list_templates(self) -> list[dict[str, str]]: for row in rows ] - async def get_template(self, template_uuid: UUID) -> TemplateDetail | None: + async def get_template(self, template_uuid: UUID, tenant_uuid: UUID) -> TemplateDetail | None: await self._ensure_schema() - statement = self.template_table.select().where(self.template_table.c.uuid == template_uuid) + statement = self.template_table.select().where( + (self.template_table.c.uuid == template_uuid) & (self.template_table.c.tenant_uuid == tenant_uuid), + ) async with self.engine.begin() as connection: result = await connection.execute(statement) @@ -306,8 +318,9 @@ async def get_template(self, template_uuid: UUID) -> TemplateDetail | None: if row is None: logger.debug( - 'No template found for uuid=%s in %s.template', + 'No template found for uuid=%s tenant_uuid=%s in %s.template', template_uuid, + tenant_uuid, self.schema_name, ) return None @@ -322,7 +335,7 @@ async def save_result( template_uuid: str, knowledge_model_uuid: str, user_uuid: str, - tenant_uuid: str, + tenant_uuid: UUID, prepolished_markdown: str, markdown: str, ) -> None: diff --git a/service/src/ai_document_plugin_service/ai/persistence/schema.py b/service/src/ai_document_plugin_service/ai/persistence/schema.py index c689bcd..65e32ed 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/schema.py +++ b/service/src/ai_document_plugin_service/ai/persistence/schema.py @@ -32,7 +32,8 @@ def create_persistence_schema(schema_name: str) -> PersistenceSchema: Column('uuid', UUID(as_uuid=True), primary_key=True), Column('title', Text, nullable=False), Column('content', JSON, nullable=False), - UniqueConstraint('title', name='uq_template_title'), + Column('tenant_uuid', UUID(as_uuid=True), nullable=False), + UniqueConstraint('title', 'tenant_uuid', name='uq_template_title_tenant_uuid'), ) assignment_table = Table( diff --git a/service/src/ai_document_plugin_service/ai/run_pipeline.py b/service/src/ai_document_plugin_service/ai/run_pipeline.py index f5ba565..ee128ee 100644 --- a/service/src/ai_document_plugin_service/ai/run_pipeline.py +++ b/service/src/ai_document_plugin_service/ai/run_pipeline.py @@ -154,6 +154,7 @@ async def run_pipeline( 'template_uuid': template_uuid, 'template_title': template_title, 'template_data': template_data, + 'tenant_uuid': tenant_uuid, }, 'dmp_generator_component': { 'replies': replies, diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index a5be9bf..ede5d3f 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -26,13 +26,15 @@ def health_check() -> dict[str, str]: @protected_router.get('/templates') -async def list_templates(database: DatabaseDI) -> list[TemplateListItem]: - return [_model_from_fields(TemplateListItem, **item) for item in await database.list_templates()] +async def list_templates(database: DatabaseDI, auth: AuthenticatedDI) -> list[TemplateListItem]: + return [ + _model_from_fields(TemplateListItem, **item) for item in await database.list_templates(auth.tenant_uuid) + ] @protected_router.get('/templates/{template_uuid}') -async def get_template(template_uuid: UUID, database: DatabaseDI) -> TemplateDetail: - template = await database.get_template(template_uuid) +async def get_template(template_uuid: UUID, database: DatabaseDI, auth: AuthenticatedDI) -> TemplateDetail: + template = await database.get_template(template_uuid, auth.tenant_uuid) if template is None: raise fastapi.HTTPException(status_code=404, detail='Template not found') @@ -41,7 +43,7 @@ async def get_template(template_uuid: UUID, database: DatabaseDI) -> TemplateDet @protected_router.post('/templates', status_code=201) -async def create_template(payload: TemplateCreateRequest, database: DatabaseDI) -> TemplateDetail: +async def create_template(payload: TemplateCreateRequest, database: DatabaseDI, auth: AuthenticatedDI) -> TemplateDetail: trimmed_title = payload.title.strip() if not trimmed_title: raise fastapi.HTTPException(status_code=400, detail='Template title is required') @@ -60,6 +62,7 @@ async def create_template(payload: TemplateCreateRequest, database: DatabaseDI) uuid=template_uuid, title=trimmed_title, content=payload.content, + tenant_uuid=auth.tenant_uuid, ) except ValueError as error: raise fastapi.HTTPException(status_code=409, detail=str(error)) from error @@ -80,7 +83,7 @@ async def start_pipeline( database: DatabaseDI, pipeline: PipelineServiceDI, ) -> PipelineRunResponse: - template = await database.get_template(payload.template_uuid) + template = await database.get_template(payload.template_uuid, auth.tenant_uuid) if template is None: raise fastapi.HTTPException(status_code=404, detail='Template not found') diff --git a/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py b/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py index 5d66292..4c69fec 100644 --- a/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py +++ b/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py @@ -49,7 +49,19 @@ def upgrade() -> None: schema=schema, ) + op.drop_constraint('uq_template_title', 'template', schema=schema, type_='unique') + op.create_unique_constraint( + 'uq_template_title_tenant_uuid', + 'template', + ['title', 'tenant_uuid'], + schema=schema, + ) + def downgrade() -> None: schema = context.get_context().version_table_schema + + op.drop_constraint('uq_template_title_tenant_uuid', 'template', schema=schema, type_='unique') + op.create_unique_constraint('uq_template_title', 'template', ['title'], schema=schema) + op.drop_column('template', 'tenant_uuid', schema=schema) 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 daa6a0d..b0f7cac 100644 --- a/service/src/ai_document_plugin_service/service/pipeline_service.py +++ b/service/src/ai_document_plugin_service/service/pipeline_service.py @@ -201,7 +201,7 @@ async def _run_pipeline( config: Config, ) -> None: run_id = run.run_id - template = await self.database.get_template(run.template_uuid) + template = await self.database.get_template(run.template_uuid, auth.tenant_uuid) if template is None: self._runs.update( run_id, From bdf32d0736afd54e0073a286cb7397c559e7b9fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Thu, 9 Jul 2026 23:30:45 +0200 Subject: [PATCH 13/20] Fix datatypes --- .../ai/persistence/database.py | 4 ++-- .../ai/persistence/saver_component.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index 76eb5bc..5fe6069 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -75,7 +75,7 @@ async def save_result( self, template_uuid: str, knowledge_model_uuid: str, - user_uuid: str, + user_uuid: UUID, tenant_uuid: UUID, prepolished_markdown: str, markdown: str, @@ -334,7 +334,7 @@ async def save_result( self, template_uuid: str, knowledge_model_uuid: str, - user_uuid: str, + user_uuid: UUID, tenant_uuid: UUID, prepolished_markdown: str, markdown: str, diff --git a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py index 0748d22..93eb5d8 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py @@ -1,5 +1,6 @@ import typing from typing import TypedDict +from uuid import UUID from haystack import component @@ -20,8 +21,8 @@ async def run_async( self, template_uuid: str, knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, debug_markdown: str, markdown: str, ) -> FileSaverComponentResult: @@ -44,8 +45,8 @@ def run( self, template_uuid: str, knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, debug_markdown: str, markdown: str, ) -> FileSaverComponentResult: From 03f3499e6febaee22ce347471db6283dd9900271 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Thu, 9 Jul 2026 23:32:52 +0200 Subject: [PATCH 14/20] Fix ruff --- .../ai/persistence/assignment_saver_component.py | 3 ++- service/src/ai_document_plugin_service/api/routes.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index c888f68..0f9c781 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -10,7 +10,6 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING, TypedDict -from uuid import UUID from haystack import component @@ -21,6 +20,8 @@ from ai_document_plugin_service.ai.common.types import AssignmentStats if TYPE_CHECKING: + from uuid import UUID + from ai_document_plugin_service.ai.persistence.database import Database logger = logging.getLogger(__name__) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index ede5d3f..bf145d6 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -43,7 +43,9 @@ async def get_template(template_uuid: UUID, database: DatabaseDI, auth: Authenti @protected_router.post('/templates', status_code=201) -async def create_template(payload: TemplateCreateRequest, database: DatabaseDI, auth: AuthenticatedDI) -> TemplateDetail: +async def create_template( + payload: TemplateCreateRequest, database: DatabaseDI, auth: AuthenticatedDI +) -> TemplateDetail: trimmed_title = payload.title.strip() if not trimmed_title: raise fastapi.HTTPException(status_code=400, detail='Template title is required') From 72af2994798dcbd3a99792ff36db6f1203a3c6de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Thu, 9 Jul 2026 23:57:37 +0200 Subject: [PATCH 15/20] Fix template id to UUID --- .../persistence/assignment_loader_component.py | 5 +++-- .../ai/persistence/assignment_saver_component.py | 10 +++++----- .../ai/persistence/database.py | 16 ++++++++-------- .../ai/persistence/saver_component.py | 4 ++-- .../ai_document_plugin_service/api/__init__.py | 3 --- .../src/ai_document_plugin_service/api/routes.py | 4 ++-- .../src/ai_document_plugin_service/api/types.py | 2 +- 7 files changed, 21 insertions(+), 23 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py index af3bfaa..70143cc 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py @@ -1,4 +1,5 @@ from typing import Any +from uuid import UUID from haystack import component @@ -14,7 +15,7 @@ def __init__(self, database: Database) -> None: assignments=JsonValue | None, found=bool, ) - async def run_async(self, knowledge_model_uuid: str, template_uuid: str) -> dict[str, Any]: + async def run_async(self, knowledge_model_uuid: str, template_uuid: UUID) -> dict[str, Any]: assignments = await self.database.get_assignments(knowledge_model_uuid, template_uuid) return { @@ -26,7 +27,7 @@ async def run_async(self, knowledge_model_uuid: str, template_uuid: str) -> dict assignments=JsonValue | None, found=bool, ) - def run(self, knowledge_model_uuid: str, template_uuid: str) -> dict[str, Any]: + def run(self, knowledge_model_uuid: str, template_uuid: UUID) -> dict[str, Any]: """Async-only component; the sync pipeline entrypoint is intentionally unsupported.""" msg = f'{type(self).__name__} is async-only; use run_async() / AsyncPipeline.run_async()' raise NotImplementedError( diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index 0f9c781..c419561 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -46,7 +46,7 @@ async def run_async( knowledge_model_uuid: str, knowledge_model_name: str, knowledge_model_version: str, - template_uuid: str, + template_uuid: UUID, template_title: str, template_data: JsonValue, tenant_uuid: UUID, @@ -82,7 +82,7 @@ def run( knowledge_model_uuid: str, knowledge_model_name: str, knowledge_model_version: str, - template_uuid: str, + template_uuid: UUID, template_title: str, template_data: JsonValue, tenant_uuid: UUID, @@ -105,7 +105,7 @@ async def save( knowledge_model_version: str, assignments: JsonValue, stats: StatsJson | None, - template_uuid: str, + template_uuid: UUID, template_title: str, template_data: JsonValue, tenant_uuid: UUID, @@ -122,7 +122,7 @@ async def save( knowledge_model_version: str, assignments: JsonValue, stats: StatsJson | None, - template_uuid: str, + template_uuid: UUID, template_title: str, template_data: JsonValue, tenant_uuid: UUID, @@ -180,7 +180,7 @@ async def save( knowledge_model_version: str, assignments: JsonValue, stats: StatsJson | None, - template_uuid: str, + template_uuid: UUID, template_title: str, template_data: JsonValue, tenant_uuid: UUID, diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index 5fe6069..d4bc86d 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -24,7 +24,7 @@ class Database(ABC): @abstractmethod async def create_template( self, - uuid: str, + uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, @@ -38,7 +38,7 @@ async def save_assignments( knowledge_model_name: str, knowledge_model_version: str, assignments: JsonValue, - template_uuid: str, + template_uuid: UUID, stats: JsonValue | None = None, created_at: datetime | None = None, ) -> None: @@ -58,7 +58,7 @@ async def save_template( async def get_assignments( self, knowledge_model_uuid: str, - template_uuid: str, + template_uuid: UUID, ) -> JsonValue | None: """Get assignments from a database backend.""" @@ -73,7 +73,7 @@ async def get_template(self, template_uuid: UUID, tenant_uuid: UUID) -> Template @abstractmethod async def save_result( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, user_uuid: UUID, tenant_uuid: UUID, @@ -197,7 +197,7 @@ async def save_assignments( async def create_template( self, - uuid: str, + uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, @@ -232,7 +232,7 @@ async def save_template( ) -> None: await self._ensure_schema() statement = postgresql_insert(self.template_table).values( - uuid=uuid, + uuid=UUID, title=title, content=content, tenant_uuid=tenant_uuid, @@ -257,7 +257,7 @@ async def save_template( async def get_assignments( self, knowledge_model_uuid: str, - template_uuid: str, + template_uuid: UUID, ) -> JsonValue | None: await self._ensure_schema() @@ -332,7 +332,7 @@ async def get_template(self, template_uuid: UUID, tenant_uuid: UUID) -> Template async def save_result( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, user_uuid: UUID, tenant_uuid: UUID, diff --git a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py index 93eb5d8..b19bae0 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py @@ -19,7 +19,7 @@ def __init__(self, database: Database) -> None: @component.output_types(markdown=str) async def run_async( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, user_uuid: UUID, tenant_uuid: UUID, @@ -43,7 +43,7 @@ async def run_async( @component.output_types(markdown=str) def run( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, user_uuid: UUID, tenant_uuid: UUID, diff --git a/service/src/ai_document_plugin_service/api/__init__.py b/service/src/ai_document_plugin_service/api/__init__.py index 4f91956..e69de29 100644 --- a/service/src/ai_document_plugin_service/api/__init__.py +++ b/service/src/ai_document_plugin_service/api/__init__.py @@ -1,3 +0,0 @@ -from ai_document_plugin_service.api.routes import protected_router, public_router - -__all__ = ['protected_router', 'public_router'] diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index bf145d6..1e21982 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -56,8 +56,8 @@ async def create_template( status_code=400, detail='Template JSON must contain a top-level "sections" array.', ) - - template_uuid = str(uuid4()) + # TODO: this should not be created here but inside database + template_uuid = uuid4() try: await database.create_template( diff --git a/service/src/ai_document_plugin_service/api/types.py b/service/src/ai_document_plugin_service/api/types.py index 4f33b26..8bfd1c3 100644 --- a/service/src/ai_document_plugin_service/api/types.py +++ b/service/src/ai_document_plugin_service/api/types.py @@ -58,7 +58,7 @@ class PipelineRunResponse(ApiModel): status: PipelineStatus run_id: str = Field(alias='runId') questionnaire_uuid: str = Field(alias='questionnaireUuid') - template_uuid: str = Field(alias='templateUuid') + template_uuid: UUID = Field(alias='templateUuid') template_title: str = Field(alias='templateTitle') From d7a7b266c09afc290b3ea63ad91c2178796d25d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 10 Jul 2026 00:06:11 +0200 Subject: [PATCH 16/20] Fix bugs with data types --- .../ai/persistence/assignment_saver_component.py | 3 +-- .../src/ai_document_plugin_service/ai/persistence/database.py | 2 +- service/src/ai_document_plugin_service/api/types.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index c419561..42389a0 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING, TypedDict +from uuid import UUID from haystack import component @@ -20,8 +21,6 @@ from ai_document_plugin_service.ai.common.types import AssignmentStats if TYPE_CHECKING: - from uuid import UUID - from ai_document_plugin_service.ai.persistence.database import Database logger = logging.getLogger(__name__) diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index d4bc86d..c4d9f2b 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -232,7 +232,7 @@ async def save_template( ) -> None: await self._ensure_schema() statement = postgresql_insert(self.template_table).values( - uuid=UUID, + uuid=uuid, title=title, content=content, tenant_uuid=tenant_uuid, diff --git a/service/src/ai_document_plugin_service/api/types.py b/service/src/ai_document_plugin_service/api/types.py index 8bfd1c3..a6a4f9b 100644 --- a/service/src/ai_document_plugin_service/api/types.py +++ b/service/src/ai_document_plugin_service/api/types.py @@ -57,7 +57,7 @@ class PipelineRunRequest(ApiModel): class PipelineRunResponse(ApiModel): status: PipelineStatus run_id: str = Field(alias='runId') - questionnaire_uuid: str = Field(alias='questionnaireUuid') + questionnaire_uuid: UUID = Field(alias='questionnaireUuid') template_uuid: UUID = Field(alias='templateUuid') template_title: str = Field(alias='templateTitle') From bdb0615f630a52c29ab42a1fa14e5943285512ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 10 Jul 2026 00:36:04 +0200 Subject: [PATCH 17/20] Fix issue with uuid imports --- .../ai/persistence/assignment_saver_component.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index 42389a0..163f11f 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -10,7 +10,8 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING, TypedDict -from uuid import UUID +# UUID must be imported outside TYPE_CHECKING block for haystack to work +from uuid import UUID # noqa: TC003 from haystack import component From 40edc500858a1af5a90945f0a8b190c3810309e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 10 Jul 2026 00:37:56 +0200 Subject: [PATCH 18/20] Refactored where uuid is created --- .../ai/persistence/assignment_saver_component.py | 1 + .../ai/persistence/database.py | 14 +++++++------- .../src/ai_document_plugin_service/api/routes.py | 5 +---- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py index 163f11f..530449f 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py +++ b/service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING, TypedDict + # UUID must be imported outside TYPE_CHECKING block for haystack to work from uuid import UUID # noqa: TC003 diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index c4d9f2b..4c7e744 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -3,7 +3,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from typing import Any -from uuid import UUID +from uuid import UUID, uuid4 from sqlalchemy import Connection, inspect from sqlalchemy.dialects.postgresql import insert as postgresql_insert @@ -24,12 +24,11 @@ class Database(ABC): @abstractmethod async def create_template( self, - uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, - ) -> None: - """Create a new template in a database backend.""" + ) -> UUID: + """Create a new template in a database backend. Return created template UUID""" @abstractmethod async def save_assignments( @@ -197,14 +196,14 @@ async def save_assignments( async def create_template( self, - uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, ) -> None: + template_uuid = uuid4() await self._ensure_schema() statement = postgresql_insert(self.template_table).values( - uuid=uuid, + uuid=template_uuid, title=title, content=content, tenant_uuid=tenant_uuid, @@ -219,9 +218,10 @@ async def create_template( logger.debug( 'Created template uuid=%s in %s.template', - uuid, + template_uuid, self.schema_name, ) + return template_uuid async def save_template( self, diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index 1e21982..a83ed14 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -56,12 +56,9 @@ async def create_template( status_code=400, detail='Template JSON must contain a top-level "sections" array.', ) - # TODO: this should not be created here but inside database - template_uuid = uuid4() try: - await database.create_template( - uuid=template_uuid, + template_uuid = await database.create_template( title=trimmed_title, content=payload.content, tenant_uuid=auth.tenant_uuid, From a6b7c294858934ffb984c3c2e396807d82244df5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 10 Jul 2026 00:43:51 +0200 Subject: [PATCH 19/20] Fix types --- .../ai_document_plugin_service/ai/persistence/database.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/service/src/ai_document_plugin_service/ai/persistence/database.py b/service/src/ai_document_plugin_service/ai/persistence/database.py index 4c7e744..3ad5682 100644 --- a/service/src/ai_document_plugin_service/ai/persistence/database.py +++ b/service/src/ai_document_plugin_service/ai/persistence/database.py @@ -46,7 +46,7 @@ async def save_assignments( @abstractmethod async def save_template( self, - uuid: str, + uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, @@ -159,7 +159,7 @@ async def save_assignments( knowledge_model_name: str, knowledge_model_version: str, assignments: JsonValue, - template_uuid: str, + template_uuid: UUID, stats: JsonValue | None = None, created_at: datetime | None = None, ) -> None: @@ -199,7 +199,7 @@ async def create_template( title: str, content: JsonValue, tenant_uuid: UUID, - ) -> None: + ) -> UUID: template_uuid = uuid4() await self._ensure_schema() statement = postgresql_insert(self.template_table).values( @@ -225,7 +225,7 @@ async def create_template( async def save_template( self, - uuid: str, + uuid: UUID, title: str, content: JsonValue, tenant_uuid: UUID, From b9cf551cb8e779cf7425679562439c9f6af7d224 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Frnka?= Date: Fri, 10 Jul 2026 00:44:02 +0200 Subject: [PATCH 20/20] Fix formatting --- service/src/ai_document_plugin_service/api/routes.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/service/src/ai_document_plugin_service/api/routes.py b/service/src/ai_document_plugin_service/api/routes.py index a83ed14..bf4b7c1 100644 --- a/service/src/ai_document_plugin_service/api/routes.py +++ b/service/src/ai_document_plugin_service/api/routes.py @@ -27,9 +27,7 @@ def health_check() -> dict[str, str]: @protected_router.get('/templates') async def list_templates(database: DatabaseDI, auth: AuthenticatedDI) -> list[TemplateListItem]: - return [ - _model_from_fields(TemplateListItem, **item) for item in await database.list_templates(auth.tenant_uuid) - ] + return [_model_from_fields(TemplateListItem, **item) for item in await database.list_templates(auth.tenant_uuid)] @protected_router.get('/templates/{template_uuid}')