Files
beaver-gateway/src/beaver_gateway/frontends/markdown/crossfront.py
T

252 lines
8.8 KiB
Python

"""Cross-frontend chat logger.
Mirrors turns completed by other frontends into the vault as ``.md``
files, matching a conversation's continuation by content-hash fingerprint.
"""
from __future__ import annotations
import asyncio
import hashlib
import logging
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
import frontmatter
from beaver_gateway.frontends.markdown import renderer
if TYPE_CHECKING:
from collections.abc import Callable, Iterable
from pathlib import Path
from anthropic.types import MessageParam
from beaver_gateway.frontends.turn_record import TurnRecord
_log = logging.getLogger("beaver_gateway.frontends.markdown.crossfront")
__all__ = [
"ChatPathFn",
"CrossFrontendLogger",
"fingerprint_messages",
"strip_trailing_user_scaffold",
]
ChatPathFn = "Callable[[str, str, Path], Path]"
"""``(title, agent, vault) -> path`` of a new chat file; relative = under the
vault. ``None`` keeps ``{vault}/{logged_subdir}/{agent}/{YYYY-MM-DD}_{slug}.md``."""
def fingerprint_messages(messages: Iterable[MessageParam]) -> str:
"""Stable, short hex hash of a conversation prefix.
Hashes ``(role, text-only content)`` pairs so differently-shaped message
histories that carry the same text still fingerprint identically.
"""
h = hashlib.sha1(usedforsecurity=False)
for msg in messages:
role = str(msg.get("role", ""))
content = msg.get("content", "")
if isinstance(content, str):
text = content
else:
parts: list[str] = []
for blk in content:
if not isinstance(blk, dict):
continue
btype = blk.get("type")
if btype == "text":
parts.append(str(blk.get("text", "")))
text = "\n".join(parts)
h.update(role.encode("utf-8"))
h.update(b"\x00")
h.update(text.strip().encode("utf-8"))
h.update(b"\x01")
return h.hexdigest()[:16]
class CrossFrontendLogger:
"""Maintains the fingerprint→file map and writes turns to disk.
The map is in-process, rebuilt by ``warm_index`` on startup; all disk
writes funnel through one lock to sidestep races.
"""
def __init__(
self,
*,
vault_path: Path,
logged_subdir: str,
chat_path: Callable[[str, str, Path], Path] | None = None,
) -> None:
self._vault = vault_path
self._root = vault_path / logged_subdir
self._index: dict[str, Path] = {}
self._lock = asyncio.Lock()
self._chat_path_fn = chat_path
self._scan_root = vault_path if chat_path is not None else self._root
def warm_index(self) -> None:
"""Scan logged files synchronously, populating the fingerprint map.
Called before any cross-frontend turn arrives; a custom ``chat_path``
forces a full-vault walk instead of scanning just ``logged_subdir``.
"""
if not self._scan_root.exists():
return
for path in self._scan_root.rglob("*.md"):
try:
post = frontmatter.load(str(path))
except Exception: # noqa: BLE001
_log.warning("could not read frontmatter from %s — skipping", path)
continue
fp = post.metadata.get("fingerprint")
if isinstance(fp, str) and fp:
self._index[fp] = path
_log.info(
"crossfront index warmed: %d logged file(s) under %s",
len(self._index),
self._scan_root,
)
async def handle(self, record: TurnRecord) -> None:
"""Append or create a logged file for ``record``.
Skips ``source == "markdown"`` records (already on disk); matches
the target file by fingerprinting ``input_messages`` sans the new turn.
"""
if record.source == "markdown":
return
async with self._lock:
prefix = record.input_messages[:-1]
prev_fp = fingerprint_messages(prefix) if prefix else None
target = self._index.get(prev_fp) if prev_fp else None
if target is None:
target = self._new_file_path(record)
assistant_msg: MessageParam = {
"role": "assistant",
"content": _flatten_text(record.output_message),
}
full_history = [*record.input_messages, assistant_msg]
new_fp = fingerprint_messages(full_history)
if target.exists():
existing = target.read_text(encoding="utf-8")
parsed = frontmatter.loads(existing)
body = strip_trailing_user_scaffold(parsed.content)
new_user = record.input_messages[-1]
new_block = renderer.render_user_param(new_user)
new_block = renderer.append_to_body(
new_block, renderer.render_assistant_message(record.output_message)
)
new_body = renderer.append_to_body(body, new_block)
metadata = dict(parsed.metadata)
else:
new_body = _render_full_history(
record.input_messages, record.output_message
)
metadata = {}
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
metadata["agent"] = record.agent_name
metadata["fingerprint"] = new_fp
metadata["source"] = record.source
self._write(target, metadata, new_body)
if prev_fp:
self._index.pop(prev_fp, None)
self._index[new_fp] = target
def _new_file_path(self, record: TurnRecord) -> Path:
"""Pick a fresh filename for a brand-new conversation.
Delegates to ``chat_path`` if set; otherwise
``{logged_subdir}/{agent}/{date}_{hex8}.md``.
"""
if self._chat_path_fn is not None:
result = self._chat_path_fn(
record.first_user_text, record.agent_name, self._vault
)
if not result.is_absolute():
result = self._vault / result
if result.suffix != ".md":
result = result.with_suffix(".md")
result.parent.mkdir(parents=True, exist_ok=True)
return result
day = datetime.now(UTC).strftime("%Y-%m-%d")
salt = fingerprint_messages(record.input_messages)[:8]
agent_dir = self._root / record.agent_name
agent_dir.mkdir(parents=True, exist_ok=True)
return agent_dir / f"{day}_{salt}.md"
def _write(self, path: Path, metadata: dict[str, Any], body: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
if metadata:
post = frontmatter.Post(content=body, **metadata)
text = frontmatter.dumps(post) + "\n"
else:
text = body if body.endswith("\n") else body + "\n"
path.write_text(text, encoding="utf-8")
def _flatten_text(message: Any) -> str:
"""Same as ``frontend._flatten_assistant_text`` but local to break a cycle."""
chunks = [
getattr(block, "text", "") or ""
for block in getattr(message, "content", ())
if getattr(block, "type", None) == "text"
]
return "\n\n".join(c for c in chunks if c)
def _render_full_history(messages: list[MessageParam], assistant: Any) -> str:
"""Render an entire conversation (used when materializing a new logged file)."""
blocks: list[str] = []
for msg in messages:
role = msg.get("role")
if role == "user":
blocks.append(renderer.render_user_param(msg))
elif role == "assistant":
content = msg.get("content", "")
text = content if isinstance(content, str) else _content_to_text(content)
blocks.append(f"### Assistant:\n\n{text.strip()}\n" if text.strip() else "")
blocks.append(renderer.render_assistant_message(assistant))
body = ""
for block in blocks:
if not block:
continue
body = renderer.append_to_body(body, block)
return body
def strip_trailing_user_scaffold(body: str) -> str:
"""Drop a trailing empty ``### User:`` block if present.
Avoids leaving two ``### User:`` headers in a row when appending a
turn that wasn't typed into the file by hand.
"""
stripped = body.rstrip()
marker = "### User:"
if not stripped.endswith(marker):
return body
head = stripped[: -len(marker)].rstrip()
if head.endswith("---"):
head = head[: -len("---")].rstrip()
return head
def _content_to_text(content: Any) -> str:
if not isinstance(content, list):
return ""
chunks = [
str(blk.get("text", ""))
for blk in content
if isinstance(blk, dict) and blk.get("type") == "text"
]
return "\n\n".join(chunks)