import random
from typing import Any, Dict, List
from contextlib import asynccontextmanager
try:
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn
except ImportError:
print("Error: FastAPI not installed")
print("Install with: pip install fastapi uvicorn")
exit(1)
from vecstore import VecStore
store: VecStore = None
@asynccontextmanager
async def lifespan(app: FastAPI):
global store
print("Starting up... Creating VecStore")
store = VecStore("./fastapi_vecstore_db")
yield
print("Shutting down...")
app = FastAPI(
title="VecStore API",
description="REST API for vector similarity search",
version="0.1.0",
lifespan=lifespan,
)
class UpsertRequest(BaseModel):
id: str
vector: List[float]
metadata: Dict[str, Any]
class QueryRequest(BaseModel):
vector: List[float]
k: int = 10
filter: str | None = None
class SearchResultResponse(BaseModel):
id: str
score: float
metadata: Dict[str, Any]
class QueryResponse(BaseModel):
results: List[SearchResultResponse]
count: int
@app.get("/")
async def root():
return {
"status": "ok",
"message": "VecStore API is running",
"version": "0.1.0",
}
@app.post("/upsert")
async def upsert(request: UpsertRequest):
try:
store.upsert(request.id, request.vector, request.metadata)
return {
"status": "success",
"id": request.id,
"message": f"Vector '{request.id}' upserted successfully",
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
try:
results = store.query(
vector=request.vector,
k=request.k,
filter=request.filter,
)
response_results = [
SearchResultResponse(
id=r.id,
score=r.score,
metadata=r.metadata,
)
for r in results
]
return QueryResponse(
results=response_results,
count=len(response_results),
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/vector/{vector_id}")
async def delete_vector(vector_id: str):
try:
store.remove(vector_id)
return {
"status": "success",
"message": f"Vector '{vector_id}' deleted",
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/stats")
async def stats():
return {
"vector_count": len(store),
"is_empty": store.is_empty(),
}
@app.post("/optimize")
async def optimize():
try:
removed_count = store.optimize()
return {
"status": "success",
"removed_count": removed_count,
"message": f"Optimized store, removed {removed_count} entries",
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
def main():
print("=" * 60)
print("Starting VecStore FastAPI Server")
print("=" * 60)
print("\nAPI Documentation:")
print(" - Interactive docs: http://localhost:8000/docs")
print(" - OpenAPI schema: http://localhost:8000/openapi.json")
print("\nPress CTRL+C to stop\n")
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
log_level="info",
)
if __name__ == "__main__":
main()