Update orbiitt_engine.py

orbiitt_engine.py

@@ -24,6 +24,7 @@ class OrbiittEngine:
         self.expected_dim = self.model.config.vision_config.hidden_size
         
         # 4. Connect to Database with Safety Logic
+        print(f"📂 Connecting to ChromaDB at: {self.db_path}")
         self.client = chromadb.PersistentClient(path=self.db_path)
         self._check_db_compatibility()
         
@@ -42,10 +43,8 @@ class OrbiittEngine:
                 existing_dim = len(sample['embeddings'][0])
                 if existing_dim != self.expected_dim:
                     print(f"⚠️  Dimension Mismatch: DB is {existing_dim}, Model is {self.expected_dim}")
-                    if input("Wipe DB and restart? (y/n): ").lower() == 'y':
-                        self.client.delete_collection(name="product_catalog")
-                    else:
-                        exit()
+                    # Auto-wipe in production to prevent crash loop, or handle gracefully
+                    print("!!! WARNING: Database dimension mismatch. Search results may be invalid.")
         except: pass
 
     def get_image_embedding(self, image_path):
@@ -85,7 +84,7 @@ class OrbiittEngine:
             except Exception as e:
                 tqdm.write(f"⚠️ Skipped {fname}: {e}")
 
-    def search(self, text_query=None, image_file=None, text_weight=0.5):
+    def search(self, text_query=None, image_file=None, text_weight=0.5, top_k=20):
         """Hybrid search blending visual and text embeddings."""
         img_vec = None
         txt_vec = None
@@ -107,9 +106,18 @@ class OrbiittEngine:
         else:
             return []
 
-        results = self.collection.query(query_embeddings=[query_emb], n_results=10)
+        # Pass dynamic top_k to ChromaDB
+        try:
+            results = self.collection.query(query_embeddings=[query_emb], n_results=top_k)
+        except Exception as e:
+            print(f"Search Query Error: {e}")
+            return []
 
         output = []
+        # Handle case where no results found
+        if not results['ids'] or not results['ids'][0]:
+            return []
+
         for i in range(len(results['ids'][0])):
             fname = results['ids'][0][i]
             # distance for 'cosine' is 1 - similarity. 
@@ -117,7 +125,8 @@ class OrbiittEngine:
             score = round((1.0 - results['distances'][0][i]) * 100)
             output.append({
                 "id": fname,
-                "url": f"http://localhost:8000/Productimages/{fname}",
+                "url": f"Productimages/{fname}", # Relative path
                 "score": score
             })
-        return output
+        return output
+