"""The streaming chat turn. Split into two halves on purpose, and the split is forced by how streaming responses actually run: * :func:`prepare_turn` runs **inside** the request, on the request's session. It validates, loads history, resolves attachments, and writes the employee's question — then commits. After it returns, the question is durable whether or not the answer ever arrives. * :func:`stream_turn` runs **after the handler has returned**, driving the ``StreamingResponse``. By then FastAPI has closed the request's session, so this half owns its own short-lived sessions and receives nothing but plain data. No ORM instance crosses the boundary — one that did would be expired and would raise ``MissingGreenlet`` on first attribute access, mid-stream, where the only symptom is a truncated answer. The gap between the two halves is also why no database connection is held across the vendor call. A thinking model's turn can run past a minute, and an idle pooled connection parked that long is how you collect MySQL 2013s under load. **Nothing is added to the employee's prompt.** No system message, no preamble, no rewriting. ``messages`` is the thread's own turns plus this one. ``test_prompt_passthrough_contract`` asserts it against the payload that actually reaches the vendor, so a later "helpful" addition fails a test rather than quietly changing the product. """ from __future__ import annotations import json import logging import time import uuid from collections.abc import AsyncIterator from dataclasses import dataclass, field from typing import Any from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.audit.service import AuditService from app.core.background import spawn from app.core.exceptions import AppException, NotFoundError from app.domains.ai_labs import attachments as attach from app.domains.ai_labs import catalog, usage from app.domains.ai_labs.models import AiLabConversation, AiLabMessage from app.domains.ai_labs.settings import AiLabsConfig from app.domains.bailian.llm_client import BailianLLMClient logger = logging.getLogger(__name__) ROLE_USER = "user" ROLE_ASSISTANT = "assistant" @dataclass class PreparedTurn: """Everything the stream needs, as plain data. Deliberately free of ORM instances — see the module docstring. """ conversation_id: uuid.UUID person_id: uuid.UUID user_message_id: uuid.UUID assistant_seq: int model_key: str vendor_model: str thinking_enabled: bool api_key: str region: str reporting_timezone: str #: The exact list handed to the vendor. vendor_messages: list[dict[str, Any]] = field(default_factory=list) #: Metadata for the attachments on this turn, already persisted. attachments: list[dict[str, Any]] = field(default_factory=list) async def _owned_conversation( db: AsyncSession, *, conversation_id: uuid.UUID, person_id: uuid.UUID, ) -> AiLabConversation: """Load a conversation, or 404 if it is not this person's. 404 rather than 403 for someone else's thread: the caller has no legitimate way to know the id exists, and confirming it does is a (small) information leak in a surface whose whole promise is privacy. """ conversation = await db.get(AiLabConversation, conversation_id) if conversation is None or conversation.person_id != person_id: raise NotFoundError("conversation not found") return conversation def _vendor_content( text: str, loaded: list[attach.Attachment], ) -> str | list[dict[str, Any]]: """One user turn's content, in the vendor's shape. A turn with no attachments stays a plain string — the simplest thing the endpoint accepts, and the shape every model handles. With attachments it becomes a parts list: images as inline data URIs (verified 2026-09-03), documents as their extracted text introduced by nothing but the filename. """ if not loaded: return text parts: list[dict[str, Any]] = [] for item in loaded: if item.kind == attach.KIND_IMAGE: parts.append({ "type": "image_url", "image_url": {"url": item.data_uri}, }) document_texts: list[str] = [] for item in loaded: if item.kind == attach.KIND_DOCUMENT and item.text: # The filename is a label, not an instruction. Anything more # would be Helix writing part of the prompt. document_texts.append(f"{item.filename}:\n{item.text}") combined = "\n\n".join([*document_texts, text]) if document_texts else text parts.append({"type": "text", "text": combined}) return parts async def _history( db: AsyncSession, *, conversation_id: uuid.UUID, before_seq: int, ) -> list[dict[str, Any]]: """Prior turns, oldest first, capped at the catalogue's window. Failed turns are skipped: an assistant row with an ``error_code`` never carried an answer, and replaying "" as the model's previous reply teaches it that empty responses are acceptable. Attachments from *earlier* turns are not re-sent. Re-inlining every image of a long thread would grow each turn's cost without bound; the text of those turns still carries the conversation. """ stmt = ( select(AiLabMessage) .where( AiLabMessage.conversation_id == conversation_id, AiLabMessage.seq < before_seq, ) .order_by(AiLabMessage.seq.desc()) .limit(catalog.MAX_HISTORY_MESSAGES) ) rows = list((await db.execute(stmt)).scalars().all()) rows.reverse() messages: list[dict[str, Any]] = [] for row in rows: if row.error_code: continue content = (row.content or "").strip() if not content: continue messages.append({"role": row.role, "content": content}) return messages async def prepare_turn( db: AsyncSession, *, conversation_id: uuid.UUID, person_id: uuid.UUID, text: str, attachment_keys: list[str], config: AiLabsConfig, attachment_names: dict[str, str] | None = None, trace_id: str = "", model_key: str | None = None, thinking: bool | None = None, client_token: str | None = None, ) -> PreparedTurn: """Validate, persist the question, and assemble the vendor payload. Commits before returning: the employee's message and its sequence number must be durable before a stream that may fail at any point. """ text = (text or "").strip() if not text: raise AppException( code="AI_LABS_MESSAGE_REQUIRED", message="A message is required.", status_code=422, field="text", ) if len(text) > catalog.MAX_MESSAGE_CHARS: raise AppException( code="AI_LABS_MESSAGE_TOO_LONG", message=( f"A message may be at most {catalog.MAX_MESSAGE_CHARS} " "characters." ), status_code=422, field="text", ) conversation = await _owned_conversation( db, conversation_id=conversation_id, person_id=person_id, ) # A model chosen for this turn sticks to the thread, so the next turn # continues with the same model instead of silently reverting. if model_key and model_key in catalog.CHAT_MODELS_BY_KEY: conversation.model_key = model_key if thinking is not None: conversation.thinking_enabled = bool(thinking) model = catalog.chat_model(conversation.model_key) if conversation.thinking_enabled and not model.supports_thinking: conversation.thinking_enabled = False loaded = await attach.load( db, object_keys=attachment_keys, owner_id=person_id, allowed_kinds={attach.KIND_IMAGE, attach.KIND_DOCUMENT}, names=attachment_names, ) if loaded and any( item.kind == attach.KIND_IMAGE for item in loaded ) and not model.supports_vision: raise AppException( code="AI_LABS_MODEL_NO_VISION", message="The selected model cannot read images.", status_code=422, field="model_key", ) user_seq = conversation.message_count + 1 assistant_seq = user_seq + 1 # History BEFORE the insert, deliberately. It reads ``seq < user_seq`` # so the new row is not in it either way — but leaving the insert # pending across a query means autoflush performs it, and a duplicate # ``seq`` then surfaces as an IntegrityError raised from the middle of # an unrelated SELECT. The failure has to happen where it is handled. history = await _history( db, conversation_id=conversation.id, before_seq=user_seq, ) stored_attachments = [item.as_stored() for item in loaded] user_message = AiLabMessage( conversation_id=conversation.id, seq=user_seq, role=ROLE_USER, content=text, attachments=stored_attachments or None, model=None, client_token=(client_token or None), ) # Two callers racing on one conversation both read the same # ``message_count`` and both claim ``user_seq``. ``uq_alm_conv_seq`` # stops the loser — that is the constraint doing its job — but a raw # IntegrityError becomes a 500, which reads as "Helix broke" when the # actual event was one person sending from two tabs. The SAVEPOINT is # what makes the clean answer possible: a failed flush aborts the # transaction, so without isolating it the 409 would be followed by # every later statement on this session failing too, including the # ones the error envelope needs. try: async with db.begin_nested(): db.add(user_message) await db.flush() except IntegrityError as exc: # Two constraints can land here and they mean different things to # the employee. ``uq_alm_conv_client`` means "this exact send was # already accepted" — a retry, so the honest answer is that the # turn exists and the client should re-read the thread rather than # pay for it twice. ``uq_alm_conv_seq`` means "another send is # mid-flight". Matched on the constraint name in the driver's # message because that is the only place MySQL puts it. if "uq_alm_conv_client" in str(getattr(exc, "orig", exc)): raise AppException( code="AI_LABS_TURN_ALREADY_SENT", message=( "This message was already sent. Reload the conversation " "to see it." ), status_code=409, ) from exc raise AppException( code="AI_LABS_TURN_IN_FLIGHT", message=( "Another message is already being sent in this " "conversation. Wait for it to finish, then try again." ), status_code=409, ) from exc # **Claim the upload grants**, in the same transaction as the message # that references them. Without this every AI Labs attachment stays # ``status="issued"`` in the platform's upload ledger, and # ``storage.orphan_upload_cleanup`` — which selects exactly that state # at 24h old — deletes the S3 object: the chip 404s the next day. It is # also a second ownership gate, re-checking that each grant belongs to # this person and carries an allowed purpose. await attach.claim( db, object_keys=attachment_keys, owner_id=person_id, resource_type="ai_labs_message", resource_id=str(user_message.id), ) if not (conversation.title or "").strip(): # First question becomes the thread's name. Trimmed on a word # boundary where there is one, so a title does not end mid-word. conversation.title = _derive_title(text) conversation.message_count = assistant_seq conversation.last_message_at = user_message.created_at user_message_id = user_message.id # The turn itself is auditable, in the shape this domain's audit rows # always take: who, when, which model, how much — and not one word of # what was asked. Written here rather than after the stream because # this is the moment the request was accepted, and because the stream # half owns no request context. await AuditService(db).log( trace_id=trace_id, actor_type="user", actor_user_id=person_id, action="ai_labs.chat.turn", status="success", resource_type="ai_labs_conversation", resource_id=str(conversation.id), detail={ "model": model.vendor_model, "thinking": bool(conversation.thinking_enabled), "seq": user_seq, "chars": len(text), "attachments": len(loaded), "history_messages": len(history), }, ) await db.commit() vendor_messages = [ *history, {"role": ROLE_USER, "content": _vendor_content(text, loaded)}, ] return PreparedTurn( conversation_id=conversation_id, person_id=person_id, user_message_id=user_message_id, assistant_seq=assistant_seq, model_key=model.key, vendor_model=config.vendor_chat_model(model.key), thinking_enabled=bool(conversation.thinking_enabled), api_key=config.api_key, region=config.region, reporting_timezone=config.reporting_timezone, vendor_messages=vendor_messages, attachments=stored_attachments, ) def _derive_title(text: str, *, limit: int = 40) -> str: """A thread title from its first message.""" flat = " ".join(text.split()) if len(flat) <= limit: return flat window = flat[:limit] if " " in window[20:]: window = window[: window.rindex(" ")] return window + "…" def _frame(event: str, payload: dict[str, Any]) -> str: """One SSE frame. ``ensure_ascii=False`` keeps CJK readable on the wire, and the double newline is the frame terminator the browser's parser needs. """ return ( f"event: {event}\n" f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" ) async def stream_turn(prepared: PreparedTurn) -> AsyncIterator[str]: """Drive one turn, yielding SSE frames. The frame protocol is small and closed: ``meta`` once, first — ids the client needs to render the pair. ``reasoning`` zero or more — the model's thinking, when it emits any. ``delta`` zero or more — answer text, in order. ``done`` once, last on success — finish reason and token counts. ``error`` once, last on failure — a code and a human sentence. Exactly one of ``done`` / ``error`` is always emitted, so a client never has to guess whether a silent stream ended or stalled. **A stream that ends early still persists what it produced.** Pressing Stop, closing the tab or losing the network cancels this generator at whichever ``yield`` it is sitting on, so neither the success path nor the failure path below is reached. Without the ``finally`` that would lose the assistant turn the employee had already read, skip the usage counter for tokens the vendor had already billed, and leave a permanent hole in the conversation's ``seq`` sequence (the user message is committed before the vendor is contacted). The partial answer is written by a **detached task**: awaiting the write inside a cancelled generator is not reliable — under ``GeneratorExit`` an ``await`` that does not complete immediately raises — whereas scheduling one does not await at all. """ started = time.monotonic() answer: list[str] = [] reasoning: list[str] = [] finish_reason: str | None = None input_tokens: int | None = None output_tokens: int | None = None #: Set by every path that has already written this turn's outcome. The #: ``finally`` uses it to tell "ended early" from "ended". settled = False # The ``meta`` frame is INSIDE the guard. A tab closed between it and # the first token is the most likely disconnect of all — it is the # window where the vendor has not answered yet — and leaving it outside # left exactly that case writing nothing and stranding the seq. try: # Constructed INSIDE the guard. It raises ``BAILIAN_LLM_NO_KEY`` on # an empty key, and outside the guard that produced neither a # ``done`` nor an ``error`` frame and persisted nothing — breaking # the invariant this docstring states absolutely. Narrow (the key # was non-empty at ``prepare_turn``), but the invariant should hold # without needing that argument. client = BailianLLMClient( prepared.api_key, model=prepared.vendor_model, region=prepared.region, purpose="ai_labs.chat") yield _frame("meta", { "conversation_id": str(prepared.conversation_id), "user_message_id": str(prepared.user_message_id), "model_key": prepared.model_key, "thinking": prepared.thinking_enabled, }) try: async for chunk in client.chat_stream( messages=prepared.vendor_messages, model=prepared.vendor_model, max_tokens=catalog.MAX_OUTPUT_TOKENS, enable_thinking=prepared.thinking_enabled, ): kind = chunk.get("type") if kind == "text": answer.append(chunk["text"]) yield _frame("delta", {"text": chunk["text"]}) elif kind == "reasoning": reasoning.append(chunk["text"]) yield _frame("reasoning", {"text": chunk["text"]}) elif kind == "done": finish_reason = chunk.get("finish_reason") input_tokens = chunk.get("input_tokens") output_tokens = chunk.get("output_tokens") except AppException as exc: # Set BEFORE the await, not after. ``_persist_failure`` catches # ``Exception``, not ``BaseException``, so a cancellation # delivered while it is awaiting propagates out with ``settled`` # still false — and the ``finally`` then writes the turn a # SECOND time and records a second usage row for one vendor # call. Being set early is the safe direction: the worst case # is a turn nobody persisted, which the row's absence makes # obvious, rather than one billed twice. settled = True await _persist_failure(prepared, code=exc.code, message=exc.message) yield _frame("error", {"code": exc.code, "message": exc.message}) return except Exception as exc: # noqa: BLE001 — the client must be told something logger.exception( "ai_labs: chat turn failed for conversation %s", prepared.conversation_id, ) settled = True # before the await — see the note above await _persist_failure( prepared, code="AI_LABS_INTERNAL", message=str(exc)[:300], ) yield _frame("error", { "code": "AI_LABS_INTERNAL", "message": "The model could not be reached. Please try again.", }) return latency_ms = int((time.monotonic() - started) * 1000) text = "".join(answer) if not text.strip(): # A finished stream that produced nothing is a failure, not an # empty answer: rendering a blank bubble tells the employee their # question was answered when it was not. settled = True # before the await — see the note above await _persist_failure( prepared, code="AI_LABS_EMPTY_REPLY", message="The model returned an empty reply.", latency_ms=latency_ms, input_tokens=input_tokens, output_tokens=output_tokens, ) yield _frame("error", { "code": "AI_LABS_EMPTY_REPLY", "message": "The model returned an empty reply. Please try again.", }) return settled = True # before the await — see the note above assistant_message_id = await _persist_answer( prepared, text=text, reasoning="".join(reasoning) or None, finish_reason=finish_reason, input_tokens=input_tokens, output_tokens=output_tokens, latency_ms=latency_ms, ) await usage.record( person_id=prepared.person_id, capability_key=usage.CAPABILITY_CHAT, model=prepared.vendor_model, timezone_name=prepared.reporting_timezone, requests=1, input_tokens=input_tokens or 0, output_tokens=output_tokens or 0, ) yield _frame("done", { "assistant_message_id": ( str(assistant_message_id) if assistant_message_id else None ), "finish_reason": finish_reason, "input_tokens": input_tokens, "output_tokens": output_tokens, "latency_ms": latency_ms, }) finally: if not settled: # Reached only when this generator was closed before any path # above recorded the turn: Stop, a closed tab, a dropped # connection. Detached on purpose (see the docstring), and # given everything it needs by value so it owns no state of # ours. spawn( _persist_interrupted( prepared, text="".join(answer), reasoning="".join(reasoning) or None, input_tokens=input_tokens, output_tokens=output_tokens, latency_ms=int((time.monotonic() - started) * 1000), ), name=f"ai_labs.interrupted.{prepared.conversation_id}", ) async def _persist_answer( prepared: PreparedTurn, *, text: str, reasoning: str | None, finish_reason: str | None, input_tokens: int | None, output_tokens: int | None, latency_ms: int, ) -> uuid.UUID | None: """Write the assistant turn on a fresh session. Returns ``None`` if the row could not be written — the employee has already read the answer off the stream by then, so this is logged rather than escalated into an error frame that would contradict what they are looking at. """ try: # Function-local import on purpose: the test harness monkeypatches # ``app.core.database.async_session_factory`` to bind the isolated # test schema, and a module-level ``from ... import`` captures the # original object at import time — every write here would silently # land in the dev database instead of the test one, and the only # symptom would be a later assertion about history failing. from app.core.database import async_session_factory async with async_session_factory() as db: message = AiLabMessage( conversation_id=prepared.conversation_id, seq=prepared.assistant_seq, role=ROLE_ASSISTANT, content=text, reasoning=reasoning, model=prepared.vendor_model, input_tokens=input_tokens, output_tokens=output_tokens, latency_ms=latency_ms, finish_reason=finish_reason, ) db.add(message) await db.flush() message_id = message.id conversation = await db.get( AiLabConversation, prepared.conversation_id, ) if conversation is not None: conversation.last_message_at = message.created_at await db.commit() return message_id except Exception: # noqa: BLE001 logger.exception( "ai_labs: could not persist the answer for conversation %s", prepared.conversation_id, ) return None async def _persist_interrupted( prepared: PreparedTurn, *, text: str, reasoning: str | None, input_tokens: int | None, output_tokens: int | None, latency_ms: int, ) -> None: """Record a turn whose stream ended before it finished. Two shapes, and which one is written matters to the employee: * **Something was produced.** Keep it, marked ``interrupted``. This is the Stop button's whole purpose — the first paragraph was the useful part — and it is also what makes a dropped mobile connection cost the tail of an answer rather than all of it. * **Nothing was produced.** Write a failed turn, so the thread does not show a question with no reply. A ``seq`` that stays empty is worse than a visible failure: the unique ``(conversation, seq)`` key means the number is spent either way, and a hole makes the next turn look like a message went missing. Never raises, and never assumes the request is still alive: it is not. """ if text.strip(): await _persist_answer( prepared, text=text, reasoning=reasoning, finish_reason="interrupted", input_tokens=input_tokens, output_tokens=output_tokens, latency_ms=latency_ms, ) # The vendor billed for the tokens it had already emitted whether or # not anyone was still reading, so this counts as a request rather # than a failure. await usage.record( person_id=prepared.person_id, capability_key=usage.CAPABILITY_CHAT, model=prepared.vendor_model, timezone_name=prepared.reporting_timezone, requests=1, input_tokens=input_tokens or 0, output_tokens=output_tokens or 0, ) return await _persist_failure( prepared, code="AI_LABS_INTERRUPTED", message="The connection closed before the model replied.", latency_ms=latency_ms, input_tokens=input_tokens, output_tokens=output_tokens, ) async def _persist_failure( prepared: PreparedTurn, *, code: str, message: str, latency_ms: int | None = None, input_tokens: int | None = None, output_tokens: int | None = None, ) -> None: """Record the failed turn, so the thread shows what happened. Without this the question would sit in the thread with no reply and no explanation, which reads as "Helix lost my message". Only the ``code`` is stored: ``ai_labs_messages`` has no ``error_message`` column, and adding one would put vendor prose (which can quote the prompt back) in the row. The client renders its own localized sentence from the code, and ``message`` is logged here so the detail is not lost to whoever is debugging. """ logger.info( "ai_labs: turn %s of conversation %s failed: %s — %s", prepared.assistant_seq, prepared.conversation_id, code, message, ) try: from app.core.database import async_session_factory async with async_session_factory() as db: db.add(AiLabMessage( conversation_id=prepared.conversation_id, seq=prepared.assistant_seq, role=ROLE_ASSISTANT, content="", model=prepared.vendor_model, error_code=code[:64], latency_ms=latency_ms, input_tokens=input_tokens, output_tokens=output_tokens, )) await db.commit() except Exception: # noqa: BLE001 logger.exception( "ai_labs: could not persist the failure for conversation %s", prepared.conversation_id, ) # Tokens spent on a turn that failed still cost the org, and a failure # rate is exactly what the admin console needs to see. await usage.record( person_id=prepared.person_id, capability_key=usage.CAPABILITY_CHAT, model=prepared.vendor_model, timezone_name=prepared.reporting_timezone, requests=1, failures=1, input_tokens=input_tokens or 0, output_tokens=output_tokens or 0, # The same code that goes on the message row, so the admin console # can say *what* failed and not only *that* something did. It is a # code, never ``message`` — see ``usage.normalize_error_code``. error_code=code, ) __all__ = [ "ROLE_ASSISTANT", "ROLE_USER", "PreparedTurn", "prepare_turn", "stream_turn", ]