FabledGrimoire/src/fabledassistant/services/tools/journal.py

"""LLM tools for journal conversations: record_moment and search_journal."""
from __future__ import annotations

import datetime
import logging
from typing import Any

from sqlalchemy import func, select

from fabledassistant.models import Note, async_session
from fabledassistant.services.journal_search import search_journal as svc_search_journal
from fabledassistant.services.moments import create_moment
from fabledassistant.services.tools._registry import tool

logger = logging.getLogger(__name__)


async def _resolve_entity_ids_by_name(
    *,
    user_id: int,
    names: list[str],
    note_type: str | None = None,
    is_task_only: bool | None = None,
) -> list[int]:
    """Case-insensitive title lookup → note IDs.

    Names with no match are silently dropped (logged at debug). Used by
    record_moment to resolve user-mentioned names without forcing the
    LLM to track tool-return IDs across calls.
    """
    if not names:
        return []
    cleaned = [n.strip() for n in names if n and n.strip()]
    if not cleaned:
        return []
    lowered = [n.lower() for n in cleaned]
    async with async_session() as session:
        stmt = select(Note.id, func.lower(Note.title).label("ltitle")).where(
            Note.user_id == user_id,
            func.lower(Note.title).in_(lowered),
        )
        if note_type is not None:
            stmt = stmt.where(Note.note_type == note_type)
        if is_task_only is True:
            stmt = stmt.where(Note.status.isnot(None))
        elif is_task_only is False:
            stmt = stmt.where(Note.status.is_(None))
        rows = (await session.execute(stmt)).all()
    # Pick first match per name; preserve input order.
    by_lower: dict[str, int] = {}
    for note_id, ltitle in rows:
        if ltitle not in by_lower:
            by_lower[ltitle] = note_id
    resolved: list[int] = []
    for low, original in zip(lowered, cleaned):
        if low in by_lower:
            resolved.append(by_lower[low])
        else:
            logger.debug(
                "record_moment: no entity match for %r (note_type=%s, task_only=%s)",
                original, note_type, is_task_only,
            )
    return resolved


@tool(
    name="record_moment",
    description=(
        "Record a meaningful beat from the conversation as a structured Moment. "
        "Use freely (no confirmation) for anything significant the user mentions: "
        "events, encounters, decisions, observations, feelings worth remembering. "
        "Each Moment is one or two sentences distilling the beat — not a full transcript. "
        "STRONGLY PREFER the *_names parameters when linking entities — the server "
        "resolves names to IDs by lookup, so you cannot accidentally invent or "
        "re-use the wrong ID. Use *_ids only when you have an exact ID returned "
        "from another tool call in this same turn."
    ),
    parameters={
        "content": {
            "type": "string",
            "description": "1-2 sentence distillation of the moment in the user's voice or third-person.",
        },
        "occurred_at": {
            "type": "string",
            "description": "ISO 8601 datetime when the moment happened. Pass 'now' for the present moment; the handler resolves it.",
        },
        "raw_excerpt": {
            "type": "string",
            "description": "Optional: literal user phrase the moment summarizes. Useful for trust UI.",
        },
        "tags": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Optional tags (no # prefix).",
        },
        "person_names": {
            "type": "array",
            "items": {"type": "string"},
            "description": "PREFERRED. Names of people mentioned (e.g. ['Victoria', 'Mom']). Server resolves to existing person notes by case-insensitive title match. Names with no match are silently skipped.",
        },
        "place_names": {
            "type": "array",
            "items": {"type": "string"},
            "description": "PREFERRED. Names of places mentioned (e.g. ['the new ramen place']). Server resolves to existing place notes by title.",
        },
        "task_titles": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Titles of tasks this moment references. Server resolves to task IDs by title match.",
        },
        "note_titles": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Titles of notes this moment references. Server resolves to note IDs by title match.",
        },
        "person_ids": {
            "type": "array",
            "items": {"type": "integer"},
            "description": "DISCOURAGED — use person_names instead. Only valid if you obtained the ID from a tool result in this same turn. Never invent IDs.",
        },
        "place_ids": {
            "type": "array",
            "items": {"type": "integer"},
            "description": "DISCOURAGED — use place_names instead. Same caveat as person_ids.",
        },
        "task_ids": {
            "type": "array",
            "items": {"type": "integer"},
            "description": "DISCOURAGED — use task_titles instead. Same caveat as person_ids.",
        },
        "note_ids": {
            "type": "array",
            "items": {"type": "integer"},
            "description": "DISCOURAGED — use note_titles instead. Same caveat as person_ids.",
        },
    },
    required=["content"],
    read_only=False,
    journal=True,
)
async def record_moment_tool(*, user_id, arguments, conv_id=None, **_ctx):
    content = arguments.get("content", "").strip()
    if not content:
        return {"success": False, "error": "content is required"}

    occurred_raw = arguments.get("occurred_at")
    now = datetime.datetime.now(datetime.timezone.utc)
    if occurred_raw and occurred_raw != "now":
        try:
            occurred_dt = datetime.datetime.fromisoformat(occurred_raw)
            if occurred_dt.tzinfo is None:
                occurred_dt = occurred_dt.replace(tzinfo=datetime.timezone.utc)
        except ValueError:
            occurred_dt = now
    else:
        occurred_dt = now

    # Start with any explicit IDs the LLM provided.
    person_ids = list(arguments.get("person_ids") or [])
    place_ids = list(arguments.get("place_ids") or [])
    task_ids = list(arguments.get("task_ids") or [])
    note_ids = list(arguments.get("note_ids") or [])

    # Resolve names → IDs and merge. Names are the preferred path because the
    # LLM is unreliable at tracking IDs across tool calls (see prior bug:
    # hallucinated person_ids=[1,2] referencing test-data notes).
    person_ids.extend(
        await _resolve_entity_ids_by_name(
            user_id=user_id,
            names=arguments.get("person_names") or [],
            note_type="person",
        )
    )
    place_ids.extend(
        await _resolve_entity_ids_by_name(
            user_id=user_id,
            names=arguments.get("place_names") or [],
            note_type="place",
        )
    )
    task_ids.extend(
        await _resolve_entity_ids_by_name(
            user_id=user_id,
            names=arguments.get("task_titles") or [],
            is_task_only=True,
        )
    )
    note_ids.extend(
        await _resolve_entity_ids_by_name(
            user_id=user_id,
            names=arguments.get("note_titles") or [],
            is_task_only=False,
        )
    )

    # Dedupe while preserving order.
    person_ids = list(dict.fromkeys(person_ids))
    place_ids = list(dict.fromkeys(place_ids))
    task_ids = list(dict.fromkeys(task_ids))
    note_ids = list(dict.fromkeys(note_ids))

    moment = await create_moment(
        user_id=user_id,
        content=content,
        occurred_at=occurred_dt,
        day_date=occurred_dt.date(),
        conversation_id=conv_id,
        raw_excerpt=arguments.get("raw_excerpt"),
        tags=arguments.get("tags") or [],
        person_ids=person_ids,
        place_ids=place_ids,
        task_ids=task_ids,
        note_ids=note_ids,
    )
    return {
        "success": True,
        "type": "moment_recorded",
        "moment": moment.to_dict(),
    }


@tool(
    name="search_journal",
    description=(
        "Search the user's journal: past Moments and (optionally) raw transcripts. "
        "Use to recall things mentioned in prior days. Three modes: "
        "(a) date filter only -> recent moments in that range; "
        "(b) person_id/place_id filter -> moments mentioning that entity; "
        "(c) query string -> semantic search, optionally constrained by date/entity."
    ),
    parameters={
        "query": {"type": "string", "description": "Optional semantic query."},
        "person_id": {"type": "integer", "description": "Filter to this person."},
        "place_id": {"type": "integer", "description": "Filter to this place."},
        "tag": {"type": "string", "description": "Filter to moments with this tag."},
        "date_from": {"type": "string", "description": "ISO date lower bound (inclusive)."},
        "date_to": {"type": "string", "description": "ISO date upper bound (inclusive)."},
        "include_transcripts": {
            "type": "boolean",
            "description": "If true, also return matching raw transcript excerpts when query is set. Default false.",
        },
        "limit": {"type": "integer", "description": "Max results, default 10."},
    },
    required=[],
    read_only=True,
    journal=True,
)
async def search_journal_tool(*, user_id, arguments, **_ctx) -> dict[str, Any]:
    df_raw = arguments.get("date_from")
    dt_raw = arguments.get("date_to")
    df = datetime.date.fromisoformat(df_raw) if df_raw else None
    dt = datetime.date.fromisoformat(dt_raw) if dt_raw else None
    results = await svc_search_journal(
        user_id=user_id,
        query=arguments.get("query"),
        person_id=arguments.get("person_id"),
        place_id=arguments.get("place_id"),
        tag=arguments.get("tag"),
        date_from=df,
        date_to=dt,
        include_transcripts=bool(arguments.get("include_transcripts", False)),
        limit=int(arguments.get("limit", 10)),
    )
    return {"success": True, "type": "journal_search_results", "count": len(results), "results": results}