Skip to content

Commit 073ad24

Browse files
authored
Show templates based on tenant (#130)
* 85: Start using Dependency injection to allow injecting the config more easily * 85: Handle usage of private function * 85: Fix formatting * 85: Refactor pipeline service - simplified * 85: Fixed ty check * 75: Removed global semaphore - moved do local semaphore for each llm client * 75: Removed global semaphore - moved do local semaphore for each llm client * 75: Fix ruff errors * 75: Improved logging * 73: Moved strings to UUIDs * Added db migration * Updated db * Fix datatypes * Fix ruff * Fix template id to UUID * Fix bugs with data types * Fix issue with uuid imports * Refactored where uuid is created * Fix types * Fix formatting
1 parent c8badfe commit 073ad24

15 files changed

Lines changed: 207 additions & 113 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818

1919
class InvalidLLMConfigError(ValueError):
20-
def __init__(self, var_name: str, tenant: str) -> None:
20+
def __init__(self, var_name: str, tenant: uuid.UUID) -> None:
2121
super().__init__(
2222
f"LLM '{var_name}' for tenant {tenant} is None, did you call `update_config` before using the client?"
2323
)
@@ -90,7 +90,7 @@ class LLMClient:
9090
This is because LLM client handles throttling to avoid spamming the LLM API.
9191
"""
9292

93-
def __init__(self, tenant_uuid: str) -> None:
93+
def __init__(self, tenant_uuid: uuid.UUID) -> None:
9494
"""
9595
Initializes the client with empty config. Call update_config before using it.
9696
"""

service/src/ai_document_plugin_service/ai/knowledgemodel/dsw_client.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from uuid import UUID
2+
13
import httpx
24

35

@@ -6,7 +8,7 @@ def __init__(self, token: str, api_url: str) -> None:
68
self.token = token
79
self.api_url = api_url.rstrip('/')
810

9-
async def get_questionnaire_detail(self, questionnaire_uuid: str) -> dict:
11+
async def get_questionnaire_detail(self, questionnaire_uuid: str | UUID) -> dict:
1012
url = f'{self.api_url}/projects/{questionnaire_uuid}/questionnaire'
1113

1214
headers: dict[str, str] = {}

service/src/ai_document_plugin_service/ai/persistence/assignment_loader_component.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from typing import Any
2+
from uuid import UUID
23

34
from haystack import component
45

@@ -14,7 +15,7 @@ def __init__(self, database: Database) -> None:
1415
assignments=JsonValue | None,
1516
found=bool,
1617
)
17-
async def run_async(self, knowledge_model_uuid: str, template_uuid: str) -> dict[str, Any]:
18+
async def run_async(self, knowledge_model_uuid: str, template_uuid: UUID) -> dict[str, Any]:
1819
assignments = await self.database.get_assignments(knowledge_model_uuid, template_uuid)
1920

2021
return {
@@ -26,7 +27,7 @@ async def run_async(self, knowledge_model_uuid: str, template_uuid: str) -> dict
2627
assignments=JsonValue | None,
2728
found=bool,
2829
)
29-
def run(self, knowledge_model_uuid: str, template_uuid: str) -> dict[str, Any]:
30+
def run(self, knowledge_model_uuid: str, template_uuid: UUID) -> dict[str, Any]:
3031
"""Async-only component; the sync pipeline entrypoint is intentionally unsupported."""
3132
msg = f'{type(self).__name__} is async-only; use run_async() / AsyncPipeline.run_async()'
3233
raise NotImplementedError(

service/src/ai_document_plugin_service/ai/persistence/assignment_saver_component.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111
from datetime import UTC, datetime
1212
from typing import TYPE_CHECKING, TypedDict
1313

14+
# UUID must be imported outside TYPE_CHECKING block for haystack to work
15+
from uuid import UUID # noqa: TC003
16+
1417
from haystack import component
1518

1619
from ai_document_plugin_service.ai.assignment.types import (
@@ -44,9 +47,10 @@ async def run_async(
4447
knowledge_model_uuid: str,
4548
knowledge_model_name: str,
4649
knowledge_model_version: str,
47-
template_uuid: str,
50+
template_uuid: UUID,
4851
template_title: str,
4952
template_data: JsonValue,
53+
tenant_uuid: UUID,
5054
assignments: list[SectionAssignment],
5155
stats: AssignmentStats | None = None,
5256
) -> AssignmentSaverComponentResult:
@@ -63,6 +67,7 @@ async def run_async(
6367
template_uuid=template_uuid,
6468
template_title=template_title,
6569
template_data=template_data,
70+
tenant_uuid=tenant_uuid,
6671
created_at=datetime.now(tz=UTC),
6772
)
6873

@@ -78,9 +83,10 @@ def run(
7883
knowledge_model_uuid: str,
7984
knowledge_model_name: str,
8085
knowledge_model_version: str,
81-
template_uuid: str,
86+
template_uuid: UUID,
8287
template_title: str,
8388
template_data: JsonValue,
89+
tenant_uuid: UUID,
8490
assignments: list[SectionAssignment],
8591
stats: AssignmentStats | None = None,
8692
) -> AssignmentSaverComponentResult:
@@ -100,9 +106,10 @@ async def save(
100106
knowledge_model_version: str,
101107
assignments: JsonValue,
102108
stats: StatsJson | None,
103-
template_uuid: str,
109+
template_uuid: UUID,
104110
template_title: str,
105111
template_data: JsonValue,
112+
tenant_uuid: UUID,
106113
created_at: datetime | None = None,
107114
) -> None:
108115
"""Persist assignments and their template."""
@@ -116,12 +123,13 @@ async def save(
116123
knowledge_model_version: str,
117124
assignments: JsonValue,
118125
stats: StatsJson | None,
119-
template_uuid: str,
126+
template_uuid: UUID,
120127
template_title: str,
121128
template_data: JsonValue,
129+
tenant_uuid: UUID,
122130
created_at: datetime | None = None,
123131
) -> None:
124-
_ = (template_uuid, template_title, template_data)
132+
_ = (template_uuid, template_title, template_data, tenant_uuid)
125133
output_name = self._build_filename(
126134
knowledge_model_uuid,
127135
knowledge_model_name,
@@ -173,15 +181,17 @@ async def save(
173181
knowledge_model_version: str,
174182
assignments: JsonValue,
175183
stats: StatsJson | None,
176-
template_uuid: str,
184+
template_uuid: UUID,
177185
template_title: str,
178186
template_data: JsonValue,
187+
tenant_uuid: UUID,
179188
created_at: datetime | None = None,
180189
) -> None:
181190
await self.database.save_template(
182191
uuid=template_uuid,
183192
title=template_title,
184193
content=template_data,
194+
tenant_uuid=tenant_uuid,
185195
)
186196
await self.database.save_assignments(
187197
knowledge_model_uuid=knowledge_model_uuid,

0 commit comments

Comments
 (0)