lorebase-keeper 0.2.3__py3-none-any.whl

lorebase_keeper/agents/feed.py

@@ -0,0 +1,259 @@
+"""Change feed of a branch: Revision events as Server-Sent Events, with a polling fallback.
+
+Events are the `revision.committed` outbox events of the branch, described by
+`lorebase_store.changefeed` (Revision metadata, changes per Kind and Relation, changed keys up
+to `agents.feed.changed_keys_max`). The SSE `id` of an event is its outbox sequence, so a
+client reconnecting with `Last-Event-ID` resumes after the last event it saw; without it the
+stream starts with the next commit. The stream sends a comment line after `heartbeat_s` of
+silence and closes after `max_connection_s`.
+
+Every `reauthorize_interval_s` the stream runs the authentication and authorization of its
+request again and closes when they fail, so a revoked API key, an expired token or a withdrawn
+grant ends it without waiting for `max_connection_s`. The check goes through the usual caches:
+a grant changed through another Keeper process is seen after at most
+`keeper.auth.authz.grants_cache_ttl_s` more, OIDC group changes after
+`keeper.auth.oidc.principal_cache_ttl_s`.
+
+`GET .../changes?since=<revision>` reads the same events from the first-parent history of the
+branch instead of the outbox: the Revisions after `since`, oldest first, a page at a time,
+with `next_since` to continue. Use it to catch up after a gap or where streams are not
+available.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from pydantic import BaseModel, Field, JsonValue
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_core.model import Schema
+from lorebase_core.tree import kind_partition
+from lorebase_keeper.agents.settings import FeedSettings
+from lorebase_keeper.errors import CONFLICT, ProblemError
+from lorebase_keeper.reading.points import branch_schema
+from lorebase_keeper.reading.problems import invalid
+from lorebase_keeper.reading.runtime import ReadRuntime, begin_read
+from lorebase_store.changefeed import EventRef, describe, first_parents, matches
+from lorebase_store.commit.pipeline import COMMITTED_EVENT
+from lorebase_store.repositories import OutboxRepository, RevisionRow, scope_to_domain
+
+logger = logging.getLogger("lorebase.keeper.feed")
+
+EVENT_STREAM = "text/event-stream"
+SSE_EVENT = "revision"
+
+
+class RevisionEventOut(BaseModel):
+    """A Revision committed on a branch."""
+
+    sequence: int | None = Field(
+        default=None, description="Outbox sequence; the SSE id. Absent in `changes` pages."
+    )
+    type: str
+    domain: str
+    branch: str
+    revision: str
+    parents: list[str]
+    generation: int
+    author: str
+    on_behalf_of: str | None = None
+    message: str
+    origin: dict[str, JsonValue]
+    stats: dict[str, JsonValue]
+    created_at: str
+    kinds: dict[str, int] = Field(description="Changed Records per Kind api name.")
+    relations: dict[str, int] = Field(description="Changed Links per Relation api name.")
+    schema_changed: bool
+    keys: dict[str, list[str]] = Field(description="Changed Record keys per Kind, capped.")
+    keys_truncated: bool
+
+
+class ChangesFeedPage(BaseModel):
+    branch: str
+    head: str
+    since: str | None = None
+    items: list[RevisionEventOut]
+    next_since: str = Field(description="Pass as `since` to continue.")
+    more: bool = Field(description="True when more Revisions follow `next_since`.")
+
+
+@dataclass(frozen=True, slots=True)
+class KindFilter:
+    """Kinds an event must touch (None: every event); keys are listed for those Kinds only."""
+
+    names: frozenset[str] | None = None
+    partitions: frozenset[str] | None = None
+
+
+def kind_filter(schema: Schema | None, values: Sequence[str] | None) -> KindFilter:
+    """The filter of a `kinds` query parameter (repeated or comma-separated api names)."""
+    if values is None:
+        return KindFilter()
+    names = [name.strip() for value in values for name in value.split(",") if name.strip()]
+    partitions = set()
+    for index, name in enumerate(names):
+        kind = None if schema is None else schema.kind_by_api_name(name)
+        if kind is None:
+            raise invalid(f"/query/kinds/{index}", f"unknown Kind {name!r}", "unknown_kind")
+        partitions.add(kind_partition(kind.id))
+    return KindFilter(frozenset(names), frozenset(partitions))
+
+
+def parse_sequence(value: str | None) -> int | None:
+    """An event id given by the client (`Last-Event-ID`); 422 when it is not one."""
+    if value is None or value == "":
+        return None
+    try:
+        sequence = int(value)
+    except ValueError:
+        sequence = -1
+    if sequence < 0:
+        raise invalid("/headers/Last-Event-ID", "is not an event id", "invalid")
+    return sequence
+
+
+async def events_after(
+    session: AsyncSession,
+    runtime: ReadRuntime,
+    domain_id: str,
+    branch: str,
+    after: int,
+    wanted: KindFilter,
+    settings: FeedSettings,
+) -> tuple[int, list[dict[str, Any]]]:
+    """Events of the branch after sequence `after` (one batch) and the last sequence read."""
+    rows = await OutboxRepository(session).after(
+        domain_id,
+        after,
+        event_types=(COMMITTED_EVENT,),
+        branch=branch,
+        limit=settings.batch_events,
+    )
+    if not rows:
+        return after, []
+    schema = await branch_schema(session, runtime, domain_id, branch)
+    documents = await describe(
+        session,
+        domain_id,
+        [EventRef(bytes.fromhex(str(r.payload["revision"])), branch, r.id) for r in rows],
+        schemas={branch: schema},
+        keys_max=settings.changed_keys_max,
+        kind_partitions=wanted.partitions,
+    )
+    return rows[-1].id, [d for d in documents if matches(d, wanted.names)]
+
+
+def sse(document: dict[str, Any]) -> bytes:
+    data = json.dumps(document, ensure_ascii=False, separators=(",", ":"))
+    return f"id: {document['sequence']}\nevent: {SSE_EVENT}\ndata: {data}\n\n".encode()
+
+
+async def stream(
+    runtime: ReadRuntime,
+    domain_id: str,
+    branch: str,
+    start: int,
+    wanted: KindFilter,
+    settings: FeedSettings,
+    disconnected: Any,
+    still_allowed: Callable[[], Awaitable[bool]],
+) -> AsyncIterator[bytes]:
+    """The SSE body: events as they are committed, heartbeats, closed after the max time or
+    once `still_allowed` (called every `reauthorize_interval_s`) answers false."""
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + settings.max_connection_s
+    next_check = loop.time() + settings.reauthorize_interval_s
+    quiet_since = loop.time()
+    cursor = start
+    yield f"retry: {settings.retry_ms}\n\n".encode()
+    while loop.time() < deadline:
+        if await disconnected():
+            return
+        if loop.time() >= next_check:
+            try:
+                allowed = await still_allowed()
+            except Exception:
+                logger.exception("change feed authorization failed", extra={"domain": domain_id})
+                return
+            if not allowed:
+                logger.info("change feed closed: access withdrawn", extra={"domain": domain_id})
+                return
+            next_check = loop.time() + settings.reauthorize_interval_s
+        try:
+            async with runtime.sessions() as session, session.begin():
+                await begin_read(session, runtime)
+                await scope_to_domain(session, domain_id)
+                cursor, documents = await events_after(
+                    session, runtime, domain_id, branch, cursor, wanted, settings
+                )
+        except Exception:
+            logger.exception("change feed poll failed", extra={"domain": domain_id})
+            return
+        for document in documents:
+            yield sse(document)
+            quiet_since = loop.time()
+        if loop.time() - quiet_since >= settings.heartbeat_s:
+            yield b": keep-alive\n\n"
+            quiet_since = loop.time()
+        await asyncio.sleep(max(min(settings.poll_interval_s, deadline - loop.time()), 0))
+
+
+async def changes_since(
+    session: AsyncSession,
+    runtime: ReadRuntime,
+    domain_id: str,
+    branch: str,
+    head: bytes,
+    since: RevisionRow | None,
+    *,
+    limit: int,
+    wanted: KindFilter,
+    settings: FeedSettings,
+) -> ChangesFeedPage:
+    """A page of the first-parent Revisions after `since` (see module doc)."""
+    if since is None:
+        segment = await first_parents(session, domain_id, head, floor=-1, max_steps=limit - 1)
+        newer = [revision for revision, _ in reversed(segment)]
+    else:
+        segment = await first_parents(
+            session, domain_id, head, floor=since.generation, max_steps=settings.max_scan
+        )
+        boundary, generation = segment[-1]
+        if boundary != since.id:
+            reason = (
+                "is not on the first-parent history of the branch"
+                if generation <= since.generation
+                else f"is more than {settings.max_scan} Revisions behind the head"
+            )
+            raise ProblemError(
+                CONFLICT,
+                f"Revision {since.id.hex()} {reason}; read the current state and follow the "
+                "branch from its head.",
+                extensions={"code": "since_unreachable"},
+            )
+        newer = [revision for revision, _ in reversed(segment[:-1])]
+    page = newer[:limit]
+    schema = await branch_schema(session, runtime, domain_id, branch)
+    documents = await describe(
+        session,
+        domain_id,
+        [EventRef(revision, branch) for revision in page],
+        schemas={branch: schema},
+        keys_max=settings.changed_keys_max,
+        kind_partitions=wanted.partitions,
+    )
+    fallback = head if since is None else since.id
+    return ChangesFeedPage(
+        branch=branch,
+        head=head.hex(),
+        since=None if since is None else since.id.hex(),
+        items=[RevisionEventOut.model_validate(d) for d in documents if matches(d, wanted.names)],
+        next_since=(page[-1] if page else fallback).hex(),
+        more=len(newer) > limit,
+    )

lorebase_keeper/agents/markdown.py

@@ -0,0 +1,219 @@
+"""Markdown rendering of context packs and Record explanations, meant to be pasted into a prompt.
+
+The output is stable (same input, same bytes) and compact: plain ASCII markup, lists for
+Records and schema, tables for neighbors and changes, empty sections left out. Revisions
+inside sections are shortened to `SHORT_REVISION` hex digits; the header carries the full id.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+from lorebase_core.model import Schema
+from lorebase_keeper.agents.budget import Draft
+from lorebase_keeper.agents.pack import links_json, schema_fragment
+
+SHORT_REVISION = 12
+SEPARATOR = "\n\n"
+
+
+def short(revision: str | None) -> str:
+    return "" if revision is None else revision[:SHORT_REVISION]
+
+
+def text(value: Any) -> str:
+    """A value on one line: strings as they are, lists comma-separated, objects as JSON."""
+    if value is None:
+        return ""
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, list):
+        return ", ".join(text(v) for v in value)
+    if isinstance(value, dict):
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
+    return " ".join(str(value).split())
+
+
+def cell(value: Any) -> str:
+    """A value inside a table cell."""
+    return text(value).replace("|", "\\|")
+
+
+def pairs(values: Mapping[str, Any] | None) -> str:
+    return "; ".join(f"{name}={text(value)}" for name, value in (values or {}).items())
+
+
+def table(header: Sequence[str], rows: Iterable[Sequence[Any]]) -> list[str]:
+    lines = ["| " + " | ".join(header) + " |", "|" + " --- |" * len(header)]
+    lines.extend("| " + " | ".join(cell(v) for v in row) + " |" for row in rows)
+    return lines
+
+
+def heading(ref: str, title: str | None) -> str:
+    return f"### {ref}" + (f": {text(title)}" if title else "")
+
+
+# Context pack sections ----------------------------------------------------------------------
+
+
+def _seeds(draft: Draft) -> str:
+    lines = ["## Records"]
+    for seed in draft.seeds:
+        lines.append(heading(seed.ref, seed.title))
+        lines.extend(f"- {name}: {text(value)}" for name, value in seed.values.items())
+        if seed.updated_revision is not None:
+            lines.append(f"- (last changed in {short(seed.updated_revision)})")
+    return "\n".join(lines)
+
+
+def _neighbors(draft: Draft) -> str:
+    if not draft.neighbors:
+        return ""
+    rows = [
+        (n.ref, n.title, n.depth, "(omitted)" if n.values is None else pairs(n.values))
+        for n in draft.neighbors
+    ]
+    return "\n".join(["## Neighbors", *table(("ref", "title", "depth", "values"), rows)])
+
+
+def _links(draft: Draft, schema: Schema) -> str:
+    if not draft.links:
+        return ""
+    lines = ["## Links"]
+    for name, links in links_json(draft.links).items():
+        found = schema.relation_by_api_name(name)
+        ends = ""
+        if found is not None:
+            relation = found[0]
+            source, target = (
+                schema.kinds.get(relation.from_kind),
+                schema.kinds.get(relation.to_kind),
+            )
+            if source is not None and target is not None:
+                ends = f" ({source.api_name} -> {target.api_name})"
+        lines.append(f"### {name}{ends}")
+        for link in links:
+            extra = f" ({pairs(link['values'])})" if link.get("values") else ""
+            lines.append(f"- {link['from']} -> {link['to']}{extra}")
+    return "\n".join(lines)
+
+
+def _attribute_line(attribute: Mapping[str, Any], indent: str = "") -> str:
+    flags = [attribute["type"]]
+    if attribute.get("required"):
+        flags.append("required")
+    if attribute.get("derived"):
+        flags.append("derived")
+    description = attribute.get("description")
+    return f"{indent}- {attribute['name']} ({', '.join(flags)})" + (
+        f": {text(description)}" if description else ""
+    )
+
+
+def schema_markdown(fragment: Mapping[str, Any]) -> str:
+    lines = ["## Schema"]
+    for kind in fragment["kinds"]:
+        lines.append(f"### Kind {kind['name']}")
+        if kind.get("description"):
+            lines.append(text(kind["description"]))
+        keys = f"key: {', '.join(kind['key'])}"
+        lines.append(keys + (f"; title: {kind['title']}" if kind.get("title") else ""))
+        lines.extend(_attribute_line(a) for a in kind["attributes"])
+    if fragment["relations"]:
+        lines.append("### Relations")
+        for relation in fragment["relations"]:
+            line = (
+                f"- {relation['name']}: {relation['from']} -> {relation['to']}, "
+                f"{relation['cardinality']}, inverse {relation['inverse']}"
+            )
+            if relation.get("description"):
+                line += f". {text(relation['description'])}"
+            lines.append(line)
+            lines.extend(_attribute_line(a, "  ") for a in relation.get("attributes", ()))
+    return "\n".join(lines)
+
+
+def _changes(draft: Draft) -> str:
+    if not draft.changes:
+        return ""
+    rows = [
+        (
+            short(c.revision),
+            c.at,
+            c.ref,
+            c.change,
+            ", ".join(c.attributes),
+            c.actor,
+            c.origin,
+            c.message,
+        )
+        for c in draft.changes
+    ]
+    header = ("revision", "at", "record", "change", "attributes", "by", "origin", "message")
+    return "\n".join(["## Recent changes", *table(header, rows)])
+
+
+def _lineage(draft: Draft) -> str:
+    if not draft.lineage:
+        return ""
+    lines = ["## Lineage"]
+    for item in draft.lineage:
+        source = " ".join(
+            part
+            for part in (
+                item.origin or "unknown origin",
+                f"by {item.actor}" if item.actor else "",
+                f"at {item.at}" if item.at else "",
+            )
+            if part
+        )
+        lines.append(f"- {item.ref}: {', '.join(item.attributes)} <- {source}")
+    return "\n".join(lines)
+
+
+def markdown_sections(draft: Draft, schema: Schema) -> dict[str, str]:
+    """Text of each section, in output order; empty for sections without content."""
+    return {
+        "seeds": _seeds(draft),
+        "neighbors": _neighbors(draft),
+        "links": _links(draft, schema),
+        "schema": schema_markdown(schema_fragment(draft, schema)) if draft.schema else "",
+        "recent_changes": _changes(draft),
+        "lineage": _lineage(draft),
+    }
+
+
+def markdown_measure(schema: Schema) -> Any:
+    """Characters of each section including its separator, a `budget.Measure`."""
+
+    def measure(draft: Draft) -> Mapping[str, int]:
+        return {
+            name: len(body) + len(SEPARATOR) if body else 0
+            for name, body in markdown_sections(draft, schema).items()
+        }
+
+    return measure
+
+
+def pack_markdown(
+    sections: Mapping[str, str],
+    *,
+    domain: str,
+    branch: str,
+    revision: str,
+    budget: Mapping[str, Any],
+    since: str | None,
+) -> str:
+    """The whole pack: a header line, then the non-empty sections."""
+    trimmed = ", ".join(f"{t['step']}({t['removed']})" for t in budget["trimmed"]) or "none"
+    header = [
+        f"# Context: {domain}/{branch}",
+        f"revision {revision}; estimated {budget['used_tokens']} of "
+        f"{budget['limit_tokens']} tokens; trimmed: {trimmed}"
+        + ("; over budget" if budget["over_budget"] else "")
+        + (f"; changes since {short(since)}" if since else ""),
+    ]
+    body = [part for part in sections.values() if part]
+    return SEPARATOR.join(["\n".join(header), *body]) + "\n"