Model/tools/knowledge_base.py

"""
Frox AI — RAG / Knowledge Base Tool

Local, self-contained document retrieval: chunk → embed (via Morph's
own model) → cosine-similarity search → return chunks with citations.
No vector DB required for this in-repo version — the production
backend architecture document specs a Qdrant-backed version with
hybrid BM25+vector search and a cross-encoder reranker (Section 6)
for when you outgrow an in-memory store; this module keeps the same
retrieval interface so swapping the backing store later doesn't
change how tools call it.
"""
from __future__ import annotations

import math
import re
import uuid
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from tools.registry import tool, ToolContext


@dataclass
class Chunk:
    id: str
    collection_id: str
    text: str
    source: str
    chunk_index: int
    embedding: List[float]


def _cosine(a: List[float], b: List[float]) -> float:
    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)


def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
    """
    Simple recursive-ish splitter: break on paragraph boundaries first,
    then fall back to sentence boundaries, packing up to chunk_size
    characters per chunk with a small overlap for context continuity.
    """
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    chunks: List[str] = []
    current = ""

    for para in paragraphs:
        if len(current) + len(para) <= chunk_size:
            current = f"{current}\n\n{para}".strip()
        else:
            if current:
                chunks.append(current)
            if len(para) <= chunk_size:
                current = para
            else:
                # paragraph itself too long — split on sentences
                sentences = re.split(r"(?<=[.!?])\s+", para)
                current = ""
                for sent in sentences:
                    if len(current) + len(sent) <= chunk_size:
                        current = f"{current} {sent}".strip()
                    else:
                        if current:
                            chunks.append(current)
                        current = sent
    if current:
        chunks.append(current)

    if overlap > 0 and len(chunks) > 1:
        overlapped = [chunks[0]]
        for i in range(1, len(chunks)):
            tail = chunks[i - 1][-overlap:]
            overlapped.append(f"{tail} {chunks[i]}".strip())
        chunks = overlapped

    return chunks


class KnowledgeBase:
    """
    In-memory (optionally per-collection) chunk store with cosine
    retrieval. One instance can hold multiple named collections so a
    single ToolContext.knowledge_base can serve several documents/
    projects without cross-contaminating search results.
    """

    def __init__(self):
        self._chunks: Dict[str, List[Chunk]] = {}   # collection_id -> chunks

    def ingest(
        self,
        collection_id: str,
        text: str,
        source: str,
        embed_fn,
        chunk_size: int = 500,
        overlap: int = 50,
    ) -> int:
        """Chunk, embed, and store a document's text. Returns chunk count."""
        pieces = chunk_text(text, chunk_size=chunk_size, overlap=overlap)
        self._chunks.setdefault(collection_id, [])

        for i, piece in enumerate(pieces):
            embedding = embed_fn(piece)
            self._chunks[collection_id].append(Chunk(
                id=str(uuid.uuid4()), collection_id=collection_id,
                text=piece, source=source, chunk_index=i, embedding=embedding,
            ))
        return len(pieces)

    def retrieve(self, collection_id: str, query_embedding: List[float], k: int = 5) -> List[Chunk]:
        chunks = self._chunks.get(collection_id, [])
        scored = sorted(chunks, key=lambda c: _cosine(c.embedding, query_embedding), reverse=True)
        return scored[:k]

    def collections(self) -> List[str]:
        return list(self._chunks.keys())

    def clear_collection(self, collection_id: str):
        self._chunks.pop(collection_id, None)


@tool(
    name="knowledge_ingest",
    description="Add a document's text to a knowledge-base collection for later retrieval",
    timeout=30.0,
)
def knowledge_ingest(ctx: ToolContext, text: str, source: str, collection_id: str = "default") -> dict:
    """
    Args:
        text: The document's raw text (already extracted — pair with
              the file_analysis tool for PDFs/DOCX/etc).
        source: A label for citations, e.g. a filename.
        collection_id: Which collection to add this document to.

    Plain `def`, not `async def`: engine.embed() is synchronous and
    GPU-bound, called once per chunk — thread-offloaded by the registry.
    """
    if ctx.knowledge_base is None:
        raise RuntimeError("No knowledge_base configured in ToolContext")
    if ctx.engine is None:
        raise RuntimeError("No engine configured in ToolContext (needed to embed chunks)")

    count = ctx.knowledge_base.ingest(
        collection_id, text, source, embed_fn=ctx.engine.embed,
    )
    return {"ingested": True, "source": source, "collection_id": collection_id, "chunks": count}


@tool(
    name="knowledge_search",
    description="Search a knowledge-base collection for relevant passages",
    timeout=15.0,
)
def knowledge_search(ctx: ToolContext, query: str, collection_id: str = "default", k: int = 5) -> dict:
    """
    Args:
        query: What to look for.
        collection_id: Which collection to search.
        k: Max number of passages to return.

    Plain `def`, not `async def`: engine.embed() is synchronous and
    GPU-bound — thread-offloaded by the registry.
    """
    if ctx.knowledge_base is None:
        raise RuntimeError("No knowledge_base configured in ToolContext")
    if ctx.engine is None:
        raise RuntimeError("No engine configured in ToolContext (needed to embed the query)")

    query_embedding = ctx.engine.embed(query)
    results = ctx.knowledge_base.retrieve(collection_id, query_embedding, k=k)

    return {
        "query": query,
        "collection_id": collection_id,
        "passages": [
            {"text": c.text, "source": c.source, "chunk_index": c.chunk_index}
            for c in results
        ],
    }