Background generation, unread replies, send/stop, PLAN.md
**Replies now run in the background.** Generation was driven by the SSE request, so navigating away or opening another chat cut the answer off mid-sentence. services/generation.py owns the work as its own task and the SSE endpoint merely follows it. Verified: attached briefly, closed the connection, went to another page -- the reply finished anyway, 832 characters, not marked stopped, auto-titled. Reattaching works because both `render` and `reasoning` frames now carry the whole block rather than a delta. A follower arriving late has no earlier fragments to append to, so deltas would leave it permanently missing the beginning. Verified: attached six seconds in and the first frame already contained 517 characters written while nobody watched. **Unread indicator.** A reply that lands with no follower attached marks its chat unread; the sidebar polls every 10s for out-of-band dot spans plus an HX-Trigger that raises a toast. Polled rather than pushed: a browser sitting on another chat has no connection to the one that finished, and an always-on channel per tab is a lot of machinery for a green dot. `unread_notified` stops the same arrival being announced every tick. Follower count is what decides "was anyone watching", so reading it as it arrives does not mark it unread -- verified both ways. **Stop is the send button.** While a reply is being written the send button becomes a red stop square, found via a MutationObserver on the thread since the composer and the streaming bubble are far apart in the document. The in-bubble Stop is gone. **Attachment border removed.** As asked -- an attachment is a picture, and the frame only ever drew at the wrong width. The anchor now shrink-wraps and the img's width/height attributes are overridden so a small image shows at its own size. Adds PLAN.md: what is built, what is not, known limits, and the decisions that look like oversights until you know the reason. 239 tests, ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -0,0 +1,285 @@
|
||||
"""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 lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Chat, Message
|
||||
from lembas.db.session import session_scope
|
||||
from lembas.services import chat as chat_service
|
||||
from lembas.services.llm.openai_client import LLMError, delta_reasoning, delta_text, 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
|
||||
|
||||
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
|
||||
# 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.
|
||||
"""
|
||||
existing = _RUNNING.get(message_id)
|
||||
if existing is not None:
|
||||
return existing
|
||||
|
||||
_prune()
|
||||
generation = Generation(chat_id=chat_id, message_id=message_id)
|
||||
_RUNNING[message_id] = generation
|
||||
_TASKS[message_id] = asyncio.create_task(_run(generation))
|
||||
return generation
|
||||
|
||||
|
||||
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."""
|
||||
splitter = ReasoningSplitter()
|
||||
started = time.monotonic()
|
||||
reasoning_started: float | None = None
|
||||
question = ""
|
||||
endpoint = model_id = None
|
||||
needs_title = False
|
||||
|
||||
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:
|
||||
generation.error = "That chat no longer exists."
|
||||
return
|
||||
|
||||
endpoint, model_id = chat_service.resolve_endpoint(db, chat)
|
||||
payload = chat_service.build_request(db, chat, upto=message)
|
||||
question = _question_from(payload)
|
||||
needs_title = not chat.title_generated
|
||||
|
||||
async for chunk in stream_chat(endpoint, payload):
|
||||
thought = delta_reasoning(chunk)
|
||||
if thought:
|
||||
if reasoning_started is None:
|
||||
reasoning_started = time.monotonic()
|
||||
generation.reasoning.append(thought)
|
||||
generation.touch()
|
||||
|
||||
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)
|
||||
generation.touch()
|
||||
|
||||
if generation.cancel:
|
||||
generation.stopped = True
|
||||
break
|
||||
|
||||
# Let followers and other tasks run between chunks.
|
||||
await asyncio.sleep(0)
|
||||
|
||||
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)
|
||||
|
||||
# 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
|
||||
)
|
||||
title = title or chat_service.fallback_title(question)
|
||||
|
||||
generation.done = True
|
||||
generation.finished_at = datetime.now(UTC)
|
||||
generation.touch()
|
||||
_persist(generation, title, time.monotonic() - started)
|
||||
|
||||
|
||||
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."""
|
||||
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.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.
|
||||
if generation.followers == 0:
|
||||
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",
|
||||
"shutdown",
|
||||
]
|
||||
Reference in New Issue
Block a user