FabledScribe/src/fabledassistant/services/journal_search.py

"""Search across Moments and (optionally) journal transcripts.

Three modes, expressed through one tool surface:
1. Pure temporal: no query, just date range -> ordered by occurred_at DESC.
2. Pure entity: person_id or place_id filter -> junction lookup.
3. Semantic: query string -> embedding similarity, optionally constrained.

This module ONLY queries Moments and (optionally) journal-conversation
messages. It MUST NOT touch notes or note_embeddings. The notes-RAG and
journal-RAG isolation is a hard invariant of the journal design.
"""
from __future__ import annotations

import datetime
import math
from typing import Sequence

from sqlalchemy import select

from fabledassistant.models import (
    Conversation,
    Message,
    Moment,
    MomentEmbedding,
    async_session,
    moment_people,
    moment_places,
)
from fabledassistant.services.embeddings import get_embedding

DEFAULT_LIMIT = 10
DEFAULT_THRESHOLD = 0.55


def _cosine(a: Sequence[float], b: Sequence[float]) -> float:
    if not a or not b:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


async def search_journal(
    *,
    user_id: int,
    query: str | None = None,
    person_id: int | None = None,
    place_id: int | None = None,
    tag: str | None = None,
    date_from: datetime.date | None = None,
    date_to: datetime.date | None = None,
    include_transcripts: bool = False,
    limit: int = DEFAULT_LIMIT,
    threshold: float = DEFAULT_THRESHOLD,
) -> list[dict]:
    """Return Moments (and optional transcript snippets) matching the filters.

    Result rows: dicts from Moment.to_dict() plus an optional `score` when
    `query` is set. Transcript snippets carry `kind='transcript'` and `id=None`
    to distinguish them.
    """
    moment_rows = await _search_moments(
        user_id=user_id,
        query=query,
        person_id=person_id,
        place_id=place_id,
        tag=tag,
        date_from=date_from,
        date_to=date_to,
        limit=limit,
        threshold=threshold,
    )

    if not include_transcripts:
        return moment_rows

    transcript_rows = await _search_transcripts(
        user_id=user_id,
        query=query,
        date_from=date_from,
        date_to=date_to,
        limit=limit,
    )
    return moment_rows + transcript_rows


async def _search_moments(
    *,
    user_id: int,
    query: str | None,
    person_id: int | None,
    place_id: int | None,
    tag: str | None,
    date_from: datetime.date | None,
    date_to: datetime.date | None,
    limit: int,
    threshold: float,
) -> list[dict]:
    async with async_session() as session:
        stmt = select(Moment).where(Moment.user_id == user_id)

        if person_id is not None:
            stmt = stmt.join(moment_people, moment_people.c.moment_id == Moment.id).where(
                moment_people.c.person_id == person_id
            )
        if place_id is not None:
            stmt = stmt.join(moment_places, moment_places.c.moment_id == Moment.id).where(
                moment_places.c.place_id == place_id
            )
        if tag is not None:
            stmt = stmt.where(Moment.tags.any(tag))
        if date_from is not None:
            stmt = stmt.where(Moment.day_date >= date_from)
        if date_to is not None:
            stmt = stmt.where(Moment.day_date <= date_to)

        if query is None:
            stmt = stmt.order_by(Moment.occurred_at.desc()).limit(limit)
            moments = (await session.execute(stmt)).scalars().all()
            return [m.to_dict() for m in moments]

        # Semantic mode. Pull a wider candidate set, then rank in Python.
        candidate_stmt = stmt.order_by(Moment.occurred_at.desc()).limit(limit * 5)
        candidates = (await session.execute(candidate_stmt)).scalars().all()
        if not candidates:
            return []

        candidate_ids = [m.id for m in candidates]
        emb_stmt = select(MomentEmbedding).where(
            MomentEmbedding.moment_id.in_(candidate_ids)
        )
        embeddings = {
            e.moment_id: e.embedding
            for e in (await session.execute(emb_stmt)).scalars()
        }

        query_vec = await get_embedding(query)
        scored = []
        for m in candidates:
            vec = embeddings.get(m.id)
            if vec is None:
                continue
            score = _cosine(query_vec, vec)
            if score >= threshold:
                row = m.to_dict()
                row["score"] = score
                scored.append(row)

        scored.sort(key=lambda r: r["score"], reverse=True)
        return scored[:limit]


async def _search_transcripts(
    *,
    user_id: int,
    query: str | None,
    date_from: datetime.date | None,
    date_to: datetime.date | None,
    limit: int,
) -> list[dict]:
    """Fallback substring search over raw journal Messages.

    Substring is intentional: transcripts catch what the LLM didn't extract
    as Moments. Semantic search over messages would require a third embedding
    index, which we deliberately don't maintain.
    """
    if query is None:
        return []
    async with async_session() as session:
        stmt = (
            select(Message, Conversation.day_date)
            .join(Conversation, Message.conversation_id == Conversation.id)
            .where(
                Conversation.user_id == user_id,
                Conversation.conversation_type == "journal",
                Message.content.ilike(f"%{query}%"),
            )
            .order_by(Message.created_at.desc())
            .limit(limit)
        )
        if date_from is not None:
            stmt = stmt.where(Conversation.day_date >= date_from)
        if date_to is not None:
            stmt = stmt.where(Conversation.day_date <= date_to)
        rows = (await session.execute(stmt)).all()
        return [
            {
                "id": None,
                "kind": "transcript",
                "message_id": msg.id,
                "conversation_id": msg.conversation_id,
                "day_date": day_date.isoformat() if day_date else None,
                "occurred_at": msg.created_at.isoformat(),
                "content": msg.content[:400],
                "raw_excerpt": None,
                "tags": [],
                "people": [],
                "places": [],
            }
            for msg, day_date in rows
        ]