Spaces:
Sleeping
Sleeping
File size: 2,551 Bytes
b16a546 322102c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
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)
|