Update main.py

main.py

@@ -1,111 +1,319 @@
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from typing import List
+
+from contextlib import asynccontextmanager
+from collections import OrderedDict
+import asyncio
 import os
 import shutil
 import uuid
 import re
-import inflect 
+import inflect
+from urllib.parse import urlparse
+from typing import List
+
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+
+import cloudinary
+import cloudinary.uploader
 import cloudinary.api
-from src.models import AIModelManager
-from src.cloud_db import CloudDB
+from pinecone import Pinecone, ServerlessSpec
+
+# ── Deferred imports so startup prints appear in order ────────────
+ai      = None   # set in lifespan
+p       = inflect.engine()
+
+# ── Semaphore: max concurrent AI inference jobs ────────────────────
+MAX_CONCURRENT_INFERENCES = int(os.getenv("MAX_CONCURRENT_INFERENCES", "6"))
+_inference_sem: asyncio.Semaphore
+
+# ── Simple LRU connection pools ───────────────────────────────────
+_pinecone_pool: OrderedDict  = OrderedDict()
+_cloudinary_pool: dict = {}
+_POOL_MAX = 64
+
+
+def _get_pinecone(api_key: str) -> Pinecone:
+    """Return a cached Pinecone client, creating one if needed."""
+    if api_key not in _pinecone_pool:
+        if len(_pinecone_pool) >= _POOL_MAX:
+            _pinecone_pool.popitem(last=False)   # evict oldest
+        _pinecone_pool[api_key] = Pinecone(api_key=api_key)
+    _pinecone_pool.move_to_end(api_key)          # refresh LRU order
+    return _pinecone_pool[api_key]
+
+
+def _configure_cloudinary(creds: dict) -> None:
+    """Configure cloudinary module only when needed, with simple caching."""
+    key = creds["cloud_name"]
+    if key not in _cloudinary_pool:
+        cloudinary.config(
+            cloud_name=creds["cloud_name"],
+            api_key=creds["api_key"],
+            api_secret=creds["api_secret"],
+        )
+        _cloudinary_pool[key] = True
 
-app = FastAPI()
+
+# ── Lifespan: load models once at startup ─────────────────────────
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global ai, _inference_sem
+    from src.models import AIModelManager
+
+    print("⏳ Loading AI models …")
+    loop = asyncio.get_event_loop()
+    ai = await loop.run_in_executor(None, AIModelManager)
+    _inference_sem = asyncio.Semaphore(MAX_CONCURRENT_INFERENCES)
+    print(f"✅ Ready! Max concurrent inference slots: {MAX_CONCURRENT_INFERENCES}")
+    yield
+    print("👋 Shutting down")
+
+
+app = FastAPI(lifespan=lifespan)
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  
+    allow_origins=["*"],        # tighten to your Vercel domain in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-print("Loading AI Models and Cloud DB...")
-ai = AIModelManager()
-db = CloudDB() 
-p = inflect.engine() 
-print("Ready!")
-
 os.makedirs("temp_uploads", exist_ok=True)
 
+
+# ── Helpers ────────────────────────────────────────────────────────
 def standardize_category_name(name: str) -> str:
-    clean_name = name.strip().lower()
-    clean_name = re.sub(r'\s+', '_', clean_name)
-    clean_name = re.sub(r'[^\w\s]', '', clean_name)
-    singular_name = p.singular_noun(clean_name)
-    return singular_name if singular_name else clean_name
+    clean = re.sub(r'\s+', '_', name.strip().lower())
+    clean = re.sub(r'[^\w]', '', clean)
+    return p.singular_noun(clean) or clean
+
 
 def sanitize_filename(filename: str) -> str:
-    clean_name = re.sub(r'\s+', '_', filename)
-    clean_name = re.sub(r'[^\w\.\-]', '', clean_name)
-    return clean_name
+    clean = re.sub(r'\s+', '_', filename)
+    return re.sub(r'[^\w.\-]', '', clean)
+
+
+def get_cloudinary_creds(env_url: str) -> dict:
+    parsed = urlparse(env_url)
+    return {
+        "api_key":    parsed.username,
+        "api_secret": parsed.password,
+        "cloud_name": parsed.hostname,
+    }
+
+
+# ══════════════════════════════════════════════════════════════════
+# 1. VERIFY KEYS & AUTO-BUILD INDEXES
+# ══════════════════════════════════════════════════════════════════
+@app.post("/api/verify-keys")
+async def verify_keys(
+    pinecone_key:    str = Form(""),
+    cloudinary_url:  str = Form(""),
+):
+    if cloudinary_url:
+        try:
+            creds = get_cloudinary_creds(cloudinary_url)
+            _configure_cloudinary(creds)
+            await asyncio.to_thread(cloudinary.api.ping)
+        except Exception:
+            raise HTTPException(400, "Invalid Cloudinary Environment URL.")
+
+    if pinecone_key:
+        try:
+            pc = _get_pinecone(pinecone_key)
+            existing = {idx.name for idx in await asyncio.to_thread(pc.list_indexes)}
+
+            tasks = []
+            if "lens-objects" not in existing:
+                tasks.append(asyncio.to_thread(
+                    pc.create_index,
+                    name="lens-objects", dimension=1536, metric="cosine",
+                    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+                ))
+            if "lens-faces" not in existing:
+                tasks.append(asyncio.to_thread(
+                    pc.create_index,
+                    name="lens-faces", dimension=512, metric="cosine",
+                    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+                ))
+
+            if tasks:
+                await asyncio.gather(*tasks)
+
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(400, f"Pinecone Error: {e}")
 
+    return {"message": "Keys verified and indexes ready!"}
+
+
+# ══════════════════════════════════════════════════════════════════
+# 2. UPLOAD (Cloudinary + Pinecone Only)
+# ══════════════════════════════════════════════════════════════════
 @app.post("/api/upload")
-async def upload_new_images(files: List[UploadFile] = File(...), folder_name: str = Form(...)):
+async def upload_new_images(
+    files:               List[UploadFile] = File(...),
+    folder_name:         str  = Form(...),
+    detect_faces:        bool = Form(True),
+    user_pinecone_key:   str  = Form(""),
+    user_cloudinary_url: str  = Form(""),
+):
+    if not user_pinecone_key or not user_cloudinary_url:
+        raise HTTPException(status_code=400, detail="Cloudinary URL and Pinecone API Key are required to upload.")
+
+    folder    = standardize_category_name(folder_name)
     uploaded_urls = []
-    standardized_folder = standardize_category_name(folder_name)
-    detect_faces: bool = Form(True)
-    
-    try:
-        for file in files:
-            safe_filename = sanitize_filename(file.filename)
-            temp_path = f"temp_uploads/{safe_filename}"
-            with open(temp_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
-                
-            image_url = db.upload_image(temp_path, standardized_folder)
-            
-            vectors_to_save = ai.process_image(temp_path, is_query=False, detect_faces=detect_faces)
-            
-            for vec_dict in vectors_to_save:
-                image_id = str(uuid.uuid4())
-                db.add_vector(vec_dict, image_url, image_id)
-            
-            os.remove(temp_path)
+
+    cld_creds = get_cloudinary_creds(user_cloudinary_url)
+    _configure_cloudinary(cld_creds)
+    pc        = _get_pinecone(user_pinecone_key)
+    idx_obj   = pc.Index("lens-objects")
+    idx_face  = pc.Index("lens-faces")
+
+    for file in files:
+        safe_name = sanitize_filename(file.filename)
+        tmp_path  = f"temp_uploads/{uuid.uuid4().hex}_{safe_name}"
+
+        try:
+            with open(tmp_path, "wb") as buf:
+                shutil.copyfileobj(file.file, buf)
+
+            # Upload image to CDN
+            result    = await asyncio.to_thread(cloudinary.uploader.upload, tmp_path, folder=folder)
+            image_url = result["secure_url"]
             uploaded_urls.append(image_url)
-            
-        return {"message": "Success!", "urls": uploaded_urls}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
 
+            # AI inference
+            async with _inference_sem:
+                vectors = await ai.process_image_async(tmp_path, is_query=False, detect_faces=detect_faces)
+
+            # Save vectors
+            face_upserts   = []
+            object_upserts = []
+
+            for v in vectors:
+                vec_list = v["vector"].tolist() if hasattr(v["vector"], "tolist") else v["vector"]
+                record   = {
+                    "id":     str(uuid.uuid4()),
+                    "values": vec_list,
+                    "metadata": {"url": image_url, "folder": folder},
+                }
+                (face_upserts if v["type"] == "face" else object_upserts).append(record)
+
+            # Fire both upserts concurrently
+            upsert_tasks = []
+            if face_upserts:
+                upsert_tasks.append(asyncio.to_thread(idx_face.upsert, vectors=face_upserts))
+            if object_upserts:
+                upsert_tasks.append(asyncio.to_thread(idx_obj.upsert, vectors=object_upserts))
+            if upsert_tasks:
+                await asyncio.gather(*upsert_tasks)
+
+        except Exception as e:
+            print(f"❌ Upload error for {file.filename}: {e}")
+            # Continue with the next file instead of aborting the whole batch
+        finally:
+            if os.path.exists(tmp_path):
+                os.remove(tmp_path)
+
+    return {"message": "Done!", "urls": uploaded_urls}
+
+
+# ══════════════════════════════════════════════════════════════════
+# 3. SEARCH (Pinecone Only)
+# ══════════════════════════════════════════════════════════════════
 @app.post("/api/search")
-async def search_database(file: UploadFile = File(...)):
+async def search_database(
+    file:                UploadFile = File(...),
+    detect_faces:        bool = Form(True),
+    user_pinecone_key:   str  = Form(""),
+    user_cloudinary_url: str  = Form(""), # Kept to match frontend form payload
+):
+    if not user_pinecone_key:
+        raise HTTPException(status_code=400, detail="Pinecone API Key is required to search.")
+
+    safe_name = sanitize_filename(file.filename)
+    tmp_path  = f"temp_uploads/query_{uuid.uuid4().hex}_{safe_name}"
+
     try:
-        safe_filename = sanitize_filename(file.filename)
-        temp_path = f"temp_uploads/query_{safe_filename}"
-        detect_faces: bool = Form(True)
-        
-        with open(temp_path, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
-            
-        vectors_to_search = ai.process_image(temp_path, is_query=True, detect_faces=detect_faces)
-        
-        all_results = []
-        for vec_dict in vectors_to_search:
-            results = db.search(vec_dict, top_k=10)
-            all_results.extend(results)
-            
-        os.remove(temp_path)
-        
-        unique_results = {}
+        with open(tmp_path, "wb") as buf:
+            shutil.copyfileobj(file.file, buf)
+
+        # AI inference
+        async with _inference_sem:
+            vectors = await ai.process_image_async(tmp_path, is_query=True, detect_faces=detect_faces)
+
+        pc       = _get_pinecone(user_pinecone_key)
+        idx_obj  = pc.Index("lens-objects")
+        idx_face = pc.Index("lens-faces")
+
+        # Fire ALL vector queries in parallel
+        async def _query_one(vec_dict: dict) -> list[dict]:
+            vec_list   = (vec_dict["vector"].tolist() if hasattr(vec_dict["vector"], "tolist") else vec_dict["vector"])
+            target_idx = idx_face if vec_dict["type"] == "face" else idx_obj
+            res = await asyncio.to_thread(
+                target_idx.query,
+                vector=vec_list, top_k=10, include_metadata=True,
+            )
+            out = []
+            for match in res.get("matches", []):
+                caption = ("👤 Verified Identity" if vec_dict["type"] == "face" else match["metadata"].get("folder", "🎯 Object Match"))
+                out.append({
+                    "url":     match["metadata"].get("url", ""),
+                    "score":   match["score"],
+                    "caption": caption,
+                })
+            return out
+
+        nested = await asyncio.gather(*[_query_one(v) for v in vectors])
+        all_results = [r for sub in nested for r in sub]
+
+        # Deduplicate, keep best score per URL
+        seen: dict[str, dict] = {}
         for r in all_results:
             url = r["url"]
-            if url not in unique_results or r["score"] > unique_results[url]["score"]:
-                unique_results[url] = r
-                
-        final_results = sorted(unique_results.values(), key=lambda x: x["score"], reverse=True)
-        
-        return {"results": final_results[:10]}
+            if url not in seen or r["score"] > seen[url]["score"]:
+                seen[url] = r
+
+        final = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:10]
+        return {"results": final}
+
     except Exception as e:
-        print(f"Production Search Error: {str(e)}") 
-        raise HTTPException(status_code=500, detail=str(e))
+        print(f"❌ Search error: {e}")
+        raise HTTPException(500, str(e))
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+
+
+# ══════════════════════════════════════════════════════════════════
+# 4. CATEGORIES (Cloudinary Folders Only)
+# ══════════════════════════════════════════════════════════════════
+@app.post("/api/categories")
+async def get_categories(user_cloudinary_url: str = Form("")):
+    if not user_cloudinary_url:
+        return {"categories": []}
     
-@app.get("/api/categories")
-async def get_categories():
     try:
-        result = cloudinary.api.root_folders()
-        folders = [folder["name"] for folder in result.get("folders", [])]
+        creds = get_cloudinary_creds(user_cloudinary_url)
+        _configure_cloudinary(creds)
+        result  = await asyncio.to_thread(cloudinary.api.root_folders)
+        folders = [f["name"] for f in result.get("folders", [])]
         return {"categories": folders}
     except Exception as e:
-        print(f"Error fetching categories from Cloudinary: {e}")
-        return {"categories": []}
+        print(f"Category fetch error: {e}")
+        return {"categories": []}
+
+
+# ══════════════════════════════════════════════════════════════════
+# 5. HEALTH CHECK
+# ══════════════════════════════════════════════════════════════════
+@app.get("/api/health")
+async def health():
+    return {
+        "status":    "ok",
+        "device":    ai.device if ai else "loading",
+        "sem_slots": _inference_sem._value if _inference_sem else 0,
+    }