"""Background reply generation. Generation used to be driven by the SSE request: the browser opening the stream was what produced the tokens, so navigating away cancelled the reply mid- sentence. Here it runs as its own task instead, and the SSE endpoint merely *follows* it. Closing the page, opening another chat, or starting a new one leaves the answer being written; coming back attaches to it and immediately receives everything produced so far. The registry is in-process, which is right for the single-worker deployment this ships with. Several workers would need the state in the database or a broker, because the request that follows a generation would not necessarily land in the process running it. """ from __future__ import annotations import asyncio import contextlib import logging import time from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from sqlalchemy import select from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Chat, Message, User from lembas.db.session import session_scope from lembas.services import chat as chat_service from lembas.services import compaction as compaction_service from lembas.services import metrics as metrics_service from lembas.services import prompts as prompts_service from lembas.services import tokens from lembas.services import tools as tools_service from lembas.services.llm.openai_client import ( LLMError, chunk_usage, delta_reasoning, delta_text, delta_tool_calls, stream_chat, ) from lembas.services.reasoning import REASONING, ReasoningSplitter log = logging.getLogger(__name__) # How often the partial answer is offered to followers. Markdown is re-rendered # whole each time -- a list or a code fence is only correct once its context # exists -- so this trades a little work for formatting that appears as the # model writes. 100ms is below the threshold where the eye reads it as stepping. RENDER_INTERVAL = 0.1 # Finished generations linger so a follower attaching at the last moment still # gets the final frames, then are pruned. KEEP_FINISHED = timedelta(minutes=5) @dataclass class Generation: """The live state of one reply being written.""" chat_id: str message_id: str content: list[str] = field(default_factory=list) reasoning: list[str] = field(default_factory=list) reasoning_ms: int = 0 # One entry per tool call made while producing this reply, in order. Shown # live as the model works and kept on the message afterwards. tool_events: list[dict] = field(default_factory=list) # --- What it cost -------------------------------------------------------- # Prompt and completion are summed across tool rounds: what the reply cost. # context_tokens is overwritten each round with that round's prompt plus # completion, because a three-round reply pays for its prompt three times # but only ever occupies the window once. prompt_tokens: int = 0 completion_tokens: int = 0 context_tokens: int = 0 context_limit: int = 0 # Filled from the assembled request before the first chunk, so a follower # has a percentage to show while the reply is still being written -- real # usage only arrives in a single chunk at the very end. prompt_estimate: int = 0 rounds: int = 0 # time.monotonic() at the start. A field rather than a local in `_run` # because `_follow` is a different function that sees only this object, and # without it there is nothing to compute a live tokens/second against. started_at: float = 0.0 elapsed_ms: int = 0 error: str = "" stopped: bool = False done: bool = False # Bumped on every change. Followers compare against it rather than being # woken individually: with a 100ms cadence a short poll is simpler than # future bookkeeping, and cannot drop a wakeup. version: int = 0 # What the reply is doing when it is not producing tokens. Shown in the # streaming bubble, because a silent multi-second pause before the first # token is what a hang looks like. status: str = "" # Number of browsers currently watching. Decides whether a finished reply # counts as unread. followers: int = 0 finished_at: datetime | None = None cancel: bool = False def touch(self) -> None: self.version += 1 @property def text(self) -> str: return "".join(self.content) @property def thinking(self) -> str: return "".join(self.reasoning) _RUNNING: dict[str, Generation] = {} _TASKS: dict[str, asyncio.Task] = {} def get(message_id: str) -> Generation | None: return _RUNNING.get(message_id) def request_stop(message_id: str) -> bool: """Ask a running generation to stop. Returns whether one was found.""" generation = _RUNNING.get(message_id) if generation is None or generation.done: return False generation.cancel = True return True def _prune() -> None: cutoff = datetime.now(UTC) - KEEP_FINISHED for message_id, generation in list(_RUNNING.items()): if generation.done and generation.finished_at and generation.finished_at < cutoff: _RUNNING.pop(message_id, None) _TASKS.pop(message_id, None) def ensure(chat_id: str, message_id: str) -> Generation: """Start generating this reply if it is not already under way. Idempotent, because more than one thing can ask for it: the route that created the message, and any page load that finds the message unfinished. `_prune` runs first, not after the lookup. Below it, a stale entry could never expire: the early return is the only path a repeated id takes, so the sweep was unreachable for exactly the message that needed it. """ _prune() existing = _RUNNING.get(message_id) if existing is not None: return existing generation = Generation(chat_id=chat_id, message_id=message_id) _RUNNING[message_id] = generation _TASKS[message_id] = asyncio.create_task(_run(generation)) return generation def restart(chat_id: str, message_id: str) -> Generation: """Produce this reply again, discarding any finished attempt at it. `ensure` is idempotent on purpose, and that is load-bearing: a page load finding an unfinished reply must attach to it rather than start a second one, and `_follow` calls it too. Regeneration is the one caller that means the opposite. It is also the one caller that reuses an existing Message row -- blanked and marked incomplete -- rather than creating a new one. The finished Generation for that id is still in the registry, because finished ones linger KEEP_FINISHED so a follower arriving at the last moment still gets the final frames. `ensure` handed that one straight back: no request was made, `_follow` replayed the previous answer, and the `done` frame re-rendered a streaming shell because the row said incomplete. That was the reconnect loop, and the Send button stuck on Stop. """ previous = _RUNNING.pop(message_id, None) task = _TASKS.pop(message_id, None) if previous is not None and not previous.done: previous.cancel = True if task is not None: task.cancel() return ensure(chat_id, message_id) async def shutdown() -> None: """Stop every running generation, keeping what each has produced.""" for task in list(_TASKS.values()): task.cancel() for task in list(_TASKS.values()): with contextlib.suppress(asyncio.CancelledError, Exception): await task async def _run(generation: Generation) -> None: """Produce one reply, then persist it. Never raises into the task. A reply is not necessarily one request. When tools are offered and the model asks to use one, the loop below runs it, appends the result to the conversation and asks again -- up to tools_service.MAX_ROUNDS times, after which the model has to answer with what it has. Text produced before a tool call is kept, so a model that narrates what it is about to look up does not lose that when the results come back. """ splitter = ReasoningSplitter() started = time.monotonic() generation.started_at = started reasoning_started: float | None = None question = "" endpoint = model_id = None needs_title = False title_prompt = "" try: # Before the request is assembled, so build_request is called once and # what goes out is the compacted conversation -- there is no second # assembly path. Here rather than in post_message because that route's # whole contract is to return immediately, and a three-second # summarisation in front of it would break exactly that. await _maybe_compact(generation) with session_scope() as db: chat = db.get(Chat, generation.chat_id) message = db.get(Message, generation.message_id) if chat is None or message is None: generation.error = "That chat no longer exists." return endpoint, model_id = chat_service.resolve_endpoint(db, chat) owner = db.get(User, chat.user_id) # Read while the session is open: everything below outlives it. # Resolved once, so that what the loop is allowed to *run* is the # same set the endpoint was *offered* -- not whatever happens to # exist by the time a call comes back. toolset = tools_service.resolve_tools(db, chat, owner) offered = toolset.schemas payload = chat_service.build_request( db, chat, upto=message, tools=offered, user=owner ) question = _question_from(payload) needs_title = not chat.title_generated # Read here, with the rest, because titling happens after this # session has closed and must not open another one. title_prompt = prompts_service.resolve(db, "task.title") tool_context = tools_service.context_for(db, owner, chat, tools=toolset) model = chat_service.model_for(db, chat) generation.context_limit = model.context_length if model is not None else 0 generation.prompt_estimate = tokens.estimate_request(payload) for round_number in range(tools_service.MAX_ROUNDS + 1): generation.rounds = round_number + 1 accumulator = tools_service.ToolCallAccumulator() # Text the model produced in *this* round, needed separately from # generation.content when echoing the assistant turn back. round_text: list[str] = [] async for chunk in stream_chat(endpoint, payload): counts = chunk_usage(chunk) if counts is not None: generation.prompt_tokens += counts.get("prompt_tokens", 0) generation.completion_tokens += counts.get("completion_tokens", 0) # Overwritten, not summed: this round's prompt already # contains every earlier round. generation.context_tokens = counts.get("prompt_tokens", 0) + counts.get( "completion_tokens", 0 ) generation.touch() thought = delta_reasoning(chunk) if thought: if reasoning_started is None: reasoning_started = time.monotonic() generation.reasoning.append(thought) generation.touch() if offered: fragments = delta_tool_calls(chunk) if fragments: accumulator.feed(fragments) text = delta_text(chunk) if text: for kind, piece in splitter.feed(text): if kind == REASONING: if reasoning_started is None: reasoning_started = time.monotonic() generation.reasoning.append(piece) else: if reasoning_started is not None and not generation.reasoning_ms: generation.reasoning_ms = int( (time.monotonic() - reasoning_started) * 1000 ) generation.content.append(piece) round_text.append(piece) generation.touch() if generation.cancel: generation.stopped = True break # Let followers and other tasks run between chunks. await asyncio.sleep(0) calls = accumulator.calls if generation.stopped or not calls: break if round_number == tools_service.MAX_ROUNDS: # Out of rounds with the model still asking for tools. Recorded # rather than silently dropped: an answer that stops here needs # to be explicable. generation.tool_events.append( { "name": calls[0]["name"], "status": "error", "error": ( f"Stopped after {tools_service.MAX_ROUNDS} rounds of tool " f"calls without an answer." ), } ) generation.touch() break messages = [ *payload["messages"], tools_service.assistant_turn(calls, "".join(round_text)), ] for call in calls: outcome = await tools_service.run_tool( tool_context, call["name"], call["arguments"] ) generation.tool_events.append(outcome.event) generation.touch() messages.append(tools_service.tool_turn(call, outcome.content)) payload = {**payload, "messages": messages} for kind, piece in splitter.flush(): (generation.reasoning if kind == REASONING else generation.content).append(piece) generation.touch() except LLMError as exc: generation.error = exc.message log.info("generation failed for chat %s: %s", generation.chat_id, exc.message) except asyncio.CancelledError: # Shutdown, not a reader navigating away -- that no longer reaches here. generation.stopped = True raise except Exception: # noqa: BLE001 - a task that dies silently is worse generation.error = "Something went wrong while generating this reply." log.exception("unexpected generation failure for chat %s", generation.chat_id) finally: if reasoning_started is not None and not generation.reasoning_ms: generation.reasoning_ms = int((time.monotonic() - reasoning_started) * 1000) generation.elapsed_ms = int((time.monotonic() - started) * 1000) if not generation.completion_tokens: # The endpoint reported nothing, so fall back to the estimate. Marked # as such everywhere it is shown -- four characters to a token is # wrong enough on code and CJK to be worth saying out loud. generation.completion_tokens = tokens.estimate( generation.text + generation.thinking ) generation.prompt_tokens = generation.prompt_estimate generation.context_tokens = generation.prompt_tokens + generation.completion_tokens # Naming the chat is a second, short completion, so it has to happen # here rather than in the synchronous persist step below. Best-effort: # a chat title is never worth surfacing an error for. title = "" if needs_title and question: if generation.error or endpoint is None: title = chat_service.fallback_title(question) else: with contextlib.suppress(Exception): title = await chat_service.generate_title( endpoint, model_id, question, generation.text, template=title_prompt, ) title = title or chat_service.fallback_title(question) # Written *before* `done`, because `_follow` breaks out of its loop the # moment it sees that flag and immediately re-renders the bubble from # the row. The other order left a window in which the finished frame # showed the previous turn's stored values. _persist(generation, title, time.monotonic() - started) generation.done = True generation.finished_at = datetime.now(UTC) generation.touch() async def _maybe_compact(generation: Generation) -> None: """Summarise the earlier turns if the window is about to be full. Never raises. A failed compaction logs and sends the uncompacted request, which either works or fails upstream with a message that says what actually happened -- refusing to answer because the summariser was unavailable would be a worse trade. The awaited call is deliberately outside any session, the same shape titling uses: read everything needed, close, ask, reopen to write. """ try: with session_scope() as db: chat = db.get(Chat, generation.chat_id) message = db.get(Message, generation.message_id) if chat is None or message is None: return pending = _pending_text(db, message) if not compaction_service.should_compact(db, chat, pending=pending): return template = prompts_service.resolve(db, "task.compact") upto = compaction_service.last_complete(db, chat) if not template.strip() or upto is None: return endpoint, model_id = chat_service.resolve_endpoint(db, chat) transcript = compaction_service.transcript(db, chat, upto=upto) previous = compaction_service.previous_summary_block(chat) upto_id = upto.id generation.status = "Summarising earlier messages…" generation.touch() summary = await chat_service.summarise_for_compaction( endpoint, model_id, transcript=transcript, previous_summary=previous, template=template, ) if not summary: return with session_scope() as db: chat = db.get(Chat, generation.chat_id) upto = db.get(Message, upto_id) if chat is None or upto is None: return compaction_service.apply(chat, summary=summary, upto=upto) db.commit() log.info("chat %s compacted automatically through %s", chat.id, upto_id) except Exception: # noqa: BLE001 - the reply matters more than the tidy-up log.exception("automatic compaction failed for chat %s", generation.chat_id) finally: generation.status = "" generation.touch() def _pending_text(db, message: Message) -> str: """The user turn this reply is answering, for the size estimate.""" previous = db.scalars( select(Message) .where(Message.chat_id == message.chat_id, Message.created_at < message.created_at) .order_by(Message.created_at.desc()) .limit(1) ).first() return previous.content if previous is not None else "" def _question_from(payload: dict) -> str: """The last thing the user said, for auto-titling.""" for entry in reversed(payload.get("messages", [])): if entry.get("role") != ROLE_USER: continue content = entry.get("content") if isinstance(content, str): return content if isinstance(content, list): return " ".join( part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" ).strip() return "" def _persist(generation: Generation, title: str, elapsed: float) -> None: """Write the finished reply, name the chat, and set the unread flag. A generation another one has replaced may not write. A regeneration cancels its predecessor, whose `finally:` then runs this on the same row -- and it would overwrite the fresh reply with the abandoned one. The test is "someone else owns this row now", not "this one is registered": an unregistered generation still writes, because that is a direct call rather than a superseded one. """ owner = _RUNNING.get(generation.message_id) if owner is not None and owner is not generation: log.debug("skipping persist for superseded generation %s", generation.message_id) return try: with session_scope() as db: message = db.get(Message, generation.message_id) chat = db.get(Chat, generation.chat_id) if message is None or chat is None: return message.content = generation.text message.reasoning = generation.thinking message.reasoning_ms = generation.reasoning_ms message.tool_calls_json = generation.tool_events message.usage_json = metrics_service.to_json( metrics_service.from_generation(generation) ) message.error = generation.error message.stopped = generation.stopped message.complete = True if title and not chat.title_generated: chat.title = title chat.title_generated = True # Nobody watching when it landed, so it is news. The chat page # clears this when it is next opened. Not for a temporary chat: # there is no sidebar row for the dot, and the toast would name a # chat nobody can navigate to. if generation.followers == 0 and not chat.temporary: chat.unread = True chat.unread_notified = False db.commit() log.debug( "chat %s finished: %d chars, %d reasoning, %.1fs", generation.chat_id, len(message.content), len(message.reasoning), elapsed, ) except Exception: # noqa: BLE001 - the task is ending either way log.exception("could not persist generation for chat %s", generation.chat_id) __all__ = [ "RENDER_INTERVAL", "Generation", "ROLE_ASSISTANT", "ensure", "get", "request_stop", "restart", "shutdown", ]