Show templates based on tenant

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.
         """

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] = {}

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(

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,