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
from prompt_budget import (
    ERROR_PROMPT_BUDGET_EXCEEDED,
    JournalBudgetError,
    day_message_safety_cap,
    estimate_tokens,
)

# 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 format_day_messages(messages: list[dict], *, with_source_ids: bool = False) -> str:
    parts = []
    for message in messages:
        role = message.get("role") or "user"
        body = message.get("body") or ""
        source_id = message.get("source_id")
        if with_source_ids and source_id:
            parts.append(f"[{source_id}] {role}: {body}")
        else:
            parts.append(f"{role}: {body}")
    return "\n".join(parts).strip()


def pair_user_priority_messages(messages: list[dict]) -> list[dict]:
    """Keep every user message. Assistant context stays inside the same conversation."""
    selected: list[dict] = []
    pending_assistant = None
    pending_conversation = None
    for message in messages:
        conversation_id = message.get("conversation_id")
        role = message.get("role") or "user"
        if pending_assistant is not None and conversation_id != pending_conversation:
            pending_assistant = None
            pending_conversation = None
        if role == "assistant":
            pending_assistant = message
            pending_conversation = conversation_id
            continue
        if role == "user":
            if pending_assistant is not None and pending_conversation == conversation_id:
                selected.append(pending_assistant)
            selected.append(message)
            pending_assistant = None
            pending_conversation = None
            continue
        pending_assistant = None
        pending_conversation = None
    return selected


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 []
    overflow = spec.get("overflow") or "limit"
    if overflow == "abort":
        cap = int(spec.get("message_cap") or day_message_safety_cap())
        fetch_limit = cap + 1
    else:
        cap = int(spec.get("message_cap") or _DAY_MESSAGE_CAP)
        fetch_limit = cap
    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, c.rowid, m.seq, m.id
                LIMIT ?
                """,
                (profile_id, journal_day_id, *conversation_ids, fetch_limit),
            ).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, c.rowid, m.seq, m.id
                LIMIT ?
                """,
                (profile_id, journal_day_id, fetch_limit),
            ).fetchall()
    messages = [row_to_dict(row) for row in rows]
    if overflow == "abort":
        if len(messages) > cap:
            raise JournalBudgetError(
                ERROR_PROMPT_BUDGET_EXCEEDED,
                diagnostics={"reason": "day_message_cap", "message_cap": cap, "fetched": len(messages)},
            )
        from journal_reconstruct import assign_source_ids

        paired = assign_source_ids(pair_user_priority_messages(messages))
        max_tokens = spec.get("max_estimated_tokens")
        if max_tokens is not None:
            used = estimate_tokens(format_day_messages(paired, with_source_ids=True))
            if used > int(max_tokens):
                raise JournalBudgetError(
                    ERROR_PROMPT_BUDGET_EXCEEDED,
                    diagnostics={
                        "reason": "day_dialogue_tokens",
                        "estimated_tokens": used,
                        "available_tokens": int(max_tokens),
                        "message_count": len(paired),
                    },
                )
        return paired
    return messages


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}]