Update main.py

main.py

@@ -1,6 +1,4 @@
-
-from contextlib import asynccontextmanager
-from collections import OrderedDict
+# main.py
 import asyncio
 import os
 import shutil
@@ -9,28 +7,30 @@ import re
 import inflect
 from urllib.parse import urlparse
 from typing import List
+from contextlib import asynccontextmanager
+from collections import OrderedDict
 
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-
 import cloudinary
 import cloudinary.uploader
 import cloudinary.api
 from pinecone import Pinecone, ServerlessSpec
 
-# ── Deferred imports so startup prints appear in order ────────────
-ai      = None   # set in lifespan
-p       = inflect.engine()
+# ── Deferred imports ─────────────────────────────────────────────
+ai = None
+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 = {}
+_pinecone_pool = OrderedDict()
+_cloudinary_pool = {}
 _POOL_MAX = 64
 
+# FIX 1: Restored your original Pinecone Index names!
+IDX_FACES = "enterprise-faces"
+IDX_OBJECTS = "enterprise-objects"
 
 def _get_pinecone(api_key: str) -> Pinecone:
     if api_key not in _pinecone_pool:
@@ -40,279 +40,198 @@ def _get_pinecone(api_key: str) -> Pinecone:
     _pinecone_pool.move_to_end(api_key)
     return _pinecone_pool[api_key]
 
-
-def _configure_cloudinary(creds: dict) -> None:
+def _configure_cloudinary(creds: dict):
     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"],
+            cloud_name=creds["cloud_name"], 
+            api_key=creds["api_key"], 
+            api_secret=creds["api_secret"]
         )
         _cloudinary_pool[key] = True
 
-
 @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}")
+    print("✅ Ready!")
     yield
-    print("👋 Shutting down")
-
 
 app = FastAPI(lifespan=lifespan)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 os.makedirs("temp_uploads", exist_ok=True)
 
-
 def standardize_category_name(name: str) -> str:
     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 = re.sub(r'\s+', '_', filename)
-    return re.sub(r'[^\w.\-]', '', clean)
-
+    return re.sub(r'[^\w.\-]', '', re.sub(r'\s+', '_', filename))
 
 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,
-    }
+    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(""),
-):
+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)
+            _configure_cloudinary(get_cloudinary_creds(cloudinary_url))
             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 IDX_OBJECTS not in existing:
+                tasks.append(asyncio.to_thread(pc.create_index, name=IDX_OBJECTS, dimension=1536, metric="cosine", spec=ServerlessSpec(cloud="aws", region="us-east-1")))
+            if IDX_FACES not in existing:
+                tasks.append(asyncio.to_thread(pc.create_index, name=IDX_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 (With Demo Fallback)
+# 2. UPLOAD (Strictly Cloud-Native + Freemium Default)
 # ══════════════════════════════════════════════════════════════════
 @app.post("/api/upload")
-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(""),
-):
-    # FALLBACK LOGIC: Use user keys if provided, otherwise use Space secrets
+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("")):
+    # FIX 2: Uses user keys if provided, otherwise falls back to HF Secrets
     actual_pc_key = user_pinecone_key or os.getenv("DEFAULT_PINECONE_KEY")
     actual_cld_url = user_cloudinary_url or os.getenv("DEFAULT_CLOUDINARY_URL")
-
+    
     if not actual_pc_key or not actual_cld_url:
-        raise HTTPException(status_code=400, detail="Cloudinary URL and Pinecone API Key are required.")
+        raise HTTPException(400, "API Keys are required. Set defaults in HF Secrets.")
 
-    folder    = standardize_category_name(folder_name)
+    folder = standardize_category_name(folder_name)
     uploaded_urls = []
-
-    cld_creds = get_cloudinary_creds(actual_cld_url)
-    _configure_cloudinary(cld_creds)
-    pc        = _get_pinecone(actual_pc_key)
-    idx_obj   = pc.Index("lens-objects")
-    idx_face  = pc.Index("lens-faces")
+    _configure_cloudinary(get_cloudinary_creds(actual_cld_url))
+    pc = _get_pinecone(actual_pc_key)
+    idx_obj = pc.Index(IDX_OBJECTS)
+    idx_face = pc.Index(IDX_FACES)
 
     for file in files:
-        safe_name = sanitize_filename(file.filename)
-        tmp_path  = f"temp_uploads/{uuid.uuid4().hex}_{safe_name}"
-
+        tmp_path = f"temp_uploads/{uuid.uuid4().hex}_{sanitize_filename(file.filename)}"
         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"]
+                
+            res = await asyncio.to_thread(cloudinary.uploader.upload, tmp_path, folder=folder)
+            image_url = res["secure_url"]
             uploaded_urls.append(image_url)
 
-            # 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 = []
-
+            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},
-                }
+                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)
 
             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)
-
+            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}")
+            print(f"❌ Upload error: {e}")
         finally:
-            if os.path.exists(tmp_path):
-                os.remove(tmp_path)
-
+            if os.path.exists(tmp_path): os.remove(tmp_path)
+            
     return {"message": "Done!", "urls": uploaded_urls}
 
 
 # ══════════════════════════════════════════════════════════════════
-# 3. SEARCH (With Demo Fallback)
+# 3. SEARCH (Strictly Cloud-Native + Freemium Default)
 # ══════════════════════════════════════════════════════════════════
 @app.post("/api/search")
-async def search_database(
-    file:                UploadFile = File(...),
-    detect_faces:        bool = Form(True),
-    user_pinecone_key:   str  = Form(""),
-    user_cloudinary_url: str  = Form(""),
-):
+async def search_database(file: UploadFile = File(...), detect_faces: bool = Form(True), user_pinecone_key: str = Form(""), user_cloudinary_url: str = Form("")):
     actual_pc_key = user_pinecone_key or os.getenv("DEFAULT_PINECONE_KEY")
-    
     if not actual_pc_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}"
+        raise HTTPException(400, "Pinecone Key is required.")
 
+    tmp_path = f"temp_uploads/query_{uuid.uuid4().hex}_{sanitize_filename(file.filename)}"
     try:
         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(actual_pc_key)
-        idx_obj  = pc.Index("lens-objects")
-        idx_face = pc.Index("lens-faces")
+        pc = _get_pinecone(actual_pc_key)
+        idx_obj = pc.Index(IDX_OBJECTS)
+        idx_face = pc.Index(IDX_FACES)
 
-        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"])
+        async def _query_one(vec_dict: 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,
-            )
+            
+            try:
+                res = await asyncio.to_thread(target_idx.query, vector=vec_list, top_k=10, include_metadata=True)
+            except Exception as e:
+                if "404" in str(e):
+                    # Graceful error if index truly doesn't exist
+                    raise HTTPException(404, f"Pinecone Index not found. Please log in and click 'Verify Keys' in Settings to build the indexes.")
+                raise e
+                
             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,
-                })
+                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]
 
-        seen: dict[str, dict] = {}
+        seen = {}
         for r in all_results:
             url = r["url"]
             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}
-
+        return {"results": sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:10]}
+    except HTTPException:
+        raise
     except Exception as e:
         print(f"❌ Search error: {e}")
         raise HTTPException(500, str(e))
     finally:
-        if os.path.exists(tmp_path):
-            os.remove(tmp_path)
+        if os.path.exists(tmp_path): os.remove(tmp_path)
 
 
 # ══════════════════════════════════════════════════════════════════
-# 4. CATEGORIES (With Demo Fallback)
+# 4. CATEGORIES (Strictly Cloud-Native + Freemium Default)
 # ══════════════════════════════════════════════════════════════════
 @app.post("/api/categories")
 async def get_categories(user_cloudinary_url: str = Form("")):
     actual_cld_url = user_cloudinary_url or os.getenv("DEFAULT_CLOUDINARY_URL")
-
     if not actual_cld_url:
         return {"categories": []}
-    
+        
     try:
-        creds = get_cloudinary_creds(actual_cld_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}
+        # FIX 3: Removed local folder scanning logic. Guests now strictly see Cloudinary folders.
+        _configure_cloudinary(get_cloudinary_creds(actual_cld_url))
+        result = await asyncio.to_thread(cloudinary.api.root_folders)
+        return {"categories": [f["name"] for f in result.get("folders", [])]}
     except Exception as e:
         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,
-    }
+    return {"status": "ok"}