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/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 738b105..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 @@ -11,6 +11,9 @@ 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 + from haystack import component from ai_document_plugin_service.ai.assignment.types import ( @@ -44,9 +47,10 @@ 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, assignments: list[SectionAssignment], stats: AssignmentStats | None = None, ) -> AssignmentSaverComponentResult: @@ -63,6 +67,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), ) @@ -78,9 +83,10 @@ 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, assignments: list[SectionAssignment], stats: AssignmentStats | None = None, ) -> AssignmentSaverComponentResult: @@ -100,9 +106,10 @@ 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, created_at: datetime | None = None, ) -> None: """Persist assignments and their template.""" @@ -116,12 +123,13 @@ 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, 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, @@ -173,15 +181,17 @@ 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, 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 d34e094..3ad5682 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, uuid4 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__) @@ -22,11 +24,11 @@ class Database(ABC): @abstractmethod async def create_template( self, - uuid: str, title: str, content: JsonValue, - ) -> None: - """Create a new template in a database backend.""" + tenant_uuid: UUID, + ) -> UUID: + """Create a new template in a database backend. Return created template UUID""" @abstractmethod async def save_assignments( @@ -35,7 +37,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: @@ -44,9 +46,10 @@ async def save_assignments( @abstractmethod async def save_template( self, - uuid: str, + uuid: UUID, title: str, content: JsonValue, + tenant_uuid: UUID, ) -> None: """Persist a template in a database backend.""" @@ -54,25 +57,25 @@ 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.""" @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: str) -> dict[str, Any] | None: + async def get_template(self, template_uuid: UUID, tenant_uuid: UUID) -> TemplateDetail | None: """Get a template record from a database backend.""" @abstractmethod async def save_result( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, prepolished_markdown: str, markdown: str, ) -> None: @@ -81,10 +84,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 +95,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.""" @@ -156,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: @@ -193,15 +196,17 @@ async def save_assignments( async def create_template( self, - uuid: str, title: str, content: JsonValue, - ) -> None: + tenant_uuid: UUID, + ) -> UUID: + 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, ) try: @@ -213,21 +218,24 @@ 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, - uuid: str, + uuid: UUID, 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], @@ -249,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() @@ -278,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) @@ -294,9 +306,11 @@ 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, 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) @@ -304,24 +318,24 @@ async def get_template(self, template_uuid: str) -> dict[str, Any] | 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 - - 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, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, prepolished_markdown: str, markdown: str, ) -> None: @@ -359,10 +373,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 +408,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/persistence/saver_component.py b/service/src/ai_document_plugin_service/ai/persistence/saver_component.py index 0748d22..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 @@ -1,5 +1,6 @@ import typing from typing import TypedDict +from uuid import UUID from haystack import component @@ -18,10 +19,10 @@ 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: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, debug_markdown: str, markdown: str, ) -> FileSaverComponentResult: @@ -42,10 +43,10 @@ async def run_async( @component.output_types(markdown=str) def run( self, - template_uuid: str, + template_uuid: UUID, knowledge_model_uuid: str, - user_uuid: str, - tenant_uuid: str, + user_uuid: UUID, + tenant_uuid: UUID, debug_markdown: str, markdown: str, ) -> FileSaverComponentResult: 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 652fa7d..ee128ee 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, @@ -153,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, @@ -198,10 +200,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/__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/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..bf4b7c1 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 @@ -26,27 +26,24 @@ 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: str, 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') - return _model_from_fields( - TemplateDetail, - uuid=template['uuid'], - title=template['title'], - content=template['content'], - ) + return template @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') @@ -58,13 +55,11 @@ async def create_template(payload: TemplateCreateRequest, database: DatabaseDI) detail='Template JSON must contain a top-level "sections" array.', ) - template_uuid = str(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, ) except ValueError as error: raise fastapi.HTTPException(status_code=409, detail=str(error)) from error @@ -85,7 +80,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') @@ -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..a6a4f9b 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') @@ -56,8 +57,8 @@ class PipelineRunRequest(ApiModel): class PipelineRunResponse(ApiModel): status: PipelineStatus run_id: str = Field(alias='runId') - questionnaire_uuid: str = Field(alias='questionnaireUuid') - template_uuid: str = Field(alias='templateUuid') + questionnaire_uuid: UUID = Field(alias='questionnaireUuid') + template_uuid: UUID = Field(alias='templateUuid') template_title: str = Field(alias='templateTitle') @@ -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/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..4c69fec --- /dev/null +++ b/service/src/ai_document_plugin_service/migrations/versions/20260709_01_add_tenant_uuid_to_template.py @@ -0,0 +1,67 @@ +"""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, + ) + + 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 1fc568c..b0f7cac 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. @@ -200,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, @@ -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,