knowledge_agent/knowledge_api.py

"""
knowledge_api.py — Public interface for the Knowledge Agent
============================================================
This module is the integration layer for the Master Orchestrator and any
other agent that needs to query your personal document store.

Usage from another agent:
    from knowledge_agent.knowledge_api import query_knowledge, index_docs

It is intentionally stateless at the call level — the KnowledgeIndexer
caches the ChromaDB connection internally for performance.
"""

from typing import List, Dict, Optional
from indexer import KnowledgeIndexer
from llm import generate_answer, stream_answer

# Module-level singleton — one DB connection shared across all calls
# (safe for single-process use; FastAPI uses one process by default)
_indexer: Optional[KnowledgeIndexer] = None


def _get_indexer() -> KnowledgeIndexer:
    global _indexer
    if _indexer is None:
        _indexer = KnowledgeIndexer()
    return _indexer


# ── Core public functions ─────────────────────────────────────────────────────

def query_knowledge(
    question: str,
    top_k: int = 5,
    return_sources: bool = True,
) -> Dict:
    """
    PRIMARY ENTRY POINT for the Master Orchestrator and other agents.

    Ask a question against your indexed personal documents.

    Args:
        question:       Natural language question.
        top_k:          Number of chunks to retrieve from ChromaDB.
        return_sources: If True, include source filenames in the response.

    Returns:
        {
            "answer":  str,          # LLM-generated, grounded answer
            "sources": [str, ...],   # source filenames (if return_sources=True)
            "chunks_used": int,      # number of relevant chunks found
            "question": str,         # echo back for logging
        }

    Example:
        >>> result = query_knowledge("What is our refund policy?")
        >>> print(result["answer"])
    """
    indexer = _get_indexer()
    chunks  = indexer.query(question, top_k=top_k)

    if not chunks:
        return {
            "answer":      "No documents have been indexed yet. "
                           "Please run `index_docs()` first.",
            "sources":     [],
            "chunks_used": 0,
            "question":    question,
        }

    result = generate_answer(question, chunks)
    result["question"] = question

    if not return_sources:
        result.pop("sources", None)

    return result


def index_docs(docs_dir: str = "./documents", force: bool = False):
    """
    Index (or re-index) all supported documents in docs_dir.
    Called by the CLI, the web UI, and can be called by the Orchestrator
    to trigger a refresh after new documents are added.

    Args:
        docs_dir: Path to the folder containing your documents.
        force:    Re-index everything even if files haven't changed.
    """
    indexer = _get_indexer()
    indexer.index_directory(docs_dir=docs_dir, force=force)


def index_single(path: str):
    """Index a single file immediately (called by file-watcher or API upload)."""
    indexer = _get_indexer()
    indexer.index_single_file(path)


def get_knowledge_stats() -> Dict:
    """Return stats about the current knowledge base (for health checks)."""
    return _get_indexer().get_stats()


def list_indexed_docs() -> List[Dict]:
    """List all documents currently indexed."""
    return _get_indexer().list_documents()


def delete_doc(filename: str):
    """Remove a document and all its chunks from the index."""
    _get_indexer().delete_document(filename)


def stream_knowledge_answer(question: str, top_k: int = 5):
    """
    Streaming variant — yields tokens one by one.
    Used by the FastAPI /stream endpoint for the web UI.
    """
    indexer = _get_indexer()
    chunks  = indexer.query(question, top_k=top_k)
    if not chunks:
        yield "No documents indexed yet."
        return
    yield from stream_answer(question, chunks)