rag_core.py

"""
rag_core.py
===========
Singleton facade over the whole `rag_pipeline` package.

Both `mcp_server.py` (stdio / MCP tools) and `main.py` (FastAPI REST) import
this module and call `await get_pipeline()` -- the first call builds every
dependency (LLM, embeddings, vector store, reranker, guardrails, cache,
compiled graph) exactly once; every later call returns the same
`RAGPipeline` instance, so the vector store / semantic cache / conversation
checkpointer stay shared no matter whether a request comes in via REST or
via an MCP tool call.
"""

from __future__ import annotations

import asyncio
import os
import uuid
from typing import Optional

from rag_pipeline.config import get_settings
from rag_pipeline.errors import make_retry_llm_call
from rag_pipeline.gateway import build_llm
from rag_pipeline.graph import build_graph
from rag_pipeline.guardrails import build_guardrails
from rag_pipeline.ingestion import load_and_split_corpus
from rag_pipeline.logging_setup import logger
from rag_pipeline.pii_masking import build_pii_engines
from rag_pipeline.reranker import build_reranker, retrieve_and_rerank
from rag_pipeline.semantic_cache import SemanticCache
from rag_pipeline.token_management import fit_docs_to_budget
from rag_pipeline.vectorstore import (
    build_embeddings,
    build_vectorstore,
    get_retriever,
    load_existing_vectorstore,
)


class RAGPipeline:
    """Thin orchestration layer: holds every built dependency and exposes a
    small, stable async API (`chat`, `search`, `get_history`, `cache_stats`)
    that callers (MCP tools, REST endpoints) use without needing to know
    about LangGraph, retrievers, or guardrails internals."""

    def __init__(self, *, llm, embeddings, retry_llm_call, base_retriever, reranker, guardrails, semantic_cache, analyzer, anonymizer, app):
        self.llm = llm
        self.embeddings = embeddings
        self.retry_llm_call = retry_llm_call
        self.base_retriever = base_retriever
        self.reranker = reranker
        self.guardrails = guardrails
        self.semantic_cache = semantic_cache
        self.analyzer = analyzer
        self.anonymizer = anonymizer
        self.app = app  # compiled LangGraph

    async def chat(self, question: str, thread_id: Optional[str] = None) -> dict:
        """Runs the full guarded, cached, reranked, hallucination-checked
        pipeline for one turn. `thread_id` ties this turn to prior turns'
        history via the LangGraph checkpointer; a new id is generated if
        none is given (i.e. a fresh conversation)."""
        thread_id = thread_id or str(uuid.uuid4())
        config = {"configurable": {"thread_id": thread_id}}

        try:
            existing = self.app.get_state(config)
            history = existing.values.get("history", []) if existing and existing.values else []

            logger.info("Pipeline invoked for query: %r (thread=%s)", question, thread_id)
            result = await self.app.ainvoke(
                {
                    "original_query": question,
                    "standalone_query": question,
                    "history": history,
                    "blocked": False,
                    "block_reason": None,
                    "cached_answer": None,
                    "context_docs": [],
                    "context_text": "",
                    "answer": "",
                    "hallucination_verdict": None,
                    "error": None,
                },
                config=config,
            )

            if result.get("error"):
                return {"thread_id": thread_id, "error": result["error"], "answer": result["answer"]}
            if result.get("blocked"):
                return {"thread_id": thread_id, "blocked": True, "reason": result["block_reason"]}
            if result.get("cached_answer"):
                return {"thread_id": thread_id, "answer": result["cached_answer"], "source": "cache"}
            return {
                "thread_id": thread_id,
                "answer": result["answer"],
                "hallucination_verdict": result["hallucination_verdict"],
                "source": "fresh",
            }
        except Exception as e:
            # Top-level safety net -- individual nodes already catch their own
            # errors via safe_node; this catches anything outside the graph
            # itself (e.g. checkpointer/state access issues).
            logger.error("Pipeline error (outside graph nodes): %s", e)
            return {"thread_id": thread_id, "error": str(e)}

    async def search(self, query: str, top_n: int = 5) -> list[dict]:
        """Retrieval + rerank only, no generation/guardrails -- for a calling
        agent that wants raw grounded context to reason over itself."""
        docs = await asyncio.to_thread(retrieve_and_rerank, query, self.base_retriever, self.reranker, top_n)
        fitted, _used_tokens = fit_docs_to_budget(docs)
        return [{"content": d.page_content, "metadata": d.metadata} for d in fitted]

    async def get_history(self, thread_id: str) -> list[list[str]]:
        config = {"configurable": {"thread_id": thread_id}}
        state = self.app.get_state(config)
        if not state or not state.values:
            return []
        return [[user, assistant] for user, assistant in state.values.get("history", [])]

    def cache_stats(self) -> dict:
        return self.semantic_cache.stats()


_pipeline: Optional[RAGPipeline] = None
_init_lock = asyncio.Lock()


async def _build_pipeline() -> RAGPipeline:
    settings = get_settings()
    settings.validate()

    llm = build_llm()
    retry_llm_call = make_retry_llm_call(llm)

    embeddings = await asyncio.to_thread(build_embeddings)

    # If the collection already exists on disk, load it instead of re-embedding
    # the whole corpus from scratch every time the process starts.
    collection_exists = os.path.isdir(settings.qdrant_path) and os.listdir(settings.qdrant_path)
    if collection_exists:
        logger.info("Existing Qdrant collection found at %s -- loading, skipping re-ingestion.", settings.qdrant_path)
        vectorstore = await asyncio.to_thread(load_existing_vectorstore, embeddings)
    else:
        logger.info("No existing Qdrant collection -- running ingestion + embedding (first run).")
        chunks = await asyncio.to_thread(load_and_split_corpus)
        vectorstore = await asyncio.to_thread(build_vectorstore, chunks, embeddings)

    base_retriever = get_retriever(vectorstore)
    reranker = await asyncio.to_thread(build_reranker)
    guardrails = build_guardrails(llm, settings.gemini_api_key_1, settings.gemini_model)
    semantic_cache = SemanticCache(embeddings, similarity_threshold=0.95, max_entries=200, ttl_seconds=3600)
    analyzer, anonymizer = build_pii_engines()

    app = build_graph(
        llm=llm,
        retry_llm_call=retry_llm_call,
        base_retriever=base_retriever,
        reranker=reranker,
        guardrails=guardrails,
        semantic_cache=semantic_cache,
        analyzer=analyzer,
        anonymizer=anonymizer,
    )

    return RAGPipeline(
        llm=llm,
        embeddings=embeddings,
        retry_llm_call=retry_llm_call,
        base_retriever=base_retriever,
        reranker=reranker,
        guardrails=guardrails,
        semantic_cache=semantic_cache,
        analyzer=analyzer,
        anonymizer=anonymizer,
        app=app,
    )


async def get_pipeline() -> RAGPipeline:
    """Lazily builds the pipeline on first call, returns the same instance
    afterwards. Guarded by a lock so two concurrent first-requests (e.g. a
    REST call and an MCP tool call landing at the same time on startup)
    don't trigger the (slow, rate-limited) build twice."""
    global _pipeline
    if _pipeline is not None:
        return _pipeline

    async with _init_lock:
        if _pipeline is None:  # re-check inside the lock
            _pipeline = await _build_pipeline()
        return _pipeline