beforeLocal

api/main.py

@@ -0,0 +1,95 @@
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from sqlalchemy.orm import Session
+from sqlalchemy import text
+from database import get_db, init_db, Email, EMBEDDING_DIMENSIONS
+from pydantic import BaseModel
+from typing import List, Optional
+from google import genai
+from google.genai import types
+import os
+import time
+
+app = FastAPI(title="Unified Email Semantic Search API")
+
+# Setup CORS for the SPA
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize Gemini client — reads GEMINI_API_KEY from environment
+gemini_client = None
+
+class SearchQuery(BaseModel):
+    query: str
+    limit: int = 10
+
+class SearchResult(BaseModel):
+    message_id: str
+    subject: str
+    sender: str
+    date: str
+    snippet: str
+    distance: float
+
+@app.on_event("startup")
+def on_startup():
+    global gemini_client
+    print("Initializing Database...")
+    time.sleep(2) # Give postgres a moment to be fully ready
+    try:
+        init_db()
+    except Exception as e:
+        print(f"Error initializing DB: {e}")
+
+    api_key = os.environ.get("GEMINI_API_KEY")
+    if api_key:
+        gemini_client = genai.Client(api_key=api_key)
+        print("Gemini client initialized.")
+    else:
+        print("WARNING: GEMINI_API_KEY not set. Embedding features disabled.")
+
+@app.post("/search", response_model=List[SearchResult])
+def search_emails(request: SearchQuery, db: Session = Depends(get_db)):
+    if not gemini_client:
+        raise HTTPException(status_code=500, detail="Gemini API Key is not configured.")
+
+    try:
+        response = gemini_client.models.embed_content(
+            model="gemini-embedding-001",
+            contents=request.query,
+            config=types.EmbedContentConfig(
+                task_type="RETRIEVAL_QUERY",
+                output_dimensionality=EMBEDDING_DIMENSIONS,
+            ),
+        )
+        query_embedding = response.embeddings[0].values
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Embedding API error: {e}")
+
+    # Use pgvector's cosine distance operator via SQLAlchemy ORM
+    results = db.query(
+        Email, 
+        Email.embedding.cosine_distance(query_embedding).label('distance')
+    ).order_by(
+        Email.embedding.cosine_distance(query_embedding)
+    ).limit(request.limit).all()
+
+    response_data = []
+    for email, distance in results:
+        # Create a snippet from the content
+        snippet = email.content[:200] + "..." if email.content and len(email.content) > 200 else (email.content or "")
+        response_data.append(SearchResult(
+            message_id=email.message_id or "",
+            subject=email.subject or "",
+            sender=email.sender or "",
+            date=email.date.isoformat() if email.date else "",
+            snippet=snippet,
+            distance=distance
+        ))
+
+    return response_data