import os from fastapi import FastAPI, Query from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles from dotenv import load_dotenv from db import get_connection from embedder import embed_text load_dotenv() PHOTOS_DIR = os.getenv("PHOTOS_DIR") FRONTEND_DIR = os.path.join(os.path.dirname(__file__), "../../pgvector-demo/frontend") app = FastAPI() app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]) app.mount("/ui", StaticFiles(directory=os.path.abspath(FRONTEND_DIR), html=True), name="ui") @app.get("/search") def search(q: str = Query(...), limit: int = Query(12)): vec = embed_text(q) conn = get_connection() cur = conn.cursor() cur.execute( """ SELECT filename, 1 - (embedding <=> %s::vector) AS score FROM images ORDER BY embedding <=> %s::vector LIMIT %s """, (vec, vec, limit), ) rows = cur.fetchall() cur.close() conn.close() return [{"filename": r[0], "score": round(r[1], 4)} for r in rows] @app.get("/stats") def stats(): conn = get_connection() cur = conn.cursor() cur.execute("SELECT COUNT(*) FROM images") count = cur.fetchone()[0] cur.close() conn.close() return {"count": count} @app.get("/photos/{filename}") def get_photo(filename: str): path = os.path.join(PHOTOS_DIR, filename) return FileResponse(path, media_type="image/jpeg")