574 lines
19 KiB
Python
574 lines
19 KiB
Python
"""Stored conversation history for the markdown frontend.
|
|
|
|
Aligns the file the user is editing against the DB-stored Anthropic-shape
|
|
message history, splicing text edits or forking on structural changes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
from sqlalchemy import delete
|
|
from sqlmodel import select
|
|
|
|
from beaver_gateway.backends.transcript import spoken_text
|
|
from beaver_gateway.storage.models import Conversation, ConversationMessage
|
|
|
|
if TYPE_CHECKING:
|
|
from anthropic.types import MessageParam
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from beaver_gateway.frontends.markdown.parser import ParsedTurn
|
|
|
|
__all__ = [
|
|
"ForkOutcome",
|
|
"diff_and_fork",
|
|
"load_conversation",
|
|
"load_messages",
|
|
"mint_conversation",
|
|
"rewrite_messages",
|
|
"set_session_id",
|
|
]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ForkOutcome:
|
|
"""Result of aligning the incoming file against stored history.
|
|
|
|
``messages`` is the backend input; ``persist_messages`` is history to
|
|
store before the new reply; ``divergence_index`` is where they diverged.
|
|
"""
|
|
|
|
messages: list[MessageParam]
|
|
persist_messages: list[dict[str, Any]]
|
|
divergence_index: int | None
|
|
edited: bool = False
|
|
"""An earlier assistant turn's prose was rewritten in the file."""
|
|
|
|
@property
|
|
def reuse_session(self) -> bool:
|
|
return self.divergence_index is None and not self.edited
|
|
|
|
|
|
async def load_conversation(
|
|
session: AsyncSession, *, frontend: str, external_id: str
|
|
) -> Conversation | None:
|
|
stmt = (
|
|
select(Conversation)
|
|
.where(Conversation.frontend == frontend)
|
|
.where(Conversation.external_id == external_id)
|
|
)
|
|
result = await session.exec(stmt)
|
|
return result.first()
|
|
|
|
|
|
async def mint_conversation(
|
|
session: AsyncSession, *, frontend: str, agent_name: str
|
|
) -> Conversation:
|
|
"""Create a fresh conversation row with a new uuid for external_id.
|
|
|
|
Caller must persist the returned ``external_id`` so future requests can find it.
|
|
"""
|
|
row = Conversation(
|
|
frontend=frontend, external_id=str(uuid.uuid4()), agent_name=agent_name
|
|
)
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return row
|
|
|
|
|
|
async def set_session_id(
|
|
session: AsyncSession, *, conversation_id: int, session_id: str | None
|
|
) -> None:
|
|
conv = await session.get(Conversation, conversation_id)
|
|
if conv is None or conv.session_id == session_id:
|
|
return
|
|
conv.session_id = session_id
|
|
session.add(conv)
|
|
await session.commit()
|
|
|
|
|
|
async def load_messages(
|
|
session: AsyncSession, *, conversation_id: int
|
|
) -> list[dict[str, Any]]:
|
|
"""Return stored messages ordered by ``seq`` ascending.
|
|
|
|
Each entry is a canonical Anthropic ``MessageParam`` dict, the same
|
|
shape fed to the backend on continuation.
|
|
"""
|
|
stmt = (
|
|
select(ConversationMessage)
|
|
.where(ConversationMessage.conversation_id == conversation_id)
|
|
.order_by(ConversationMessage.seq.asc()) # ty: ignore[unresolved-attribute]
|
|
)
|
|
result = await session.exec(stmt)
|
|
rows = result.all()
|
|
return [
|
|
{"role": r.role, "content": _sanitize_content(json.loads(r.content_json))}
|
|
for r in rows
|
|
]
|
|
|
|
|
|
def _sanitize_content(content: Any) -> Any:
|
|
"""Strip wire-illegal fields from stored Anthropic content blocks.
|
|
|
|
Drops ``tool_result.is_error: null`` — the API rejects null there
|
|
though the field is optional.
|
|
"""
|
|
if not isinstance(content, list):
|
|
return content
|
|
cleaned: list[Any] = []
|
|
for blk in content:
|
|
out_blk = blk
|
|
if (
|
|
isinstance(blk, dict)
|
|
and blk.get("type") == "tool_result"
|
|
and blk.get("is_error") is None
|
|
and "is_error" in blk
|
|
):
|
|
out_blk = {k: v for k, v in blk.items() if k != "is_error"}
|
|
cleaned.append(out_blk)
|
|
return cleaned
|
|
|
|
|
|
async def rewrite_messages(
|
|
session: AsyncSession, *, conversation_id: int, messages: list[dict[str, Any]]
|
|
) -> None:
|
|
"""Replace the conversation's stored messages (full overwrite, no branch history).
|
|
|
|
Deletes and flushes before inserting — SQLAlchemy's default INSERT-before-DELETE
|
|
flush order would otherwise collide with ``uq_msg_conv_seq``.
|
|
"""
|
|
await session.execute( # ty: ignore[deprecated]
|
|
delete(ConversationMessage).where(
|
|
ConversationMessage.conversation_id == conversation_id # ty: ignore[invalid-argument-type]
|
|
)
|
|
)
|
|
await session.flush()
|
|
for seq, m in enumerate(messages):
|
|
session.add(
|
|
ConversationMessage(
|
|
conversation_id=conversation_id,
|
|
seq=seq,
|
|
role=str(m["role"]),
|
|
content_json=json.dumps(m["content"], separators=(",", ":")),
|
|
)
|
|
)
|
|
conv = await session.get(Conversation, conversation_id)
|
|
if conv is not None:
|
|
from datetime import UTC, datetime
|
|
|
|
conv.updated_at = datetime.now(UTC)
|
|
session.add(conv)
|
|
await session.commit()
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _StoredDisplayTurn:
|
|
"""A "display turn" reconstructed from stored raw messages.
|
|
|
|
Parser-equivalent view used to diff the file against the DB.
|
|
"""
|
|
|
|
role: str
|
|
messages: tuple[dict[str, Any], ...]
|
|
spoken_text: str
|
|
skeleton: tuple[str, ...]
|
|
text_segment_count: int
|
|
|
|
|
|
def _group_display_turns(
|
|
stored: list[dict[str, Any]], *, thinking_is_progress: bool
|
|
) -> list[_StoredDisplayTurn]:
|
|
"""Walk raw stored messages, group them into Obsidian-visible turns.
|
|
|
|
A user-prompt message (no ``tool_result`` blocks) opens a new turn;
|
|
tool-result-only messages roll into the current assistant turn.
|
|
"""
|
|
out: list[_StoredDisplayTurn] = []
|
|
i = 0
|
|
while i < len(stored):
|
|
msg = stored[i]
|
|
role = msg["role"]
|
|
if role == "user" and _is_user_prompt(msg.get("content")):
|
|
out.append(
|
|
_StoredDisplayTurn(
|
|
role="user",
|
|
messages=(msg,),
|
|
spoken_text=_user_prompt_text(msg.get("content")),
|
|
skeleton=(),
|
|
text_segment_count=0,
|
|
)
|
|
)
|
|
i += 1
|
|
continue
|
|
group: list[dict[str, Any]] = []
|
|
while i < len(stored):
|
|
m = stored[i]
|
|
if m["role"] == "user" and _is_user_prompt(m.get("content")):
|
|
break
|
|
group.append(m)
|
|
i += 1
|
|
spoken, skeleton, text_count = _summarize_assistant_group(
|
|
group, thinking_is_progress=thinking_is_progress
|
|
)
|
|
out.append(
|
|
_StoredDisplayTurn(
|
|
role="assistant",
|
|
messages=tuple(group),
|
|
spoken_text=spoken,
|
|
skeleton=tuple(skeleton),
|
|
text_segment_count=text_count,
|
|
)
|
|
)
|
|
return out
|
|
|
|
|
|
def _is_user_prompt(content: Any) -> bool:
|
|
"""A user message is a *prompt* unless its content carries tool_result blocks.
|
|
|
|
Unknown content shapes are conservatively treated as a prompt.
|
|
"""
|
|
if isinstance(content, str):
|
|
return True
|
|
if isinstance(content, list):
|
|
return not any(
|
|
isinstance(b, dict) and b.get("type") == "tool_result" for b in content
|
|
)
|
|
return True
|
|
|
|
|
|
def _user_prompt_text(content: Any) -> str:
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
chunks = [
|
|
str(b.get("text", ""))
|
|
for b in content
|
|
if isinstance(b, dict) and b.get("type") == "text"
|
|
]
|
|
return "\n\n".join(c for c in chunks if c)
|
|
return ""
|
|
|
|
|
|
def _summarize_assistant_group(
|
|
group: list[dict[str, Any]], *, thinking_is_progress: bool
|
|
) -> tuple[str, list[str], int]:
|
|
"""Compute (spoken_text, tool_skeleton, text_segment_count) for a display group.
|
|
|
|
Must mirror what the renderer writes and the parser reads back: text and
|
|
progress thinking collapse into one segment, tool_use becomes a skeleton
|
|
entry, tool_result, reasoning and empty thinking are invisible.
|
|
"""
|
|
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
|
|
|
segments: list[TextSegment | ToolSegment] = []
|
|
pending: list[str] = []
|
|
|
|
def _flush() -> None:
|
|
if not pending:
|
|
return
|
|
joined = "\n\n".join(p for p in pending if p)
|
|
pending.clear()
|
|
cleaned = joined.strip()
|
|
if cleaned:
|
|
segments.append(TextSegment(text=cleaned))
|
|
|
|
for msg in group:
|
|
if msg["role"] == "user":
|
|
_flush()
|
|
continue
|
|
content = msg.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
for blk in content:
|
|
if not isinstance(blk, dict):
|
|
continue
|
|
text = spoken_text(blk, thinking_is_progress=thinking_is_progress)
|
|
if text:
|
|
pending.append(text)
|
|
elif blk.get("type") == "tool_use":
|
|
_flush()
|
|
segments.append(ToolSegment(name=str(blk.get("name", ""))))
|
|
_flush()
|
|
spoken_chunks = [s.text for s in segments if isinstance(s, TextSegment)]
|
|
spoken = "\n\n".join(c for c in spoken_chunks if c).strip()
|
|
skeleton = [s.name for s in segments if isinstance(s, ToolSegment)]
|
|
text_count = sum(1 for s in segments if isinstance(s, TextSegment))
|
|
return spoken, skeleton, text_count
|
|
|
|
|
|
def diff_and_fork(
|
|
*,
|
|
stored: list[dict[str, Any]],
|
|
incoming: list[ParsedTurn],
|
|
thinking_is_progress: bool,
|
|
) -> ForkOutcome:
|
|
"""Align the incoming parsed file against stored history.
|
|
|
|
``incoming`` must end with a user turn (the new prompt); raises otherwise.
|
|
Segment-class imports below are deferred to avoid a cycle with ``parser``.
|
|
"""
|
|
if not incoming or incoming[-1].role != "user":
|
|
msg = (
|
|
"diff_and_fork expects incoming to end with a user turn "
|
|
"(the new prompt); got "
|
|
f"{incoming[-1].role if incoming else 'empty'}"
|
|
)
|
|
raise ValueError(msg)
|
|
|
|
stored_groups = _group_display_turns(
|
|
stored, thinking_is_progress=thinking_is_progress
|
|
)
|
|
new_user_turn = incoming[-1]
|
|
prior_incoming = incoming[:-1]
|
|
|
|
spliced_groups, divergence, edited = _walk_prefix(
|
|
prior_incoming, stored_groups, thinking_is_progress=thinking_is_progress
|
|
)
|
|
|
|
if divergence is None and len(prior_incoming) < len(stored_groups):
|
|
if _file_lags_store(stored_groups, len(prior_incoming), new_user_turn):
|
|
spliced_groups.extend(
|
|
list(g.messages) for g in stored_groups[len(prior_incoming) :]
|
|
)
|
|
else:
|
|
divergence = len(prior_incoming)
|
|
|
|
backend_msgs, persist_msgs = _assemble_tail(
|
|
spliced_groups=spliced_groups,
|
|
prior_incoming=prior_incoming,
|
|
divergence=divergence,
|
|
new_user_turn=new_user_turn,
|
|
)
|
|
return ForkOutcome(
|
|
messages=backend_msgs,
|
|
persist_messages=persist_msgs,
|
|
divergence_index=divergence,
|
|
edited=edited,
|
|
)
|
|
|
|
|
|
def _file_lags_store(
|
|
stored_groups: list[_StoredDisplayTurn], prior_len: int, new_user_turn: ParsedTurn
|
|
) -> bool:
|
|
"""Is the shorter incoming file a stale view rather than a deletion?
|
|
|
|
Tell: if the DB already holds the submitted prompt at the position the
|
|
file stops at, it's a lagging render, not a deletion — misjudging this
|
|
forks history, breaks the fingerprint, and respawns the backend session.
|
|
"""
|
|
if prior_len >= len(stored_groups):
|
|
return False
|
|
candidate = stored_groups[prior_len]
|
|
return candidate.role == "user" and candidate.spoken_text == new_user_turn.text
|
|
|
|
|
|
def _walk_prefix(
|
|
prior_incoming: list[ParsedTurn],
|
|
stored_groups: list[_StoredDisplayTurn],
|
|
*,
|
|
thinking_is_progress: bool,
|
|
) -> tuple[list[list[dict[str, Any]]], int | None, bool]:
|
|
"""Walk incoming vs stored side-by-side until first divergence.
|
|
|
|
An empty incoming skeleton (no tool callouts rendered) means prose
|
|
alone decides the match; a mismatch forces a fresh backend session.
|
|
"""
|
|
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
|
|
|
spliced_groups: list[list[dict[str, Any]]] = []
|
|
edited = False
|
|
for i, inc in enumerate(prior_incoming):
|
|
if i >= len(stored_groups):
|
|
return spliced_groups, i, edited
|
|
st = stored_groups[i]
|
|
if inc.role != st.role:
|
|
return spliced_groups, i, edited
|
|
if inc.role == "user":
|
|
if inc.text != st.spoken_text:
|
|
return spliced_groups, i, edited
|
|
spliced_groups.append(list(st.messages))
|
|
continue
|
|
inc_skeleton = tuple(
|
|
s.name for s in inc.structure if isinstance(s, ToolSegment)
|
|
)
|
|
inc_text_count = sum(1 for s in inc.structure if isinstance(s, TextSegment))
|
|
if inc_skeleton and inc_skeleton != st.skeleton:
|
|
return spliced_groups, i, edited
|
|
if inc.text == st.spoken_text:
|
|
spliced_groups.append(list(st.messages))
|
|
continue
|
|
if inc_skeleton and inc_text_count != st.text_segment_count:
|
|
return spliced_groups, i, edited
|
|
spliced = _splice_assistant_group(
|
|
stored_group=st, incoming=inc, thinking_is_progress=thinking_is_progress
|
|
)
|
|
if spliced is None:
|
|
return spliced_groups, i, edited
|
|
spliced_groups.append(spliced)
|
|
edited = True
|
|
return spliced_groups, None, edited
|
|
|
|
|
|
def _assemble_tail(
|
|
*,
|
|
spliced_groups: list[list[dict[str, Any]]],
|
|
prior_incoming: list[ParsedTurn],
|
|
divergence: int | None,
|
|
new_user_turn: ParsedTurn,
|
|
) -> tuple[list[MessageParam], list[dict[str, Any]]]:
|
|
"""Build the (backend, persist) lists from aligned + post-divergence tail."""
|
|
backend_msgs: list[MessageParam] = []
|
|
persist_msgs: list[dict[str, Any]] = []
|
|
for spliced in spliced_groups:
|
|
for m in spliced:
|
|
entry: dict[str, Any] = {"role": m["role"], "content": m["content"]}
|
|
backend_msgs.append(cast("MessageParam", entry))
|
|
persist_msgs.append(entry)
|
|
if divergence is not None:
|
|
for inc in prior_incoming[divergence:]:
|
|
if not inc.text:
|
|
continue
|
|
entry = {"role": inc.role, "content": inc.text}
|
|
backend_msgs.append(cast("MessageParam", entry))
|
|
persist_msgs.append(entry)
|
|
backend_msgs.append({"role": "user", "content": new_user_turn.text})
|
|
return backend_msgs, persist_msgs
|
|
|
|
|
|
def _splice_assistant_group(
|
|
*,
|
|
stored_group: _StoredDisplayTurn,
|
|
incoming: ParsedTurn,
|
|
thinking_is_progress: bool,
|
|
) -> list[dict[str, Any]] | None:
|
|
"""Apply the file's prose to a stored assistant turn.
|
|
|
|
Prefers substituting text in place, which keeps everything the
|
|
markdown can't express. Falls back to rebuilding from the file's
|
|
structure when the text blocks don't line up one-to-one.
|
|
"""
|
|
preserved = _splice_in_place(
|
|
stored_group=stored_group,
|
|
incoming=incoming,
|
|
thinking_is_progress=thinking_is_progress,
|
|
)
|
|
if preserved is not None:
|
|
return preserved
|
|
return _splice_by_rebuild(stored_group=stored_group, incoming=incoming)
|
|
|
|
|
|
def _splice_in_place(
|
|
*,
|
|
stored_group: _StoredDisplayTurn,
|
|
incoming: ParsedTurn,
|
|
thinking_is_progress: bool,
|
|
) -> list[dict[str, Any]] | None:
|
|
"""Copy the stored messages, swapping only their text block contents.
|
|
|
|
Preserves message/thinking-block boundaries that a rebuild would lose
|
|
(and that a cache-hit continuation depends on); returns ``None`` if
|
|
text blocks aren't 1:1 with incoming segments or the prose runs through
|
|
signed thinking blocks, so the caller rebuilds instead.
|
|
"""
|
|
from beaver_gateway.frontends.markdown.parser import TextSegment
|
|
|
|
new_texts = [
|
|
s.text for s in incoming.structure if isinstance(s, TextSegment) and s.text
|
|
]
|
|
positions: list[tuple[int, int]] = []
|
|
for mi, msg in enumerate(stored_group.messages):
|
|
content = msg.get("content")
|
|
if msg["role"] != "assistant" or not isinstance(content, list):
|
|
continue
|
|
for bi, blk in enumerate(cast("list[dict[str, Any]]", content)):
|
|
if not spoken_text(blk, thinking_is_progress=thinking_is_progress):
|
|
continue
|
|
if blk["type"] == "thinking":
|
|
return None
|
|
positions.append((mi, bi))
|
|
if len(positions) != len(new_texts):
|
|
return None
|
|
|
|
out: list[dict[str, Any]] = [
|
|
{
|
|
"role": m["role"],
|
|
"content": list(m["content"])
|
|
if isinstance(m.get("content"), list)
|
|
else m.get("content"),
|
|
}
|
|
for m in stored_group.messages
|
|
]
|
|
for (mi, bi), text in zip(positions, new_texts, strict=True):
|
|
out[mi]["content"][bi] = {**out[mi]["content"][bi], "text": text}
|
|
return out
|
|
|
|
|
|
def _splice_by_rebuild(
|
|
*, stored_group: _StoredDisplayTurn, incoming: ParsedTurn
|
|
) -> list[dict[str, Any]] | None:
|
|
"""Rebuild an assistant display turn with new text + stored tool_use blocks.
|
|
|
|
Matches tool_use blocks to incoming ``ToolSegment``s by position; returns
|
|
``None`` (caller forks) if a matching stored ``tool_result`` is missing.
|
|
"""
|
|
from beaver_gateway.frontends.markdown.parser import TextSegment
|
|
|
|
tool_uses, tool_results_by_id = _harvest_tool_blocks(stored_group)
|
|
|
|
spliced: list[dict[str, Any]] = []
|
|
current_asst: list[dict[str, Any]] = []
|
|
next_tool = 0
|
|
for seg in incoming.structure:
|
|
if isinstance(seg, TextSegment):
|
|
if seg.text:
|
|
current_asst.append({"type": "text", "text": seg.text})
|
|
continue
|
|
if next_tool >= len(tool_uses):
|
|
return None
|
|
tu = tool_uses[next_tool]
|
|
next_tool += 1
|
|
current_asst.append(tu)
|
|
spliced.append({"role": "assistant", "content": current_asst})
|
|
current_asst = []
|
|
tr = tool_results_by_id.get(str(tu.get("id", "")))
|
|
if tr is None:
|
|
return None
|
|
spliced.append({"role": "user", "content": [tr]})
|
|
if current_asst:
|
|
spliced.append({"role": "assistant", "content": current_asst})
|
|
elif not spliced:
|
|
return None
|
|
return spliced
|
|
|
|
|
|
def _harvest_tool_blocks(
|
|
stored_group: _StoredDisplayTurn,
|
|
) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]]]:
|
|
"""Pull stored ``tool_use`` blocks (ordered) and ``tool_result`` blocks (by id)."""
|
|
tool_uses: list[dict[str, Any]] = []
|
|
tool_results_by_id: dict[str, dict[str, Any]] = {}
|
|
for msg in stored_group.messages:
|
|
content = msg.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
if msg["role"] == "assistant":
|
|
tool_uses.extend(
|
|
blk
|
|
for blk in content
|
|
if isinstance(blk, dict) and blk.get("type") == "tool_use"
|
|
)
|
|
continue
|
|
for blk in content:
|
|
if not isinstance(blk, dict) or blk.get("type") != "tool_result":
|
|
continue
|
|
tid = blk.get("tool_use_id")
|
|
if isinstance(tid, str):
|
|
tool_results_by_id[tid] = blk
|
|
return tool_uses, tool_results_by_id
|