localRAG/server.py

"""
FastAPI server for Local RAG with chat GUI.
Run with: uvicorn server:app --reload
"""
import os
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from local_rag import LocalRAG
from personas import DEFAULT_PERSONA, PERSONAS, PUBLIC_FIELDS

OLLAMA_MODEL = "gemma4:12b"
OPENAI_MODEL = "gpt-6-luna"
# One RAG instance per (LLM provider, vector store), for on-the-fly provider switching.
# Each persona names its vector store in personas.py; personas sharing one share the instance.
rags = {}
for vectorstore in sorted({p["vectorstore"] for p in PERSONAS.values()}):
    rags[("ollama", vectorstore)] = LocalRAG(
        vectorstore_path=vectorstore,
        llm_provider="ollama",
        ollama_model=OLLAMA_MODEL,
        openai_model=OPENAI_MODEL,
    )
    if os.environ.get("OPENAI_API_KEY"):
        try:
            rags[("openai", vectorstore)] = LocalRAG(
                vectorstore_path=vectorstore,
                llm_provider="openai",
                ollama_model=OLLAMA_MODEL,
                openai_model=OPENAI_MODEL,
            )
        except Exception as e:
            print(f"OpenAI RAG not available for {vectorstore}: {e}")
openai_available = any(provider == "openai" for provider, _ in rags)

app = FastAPI(title="Local RAG Chat", version="1.0.0")
# Images for the chat GUI (researcher portrait, avatar)
app.mount("/visuals", StaticFiles(directory=Path(__file__).parent / "templates" / "visuals"), name="visuals")


class ChatMessage(BaseModel):
    role: str  # "user" | "assistant"
    content: str


class ChatRequest(BaseModel):
    message: str
    history: list[ChatMessage] = []  # previous turns for conversation context
    llm_provider: str = "ollama"  # "ollama" | "openai"
    persona: str = DEFAULT_PERSONA  # key in personas.PERSONAS


class RetrievedChunk(BaseModel):
    content: str
    source: str
    page: int | None
    score: float | None = None  # L2 distance from FAISS (lower = more similar)


class ChatResponse(BaseModel):
    answer: str
    error: str | None = None
    retrieved: list[RetrievedChunk] | None = None
    searches: list[str] | None = None  # queries the LLM sent to the corpus search tool


@app.get("/", response_class=HTMLResponse)
def chat_view():
    """Serve the chat GUI."""
    html_path = Path(__file__).parent / "templates" / "chat.html"
    if not html_path.exists():
        raise HTTPException(status_code=500, detail="Chat template not found")
    return HTMLResponse(content=html_path.read_text(encoding="utf-8"))


def _get_rag(provider: str, persona: str):
    """Return the RAG instance for the provider and the persona's vector store.
    Fall back to Ollama if OpenAI is unavailable."""
    vectorstore = PERSONAS[persona]["vectorstore"]
    return rags.get((provider, vectorstore)) or rags[("ollama", vectorstore)]


@app.post("/api/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    """Handle a chat message and return the RAG answer."""
    if not request.message or not request.message.strip():
        return ChatResponse(answer="", error="Message cannot be empty")
    if request.llm_provider == "openai" and not openai_available:
        return ChatResponse(answer="", error="OpenAI not configured. Set OPENAI_API_KEY.")
    if request.persona not in PERSONAS:
        return ChatResponse(answer="", error=f"Unknown persona: {request.persona}")
    persona = PERSONAS[request.persona]
    rag = _get_rag(request.llm_provider, request.persona)
    try:
        chat_history = [{"role": m.role, "content": m.content} for m in request.history]
        result = rag.query_with_history(
            request.message.strip(),
            chat_history=chat_history,
            system_prompt=persona["system_prompt"],
        )
        answer = result["answer"]
        retrieved = result.get("retrieved", [])
        searches = result.get("searches", [])

        # Server-side console trace: shorter chunk logs + raw LLM response
        if retrieved:
            print(f"\n[RAG] Retrieved {len(retrieved)} chunk(s)")
            for i, chunk in enumerate(retrieved):
                content = chunk.get("content", "")
                preview = (content[:150] + "...") if len(content) > 150 else content
                print(f"  [{i + 1}] {chunk.get('source', '')} p.{chunk.get('page', '?')} s={chunk.get('score')} | {preview!r}")
        else:
            print(f"\n[RAG] No corpus search for this message")
        provider_label = "OpenAI" if request.llm_provider == "openai" else "Ollama"
        model_name = OPENAI_MODEL if request.llm_provider == "openai" else OLLAMA_MODEL
        print(f"[RAG] LLM response ({provider_label} / {model_name}, {request.persona}):\n{answer}")

        return ChatResponse(answer=answer, retrieved=retrieved, searches=searches)
    except Exception as e:
        return ChatResponse(answer="", error=str(e))


@app.get("/api/health")
def health():
    """Health check and vector store status."""
    has_docs = all(rag.vectorstore is not None for rag in rags.values())
    return {"status": "ok", "vectorstore_loaded": has_docs}


@app.get("/api/providers")
def providers():
    """Return which LLM providers are available."""
    return {"ollama": True, "openai": openai_available}


@app.get("/api/personas")
def personas():
    """Return the selectable conversation partners (texts and images for the GUI)."""
    return {
        "default": DEFAULT_PERSONA,
        "personas": [{"id": key, **{f: p[f] for f in PUBLIC_FIELDS}} for key, p in PERSONAS.items()],
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)