595 lines
22 KiB
Python
595 lines
22 KiB
Python
"""The conversation rows: create, find, bind to windows, flags, status, history."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import String, func
|
|
from sqlalchemy import cast as sa_cast
|
|
from sqlmodel import col, select
|
|
|
|
from beaver_gateway.backends.transcript import (
|
|
messages_from_entries,
|
|
render_messages,
|
|
text_of,
|
|
)
|
|
from beaver_gateway.conversations.kinds import KINDS, Kind
|
|
from beaver_gateway.conversations.state import State, iso
|
|
from beaver_gateway.conversations.texts import split_service
|
|
from beaver_gateway.frontends.markdown.history import load_messages
|
|
from beaver_gateway.storage.models import (
|
|
Conversation,
|
|
ConversationBinding,
|
|
ConversationMessage,
|
|
RateLimit,
|
|
TranscriptEntry,
|
|
Usage,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable, Iterable
|
|
|
|
from beaver_gateway.frontends.base import Frontend
|
|
|
|
__all__ = [
|
|
"MASTER_ALIAS",
|
|
"PARENT_ALIAS",
|
|
"STATUSES",
|
|
"TITLE_MAX",
|
|
"Rows",
|
|
"context_of",
|
|
"implied_title",
|
|
]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
MASTER_ALIAS = "master"
|
|
PARENT_ALIAS = "parent"
|
|
STATUSES = ("open", "merged", "closed", "archived")
|
|
TITLE_MAX = 80
|
|
|
|
|
|
class Rows(State):
|
|
async def create(
|
|
self,
|
|
*,
|
|
kind: Kind,
|
|
agent: str,
|
|
parent: Conversation | None = None,
|
|
title: str | None = None,
|
|
origin: str = "api",
|
|
session_id: str | None = None,
|
|
flags: dict[str, Any] | None = None,
|
|
) -> Conversation:
|
|
if kind not in KINDS:
|
|
msg = f"unknown conversation kind {kind!r}"
|
|
raise ValueError(msg)
|
|
if not self._claude_agent(agent).serves(kind):
|
|
msg = f"agent {agent!r} does not serve kind {kind!r}"
|
|
raise ValueError(msg)
|
|
now = datetime.now(UTC)
|
|
row = Conversation(
|
|
frontend=origin,
|
|
external_id=str(uuid4()),
|
|
agent_name=agent,
|
|
kind=kind,
|
|
parent_id=parent.id if parent is not None else None,
|
|
title=title,
|
|
session_id=session_id,
|
|
flags=dict(flags or {}),
|
|
last_activity_at=now,
|
|
)
|
|
async with self._db.session() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
self._bus.publish("conversation.created", **self.public(row))
|
|
return row
|
|
|
|
async def get(self, public_id: str) -> Conversation | None:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(Conversation).where(Conversation.external_id == public_id)
|
|
)
|
|
return result.first()
|
|
|
|
async def get_row(self, row_id: int) -> Conversation | None:
|
|
async with self._db.session() as session:
|
|
return await session.get(Conversation, row_id)
|
|
|
|
async def resolve(
|
|
self, key: str, *, origin: Conversation | None = None
|
|
) -> Conversation | None:
|
|
"""By public id, or ``master`` / ``parent`` relative to ``origin``."""
|
|
key = key.strip()
|
|
if key == MASTER_ALIAS:
|
|
return await self.open_master()
|
|
if key == PARENT_ALIAS:
|
|
if origin is None or origin.parent_id is None:
|
|
return None
|
|
return await self.get_row(origin.parent_id)
|
|
return await self.get(key)
|
|
|
|
async def open_master(self) -> Conversation | None:
|
|
masters = await self.find(kind="master", status="open", limit=1)
|
|
return masters[0] if masters else None
|
|
|
|
async def find(
|
|
self,
|
|
*,
|
|
status: str | None = None,
|
|
kind: str | None = None,
|
|
parent: Conversation | None = None,
|
|
limit: int = 200,
|
|
) -> list[Conversation]:
|
|
stmt = select(Conversation).order_by(col(Conversation.id).desc()).limit(limit)
|
|
if status is not None:
|
|
stmt = stmt.where(Conversation.status == status)
|
|
if kind is not None:
|
|
stmt = stmt.where(Conversation.kind == kind)
|
|
if parent is not None:
|
|
stmt = stmt.where(Conversation.parent_id == parent.id)
|
|
async with self._db.session() as session:
|
|
return list((await session.exec(stmt)).all())
|
|
|
|
async def search(self, query: str, *, limit: int = 20) -> list[Conversation]:
|
|
needle = f"%{query}%"
|
|
by_title = select(Conversation.id).where(col(Conversation.title).ilike(needle))
|
|
by_first = select(ConversationMessage.conversation_id).where(
|
|
ConversationMessage.seq == 0,
|
|
col(ConversationMessage.content_json).ilike(needle),
|
|
)
|
|
by_entry = (
|
|
select(Conversation.id)
|
|
.join(
|
|
TranscriptEntry,
|
|
col(TranscriptEntry.session_id) == col(Conversation.session_id),
|
|
)
|
|
.where(sa_cast(col(TranscriptEntry.entry), String).ilike(needle))
|
|
)
|
|
stmt = (
|
|
select(Conversation)
|
|
.where(col(Conversation.id).in_(by_title.union(by_first, by_entry)))
|
|
.order_by(col(Conversation.id).desc())
|
|
.limit(limit)
|
|
)
|
|
async with self._db.session() as session:
|
|
return list((await session.exec(stmt)).all())
|
|
|
|
async def bindings(self, conv: Conversation) -> list[ConversationBinding]:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(ConversationBinding)
|
|
.where(ConversationBinding.conversation_id == conv.id)
|
|
.order_by(col(ConversationBinding.id))
|
|
)
|
|
return list(result.all())
|
|
|
|
async def bind(
|
|
self,
|
|
conv: Conversation,
|
|
*,
|
|
frontend: str,
|
|
external_id: str,
|
|
visible: bool = True,
|
|
) -> ConversationBinding:
|
|
if conv.kind not in self.frontend(frontend).kinds:
|
|
msg = f"frontend {frontend!r} does not show kind {conv.kind!r}"
|
|
raise ValueError(msg)
|
|
async with self._db.session() as session:
|
|
existing = list(
|
|
(
|
|
await session.exec(
|
|
select(ConversationBinding).where(
|
|
ConversationBinding.conversation_id == conv.id,
|
|
ConversationBinding.frontend == frontend,
|
|
)
|
|
)
|
|
).all()
|
|
)
|
|
row = next((b for b in existing if b.external_id == external_id), None)
|
|
if visible:
|
|
for other in existing:
|
|
if other is not row and other.visible:
|
|
other.visible = False
|
|
session.add(other)
|
|
same_window = await session.exec(
|
|
select(ConversationBinding).where(
|
|
ConversationBinding.frontend == frontend,
|
|
ConversationBinding.external_id == external_id,
|
|
ConversationBinding.conversation_id != conv.id,
|
|
col(ConversationBinding.visible).is_(True),
|
|
)
|
|
)
|
|
for other in same_window.all():
|
|
other.visible = False
|
|
session.add(other)
|
|
if row is None:
|
|
row = ConversationBinding(
|
|
conversation_id=cast("int", conv.id),
|
|
frontend=frontend,
|
|
external_id=external_id,
|
|
visible=visible,
|
|
)
|
|
else:
|
|
row.visible = visible
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
self._bus.publish(
|
|
"conversation.bound",
|
|
conversation_id=conv.external_id,
|
|
frontend=frontend,
|
|
external_id=external_id,
|
|
visible=visible,
|
|
)
|
|
return row
|
|
|
|
async def find_bound(
|
|
self, *, frontend: str, external_id: str
|
|
) -> Conversation | None:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(Conversation)
|
|
.join(
|
|
ConversationBinding,
|
|
col(ConversationBinding.conversation_id) == col(Conversation.id),
|
|
)
|
|
.where(
|
|
ConversationBinding.frontend == frontend,
|
|
ConversationBinding.external_id == external_id,
|
|
col(ConversationBinding.visible).is_(True),
|
|
)
|
|
.order_by(col(Conversation.id).desc())
|
|
)
|
|
return result.first()
|
|
|
|
async def last_binding(
|
|
self, *, frontend: str, kind: str
|
|
) -> ConversationBinding | None:
|
|
"""The window ``frontend`` last used for ``kind``; outlives a rotation."""
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(ConversationBinding)
|
|
.join(
|
|
Conversation,
|
|
col(Conversation.id) == col(ConversationBinding.conversation_id),
|
|
)
|
|
.where(
|
|
ConversationBinding.frontend == frontend, Conversation.kind == kind
|
|
)
|
|
.order_by(col(ConversationBinding.id).desc())
|
|
)
|
|
return result.first()
|
|
|
|
async def window_of(self, conv: Conversation) -> str | None:
|
|
for binding in await self.bindings(conv):
|
|
if binding.visible:
|
|
return binding.external_id
|
|
return None
|
|
|
|
async def set_flags(
|
|
self, conv: Conversation, flags: dict[str, Any]
|
|
) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.flags = {**row.flags, **flags}
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def set_status(self, conv: Conversation, status: str) -> Conversation:
|
|
if status not in STATUSES:
|
|
msg = f"unknown status {status!r}"
|
|
raise ValueError(msg)
|
|
|
|
async def apply(row: Conversation) -> None:
|
|
row.status = status
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def set_title(self, conv: Conversation, title: str) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.title = title
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def touch_user(self, conv: Conversation) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.last_user_activity_at = datetime.now(UTC)
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def reparent(self, conv: Conversation, parent: Conversation) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.parent_id = parent.id
|
|
|
|
return await self._update(conv, apply)
|
|
|
|
async def _update(
|
|
self, conv: Conversation, apply: Callable[[Conversation], Awaitable[None]]
|
|
) -> Conversation:
|
|
async with self._db.session() as session:
|
|
row = await session.get(Conversation, conv.id)
|
|
if row is None:
|
|
msg = f"conversation {conv.external_id} vanished"
|
|
raise LookupError(msg)
|
|
await apply(row)
|
|
row.updated_at = datetime.now(UTC)
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
self._bus.publish("conversation.updated", **self.public(row))
|
|
return row
|
|
|
|
def public(self, conv: Conversation) -> dict[str, Any]:
|
|
return {
|
|
"id": conv.external_id,
|
|
"kind": conv.kind,
|
|
"agent": conv.agent_name,
|
|
"title": conv.title,
|
|
"status": conv.status,
|
|
"parent_row": conv.parent_id,
|
|
"session_id": conv.session_id,
|
|
"running_turn": conv.running_turn,
|
|
"pending_question": conv.pending_question,
|
|
"flags": conv.flags,
|
|
"origin": conv.frontend,
|
|
"created_at": iso(conv.created_at),
|
|
"last_user_activity_at": iso(conv.last_user_activity_at),
|
|
"last_activity_at": iso(conv.last_activity_at),
|
|
}
|
|
|
|
async def describe(self, conv: Conversation) -> dict[str, Any]:
|
|
out = self.public(conv)
|
|
out["title"] = await self.implied_title(conv)
|
|
out["context_tokens"] = await self.context_tokens(conv)
|
|
parent = await self.get_row(conv.parent_id) if conv.parent_id else None
|
|
out["parent"] = parent.external_id if parent is not None else None
|
|
out["bindings"] = [
|
|
{"frontend": b.frontend, "external_id": b.external_id, "visible": b.visible}
|
|
for b in await self.bindings(conv)
|
|
]
|
|
live = self._pool.get(conv.external_id)
|
|
agent = self._agents.get(conv.agent_name)
|
|
out["thinking_is_progress"] = agent is not None and agent.thinking_is_progress
|
|
out["live"] = live is not None
|
|
out["busy"] = live.busy if live is not None else False
|
|
runner = self._runners.get(cast("int", conv.id))
|
|
out["turn"] = runner.snapshot() if runner is not None else None
|
|
pending = self.pending_question(conv.external_id)
|
|
out["question"] = (
|
|
{"id": pending[0], "questions": pending[1]} if pending else None
|
|
)
|
|
return out
|
|
|
|
def pending_question(self, key: str) -> tuple[str, list[dict[str, Any]]] | None:
|
|
for question_id, pending in self._questions.items():
|
|
if pending.conversation_id == key and not pending.answer.done():
|
|
return question_id, pending.questions
|
|
return None
|
|
|
|
async def rate_limits(self, *, limit: int = 100) -> list[RateLimit]:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(RateLimit).order_by(col(RateLimit.id).desc()).limit(limit)
|
|
)
|
|
return list(result.all())
|
|
|
|
async def context_tokens(self, conv: Conversation) -> int:
|
|
async with self._db.session() as session:
|
|
row = (
|
|
await session.exec(
|
|
select(Usage)
|
|
.where(Usage.conversation_id == conv.external_id)
|
|
.order_by(col(Usage.id).desc())
|
|
.limit(1)
|
|
)
|
|
).first()
|
|
return context_of(row)
|
|
|
|
async def context_tokens_by_id(self, ids: Iterable[str]) -> dict[str, int]:
|
|
wanted = list(ids)
|
|
if not wanted:
|
|
return {}
|
|
latest = (
|
|
select(func.max(Usage.id))
|
|
.where(col(Usage.conversation_id).in_(wanted))
|
|
.group_by(Usage.conversation_id)
|
|
)
|
|
async with self._db.session() as session:
|
|
rows = (
|
|
await session.exec(select(Usage).where(col(Usage.id).in_(latest)))
|
|
).all()
|
|
return {
|
|
cast("str", row.conversation_id): context_of(row)
|
|
for row in rows
|
|
if row.conversation_id
|
|
}
|
|
|
|
async def usage_tokens(self, since: datetime) -> int:
|
|
async with self._db.session() as session:
|
|
rows = (
|
|
await session.exec(
|
|
select(Usage).where(
|
|
col(Usage.ts) >= since.astimezone(UTC).replace(tzinfo=None)
|
|
)
|
|
)
|
|
).all()
|
|
return sum(
|
|
r.input_tokens + r.output_tokens + r.cache_creation_tokens for r in rows
|
|
)
|
|
|
|
async def busy(self, conv: Conversation) -> bool:
|
|
row = await self.get_row(cast("int", conv.id)) or conv
|
|
if row.running_turn or row.pending_question:
|
|
return True
|
|
live = self._pool.get(row.external_id)
|
|
if live is not None and live.busy:
|
|
return True
|
|
pending = await self._queue.pending(cast("int", row.id))
|
|
return any(i.priority in ("user", "urgent", "wake") for i in pending)
|
|
|
|
@property
|
|
def frontends(self) -> list[Frontend]:
|
|
return list(self._frontends)
|
|
|
|
def frontend(self, name: str) -> Frontend:
|
|
for fe in self._frontends:
|
|
if fe.name == name:
|
|
return fe
|
|
msg = f"unknown frontend {name!r}"
|
|
raise ValueError(msg)
|
|
|
|
def default_agent(self, kind: Kind) -> str | None:
|
|
for fe in self._frontends:
|
|
if kind in fe.kinds and (agent := fe.agent_for(kind)):
|
|
return agent
|
|
return None
|
|
|
|
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
|
|
for fe in self._frontends:
|
|
if conv.kind not in fe.kinds:
|
|
continue
|
|
binding = await fe.materialize(conv)
|
|
if binding is not None:
|
|
return binding
|
|
return None
|
|
|
|
async def mark_closed(self, conv: Conversation) -> bool:
|
|
marked = False
|
|
for fe in self._frontends:
|
|
if conv.kind in fe.kinds:
|
|
try:
|
|
marked = await fe.mark_closed(conv) or marked
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("%s could not mark %s", fe.name, conv.external_id)
|
|
return marked
|
|
|
|
async def read(self, conv: Conversation, *, window: int | None = None) -> str:
|
|
return render_messages(
|
|
await self.history(conv),
|
|
thinking_is_progress=self._claude_agent(
|
|
conv.agent_name
|
|
).thinking_is_progress,
|
|
window=window,
|
|
)
|
|
|
|
async def history(self, conv: Conversation) -> list[dict[str, Any]]:
|
|
if conv.session_id is None:
|
|
async with self._db.session() as session:
|
|
return await load_messages(
|
|
session, conversation_id=cast("int", conv.id)
|
|
)
|
|
messages = messages_from_entries(
|
|
cast("Any", await self.entries(conv)), stamped=True
|
|
)
|
|
return [self._unwrapped(message) for message in messages]
|
|
|
|
def _unwrapped(self, message: dict[str, Any]) -> dict[str, Any]:
|
|
if message["role"] != "user":
|
|
return message
|
|
content = message["content"]
|
|
if isinstance(content, str):
|
|
text, service = split_service(content, self._texts)
|
|
return (
|
|
message
|
|
if service is None
|
|
else {**message, "content": text, "service": service}
|
|
)
|
|
for index, block in enumerate(content):
|
|
if block.get("type") != "text":
|
|
continue
|
|
text, service = split_service(str(block.get("text") or ""), self._texts)
|
|
if service is not None:
|
|
blocks = [*content]
|
|
blocks[index] = {**block, "text": text}
|
|
return {**message, "content": blocks, "service": service}
|
|
return message
|
|
|
|
async def entries(self, conv: Conversation, *, subpath: str = "") -> list[Any]:
|
|
if conv.session_id is None:
|
|
return []
|
|
key = {**self._store_key(conv), "subpath": subpath}
|
|
return list(await self._store.load(cast("Any", key)) or [])
|
|
|
|
async def subpaths(self, conv: Conversation) -> list[str]:
|
|
if conv.session_id is None:
|
|
return []
|
|
return list(await self._store.list_subkeys(cast("Any", self._store_key(conv))))
|
|
|
|
async def first_user_texts(self, ids: Iterable[int]) -> dict[int, str]:
|
|
wanted = list(ids)
|
|
if not wanted:
|
|
return {}
|
|
async with self._db.session() as session:
|
|
rows = (
|
|
await session.exec(
|
|
select(ConversationMessage).where(
|
|
col(ConversationMessage.conversation_id).in_(wanted),
|
|
ConversationMessage.seq == 0,
|
|
ConversationMessage.role == "user",
|
|
)
|
|
)
|
|
).all()
|
|
return {
|
|
r.conversation_id: text_of(json.loads(r.content_json)).strip() for r in rows
|
|
}
|
|
|
|
async def implied_title(self, conv: Conversation) -> str | None:
|
|
if conv.title:
|
|
return conv.title
|
|
text = (await self.first_user_texts([cast("int", conv.id)])).get(
|
|
cast("int", conv.id)
|
|
)
|
|
return implied_title(text)
|
|
|
|
async def chat_name(self, conv: Conversation) -> str:
|
|
"""What a ``[[wikilink]]`` to the chat says: the file's stem when it has one."""
|
|
window = await self.window_of(conv)
|
|
if window and window.endswith(".md"):
|
|
return window.rsplit("/", 1)[-1][: -len(".md")]
|
|
return await self.implied_title(conv) or conv.external_id
|
|
|
|
async def adopt(self, *, kind: Kind, first_user_text: str) -> Conversation | None:
|
|
"""The one unbound, session-less conversation whose history starts here."""
|
|
text = first_user_text.strip()
|
|
if not text:
|
|
return None
|
|
bound = select(ConversationBinding.conversation_id).where(
|
|
col(ConversationBinding.visible).is_(True)
|
|
)
|
|
async with self._db.session() as session:
|
|
rows = (
|
|
await session.exec(
|
|
select(Conversation).where(
|
|
Conversation.kind == kind,
|
|
Conversation.status == "open",
|
|
col(Conversation.session_id).is_(None),
|
|
col(Conversation.id).not_in(bound),
|
|
)
|
|
)
|
|
).all()
|
|
firsts = await self.first_user_texts(cast("int", r.id) for r in rows)
|
|
hits = [r for r in rows if firsts.get(cast("int", r.id)) == text]
|
|
return hits[0] if len(hits) == 1 else None
|
|
|
|
|
|
def implied_title(text: str | None) -> str | None:
|
|
if not text:
|
|
return None
|
|
line = text.strip().splitlines()[0].strip()
|
|
return line if len(line) <= TITLE_MAX else line[: TITLE_MAX - 1] + "…"
|
|
|
|
|
|
def context_of(row: Usage | None) -> int:
|
|
"""The last API call's input, or the per-call average for older rows."""
|
|
if row is None:
|
|
return 0
|
|
if row.context_tokens:
|
|
return row.context_tokens
|
|
total = row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
|
|
return round(total / max(row.num_turns or 1, 1))
|