Files
LLeMbas/src/lembas/services/generation.py
T
Jaroslav Beneš ecb52e9978 MCP servers, over streamable HTTP
A server is a row with a URL; its tools are discovered by a button and
cached, then offered beside the built-in ones. Written by hand rather than
taken from the reference SDK, because that SDK's transport does its own
connecting -- and the one thing that must not be bypassed is check_url on
every hop. Owning the transport is the point; the framing beside it is the
small part.

Sessions are per call: initialize, initialized, the call, a best-effort
DELETE. Caching one wants an owner, a TTL, eviction, a lock and a shutdown
hook, and the server may expire it under all of that anyway -- ToolContext
is a session-free snapshot precisely so nothing in a tool holds live state.

A server's names and descriptions reach the model as instructions and are
bounded before they do; what it returns is escaped preformatted text, never
markdown. Tools are namespaced per server, so two servers exposing "search"
do not collide and neither shadows a built-in.

Also: a round's calls now run together under a semaphore, results indexed
so each tool turn stays paired with its call, and generation.status names
what is running -- a remote tool is latency-bound, and a silent pause is
what a hang looks like.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-01 16:44:29 +02:00

613 lines
24 KiB
Python

"""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)),
]
generation.status = _tool_status(calls)
generation.touch()
try:
outcomes = await _run_calls(tool_context, calls)
finally:
generation.status = ""
generation.touch()
for call, outcome in zip(calls, outcomes, strict=True):
generation.tool_events.append(outcome.event)
messages.append(tools_service.tool_turn(call, outcome.content))
generation.touch()
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()
# How many of a round's tool calls may be in flight at once. A bound rather
# than none: a model that asks for eight would otherwise open eight sockets and
# eight database sessions at the same moment.
MAX_PARALLEL_TOOLS = 4
def _tool_status(calls: list[dict]) -> str:
"""What to show while tools run.
A remote tool -- an HTTP endpoint, an MCP server -- can take seconds with
nothing streaming, and a silent pause is exactly what a hang looks like.
"""
if len(calls) == 1:
return f"Running {calls[0]['name']}…"
return f"Running {len(calls)} tools…"
async def _run_calls(context, calls: list[dict]) -> list:
"""Run one round's calls together, results in call order.
Sequential was right when every tool was a local database read. A remote one
is latency-bound, and three two-second calls in a row are six seconds of a
reply looking hung -- while the model has already been told it may ask for
several at once.
The results are indexed rather than appended as they finish, because each
tool turn has to line up with the assistant turn's `tool_calls`: an endpoint
matching on `tool_call_id` would otherwise pair the right id with the wrong
content the moment two calls came back out of order.
Safe to run together because `run_tool` never raises, so no failure cancels
its siblings, and each runner opens its own `session_scope()` against a
database in WAL mode with a busy timeout.
"""
limit = asyncio.Semaphore(MAX_PARALLEL_TOOLS)
async def one(call: dict):
async with limit:
return await tools_service.run_tool(context, call["name"], call["arguments"])
return list(await asyncio.gather(*(one(call) for call in calls)))
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",
]