chat/chat/state/memory.py

from __future__ import annotations
from sqlite3 import Connection
from chat.eventlog.projector import on
from chat.eventlog.log import Event

_VALID_WITNESS_ROLES = {"you", "host", "guest"}


def _row_to_dict(conn: Connection, row: tuple) -> dict:
    cols = [c[1] for c in conn.execute("PRAGMA table_info(memories)").fetchall()]
    return dict(zip(cols, row))


@on("memory_written")
def _apply_memory_written(conn: Connection, e: Event) -> None:
    p = e.payload
    conn.execute(
        "INSERT INTO memories ("
        "owner_id, chat_id, scene_id, pov_summary, "
        "witness_you, witness_host, witness_guest, "
        "chat_clock_at, source, reliability, significance, pinned, auto_pinned"
        ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (
            p["owner_id"],
            p["chat_id"],
            p.get("scene_id"),
            p["pov_summary"],
            int(p["witness_you"]),
            int(p["witness_host"]),
            int(p["witness_guest"]),
            p.get("chat_clock_at"),
            p.get("source", "direct"),
            float(p.get("reliability", 1.0)),
            int(p.get("significance", 1)),
            int(p.get("pinned", 0)),
            int(p.get("auto_pinned", 0)),
        ),
    )


@on("memory_significance_set")
def _apply_memory_significance_set(conn: Connection, e: Event) -> None:
    """Update an existing memory's significance score (T22).

    Emitted by the async significance worker after it scores the turn.
    """
    p = e.payload
    conn.execute(
        "UPDATE memories SET significance = ? WHERE id = ?",
        (int(p["significance"]), int(p["memory_id"])),
    )


@on("memory_pin_changed")
def _apply_memory_pin_changed(conn: Connection, e: Event) -> None:
    """Toggle a memory's pin state (T22, §8.5).

    Used both for auto-pinning a pivotal turn and for evicting the oldest
    auto-pin when the per-owner soft cap is exceeded. Manual pins use the
    same handler; the ``auto_pinned`` flag distinguishes them so the
    eviction query can leave manual pins alone.
    """
    p = e.payload
    conn.execute(
        "UPDATE memories SET pinned = ?, auto_pinned = ? WHERE id = ?",
        (int(p["pinned"]), int(p["auto_pinned"]), int(p["memory_id"])),
    )


def get_memory(conn: Connection, memory_id: int) -> dict | None:
    row = conn.execute(
        "SELECT * FROM memories WHERE id = ?", (memory_id,)
    ).fetchone()
    if not row:
        return None
    return _row_to_dict(conn, row)


def get_pinned(conn: Connection, owner_id: str) -> list[dict]:
    cur = conn.execute(
        "SELECT * FROM memories WHERE owner_id = ? AND pinned = 1 "
        "ORDER BY created_at DESC, id DESC",
        (owner_id,),
    )
    rows = cur.fetchall()
    cols = [c[1] for c in conn.execute("PRAGMA table_info(memories)").fetchall()]
    return [dict(zip(cols, row)) for row in rows]


# Composite-score weights used by ``search_memories`` (T23, §8 retrieval).
# FTS5 BM25 ``rank`` is *more negative* for better matches, so subtracting a
# positive boost from it drives stronger candidates further down (i.e. earlier
# in an ascending sort). Hardcoded for v1 — tunable in a later pass.
_SIGNIFICANCE_WEIGHT = 0.3
_RECENCY_WEIGHT = 0.5


def search_memories(
    conn: Connection,
    owner_id: str,
    witness_role: str,
    query: str,
    k: int = 4,
) -> list[dict]:
    """FTS5 search over pov_summary, scoped by owner and witness role.

    witness_role must be one of {"you", "host", "guest"} per the witness flags
    on each memory row. Returns up to ``k`` rows ranked by a composite score
    that combines the FTS5 BM25 rank with two boosts (§8 retrieval rules):

    * **significance boost** — ``0.3 * significance`` (0..3 per §11.1).
    * **recency boost** — ``0.5 * (id / max_id)``, using the row id as a
      monotonic recency proxy. Newer memories therefore tilt above older ones
      when the BM25 rank and significance are otherwise tied.

    BM25 returns negative scores (lower = better). Both boosts are subtracted
    so that stronger candidates yield smaller composite scores; the result is
    sorted ascending and truncated to ``k``. The unmodified ``fts_rank`` and a
    debug-friendly ``composite_score`` are kept on each returned dict.
    """
    if witness_role not in _VALID_WITNESS_ROLES:
        raise ValueError(
            f"witness_role must be one of {sorted(_VALID_WITNESS_ROLES)}, "
            f"got {witness_role!r}"
        )
    if not query.strip():
        return []
    witness_col = f"witness_{witness_role}"
    cols = [c[1] for c in conn.execute("PRAGMA table_info(memories)").fetchall()]
    select_list = ", ".join(f"m.{c}" for c in cols)
    # Over-fetch from FTS so the Python-side re-rank has room to reorder
    # results that BM25 alone would have demoted past the top-k boundary.
    over_fetch = max(k * 4, 20)
    sql = (
        f"SELECT {select_list}, memories_fts.rank AS fts_rank "
        "FROM memories_fts "
        "JOIN memories m ON m.id = memories_fts.rowid "
        f"WHERE m.owner_id = ? AND m.{witness_col} = 1 "
        "AND memories_fts MATCH ? "
        "ORDER BY memories_fts.rank "
        "LIMIT ?"
    )
    cur = conn.execute(sql, (owner_id, query, over_fetch))
    rows = cur.fetchall()
    if not rows:
        return []

    # Recency normalises against the current max id for this owner so the
    # boost magnitude is bounded regardless of dataset size.
    max_id_row = conn.execute(
        "SELECT MAX(id) FROM memories WHERE owner_id = ?", (owner_id,)
    ).fetchone()
    max_id = max_id_row[0] if max_id_row and max_id_row[0] else 1

    result_cols = cols + ["fts_rank"]
    enriched: list[dict] = []
    for row in rows:
        d = dict(zip(result_cols, row))
        fts_rank = d.get("fts_rank") or 0.0
        sig_boost = _SIGNIFICANCE_WEIGHT * (d.get("significance") or 0)
        recency_boost = _RECENCY_WEIGHT * ((d.get("id") or 0) / max_id)
        d["composite_score"] = fts_rank - sig_boost - recency_boost
        enriched.append(d)

    enriched.sort(key=lambda x: x["composite_score"])
    return enriched[:k]