Kansho/backend/retrieval.py

"""Swappable retrieval interface. Specs name what to select; recency/SQL is implementation."""
from __future__ import annotations

from typing import Any

from db import get_db, row_to_dict
from journal_body import plain_text

# Implementation caps. Not part of the selection-spec contract.
_DAY_MESSAGE_CAP = 80
_SPACE_ENTRY_LIMIT = 5
_SPACE_ENTRY_EXCERPT_CHARS = 400
_SPACE_SOURCE_CONV_LIMIT = 3
_SPACE_SOURCE_EXCERPT_CHARS = 400
_SPACE_SOURCE_MESSAGE_CAP = 8


def retrieve(profile_id: str, spec: dict[str, Any]) -> list[dict]:
    kind = spec.get("kind")
    if kind == "day_messages":
        return _day_messages(profile_id, spec)
    if kind == "space_entries":
        return _space_entries(profile_id, spec)
    if kind == "space_recent_sources":
        return _space_recent_sources(profile_id, spec)
    if kind == "writing_profile":
        return _writing_profile(profile_id)
    return []


def _day_messages(profile_id: str, spec: dict[str, Any]) -> list[dict]:
    journal_day_id = spec.get("journal_day_id")
    conversation_ids = spec.get("conversation_ids") or []
    if not journal_day_id:
        return []
    with get_db() as conn:
        if conversation_ids:
            placeholders = ",".join("?" * len(conversation_ids))
            rows = conn.execute(
                f"""
                SELECT m.*, c.title AS conversation_title
                FROM messages m
                JOIN conversations c ON c.id = m.conversation_id
                WHERE m.profile_id = ? AND c.journal_day_id = ? AND m.conversation_id IN ({placeholders})
                ORDER BY c.created, m.seq
                LIMIT ?
                """,
                (profile_id, journal_day_id, *conversation_ids, _DAY_MESSAGE_CAP),
            ).fetchall()
        else:
            rows = conn.execute(
                """
                SELECT m.*, c.title AS conversation_title
                FROM messages m
                JOIN conversations c ON c.id = m.conversation_id
                WHERE m.profile_id = ? AND c.journal_day_id = ?
                ORDER BY c.created, m.seq
                LIMIT ?
                """,
                (profile_id, journal_day_id, _DAY_MESSAGE_CAP),
            ).fetchall()
    return [row_to_dict(row) for row in rows]


def _space_entries(profile_id: str, spec: dict[str, Any]) -> list[dict]:
    space_id = spec.get("space_id")
    exclude_day_id = spec.get("exclude_day_id")
    if not space_id:
        return []
    params: list[Any] = [profile_id, space_id]
    exclude_sql = ""
    if exclude_day_id:
        exclude_sql = "AND e.journal_day_id != ?"
        params.append(exclude_day_id)
    params.append(_SPACE_ENTRY_LIMIT)
    with get_db() as conn:
        rows = conn.execute(
            f"""
            SELECT e.id AS entry_id, e.journal_day_id, d.calendar_date, v.title, v.body, v.origin
            FROM journal_entries e
            JOIN journal_days d ON d.id = e.journal_day_id
            JOIN journal_entry_versions v ON v.id = e.current_version_id
            WHERE e.profile_id = ? AND e.space_id = ? AND e.deleted_at IS NULL {exclude_sql}
            ORDER BY d.calendar_date DESC, e.updated DESC
            LIMIT ?
            """,
            params,
        ).fetchall()
    result = []
    for row in rows:
        item = row_to_dict(row)
        item["excerpt"] = plain_text(item.get("body") or "")[:_SPACE_ENTRY_EXCERPT_CHARS]
        result.append(item)
    return result


def _space_recent_sources(profile_id: str, spec: dict[str, Any]) -> list[dict]:
    """Recent original conversations in the same space. Recency is the current heuristic."""
    space_id = spec.get("space_id")
    exclude_conversation_id = spec.get("exclude_conversation_id")
    exclude_journal_day_id = spec.get("exclude_journal_day_id")
    if not space_id:
        return []
    params: list[Any] = [profile_id, space_id]
    extra = ""
    if exclude_conversation_id:
        extra += " AND c.id != ?"
        params.append(exclude_conversation_id)
    if exclude_journal_day_id:
        extra += " AND (c.journal_day_id IS NULL OR c.journal_day_id != ?)"
        params.append(exclude_journal_day_id)
    params.append(_SPACE_SOURCE_CONV_LIMIT)
    with get_db() as conn:
        conversations = conn.execute(
            f"""
            SELECT c.id, c.updated, c.journal_day_id
            FROM conversations c
            WHERE c.profile_id = ? AND c.space_id = ? {extra}
            ORDER BY c.updated DESC, c.created DESC
            LIMIT ?
            """,
            params,
        ).fetchall()
        result = []
        for conv in conversations:
            messages = conn.execute(
                """
                SELECT body FROM messages
                WHERE profile_id = ? AND conversation_id = ? AND role = 'user'
                ORDER BY seq
                LIMIT ?
                """,
                (profile_id, conv["id"], _SPACE_SOURCE_MESSAGE_CAP),
            ).fetchall()
            excerpt = " ".join((row["body"] or "").strip() for row in messages if (row["body"] or "").strip())
            excerpt = excerpt[:_SPACE_SOURCE_EXCERPT_CHARS]
            if not excerpt:
                continue
            result.append(
                {
                    "conversation_id": conv["id"],
                    "journal_day_id": conv["journal_day_id"],
                    "excerpt": excerpt,
                }
            )
    return result


def _writing_profile(profile_id: str) -> list[dict]:
    with get_db() as conn:
        row = row_to_dict(
            conn.execute(
                "SELECT compiled_brief, updated FROM writing_profiles WHERE profile_id = ?",
                (profile_id,),
            ).fetchone()
        )
        sources = [
            row_to_dict(item)
            for item in conn.execute(
                """
                SELECT id, kind, body, entry_id, weight, created
                FROM writing_profile_sources
                WHERE profile_id = ?
                ORDER BY CASE kind
                  WHEN 'journal_entry' THEN 0
                  WHEN 'imported_text' THEN 1
                  ELSE 2
                END, created DESC
                """,
                (profile_id,),
            ).fetchall()
        ]
    if not row and not sources:
        return []
    return [{"compiled_brief": (row or {}).get("compiled_brief") or "", "sources": sources}]