FastAPI_server.py

"""
FastAPI_server.py — FastAPI server exposing tools for a local RAG index.

Purpose
-------
Provide a minimal "tool server" that a client (CLI/Streamlit/agent)
can call to:
  1) Ingest PDFs and CSVs into a persistent vector DB (Chroma)
  2) Search the index by semantic similarity and return top-k hits with metadata

Endpoints
---------
POST /tools/ingest_pdf    -> Body: { "path": "<local file>", "doc_type": "pdf|qna|..." }
POST /tools/ingest_csv    -> Body: { "path": "<local file>", "text_col": "optional", "doc_type": "csv|qna|..." }
POST /tools/search_docs   -> Body: { "query": "<text>", "k": 5 }  # returns content, score, source metadata
GET  /                    -> Simple health check

Run
---
pip install -r requirements.txt
uvicorn FastAPI_server:app --port 8799 --reload

Environment variables (optional)
--------------------------------
DB_DIR=./db
EMBED_MODEL=sentence-transformers/all-MiniLM-L6-v2

Notes
-----
- This server keeps secrets out of the client. If you later call external APIs,
  do it here and apply guardrails (validation, allowlists, rate limits).
- File *paths* must be reachable by this process. If the UI runs elsewhere,
  previews may not work locally, but ingestion/search still will.
- This is intentionally simple: no auth, basic CORS for Streamlit, no quotas.
"""

from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())  # loads ${workspaceFolder}/.env by default

import os
from typing import Optional, List

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, conint

# LangChain: loaders (PDF/CSV), splitter, embeddings, vector store (Chroma)
from langchain_community.document_loaders import PyPDFLoader, CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma


# =============================================================================
# Configuration
# =============================================================================

# Where Chroma will persist its data (collection + index). Delete this dir to reset.
DB_DIR = os.getenv("DB_DIR", "./db")

# Local embedding model (small, fast, downloaded on first use). Swap as needed.
EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")


def get_vector_store() -> Chroma:
    """
    Create (or reopen) the persistent Chroma collection.

    We build the embeddings object once per call; for higher throughput you can
    cache it globally. For local demos this is fine.
    """
    embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
    return Chroma(
        collection_name="insights",            # collection name (namespace)
        embedding_function=embeddings,         # how documents are embedded
        persist_directory=DB_DIR,              # on-disk storage for persistence
    )


# =============================================================================
# Pydantic Schemas (validate inputs / shape responses)
# =============================================================================

class IngestPDFReq(BaseModel):
    path: str = Field(..., description="Local path to a PDF file.")
    doc_type: str = Field(
        default="pdf",
        description="Tag for filtering in clients (e.g., qna, whitepaper, ticket, pdf)."
    )


class IngestCSVReq(BaseModel):
    path: str = Field(..., description="Local path to a CSV file.")
    text_col: Optional[str] = Field(
        default=None,
        description="Optional column containing text; if None, the loader concatenates row cells."
    )
    doc_type: str = Field(
        default="csv",
        description="Tag for filtering in clients (e.g., qna, whitepaper, ticket, csv)."
    )


class SearchReq(BaseModel):
    query: str = Field(..., description="Natural language query.")
    k: conint(ge=1, le=50) = Field(5, description="Top-k results to return (1..50).")


class SearchHit(BaseModel):
    # What we return to the client for each result
    content: str = Field(..., description="Text chunk content.")
    score: float = Field(..., description="Relevance score (lower is more similar in Chroma).")
    source: dict = Field(..., description="Normalized metadata (source path, page/row, doc_type).")


class SearchResp(BaseModel):
    hits: List[SearchHit]


# =============================================================================
# FastAPI App
# =============================================================================

app = FastAPI(
    title="FastAPI RAG Server",
    version="1.0.0",
    description="Minimal tool server exposing ingest/search over a local Chroma index."
)

# --- CORS for Streamlit (and local dev UIs) ---
# Allow Streamlit origins by default; override by setting FASTAPI_CORS_ORIGINS as a comma-separated list.
_default_origins = ["http://localhost:8501", "http://127.0.0.1:8501"]
env_origins = os.getenv("FASTAPI_CORS_ORIGINS")
allow_origins = [o.strip() for o in env_origins.split(",")] if env_origins else _default_origins

app.add_middleware(
    CORSMiddleware,
    allow_origins=allow_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Reusable chunker. Choose sizes that fit your model context window and typical doc structure.
_SPLITTER = RecursiveCharacterTextSplitter(
    chunk_size=250,     # ~1000 chars per chunk
    chunk_overlap=100   # keep overlap to preserve context across boundaries
)


def _normalize_metadata(md: dict | None) -> dict:
    """
    Ensure common metadata keys exist so client UIs can render consistently.

    Keys:
      - source: path to original file (string or None)
      - page: PDF page index (int or None)
      - row: CSV row index (int or None)
      - doc_type: tag set at ingest (e.g., 'qna', 'pdf', 'csv')
    """
    md = md or {}
    for key in ("source", "page", "row", "doc_type"):
        md.setdefault(key, None)
    return md


# =============================================================================
# Endpoints
# =============================================================================

@app.post("/tools/ingest_pdf")
def ingest_pdf(req: IngestPDFReq):
    """
    Ingest a local PDF into the vector store.

    - Uses PyPDFLoader to extract page text (no OCR).
    - Normalizes metadata and tags chunks with 'doc_type'.
    - Splits pages into chunks before adding to Chroma.
    """
    if not os.path.exists(req.path):
        # 400 = client error (bad path). Clients can fix and retry.
        raise HTTPException(status_code=400, detail=f"PDF not found: {req.path}")

    # Load the PDF; loader adds page numbers and source path in metadata
    docs = PyPDFLoader(req.path).load()

    # Normalize + tag metadata so clients (CLI/Streamlit) can show consistent badges
    for d in docs:
        d.metadata = _normalize_metadata(d.metadata)
        d.metadata["doc_type"] = req.doc_type

    # Split -> embed -> add to Chroma (persist to disk)
    chunks = _SPLITTER.split_documents(docs)
    vs = get_vector_store()
    vs.add_documents(chunks)
    vs.persist()

    return {"ok": True, "chunks_added": len(chunks)}


@app.post("/tools/ingest_csv")
def ingest_csv(req: IngestCSVReq):
    """
    Ingest a local CSV into the vector store.

    - If 'text_col' is provided, each row uses that column as its text.
    - If not, the loader concatenates row values into a single text string.
    - Metadata will contain 'row' indices and 'source' path if available.
    """
    if not os.path.exists(req.path):
        raise HTTPException(status_code=400, detail=f"CSV not found: {req.path}")

    # CSVLoader handles both text_col and "concat-row" modes via source_column=None
    loader = CSVLoader(file_path=req.path, csv_args={"delimiter": ","}, source_column=req.text_col)
    docs = loader.load()

    # Normalize + tag metadata
    for d in docs:
        d.metadata = _normalize_metadata(d.metadata)
        d.metadata["doc_type"] = req.doc_type

    # Split -> embed -> add to Chroma (persist to disk)
    chunks = _SPLITTER.split_documents(docs)
    vs = get_vector_store()
    vs.add_documents(chunks)
    vs.persist()

    return {"ok": True, "chunks_added": len(chunks)}


@app.post("/tools/search_docs", response_model=SearchResp)
def search_docs(req: SearchReq):
    """
    Semantic search with relevance scores.

    Returns a list of hits (content + score + normalized metadata).
    Client UIs can show:
      - badges (doc_type / page / row)
      - a preview pane (e.g., PDF page text, CSV row)
    """
    vs = get_vector_store()

    # Chroma returns (Document, score). Score semantics:
    # For similarity_search_with_relevance_scores, *lower is closer* (cosine distance).
    results = vs.similarity_search_with_relevance_scores(req.query, k=req.k)

    hits: list[SearchHit] = []
    for doc, score in results:
        hits.append(
            SearchHit(
                content=doc.page_content,
                score=float(score),
                source={
                    "source": doc.metadata.get("source"),
                    "page": doc.metadata.get("page"),
                    "row": doc.metadata.get("row"),
                    "doc_type": doc.metadata.get("doc_type"),
                },
            )
        )

    return SearchResp(hits=hits)


@app.get("/")
def health():
    """
    Basic health check so you can 'ping' the server or use it in a readiness probe.
    """
    return {"ok": True, "message": "FastAPI RAG server is running", "db_dir": DB_DIR, "embed_model": EMBED_MODEL}