Separate llm clients based on tenant

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

@@ -15,7 +15,12 @@
 
 logger = logging.getLogger(__name__)
 
-semaphore = DynamicSemaphore(1)
+
+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):
@@ -77,42 +82,104 @@ 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: AsyncOpenAI | None = 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,
         )
 
     def get_max_workers(self) -> int:
+        if self.max_workers is None:
+            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:
+            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(
         self,
         *args: Any,  # noqa: ANN401
         **kwargs: Any,  # noqa: ANN401
     ) -> ChatCompletion:
+        if self.model is None:
+            raise InvalidLLMConfigError('model', self.tenant_uuid)  # noqa: EM101
+        if self.max_workers is None:
+            raise InvalidLLMConfigError('max_workers', self.tenant_uuid)  # noqa: EM101
+        if self.api_url is None:
+            raise InvalidLLMConfigError('api_url', self.tenant_uuid)  # noqa: EM101
+        if self.api_key is None:
+            raise InvalidLLMConfigError('api_key', self.tenant_uuid)  # noqa: EM101
+        if self.client is None:
+            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)
-        async with semaphore:
+        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,
-                semaphore.limit,
+                self.semaphore.limit,
             )
             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,
             )

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)

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'],
     )
@@ -128,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)

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

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,12 @@ 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 +159,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 +177,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 +195,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 +214,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 +231,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(