Update server.py

server.py

@@ -7,6 +7,7 @@ import unicodedata
 import threading
 import anyio
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from huggingface_hub import snapshot_download
@@ -17,7 +18,7 @@ logger = logging.getLogger(__name__)
 
 DATASET_REPO = "aniketkumar1106/orbit-data" 
 IMAGE_DIR = "Productimages"
-DB_TARGET = "orbiitt.db" 
+DB_TARGET_FOLDER = "orbiitt_db" # The folder ChromaDB expects
 MIN_CONFIDENCE_THRESHOLD = 0.1
 
 # BOOTSTRAP: Mandatory folder creation
@@ -48,6 +49,7 @@ def background_sync():
     global engine, loading_status
     token = os.environ.get("HF_TOKEN")
     
+    # Cleanup old images
     for f in os.listdir(IMAGE_DIR):
         p = os.path.join(IMAGE_DIR, f)
         try:
@@ -60,18 +62,33 @@ def background_sync():
         snapshot_download(repo_id=DATASET_REPO, repo_type="dataset", token=token, local_dir=".")
 
         if os.path.exists("orbiitt_db.zip"):
+            loading_status = "Extracting Database..."
             with zipfile.ZipFile("orbiitt_db.zip", 'r') as z:
                 z.extractall("temp_extract")
             
-            for root, _, files in os.walk("temp_extract"):
+            # Smart Extraction: Find the ChromaDB folder logic
+            db_found = False
+            for root, dirs, files in os.walk("temp_extract"):
+                # Identifying ChromaDB by its signature file
+                if "chroma.sqlite3" in files:
+                    if os.path.exists(DB_TARGET_FOLDER):
+                        shutil.rmtree(DB_TARGET_FOLDER)
+                    # Move the directory containing the sqlite3 file to our target location
+                    shutil.move(root, DB_TARGET_FOLDER)
+                    db_found = True
+                
+                # Move images
                 for f in files:
-                    src = os.path.join(root, f)
                     if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+                        src = os.path.join(root, f)
                         clean_name = normalize_filename(f)
                         shutil.copy(src, os.path.join(IMAGE_DIR, clean_name))
-                    elif "orbiitt" in f.lower() and f.endswith(".db"):
-                        shutil.copy(src, f"./{DB_TARGET}")
+
             shutil.rmtree("temp_extract")
+            
+            # If no DB folder was moved (maybe zip structure was flat?), ensure folder exists
+            if not db_found and not os.path.exists(DB_TARGET_FOLDER):
+                os.makedirs(DB_TARGET_FOLDER, exist_ok=True)
 
         # LOGGING FILE COUNT FOR VALIDATION
         final_count = len(os.listdir(IMAGE_DIR))
@@ -80,21 +97,31 @@ def background_sync():
         loading_status = "Loading AI Engine..."
         try:
             from orbiitt_engine import OrbiittEngine
-            engine = OrbiittEngine()
+            # Initialize with the CORRECT folder path
+            engine = OrbiittEngine(db_path=f"./{DB_TARGET_FOLDER}")
             loading_status = "Ready"
             logger.info(">>> ENGINE ONLINE <<<")
         except Exception as e:
             loading_status = f"Engine Error: {str(e)}"
+            logger.error(f"Engine Failed: {e}")
+            
     except Exception as e:
         loading_status = f"Sync Error: {str(e)}"
+        logger.error(f"Sync Failed: {e}")
 
 @app.on_event("startup")
 async def startup_event():
     thread = threading.Thread(target=background_sync, daemon=True)
     thread.start()
 
+# Mount Static Images
 app.mount("/Productimages", StaticFiles(directory=IMAGE_DIR), name="Productimages")
 
+# Serve UI (Root Endpoint)
+@app.get("/")
+async def read_index():
+    return FileResponse('index.html')
+
 @app.get("/health")
 def health():
     return {"status": loading_status, "ready": engine is not None}
@@ -117,13 +144,13 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
             async with await anyio.open_file(t_path, "wb") as f:
                 await f.write(content)
         
-        # THE FIX: Use lambda to ensure keyword arguments are passed correctly
+        # CORRECTED: Calling engine with top_k
         results = await anyio.to_thread.run_sync(
             lambda: engine.search(
                 text_query=text, 
                 image_file=t_path, 
-                text_weight=actual_weight, 
-                top_k=50
+                text_weight=actual_weight,
+                top_k=50 
             )
         )
         
@@ -138,13 +165,15 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
             if score < MIN_CONFIDENCE_THRESHOLD or pid in seen_ids:
                 continue
 
-            raw_path = r.get('url') or r.get('path') or ""
-            fname = normalize_filename(os.path.basename(raw_path))
+            # The engine returns a clean ID/path, let's match it to disk
+            fname_from_db = os.path.basename(r.get('id', ''))
+            fname = normalize_filename(fname_from_db)
             
             match = None
             if fname in all_files:
                 match = fname
             else:
+                # Fuzzy fallback if exact match fails
                 for disk_f in all_files:
                     if fname[:15].lower() in disk_f.lower():
                         match = disk_f
@@ -153,6 +182,7 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
             if match:
                 final_list.append({
                     "id": pid,
+                    # Ensure URL is properly encoded for web
                     "url": f"Productimages/{urllib.parse.quote(match)}",
                     "score": round(float(score), 4)
                 })