app.py

"""
app.py — FastAPI server for the Vector DB Explorer UI.

Run:  python app.py
Then open http://127.0.0.1:8000
"""

from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
import uvicorn

from vector_db import VectorDB

DB_PATH = "data.json"
db = VectorDB(db_path=DB_PATH)


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Pre-load the embedding model at startup so the first search is instant.
    await run_in_threadpool(db._get_model)
    yield


app = FastAPI(title="Vector DB Explorer", lifespan=lifespan)


# ── Request / Response models ─────────────────────────────────────────

class SearchRequest(BaseModel):
    query: str
    top_k: int = 3


class AddRequest(BaseModel):
    text: str
    category: str = "general"
    topic: str = ""


# ── Routes ────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def root():
    # Read on every request so the file can be edited without restarting.
    return Path("static/index.html").read_text(encoding="utf-8")


@app.get("/api/documents")
async def get_documents():
    return [
        {"id": r["id"], "text": r["text"], "metadata": r["metadata"]}
        for r in db._records
    ]


@app.post("/api/search")
async def search(req: SearchRequest):
    if not req.query.strip():
        raise HTTPException(status_code=400, detail="Query cannot be empty")
    # Embedding is CPU-bound; run in a thread so the event loop stays free.
    results = await run_in_threadpool(db.search, req.query.strip(), req.top_k)
    return results


@app.post("/api/add")
async def add_document(req: AddRequest):
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    await run_in_threadpool(db.add, req.text.strip(), {"category": req.category, "topic": req.topic})
    db.save()
    rec = db._records[-1]
    return {"id": rec["id"], "text": rec["text"], "metadata": rec["metadata"]}


if __name__ == "__main__":
    import os
    port   = int(os.environ.get("PORT", 8000))
    is_dev = os.environ.get("ENV") != "production"
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=is_dev)