opusdev/vector-similarity-api

1import os2 3from dotenv import load_dotenv4from fastapi import FastAPI, HTTPException5from pydantic import BaseModel6from sentence_transformers import SentenceTransformer7from typing import List8import uvicorn9 10load_dotenv()11 12app = FastAPI(13    title="Embedding Service",14    description="Microservice for generating text embeddings",15    version="1.0.0"16)17 18# Load model at startup19model_name = os.getenv('EMBEDDING_MODEL', 'all-MiniLM-L6-v2')20model_dir = os.getenv('EMBEDDING_MODEL_DIR', '')21cache_dir = os.getenv('HUGGINGFACE_CACHE', os.path.join(os.getcwd(), '.cache', 'huggingface'))22 23os.environ.setdefault('TRANSFORMERS_CACHE', cache_dir)24os.environ.setdefault('HF_HOME', cache_dir)25 26model_source = model_dir if model_dir else model_name27print(f"Loading embedding model from: {model_source}")28 29try:30    model = SentenceTransformer(model_source)31    print("model loaded successfully")32except Exception as e:33    print("Failed to load embedding model:", e)34    print("If your network blocks Hugging Face, download the model locally and set EMBEDDING_MODEL_DIR to that path.")35    raise36 37class EmbedRequest(BaseModel):38    text: str39 40class EmbedResponse(BaseModel):41    embedding: List[float]42    dimension: int43 44@app.post("/embed", response_model=EmbedResponse)45def generate_embedding(request: EmbedRequest):46    """Generate embedding for given text"""47    try:48        if not request.text.strip():49            raise HTTPException(status_code=400, detail="Text cannot be empty")50        51        embedding = model.encode(request.text).tolist()52        53        return EmbedResponse(54            embedding=embedding,55            dimension=len(embedding)56        )57    except Exception as e:58        raise HTTPException(status_code=500, detail=str(e))59 60@app.get("/health")61def health_check():62    """Health check endpoint"""63    return {64        "status": "healthy",65        "model": "all-MiniLM-L6-v2",66        "dimension": 38467    }68 69@app.get("/")70def root():71    return {72        "service": "Embedding Service",73        "model": "all-MiniLM-L6-v2",74        "endpoints": {75            "POST /embed": "Generate embedding for text",76            "GET /health": "Health check"77        }78    }79 80if __name__ == "__main__":81    print("\n" + "="*60)82    print(" Starting Embedding Service")83    print("="*60)84    print(" Server: http://localhost:8001")85    print(" Health: http://localhost:8001/health")86    print("="*60 + "\n")87    88    uvicorn.run(app, host="0.0.0.0", port=8001)89 90 91 92 93 94 95 96