"""AI Labs service layer: conversations, generations, translation, stats. Every read in here is self-scoped, and the scoping is done with the caller's person id in the ``WHERE`` clause rather than by filtering after the fetch — there is no code path that loads someone else's row and then decides not to show it. Deletion is a **hard** delete, and it takes the S3 objects with it. That is the shape the product promise requires: an employee told that Helix does not keep what they typed here must not discover a ``deleted_at`` column holding it. The stored artifacts of a deleted generation, and the attachments of a deleted conversation, are removed on a best-effort basis and the row goes regardless — a bucket hiccup must not be the reason someone cannot delete their own data. """ from __future__ import annotations import logging import uuid from datetime import UTC, date, datetime, timedelta from typing import Any, ClassVar from sqlalchemy import delete, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql.elements import ColumnElement from app.audit.service import AuditService from app.core.background import spawn from app.core.exceptions import AppException, NotFoundError from app.core.storage import ( StorageNotConfiguredError, delete_object, presign_get, ) from app.core.storage_uploads import forget_upload_grants from app.domains.ai_labs import attachments as attach from app.domains.ai_labs import catalog, executor, usage from app.domains.ai_labs.models import ( GEN_CANCELLED, GEN_QUEUED, GEN_RUNNING, GEN_TERMINAL_STATUSES, AiLabConversation, AiLabFailureDaily, AiLabGeneration, AiLabMessage, AiLabUsageDaily, ) from app.domains.ai_labs.settings import AiLabsConfig, resolve_config from app.domains.bailian.llm_client import BailianLLMClient from app.identity.person_models import IdentityOrgMembership, IdentityPerson logger = logging.getLogger(__name__) #: How long an artifact download link is valid. Long enough to click and #: for a video to start playing; short enough that a copied URL is not a #: durable handle to private content. ARTIFACT_URL_TTL_SECONDS = 900 #: Virtual value for the gallery's status filter, covering both #: pre-terminal states. Not a stored status. STATUS_FILTER_ACTIVE = "active" def _artifact_download_name( row: AiLabGeneration, *, index: int, object_key: str, ) -> str: """A filename for a saved artifact. Built from the capability and the job's short id rather than the prompt: a prompt can be a paragraph, can contain path separators, and is the employee's content — putting it in a header that travels to S3 and into their Downloads folder is more exposure than a name needs. The extension comes from the stored key. """ _, _, extension = object_key.rpartition(".") suffix = f".{extension}" if extension and len(extension) <= 8 else "" part = f"-{index + 1}" if index else "" return f"helix-{row.capability}-{str(row.id)[:8]}{part}{suffix}" def _iso(value: datetime | None) -> str | None: if value is None: return None aware = value.replace(tzinfo=UTC) if value.tzinfo is None else value return aware.isoformat() class AiLabsService: """Everything a request does, other than streaming a chat turn.""" def __init__(self, db: AsyncSession) -> None: self.db = db self.audit = AuditService(db) # ── conversations ────────────────────────────────────────── async def list_conversations( self, *, person_id: uuid.UUID, page: int, page_size: int, ) -> tuple[list[dict[str, Any]], int]: base = select(AiLabConversation).where( AiLabConversation.person_id == person_id, ) total = int( ( await self.db.execute( select(func.count()).select_from(base.subquery()), ) ).scalar() or 0, ) stmt = ( base # Pinned first, then most recent activity. ``last_message_at`` # is NULL on a thread nobody has written in, so ``created_at`` # is the tiebreak that keeps a fresh empty thread visible. .order_by( AiLabConversation.pinned.desc(), # MySQL/TiDB sort NULL as the lowest value, so a DESC sort # already puts never-used threads last. ``NULLS LAST`` is not # valid MySQL syntax and would fail at execution, not parse. AiLabConversation.last_message_at.desc(), AiLabConversation.created_at.desc(), ) .offset((page - 1) * page_size) .limit(page_size) ) rows = list((await self.db.execute(stmt)).scalars().all()) return [self._conversation_summary(row) for row in rows], total @staticmethod def _conversation_summary(row: AiLabConversation) -> dict[str, Any]: return { "id": str(row.id), "title": row.title or "", "model_key": row.model_key, "thinking_enabled": bool(row.thinking_enabled), "pinned": bool(row.pinned), "message_count": int(row.message_count or 0), "last_message_at": _iso(row.last_message_at), "created_at": _iso(row.created_at), "updated_at": _iso(row.updated_at), } async def create_conversation( self, *, person_id: uuid.UUID, model_key: str | None, thinking: bool | None, ) -> dict[str, Any]: count = int( ( await self.db.execute( select(func.count()) .select_from(AiLabConversation) .where(AiLabConversation.person_id == person_id), ) ).scalar() or 0, ) if count >= catalog.MAX_CONVERSATIONS_PER_PERSON: raise AppException( code="AI_LABS_TOO_MANY_CONVERSATIONS", message=( "You have reached the maximum number of saved " f"conversations ({catalog.MAX_CONVERSATIONS_PER_PERSON}). " "Delete one to start a new thread." ), status_code=409, ) model = catalog.chat_model(model_key) conversation = AiLabConversation( person_id=person_id, title="", model_key=model.key, thinking_enabled=bool(thinking) and model.supports_thinking, ) self.db.add(conversation) await self.db.flush() return self._conversation_summary(conversation) async def get_conversation( self, *, person_id: uuid.UUID, conversation_id: uuid.UUID, ) -> dict[str, Any]: conversation = await self._owned_conversation( person_id=person_id, conversation_id=conversation_id, ) stmt = ( select(AiLabMessage) .where(AiLabMessage.conversation_id == conversation.id) .order_by(AiLabMessage.seq) ) messages = list((await self.db.execute(stmt)).scalars().all()) return { **self._conversation_summary(conversation), "messages": [self._message_out(row) for row in messages], } @staticmethod def _message_out(row: AiLabMessage) -> dict[str, Any]: return { "id": str(row.id), "seq": int(row.seq), "role": row.role, "content": row.content or "", "reasoning": row.reasoning, "attachments": row.attachments or [], "model": row.model, "input_tokens": row.input_tokens, "output_tokens": row.output_tokens, "latency_ms": row.latency_ms, "finish_reason": row.finish_reason, "error_code": row.error_code, # The token the CLIENT minted for this send, echoed back. # # It is what lets the thread recognise its own in-flight turn # without waiting for the ``meta`` frame. Matching on the text # instead could not distinguish "the server already stored this # turn" from "the employee re-sent the same words", and those # need opposite renderings — so the ambiguity is removed rather # than guessed at. Only ever returned to the employee whose row # it is, and it carries no content. "client_token": row.client_token, "created_at": _iso(row.created_at), } async def _owned_conversation( self, *, person_id: uuid.UUID, conversation_id: uuid.UUID, ) -> AiLabConversation: conversation = await self.db.get(AiLabConversation, conversation_id) if conversation is None or conversation.person_id != person_id: raise NotFoundError("conversation not found") return conversation async def update_conversation( self, *, person_id: uuid.UUID, conversation_id: uuid.UUID, title: str | None = None, pinned: bool | None = None, model_key: str | None = None, thinking: bool | None = None, ) -> dict[str, Any]: conversation = await self._owned_conversation( person_id=person_id, conversation_id=conversation_id, ) if title is not None: cleaned = " ".join(title.split())[:200] if not cleaned: raise AppException( code="AI_LABS_TITLE_REQUIRED", message="A title cannot be blank.", status_code=422, field="title", ) conversation.title = cleaned if pinned is not None: conversation.pinned = bool(pinned) if model_key is not None: if model_key not in catalog.CHAT_MODELS_BY_KEY: raise AppException( code="AI_LABS_MODEL_UNKNOWN", message="That model is not available.", status_code=422, field="model_key", ) conversation.model_key = model_key if thinking is not None: model = catalog.chat_model(conversation.model_key) conversation.thinking_enabled = ( bool(thinking) and model.supports_thinking ) await self.db.flush() return self._conversation_summary(conversation) async def delete_conversation( self, *, person_id: uuid.UUID, conversation_id: uuid.UUID, trace_id: str, ) -> None: """Hard-delete a thread, its messages and its uploaded attachments.""" conversation = await self._owned_conversation( person_id=person_id, conversation_id=conversation_id, ) keys: list[str] = [] stmt = select(AiLabMessage.attachments).where( AiLabMessage.conversation_id == conversation.id, ) for attachments in (await self.db.execute(stmt)).scalars().all(): for item in attachments or []: key = (item or {}).get("object_key") if isinstance(key, str) and key: keys.append(key) # ORM cascade needs the children loaded; a bulk DELETE on the # parent would leave the message rows behind on a database whose # FK has no ON DELETE (TiDB configurations vary), so delete the # children explicitly and then the parent. await self.db.execute( delete(AiLabMessage).where( AiLabMessage.conversation_id == conversation.id, ), ) await self.db.delete(conversation) # The grant ledger row is part of what "hard delete" has to reach. # It keeps a full ``checksum_sha256`` of the file plus the exact # size, type and minute — enough for someone holding a copy of a # document to prove a named employee put it into AI Labs, months # after they deleted the conversation. await forget_upload_grants( self.db, object_keys=keys, person_id=person_id, ) await self.audit.log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.conversation.delete", status="success", resource_type="ai_labs_conversation", resource_id=str(conversation_id), # Metadata only. The audit trail records that a thread was # deleted and how big it was — never a word of what was in it. detail={"attachments": len(keys)}, ) await self.db.commit() await self._forget_objects(keys) async def _forget_objects(self, keys: list[str]) -> None: """Best-effort S3 cleanup after the rows are already gone. Runs after the commit and never raises: the employee asked for the row to be gone, and it is. A key that survives is stranded **for good**, and the earlier version of this docstring was wrong to hand it to "the orphan sweep": that sweep walks the ``storage_upload_grants`` ledger, and neither of these key families is in it by the time we get here — ``ai_labs_output/`` was never granted at all, and the input grants are deleted in the same transaction as the row. So a failure is logged at ``error`` for an operator to act on, and the row still goes, because the employee asked for it to. """ # Call-time lookup so the test harness's patched factory applies. from app.core.database import async_session_factory for key in keys: try: async with async_session_factory() as own: await delete_object(own, object_key=key) except Exception as exc: # noqa: BLE001 logger.error( "ai_labs: could not delete object %s (%s) — it is now " "stranded in the bucket with nothing referencing it and " "no sweep that can find it; remove it by hand", key, exc, ) # ── generations ──────────────────────────────────────────── async def submit_generation( self, *, person_id: uuid.UUID, capability_key: str, prompt: str, params: dict[str, Any], input_object_keys: list[str], config: AiLabsConfig, trace_id: str, ) -> dict[str, Any]: """Validate, persist a ``queued`` job, and start it. The row is committed **before** the runner is spawned, so a job always exists to resume: spawning first would create a window where the vendor is working on something Helix has no record of. """ capability = catalog.capability(capability_key) if capability is None: raise AppException( code="AI_LABS_CAPABILITY_UNKNOWN", message="That capability is not available.", status_code=422, field="capability", ) prompt = (prompt or "").strip() if capability.prompt_required and not prompt: raise AppException( code="AI_LABS_PROMPT_REQUIRED", message="A prompt is required.", status_code=422, field="prompt", ) if prompt and capability.prompt_max_chars == 0: # A capability whose vendor request has no slot for a prompt. # Storing one and dropping it would show the employee their # own words on the tile as though they had been used. raise AppException( code="AI_LABS_PROMPT_NOT_ACCEPTED", message=( "This capability does not take a prompt — it works from " "the file you attached." ), status_code=422, field="prompt", ) if len(prompt) > capability.prompt_max_chars: raise AppException( code="AI_LABS_PROMPT_TOO_LONG", message=( "The prompt may be at most " f"{capability.prompt_max_chars} characters." ), status_code=422, field="prompt", ) keys = list(input_object_keys or []) if len(keys) < capability.min_inputs: raise AppException( code="AI_LABS_INPUT_REQUIRED", message=( f"This capability needs at least {capability.min_inputs} " "input file(s)." ), status_code=422, field="input_object_keys", ) if len(keys) > capability.max_inputs: raise AppException( code="AI_LABS_TOO_MANY_INPUTS", message=( f"At most {capability.max_inputs} input file(s) are " "accepted here." ), status_code=422, field="input_object_keys", ) # Verify ownership and kind now, at submit time, so a bad key is a # clear 4xx on the click rather than a failed job minutes later. expected_kind = ( attach.KIND_IMAGE if capability.input_kind == catalog.INPUT_IMAGE else attach.KIND_AUDIO if capability.input_kind == catalog.INPUT_AUDIO else None ) for key in keys: kind = attach.classify(key, owner_id=person_id) if expected_kind and kind != expected_kind: raise AppException( code="AI_LABS_INPUT_WRONG_KIND", message=f"This capability expects a {expected_kind} file.", status_code=422, field="input_object_keys", ) clean_params = self._validate_params(capability, params) await usage.assert_within_quota( self.db, person_id=person_id, config=config, capability_key=capability.key, video_seconds=int(clean_params.get("duration") or 0), ) generation = AiLabGeneration( person_id=person_id, capability=capability.key, model=config.vendor_capability_model(capability), status=GEN_QUEUED, prompt=prompt, params=clean_params or None, input_object_keys=keys or None, ) self.db.add(generation) await self.db.flush() generation_id = generation.id # Claim the upload grants — see the note in ``chat.prepare_turn``. # Without it the hourly orphan sweep deletes these input objects 24 # hours after upload, and a job the reaper resumes after that dies # on a file the employee can still see listed on their own row. await attach.claim( self.db, object_keys=keys, owner_id=person_id, resource_type="ai_labs_generation", resource_id=str(generation_id), ) await self.audit.log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.generation.submit", status="success", resource_type="ai_labs_generation", resource_id=str(generation_id), # Which capability, which model, how many inputs, how long the # prompt was — never the prompt itself. detail={ "capability": capability.key, "model": generation.model, "params": clean_params, "inputs": len(keys), "prompt_chars": len(prompt), }, ) payload = self._generation_out(generation) await self.db.commit() spawn( executor.run_job(generation_id), name=f"ai_labs.generation.{generation_id}", ) return payload @staticmethod def _validate_params( capability: catalog.GenerationCapability, raw: dict[str, Any] | None, ) -> dict[str, Any]: """Keep only catalogue-declared params with catalogue-allowed values. Unknown keys are dropped rather than rejected — a slightly stale frontend sending an extra field should not fail an employee's generation. A *known* key with a disallowed value IS rejected: that is a real request for something the model was never verified to accept. """ clean: dict[str, Any] = {} provided = raw or {} for param in capability.params: if param.name in provided: value = provided[param.name] if value not in param.allowed_values(): raise AppException( code="AI_LABS_PARAM_INVALID", message=( f"'{value}' is not an accepted value for " f"{param.name}." ), status_code=422, field=param.name, ) clean[param.name] = value else: clean[param.name] = param.default return clean @staticmethod def _generation_out(row: AiLabGeneration) -> dict[str, Any]: artifacts = row.artifacts or [] return { "id": str(row.id), "capability": row.capability, "model": row.model, "status": row.status, "prompt": row.prompt or "", "params": row.params or {}, "input_object_keys": row.input_object_keys or [], # Object keys are useful to the owner (they can presign them), # and the artifact bytes are fetched through this domain's own # endpoint, which re-checks ownership. "artifact_count": len(artifacts), "artifacts": [ { "index": index, "content_type": item.get("content_type"), "size_bytes": item.get("size_bytes"), } for index, item in enumerate(artifacts) ], "output_text": row.output_text, "error_code": row.error_code, "error_message": row.error_message, "attempts": int(row.attempts or 0), "created_at": _iso(row.created_at), "started_at": _iso(row.started_at), "finished_at": _iso(row.finished_at), } async def list_generations( self, *, person_id: uuid.UUID, capability_key: str | None, status: str | None, page: int, page_size: int, ) -> tuple[list[dict[str, Any]], int]: base = select(AiLabGeneration).where( AiLabGeneration.person_id == person_id, ) if capability_key: base = base.where(AiLabGeneration.capability == capability_key) if status == STATUS_FILTER_ACTIVE: # "In progress" has to mean both, or a just-submitted job is # invisible under the filter whose label says it is there: a # row is ``queued`` until a runner leases it, which for a video # can be the whole first tick. base = base.where( AiLabGeneration.status.in_((GEN_QUEUED, GEN_RUNNING)), ) elif status: base = base.where(AiLabGeneration.status == status) total = int( ( await self.db.execute( select(func.count()).select_from(base.subquery()), ) ).scalar() or 0, ) stmt = ( base.order_by(AiLabGeneration.created_at.desc()) .offset((page - 1) * page_size) .limit(page_size) ) rows = list((await self.db.execute(stmt)).scalars().all()) return [self._generation_out(row) for row in rows], total async def get_generation( self, *, person_id: uuid.UUID, generation_id: uuid.UUID, ) -> dict[str, Any]: return self._generation_out( await self._owned_generation( person_id=person_id, generation_id=generation_id, ), ) async def _owned_generation( self, *, person_id: uuid.UUID, generation_id: uuid.UUID, ) -> AiLabGeneration: row = await self.db.get(AiLabGeneration, generation_id) if row is None or row.person_id != person_id: raise NotFoundError("generation not found") return row async def artifact_url( self, *, person_id: uuid.UUID, generation_id: uuid.UUID, index: int, download: bool = False, ) -> dict[str, Any]: """A short-lived download link for one artifact of one's own job. Ownership is re-checked here rather than trusting a key the client holds: minting a presigned URL is an authorisation decision, and it has to be made where the owner is known. ``download`` signs a ``Content-Disposition: attachment``, which is what makes a button labelled "Download" actually save the file. Without it the object serves with its real content type and an image or clip simply opens in a tab — and inside the Feishu webview that tab often does not open at all. The inline players use the same endpoint **without** the flag, because a disposition there would stop them playing. """ row = await self._owned_generation( person_id=person_id, generation_id=generation_id, ) artifacts = row.artifacts or [] if index < 0 or index >= len(artifacts): raise NotFoundError("artifact not found") object_key = (artifacts[index] or {}).get("object_key") if not isinstance(object_key, str) or not object_key: raise NotFoundError("artifact not found") try: url = await presign_get( self.db, object_key=object_key, expires_in=ARTIFACT_URL_TTL_SECONDS, download_filename=( _artifact_download_name( row, index=index, object_key=object_key, ) if download else None ), ) except StorageNotConfiguredError as exc: raise AppException( code="STORAGE_NOT_CONFIGURED", message="File storage is not configured.", status_code=503, ) from exc return { "url": url, "expires_in": ARTIFACT_URL_TTL_SECONDS, "content_type": (artifacts[index] or {}).get("content_type"), } async def cancel_generation( self, *, person_id: uuid.UUID, generation_id: uuid.UUID, trace_id: str = "", ) -> dict[str, Any]: """Mark one's own unfinished job cancelled. Cooperative, not forceful: the vendor has no cancel API for these families, so this stops the reaper from resuming the job and stops Helix from *recording* an outcome for it. A clip the vendor is already rendering still renders; nobody is billed differently for looking away. It does not necessarily stop Helix waiting. An async job notices within one poll (~6s) because its loop heartbeats. A **sync** job is one long ``await`` with no in-loop check, so its runner waits out the whole vendor call and its download, then sees ``keeper.lost`` and discards what it stored. The row goes ``cancelled`` immediately either way, which is what the employee sees. """ row = await self._owned_generation( person_id=person_id, generation_id=generation_id, ) if row.status in GEN_TERMINAL_STATUSES: raise AppException( code="AI_LABS_ALREADY_FINISHED", message="That job has already finished.", status_code=409, ) previous_status = row.status row.status = GEN_CANCELLED row.finished_at = datetime.now(UTC) row.lease_owner = None row.lease_expires_at = None row.error_code = "AI_LABS_CANCELLED" row.error_message = "Cancelled by the requester." await self.db.flush() payload = self._generation_out(row) await self.audit.log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.generation.cancel", status="success", resource_type="ai_labs_generation", resource_id=str(generation_id), # Which capability and how far it got. Not the prompt: this # domain's audit rows record that an action happened, never # what the employee asked for. detail={ "capability": row.capability, "attempts": int(row.attempts or 0), "was_status": previous_status, }, ) await self.db.commit() return payload async def delete_generation( self, *, person_id: uuid.UUID, generation_id: uuid.UUID, trace_id: str, ) -> None: row = await self._owned_generation( person_id=person_id, generation_id=generation_id, ) if row.status in (GEN_QUEUED, GEN_RUNNING): raise AppException( code="AI_LABS_STILL_RUNNING", message=( "Cancel this job before deleting it, so it is not " "resumed after the row is gone." ), status_code=409, ) keys = [ key for item in (row.artifacts or []) if isinstance(key := (item or {}).get("object_key"), str) and key ] keys += [ key for key in (row.input_object_keys or []) if isinstance(key, str) ] await self.db.delete(row) # Only the INPUT keys have grant rows — ``ai_labs_output/`` is # written server-side and was never granted — and passing the # artifact keys too is harmless because the delete is scoped to # both the key and this person. await forget_upload_grants( self.db, object_keys=keys, person_id=person_id, ) await self.audit.log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.generation.delete", status="success", resource_type="ai_labs_generation", resource_id=str(generation_id), detail={"capability": row.capability, "objects": len(keys)}, ) await self.db.commit() await self._forget_objects(keys) # ── translation ──────────────────────────────────────────── async def translate( self, *, person_id: uuid.UUID, text: str, target_language: str, config: AiLabsConfig, trace_id: str = "", ) -> dict[str, Any]: """One synchronous translation. Uses the vendor's dedicated translation model with the target as a real parameter, rather than a chat model told to translate. That is not a performance choice: an instruction would be Helix writing part of the prompt, and the whole surface is built on not doing that. """ text = (text or "").strip() if not text: raise AppException( code="AI_LABS_TEXT_REQUIRED", message="Text to translate is required.", status_code=422, field="text", ) if len(text) > catalog.TRANSLATE_MAX_CHARS: raise AppException( code="AI_LABS_TEXT_TOO_LONG", message=( "At most " f"{catalog.TRANSLATE_MAX_CHARS} characters per " "translation." ), status_code=422, field="text", ) if target_language not in catalog.TRANSLATE_LANGUAGE_VALUES: raise AppException( code="AI_LABS_LANGUAGE_UNKNOWN", message="That target language is not available.", status_code=422, field="target_language", ) await usage.assert_within_quota( self.db, person_id=person_id, config=config, capability_key=usage.CAPABILITY_TRANSLATE, ) model = config.model_overrides.get( "translate", catalog.TRANSLATE_MODEL, ) client = BailianLLMClient( config.api_key, model=model, region=config.region, purpose="ai_labs.media") try: result = await client.chat( messages=[{"role": "user", "content": text}], translation_options={ "source_lang": "auto", "target_lang": target_language, }, # A translation is roughly as long as its source, and the # client's 1500-token default would silently truncate an # 8,000-character input at about a fifth of the way # through. max_tokens=8000, temperature=0.0, ) except AppException as exc: # Translation is the one capability whose failures were counted # nowhere: chat books its own through ``_persist_failure`` and # generations through ``_note_outcome``, but this call simply # raised past the recorder. "Translate is broken for everyone" # was therefore invisible on the admin console — no request, no # failure, no reason. Counted here, then re-raised unchanged. # # Only ``exc.code`` travels: the vendor's prose can quote the # submitted text back, and ``usage.record`` would refuse it # anyway (see ``normalize_error_code``). await usage.record( person_id=person_id, capability_key=usage.CAPABILITY_TRANSLATE, model=model, timezone_name=config.reporting_timezone, requests=1, failures=1, error_code=exc.code, ) raise await usage.record( person_id=person_id, capability_key=usage.CAPABILITY_TRANSLATE, model=model, timezone_name=config.reporting_timezone, requests=1, input_tokens=result.get("input_tokens") or 0, output_tokens=result.get("output_tokens") or 0, ) await self.audit.log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.translate", status="success", resource_type="ai_labs_translation", resource_id=str(person_id), # Shape and cost only. ``chars`` is a length, not a sample. detail={ "target_language": target_language, "model": model, "chars": len(text or ""), }, ) await self.db.commit() return { "text": result.get("content") or "", "target_language": target_language, "model": model, } # ── admin: counters only ─────────────────────────────────── async def usage_overview(self, *, days: int) -> dict[str, Any]: """Aggregate usage for the admin console. Reads :class:`AiLabUsageDaily` and nothing else, which is what makes this endpoint safe to expose: the table it queries has no column that could contain a prompt, a reply or an object key. ``days`` is **inclusive of today** and counted in the reporting timezone, because that is the only calendar ``usage_date`` is written in (see :func:`usage.local_today`). Deriving the window from ``datetime.now(UTC).date()`` instead would be wrong twice over: it shifts the boundary by the zone's offset for eight hours of every Beijing day, and ``now - days`` spans ``days + 1`` calendar dates, so "7 days" summed eight of them. """ config = await resolve_config(self.db) today = usage.local_today(config.reporting_timezone) since = today - timedelta(days=max(1, days) - 1) totals_stmt = select( func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce(func.sum(AiLabUsageDaily.input_tokens), 0), func.coalesce(func.sum(AiLabUsageDaily.output_tokens), 0), func.coalesce(func.sum(AiLabUsageDaily.artifacts), 0), func.coalesce(func.sum(AiLabUsageDaily.media_seconds), 0), func.count(func.distinct(AiLabUsageDaily.person_id)), ).where(AiLabUsageDaily.usage_date.between(since, today)) totals = (await self.db.execute(totals_stmt)).one() by_capability_stmt = ( select( AiLabUsageDaily.capability, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce(func.sum(AiLabUsageDaily.artifacts), 0), func.count(func.distinct(AiLabUsageDaily.person_id)), ) .where(AiLabUsageDaily.usage_date.between(since, today)) .group_by(AiLabUsageDaily.capability) .order_by(func.sum(AiLabUsageDaily.requests).desc()) ) by_capability = [ { "capability": row[0], "requests": int(row[1] or 0), "failures": int(row[2] or 0), "artifacts": int(row[3] or 0), "people": int(row[4] or 0), } for row in (await self.db.execute(by_capability_stmt)).all() ] by_model_stmt = ( select( AiLabUsageDaily.model, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), ) .where(AiLabUsageDaily.usage_date.between(since, today)) .group_by(AiLabUsageDaily.model) .order_by(func.sum(AiLabUsageDaily.requests).desc()) .limit(20) ) by_model = [ { "model": row[0], "requests": int(row[1] or 0), "tokens": int(row[2] or 0), } for row in (await self.db.execute(by_model_stmt)).all() ] daily_stmt = ( select( AiLabUsageDaily.usage_date, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.count(func.distinct(AiLabUsageDaily.person_id)), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), ) .where(AiLabUsageDaily.usage_date.between(since, today)) .group_by(AiLabUsageDaily.usage_date) .order_by(AiLabUsageDaily.usage_date) ) daily = [ { "date": row[0].isoformat() if row[0] else None, "requests": int(row[1] or 0), "people": int(row[2] or 0), "tokens": int(row[3] or 0), } for row in (await self.db.execute(daily_stmt)).all() ] # Live job health, so an operator can see a stuck queue without # reading anyone's prompt. queue_stmt = ( select(AiLabGeneration.status, func.count()) .group_by(AiLabGeneration.status) ) queue = { str(row[0]): int(row[1] or 0) for row in (await self.db.execute(queue_stmt)).all() } return { "window_days": max(1, days), "since": since.isoformat(), "through": today.isoformat(), "timezone": config.reporting_timezone, "totals": { "requests": int(totals[0] or 0), "failures": int(totals[1] or 0), "input_tokens": int(totals[2] or 0), "output_tokens": int(totals[3] or 0), "artifacts": int(totals[4] or 0), "media_seconds": int(totals[5] or 0), "people": int(totals[6] or 0), }, "by_capability": by_capability, "by_model": by_model, "daily": daily, "queue": queue, "catalog_version": catalog.CATALOG_VERSION, } # ── admin: drill-downs (still counters only) ─────────────── # # Everything below answers "who" and "why" without widening what an # admin can read. ``ai_labs_usage_daily`` already carried # ``person_id`` — the per-person question was never blocked by a # privacy boundary, only by a missing GROUP BY — and the reason # dimension comes from ``ai_labs_failures_daily``, whose only # non-numeric column is a normalised error code. Neither table has a # column that could hold a prompt, a reply or an object key, so the # guarantee the console advertises stays structural. #: Sort keys the people list accepts, mapped to their aggregate. A #: caller-supplied column name is never interpolated — it is a lookup #: in this map, and an unknown key falls back to ``requests``. #: #: ``ClassVar`` and built once at import: SQLAlchemy expressions are #: immutable constructs, so sharing them across instances and requests #: is safe (each ``select()`` copies what it is handed). _PEOPLE_SORTS: ClassVar[dict[str, ColumnElement[Any]]] = { "requests": func.sum(AiLabUsageDaily.requests), "failures": func.sum(AiLabUsageDaily.failures), "tokens": func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), "last_used": func.max(AiLabUsageDaily.usage_date), } def _window(self, config: AiLabsConfig, days: int) -> tuple[date, date]: """``(since, through)`` in the reporting timezone, inclusive. The same arithmetic as :meth:`usage_overview`, and for the same reason: ``usage_date`` is only ever written in this calendar, and ``today - days`` would span ``days + 1`` of them. """ today = usage.local_today(config.reporting_timezone) return today - timedelta(days=max(1, days) - 1), today async def _hydrate_people( self, person_ids: list[uuid.UUID], ) -> dict[uuid.UUID, dict[str, Any]]: """Directory labels for a page of person ids. A **second** query rather than a join on the aggregate, on purpose: a person can hold more than one primary org membership (the multi-instance HRIS case), and joining that table into a ``GROUP BY person_id`` would multiply their counters. Aggregate first, label afterwards. ``person_id`` on the counter tables carries no FK — a counter outliving its person row is correct for a spend statistic — so a missing directory row is an expected outcome, not an error. The caller renders the id in that case instead of dropping the row, because a leaver's usage is still usage. """ if not person_ids: return {} rows = ( await self.db.execute( select( IdentityPerson.id, IdentityPerson.display_name, IdentityPerson.work_email, IdentityPerson.employee_no, IdentityPerson.job_title, IdentityPerson.employment_status, ).where(IdentityPerson.id.in_(person_ids)), ) ).all() people = { row[0]: { "person_id": str(row[0]), # Kept as separate fields, never pre-joined into one # "Name · email" string: the frontend needs the halves # apart to render a name column and an email column, and # splitting a concatenation back open fails for anyone # whose name contains the separator. "display_name": row[1], "work_email": row[2], "employee_no": row[3], "job_title": row[4], "employment_status": row[5], "department": None, } for row in rows } # Primary department, one row per person. ``GROUP BY`` + ``MIN`` # rather than a bare select so a duplicate primary membership # yields one deterministic label instead of two rows. dept_rows = ( await self.db.execute( select( IdentityOrgMembership.person_id, func.min(IdentityOrgMembership.org_node_name), ) .where( IdentityOrgMembership.person_id.in_(person_ids), IdentityOrgMembership.is_primary.is_(True), IdentityOrgMembership.deleted_at.is_(None), ) .group_by(IdentityOrgMembership.person_id), ) ).all() for person_id, name in dept_rows: if person_id in people: people[person_id]["department"] = name return people def _person_row( self, person_id: uuid.UUID, hydrated: dict[uuid.UUID, dict[str, Any]], ) -> dict[str, Any]: """Directory fields for one counter row, present or not.""" known = hydrated.get(person_id) if known is not None: return dict(known) return { "person_id": str(person_id), "display_name": None, "work_email": None, "employee_no": None, "job_title": None, "employment_status": None, "department": None, } async def usage_by_person( self, *, days: int, page: int = 1, page_size: int = 20, sort: str = "requests", search: str | None = None, ) -> dict[str, Any]: """Who used AI Labs in the window, and how much. Counters per person, newest-heaviest first. ``search`` filters on name / email / employee number through the directory, so it is resolved to a person-id set *before* the aggregate runs rather than joined into it. """ config = await resolve_config(self.db) since, today = self._window(config, days) order_expr = self._PEOPLE_SORTS.get( sort, self._PEOPLE_SORTS["requests"], ) conditions: list[ColumnElement[bool]] = [ AiLabUsageDaily.usage_date.between(since, today), ] term = (search or "").strip() if term: like = f"%{term}%" matched = ( await self.db.execute( select(IdentityPerson.id).where( or_( IdentityPerson.display_name.ilike(like), IdentityPerson.work_email.ilike(like), IdentityPerson.employee_no.ilike(like), ), ).limit(2000), ) ).scalars().all() if not matched: # No directory match means no rows, and it must not mean # "no filter" — an empty ``IN ()`` would return everyone. return { "window_days": max(1, days), "since": since.isoformat(), "through": today.isoformat(), "timezone": config.reporting_timezone, "sort": sort, "items": [], "meta": { "page": page, "page_size": page_size, "total": 0, "total_pages": 0, }, } conditions.append(AiLabUsageDaily.person_id.in_(matched)) total = int( ( await self.db.execute( select( func.count(func.distinct(AiLabUsageDaily.person_id)), ).where(*conditions), ) ).scalar() or 0, ) page = max(1, page) page_size = max(1, min(100, page_size)) rows = ( await self.db.execute( select( AiLabUsageDaily.person_id, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), func.coalesce(func.sum(AiLabUsageDaily.artifacts), 0), func.coalesce(func.sum(AiLabUsageDaily.media_seconds), 0), func.max(AiLabUsageDaily.usage_date), func.count(func.distinct(AiLabUsageDaily.capability)), ) .where(*conditions) .group_by(AiLabUsageDaily.person_id) .order_by(order_expr.desc(), AiLabUsageDaily.person_id) .offset((page - 1) * page_size) .limit(page_size), ) ).all() hydrated = await self._hydrate_people([row[0] for row in rows]) items = [] for row in rows: item = self._person_row(row[0], hydrated) item.update({ "requests": int(row[1] or 0), "failures": int(row[2] or 0), "tokens": int(row[3] or 0), "artifacts": int(row[4] or 0), "media_seconds": int(row[5] or 0), "last_used_date": row[6].isoformat() if row[6] else None, "capabilities_used": int(row[7] or 0), }) items.append(item) return { "window_days": max(1, days), "since": since.isoformat(), "through": today.isoformat(), "timezone": config.reporting_timezone, "sort": sort, "items": items, "meta": { "page": page, "page_size": page_size, "total": total, "total_pages": (total + page_size - 1) // page_size, }, } async def person_usage_detail( self, *, person_id: uuid.UUID, days: int, ) -> dict[str, Any]: """One person's counters, broken out four ways. By day, by capability, by model, and by failure reason. Still counters: this endpoint cannot show what they asked, only how much they asked and what came back broken. """ config = await resolve_config(self.db) since, today = self._window(config, days) in_window = AiLabUsageDaily.usage_date.between(since, today) mine = AiLabUsageDaily.person_id == person_id totals = ( await self.db.execute( select( func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), func.coalesce(func.sum(AiLabUsageDaily.artifacts), 0), func.coalesce(func.sum(AiLabUsageDaily.media_seconds), 0), func.max(AiLabUsageDaily.usage_date), ).where(mine, in_window), ) ).one() by_capability = [ { "capability": row[0], "requests": int(row[1] or 0), "failures": int(row[2] or 0), "tokens": int(row[3] or 0), "artifacts": int(row[4] or 0), } for row in ( await self.db.execute( select( AiLabUsageDaily.capability, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), func.coalesce(func.sum(AiLabUsageDaily.artifacts), 0), ) .where(mine, in_window) .group_by(AiLabUsageDaily.capability) .order_by(func.sum(AiLabUsageDaily.requests).desc()), ) ).all() ] by_model = [ { "model": row[0], "requests": int(row[1] or 0), "tokens": int(row[2] or 0), } for row in ( await self.db.execute( select( AiLabUsageDaily.model, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), ) .where(mine, in_window) .group_by(AiLabUsageDaily.model) .order_by(func.sum(AiLabUsageDaily.requests).desc()) .limit(20), ) ).all() ] daily = [ { "date": row[0].isoformat() if row[0] else None, "requests": int(row[1] or 0), "failures": int(row[2] or 0), "tokens": int(row[3] or 0), } for row in ( await self.db.execute( select( AiLabUsageDaily.usage_date, func.coalesce(func.sum(AiLabUsageDaily.requests), 0), func.coalesce(func.sum(AiLabUsageDaily.failures), 0), func.coalesce( func.sum( AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens, ), 0, ), ) .where(mine, in_window) .group_by(AiLabUsageDaily.usage_date) .order_by(AiLabUsageDaily.usage_date), ) ).all() ] by_error = await self._failures_by_error( since, today, person_id=person_id, ) hydrated = await self._hydrate_people([person_id]) person = self._person_row(person_id, hydrated) return { "window_days": max(1, days), "since": since.isoformat(), "through": today.isoformat(), "timezone": config.reporting_timezone, "person": person, "totals": { "requests": int(totals[0] or 0), "failures": int(totals[1] or 0), "tokens": int(totals[2] or 0), "artifacts": int(totals[3] or 0), "media_seconds": int(totals[4] or 0), "last_used_date": ( totals[5].isoformat() if totals[5] else None ), }, "by_capability": by_capability, "by_model": by_model, "daily": daily, "by_error": by_error, # The same reconciliation the failure card does, scoped to one # person — see :meth:`failure_overview`. "unclassified_failures": max( 0, int(totals[1] or 0) - await self._classified_failures( since, today, person_id=person_id, ), ), } async def _classified_failures( self, since: date, through: date, *, person_id: uuid.UUID | None = None, ) -> int: """Total failures that carry a reason, over **every** code. Its own SUM rather than adding up :meth:`_failures_by_error`: that one is a display list and is ``LIMIT``-ed. Deriving the classified total from a truncated list would inflate ``unclassified_failures`` by whatever fell off the end — the reconciliation number would then misreport the very thing it exists to report honestly. """ conditions: list[ColumnElement[bool]] = [ AiLabFailureDaily.usage_date.between(since, through), ] if person_id is not None: conditions.append(AiLabFailureDaily.person_id == person_id) return int( ( await self.db.execute( select( func.coalesce(func.sum(AiLabFailureDaily.failures), 0), ).where(*conditions), ) ).scalar() or 0, ) async def _failures_by_error( self, since: date, through: date, *, person_id: uuid.UUID | None = None, ) -> list[dict[str, Any]]: """Failure counts grouped by error code, heaviest first.""" conditions: list[ColumnElement[bool]] = [ AiLabFailureDaily.usage_date.between(since, through), ] if person_id is not None: conditions.append(AiLabFailureDaily.person_id == person_id) rows = ( await self.db.execute( select( AiLabFailureDaily.error_code, func.coalesce(func.sum(AiLabFailureDaily.failures), 0), func.count(func.distinct(AiLabFailureDaily.person_id)), func.max(AiLabFailureDaily.last_seen_at), ) .where(*conditions) .group_by(AiLabFailureDaily.error_code) .order_by(func.sum(AiLabFailureDaily.failures).desc()) .limit(50), ) ).all() return [ { "error_code": row[0], "failures": int(row[1] or 0), "people": int(row[2] or 0), "last_seen_at": row[3].isoformat() if row[3] else None, } for row in rows ] async def failure_overview(self, *, days: int) -> dict[str, Any]: """What is failing, and for whom. The one number here that needs explaining is ``unclassified_failures``: ``ai_labs_usage_daily.failures`` is the authoritative failure count and predates the reason table, so a window that reaches back before this shipped — or that contains a failure booked on a path that had no code in hand — has more counted failures than classified ones. Reporting the difference explicitly is the point: a reason breakdown that silently covers two of eight failures, presented as *the* breakdown, is worse than no breakdown at all. """ config = await resolve_config(self.db) since, today = self._window(config, days) counted = int( ( await self.db.execute( select( func.coalesce(func.sum(AiLabUsageDaily.failures), 0), ).where(AiLabUsageDaily.usage_date.between(since, today)), ) ).scalar() or 0, ) by_error = await self._failures_by_error(since, today) classified = await self._classified_failures(since, today) by_capability_error = [ { "capability": row[0], "error_code": row[1], "failures": int(row[2] or 0), "last_seen_at": row[3].isoformat() if row[3] else None, } for row in ( await self.db.execute( select( AiLabFailureDaily.capability, AiLabFailureDaily.error_code, func.coalesce(func.sum(AiLabFailureDaily.failures), 0), func.max(AiLabFailureDaily.last_seen_at), ) .where(AiLabFailureDaily.usage_date.between(since, today)) .group_by( AiLabFailureDaily.capability, AiLabFailureDaily.error_code, ) .order_by(func.sum(AiLabFailureDaily.failures).desc()) .limit(50), ) ).all() ] by_model_error = [ { "model": row[0], "error_code": row[1], "failures": int(row[2] or 0), } for row in ( await self.db.execute( select( AiLabFailureDaily.model, AiLabFailureDaily.error_code, func.coalesce(func.sum(AiLabFailureDaily.failures), 0), ) .where(AiLabFailureDaily.usage_date.between(since, today)) .group_by( AiLabFailureDaily.model, AiLabFailureDaily.error_code, ) .order_by(func.sum(AiLabFailureDaily.failures).desc()) .limit(50), ) ).all() ] daily = [ { "date": row[0].isoformat() if row[0] else None, "failures": int(row[1] or 0), } for row in ( await self.db.execute( select( AiLabFailureDaily.usage_date, func.coalesce(func.sum(AiLabFailureDaily.failures), 0), ) .where(AiLabFailureDaily.usage_date.between(since, today)) .group_by(AiLabFailureDaily.usage_date) .order_by(AiLabFailureDaily.usage_date), ) ).all() ] people_rows = ( await self.db.execute( select( AiLabFailureDaily.person_id, func.coalesce(func.sum(AiLabFailureDaily.failures), 0), func.count(func.distinct(AiLabFailureDaily.error_code)), func.max(AiLabFailureDaily.last_seen_at), ) .where(AiLabFailureDaily.usage_date.between(since, today)) .group_by(AiLabFailureDaily.person_id) .order_by(func.sum(AiLabFailureDaily.failures).desc()) .limit(10), ) ).all() hydrated = await self._hydrate_people( [row[0] for row in people_rows], ) top_people = [] for row in people_rows: entry = self._person_row(row[0], hydrated) entry.update({ "failures": int(row[1] or 0), "error_codes": int(row[2] or 0), "last_seen_at": row[3].isoformat() if row[3] else None, }) top_people.append(entry) return { "window_days": max(1, days), "since": since.isoformat(), "through": today.isoformat(), "timezone": config.reporting_timezone, "counted_failures": counted, "classified_failures": classified, "unclassified_failures": max(0, counted - classified), "by_error": by_error, "by_capability_error": by_capability_error, "by_model_error": by_model_error, "daily": daily, "top_people": top_people, } async def my_usage( self, *, person_id: uuid.UUID, config: AiLabsConfig, ) -> dict[str, Any]: """One person's own totals for today, plus the ceilings in force.""" today = usage.local_today(config.reporting_timezone) totals = await usage.daily_totals( self.db, person_id=person_id, on_date=today, ) quotas = config.quotas return { "date": today.isoformat(), "timezone": config.reporting_timezone, "used": { "tokens": totals.tokens, "images": totals.images, "video_seconds": totals.video_seconds, "speech": totals.speech, "transcribe": totals.transcribe, }, # ``0`` means unlimited, and the frontend renders that as "no # limit" rather than as a full bar at zero. "limits": { "tokens": quotas.daily_tokens, "images": quotas.daily_images, "video_seconds": quotas.daily_video_seconds, "speech": quotas.daily_speech, "transcribe": quotas.daily_transcribe, }, "enforced": quotas.any_enforced, } __all__ = ["ARTIFACT_URL_TTL_SECONDS", "AiLabsService"]