FabledGrimoire/src/scribe/services/knowledge.py

"""Knowledge service — unified query across notes, tasks, plans, and processes."""
import logging

from sqlalchemy import func, select

from scribe.models import async_session
from scribe.models.note import Note

logger = logging.getLogger(__name__)

_SNIPPET_LEN = 200


def _note_to_item(note: Note) -> dict:
    item: dict = {
        "id": note.id,
        "note_type": note.note_type or "note",
        "title": note.title,
        "snippet": (note.body or "")[:_SNIPPET_LEN],
        "tags": note.tags or [],
        "project_id": note.project_id,
        "created_at": note.created_at.isoformat(),
        "updated_at": note.updated_at.isoformat(),
    }
    # Task fields — override note_type and add status/priority/due_date
    if note.is_task:
        item["note_type"] = "task"
        item["task_kind"] = note.task_kind
        item["status"] = note.status
        item["priority"] = note.priority
        item["due_date"] = note.due_date.isoformat() if note.due_date else None

    return item


def _apply_type_filter(stmt, note_type: str | None):
    """Apply the type facet to a Note select.

    'task' = any task (status not null); 'plan' = a task with task_kind='plan';
    any other non-empty type = a non-task note of that note_type; None = all.

    Trashed rows (deleted_at set) are always excluded.
    """
    stmt = stmt.where(Note.deleted_at.is_(None))
    if note_type == "task":
        return stmt.where(Note.status.isnot(None))
    if note_type == "plan":
        return stmt.where(Note.status.isnot(None)).where(Note.task_kind == "plan")
    if note_type:
        return stmt.where(Note.note_type == note_type).where(Note.status.is_(None))
    return stmt


async def query_knowledge(
    user_id: int,
    note_type: str | None,
    tags: list[str],
    sort: str,
    q: str | None,
    limit: int,
    offset: int,
) -> tuple[list[dict], int]:
    """Query knowledge objects (non-task notes) with filters.

    Returns (items, total_count).
    """
    # Semantic search path — scores take priority over sort
    if q:
        return await _semantic_knowledge_search(
            user_id, q, note_type=note_type, tags=tags, limit=limit, offset=offset
        )

    async with async_session() as session:
        base = select(Note).where(Note.user_id == user_id)

        base = _apply_type_filter(base, note_type)

        for tag in tags:
            base = base.where(Note.tags.contains([tag]))

        # Count before pagination
        count_stmt = select(func.count()).select_from(base.subquery())
        total: int = (await session.execute(count_stmt)).scalar_one()

        # Apply sort
        if sort == "created":
            base = base.order_by(Note.created_at.desc())
        elif sort == "alpha":
            base = base.order_by(Note.title.asc())
        elif sort == "type":
            base = base.order_by(Note.note_type.asc(), Note.updated_at.desc())
        else:  # modified (default)
            base = base.order_by(Note.updated_at.desc())

        rows = list((await session.execute(base.limit(limit).offset(offset))).scalars().all())

    return [_note_to_item(n) for n in rows], total


async def _semantic_knowledge_search(
    user_id: int,
    q: str,
    note_type: str | None,
    tags: list[str],
    limit: int,
    offset: int,
) -> tuple[list[dict], int]:
    """Hybrid search: keyword matches first (title/body ILIKE), then semantic results.

    Exact keyword matches always rank above semantic-only matches so that
    searching for a name like "Weston" surfaces the note with that title
    before conceptually related notes.

    BEST-EFFORT TOP-N, not exhaustive pagination: the ranked candidate set is
    capped (keyword limit*2 + up to ~200 semantic), so `total` is the size of
    that window, NOT the true match count, and matches beyond the cap are not
    reachable by paging. Each page also recomputes the full merge (O(corpus)
    per page). Acceptable for an interactive "best results" feed; a cached
    ranked-id list or pgvector ORDER BY/LIMIT is the fix if exhaustive,
    cheap pagination is ever needed.
    """
    # 1. Keyword search — title and body ILIKE
    keyword_notes: list[Note] = []
    try:
        async with async_session() as session:
            pattern = f"%{q}%"
            base = (
                select(Note)
                .where(Note.user_id == user_id)
                .where(Note.title.ilike(pattern) | Note.body.ilike(pattern))
            )
            base = _apply_type_filter(base, note_type)
            for tag in tags:
                base = base.where(Note.tags.contains([tag]))
            # Title matches first, then body-only matches, newest first within each
            base = base.order_by(
                Note.title.ilike(pattern).desc(),
                Note.updated_at.desc(),
            ).limit(limit * 2)
            keyword_notes = list((await session.execute(base)).scalars().all())
    except Exception:
        logger.warning("Keyword search failed", exc_info=True)

    # 2. Semantic search — conceptual similarity
    semantic_notes: list[Note] = []
    try:
        from scribe.services.embeddings import semantic_search_notes
        is_task_filter = True if note_type in ("task", "plan") else (False if note_type else None)
        candidates = await semantic_search_notes(
            user_id=user_id,
            query=q,
            limit=min(200, limit * 4),
            threshold=0.3,
            is_task=is_task_filter,
        )
        for _score, note in candidates:
            if note.deleted_at is not None:
                continue
            if note_type == "task" and not note.is_task:
                continue
            elif note_type == "plan" and (not note.is_task or note.task_kind != "plan"):
                continue
            elif note_type and note_type not in ("task", "plan") and note.note_type != note_type:
                continue
            if tags and not all(t in (note.tags or []) for t in tags):
                continue
            semantic_notes.append(note)
    except Exception:
        logger.warning("Semantic search unavailable, using keyword results only", exc_info=True)

    # 3. Merge — keyword matches first, then semantic (deduplicated)
    seen_ids: set[int] = set()
    merged: list[Note] = []
    for note in keyword_notes:
        if note.id not in seen_ids:
            seen_ids.add(note.id)
            merged.append(note)
    for note in semantic_notes:
        if note.id not in seen_ids:
            seen_ids.add(note.id)
            merged.append(note)

    total = len(merged)
    page_items = merged[offset: offset + limit]
    return [_note_to_item(n) for n in page_items], total


async def get_knowledge_tags(user_id: int, note_type: str | None = None) -> list[str]:
    """Return all distinct tags used across knowledge objects for this user."""
    async with async_session() as session:
        base = (
            select(func.unnest(Note.tags).label("tag"))
            .where(Note.user_id == user_id)
        )
        base = _apply_type_filter(base, note_type)
        stmt = base.distinct().order_by("tag")
        rows = list((await session.execute(stmt)).scalars().all())
    return [r for r in rows if r]


async def get_knowledge_counts(user_id: int, tags: list[str] | None = None) -> dict[str, int]:
    """Return per-type count of knowledge objects for the sidebar display."""
    async with async_session() as session:
        # Count non-task types
        stmt = (
            select(Note.note_type, func.count(Note.id))
            .where(Note.user_id == user_id)
            .where(Note.status.is_(None))
            .where(Note.deleted_at.is_(None))
            .where(Note.note_type.in_(["note", "process"]))
            .group_by(Note.note_type)
        )
        if tags:
            for tag in tags:
                stmt = stmt.where(Note.tags.contains([tag]))
        rows = list((await session.execute(stmt)).all())
        counts = {row[0]: row[1] for row in rows}

        # Count tasks separately (is_task = status IS NOT NULL)
        task_stmt = (
            select(func.count(Note.id))
            .where(Note.user_id == user_id)
            .where(Note.status.isnot(None))
            .where(Note.deleted_at.is_(None))
        )
        if tags:
            for tag in tags:
                task_stmt = task_stmt.where(Note.tags.contains([tag]))
        task_count: int = (await session.execute(task_stmt)).scalar_one()
        counts["task"] = task_count

        # Plans are a subset of tasks (task_kind='plan'); counted for the facet
        # but NOT added to total to avoid double-counting against "task".
        plan_stmt = (
            select(func.count(Note.id))
            .where(Note.user_id == user_id)
            .where(Note.status.isnot(None))
            .where(Note.task_kind == "plan")
            .where(Note.deleted_at.is_(None))
        )
        if tags:
            for tag in tags:
                plan_stmt = plan_stmt.where(Note.tags.contains([tag]))
        counts["plan"] = (await session.execute(plan_stmt)).scalar_one()

    for t in ("note", "task", "plan", "process"):
        counts.setdefault(t, 0)
    counts["total"] = sum(counts[t] for t in ("note", "task", "process"))
    return counts


async def query_knowledge_ids(
    user_id: int,
    note_type: str | None,
    tags: list[str],
    sort: str,
    q: str | None,
    limit: int = 100,
    offset: int = 0,
) -> tuple[list[int], int]:
    """Return note IDs only — cheap query for the two-tier pagination feed."""
    if q:
        # Re-use semantic search, extract IDs in rank order
        items, total = await _semantic_knowledge_search(
            user_id, q, note_type=note_type, tags=tags,
            limit=limit, offset=offset,
        )
        return [item["id"] for item in items], total

    async with async_session() as session:
        base = select(Note.id).where(Note.user_id == user_id)

        base = _apply_type_filter(base, note_type)
        for tag in tags:
            base = base.where(Note.tags.contains([tag]))

        count_stmt = select(func.count()).select_from(base.subquery())
        total: int = (await session.execute(count_stmt)).scalar_one()

        if sort == "created":
            base = base.order_by(Note.created_at.desc())
        elif sort == "alpha":
            base = base.order_by(Note.title.asc())
        elif sort == "type":
            base = base.order_by(Note.note_type.asc(), Note.updated_at.desc())
        else:
            base = base.order_by(Note.updated_at.desc())

        ids = list((await session.execute(base.limit(limit).offset(offset))).scalars().all())

    return ids, total


async def get_knowledge_by_ids(user_id: int, ids: list[int]) -> list[dict]:
    """Fetch full items for the given IDs, preserving the requested order."""
    if not ids:
        return []
    async with async_session() as session:
        stmt = (
            select(Note)
            .where(Note.user_id == user_id)
            .where(Note.id.in_(ids))
            .where(Note.deleted_at.is_(None))
        )
        rows = list((await session.execute(stmt)).scalars().all())
    by_id = {n.id: n for n in rows}
    return [_note_to_item(by_id[i]) for i in ids if i in by_id]