api.py

"""
FastAPI Backend + Gradio UI
Servește RAG pipeline-ul ca API REST și interfață web.
"""

import shutil
from pathlib import Path
from contextlib import asynccontextmanager

import gradio as gr
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from config import DOCS_DIR, COLLECTION_NAME, CHROMA_DIR
from rag import RAGPipeline
from ingest import run_ingestion

# ── Global state ──────────────────────────────────────────

rag_pipeline: RAGPipeline | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Inițializează RAG pipeline-ul la pornirea serverului."""
    global rag_pipeline
    try:
        rag_pipeline = RAGPipeline()
    except Exception as e:
        print(f"⚠ RAG pipeline nu s-a putut inițializa: {e}")
        print("  Rulează mai întâi: python ingest.py")
    yield


app = FastAPI(
    title="RAG Tech Docs",
    description="Retrieval-Augmented Generation pentru documentație tehnică",
    version="1.0.0",
    lifespan=lifespan,
)


# ── Pydantic models ──────────────────────────────────────

class QueryRequest(BaseModel):
    question: str
    top_k: int = 5


class QueryResponse(BaseModel):
    question: str
    answer: str
    sources: list[dict]
    timing: dict


# ── API Endpoints ─────────────────────────────────────────

@app.get("/health")
def health_check():
    return {
        "status": "ok",
        "pipeline_ready": rag_pipeline is not None,
    }


@app.post("/query", response_model=QueryResponse)
def query_docs(req: QueryRequest):
    if rag_pipeline is None:
        raise HTTPException(status_code=503, detail="Pipeline-ul nu e inițializat. Rulează ingestia mai întâi.")
    result = rag_pipeline.query(req.question, req.top_k)
    return result


@app.post("/ingest")
def ingest_docs():
    """Re-indexează toate documentele din directorul docs/."""
    global rag_pipeline
    try:
        run_ingestion()
        rag_pipeline = RAGPipeline()
        return {"status": "ok", "message": "Documente re-indexate cu succes."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/upload")
async def upload_document(file: UploadFile = File(...)):
    """Încarcă un document nou în directorul docs/."""
    allowed = {".txt", ".md", ".pdf"}
    ext = Path(file.filename).suffix.lower()
    if ext not in allowed:
        raise HTTPException(status_code=400, detail=f"Format nesuportat: {ext}. Permis: {allowed}")

    DOCS_DIR.mkdir(parents=True, exist_ok=True)
    dest = DOCS_DIR / file.filename
    with open(dest, "wb") as f:
        shutil.copyfileobj(file.file, f)

    return {"status": "ok", "message": f"Fișierul {file.filename} a fost încărcat. Rulează /ingest pentru indexare."}


# ── Gradio Chat UI ────────────────────────────────────────

def chat_fn(message: str, history: list[dict]) -> str:
    """Funcția de chat pentru Gradio."""
    if rag_pipeline is None:
        return "⚠ Pipeline-ul nu e inițializat. Rulează mai întâi `python ingest.py` pentru a indexa documente."

    result = rag_pipeline.query(message)

    sources_text = "\n".join(
        f"  - **{s['source']}** (relevanță: {s['score']:.0%})"
        for s in result["sources"][:3]
    )

    response = f"{result['answer']}\n\n---\n📚 **Surse:**\n{sources_text}\n⏱ *{result['timing']['total_ms']}ms*"
    return response


gradio_app = gr.ChatInterface(
    fn=chat_fn,
    type="messages",
    title="📖 RAG Tech Docs",
    description="Pune întrebări despre documentația tehnică. Răspunsurile sunt generate din documente reale.",
    examples=[
        "How do I create a FastAPI application?",
        "What is dependency injection?",
        "How to handle errors?",
    ],
    theme=gr.themes.Soft(),
)

app = gr.mount_gradio_app(app, gradio_app, path="/chat")


# ── Main ──────────────────────────────────────────────────

if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("  RAG Tech Docs Server")
    print("  API:  http://localhost:8000/docs")
    print("  Chat: http://localhost:8000/chat")
    print("=" * 60 + "\n")
    uvicorn.run(app, host="0.0.0.0", port=8000)