Jaroslav Beneš
d4cefb066a
A row in custom_tools becomes a ToolDef like any built-in, offered beside the thirteen. The registry had to stop being an import-time constant for that: `resolve_tools` now returns the schemas *and* the runners together, carried to the loop on the ToolContext. That closes a hole on the way. `run_tool` looked names up in the global REGISTRY with no reference to what had been offered, so a model naming a tool its chat was gated out of -- a family switched off, a permission the reader lacks -- had it run anyway. The resolved set is now authoritative. Arguments come from a model, so an argument may fill a hole but never move the target: the scheme and host of a URL template are literal, values are escaped for where they land, and the origin is pinned afterwards. Every redirect hop is checked the way services/fetch.py checks one, and the secret is dropped if a hop leaves the origin it was issued for. Also fixes the tool-activity block claiming every library tool had "searched the web", which it has done since the second family landed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
564 lines
22 KiB
Python
564 lines
22 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)),
|
|
]
|
|
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",
|
|
]
|