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..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 @@ -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,20 +82,69 @@ 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( @@ -98,21 +152,34 @@ async def completion( *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, ) 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..101b87f 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'], ) @@ -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) 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..1fc568c 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,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(