Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -23,12 +23,11 @@ class PolicyVectorDB:
         self.collection_name = "neepco_dop_policies"
         
         # Using a powerful open-source embedding model.
-        # Change 'cpu' to 'cuda' if a GPU is available for significantly faster embedding.
         logger.info("Loading embedding model 'BAAI/bge-large-en-v1.5'. This may take a moment...")
         self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cpu')
         logger.info("Embedding model loaded successfully.")
         
-        self.collection = None  # Initialize collection as None for lazy loading
+        self.collection = None # Initialize collection as None for lazy loading
         self.top_k_default = top_k_default
         self.relevance_threshold = relevance_threshold
 
@@ -82,7 +81,7 @@ class PolicyVectorDB:
         logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
 
         # Process in batches for efficiency
-        batch_size = 32  # Reduced batch size for potentially large embeddings
+        batch_size = 32 # Reduced batch size for potentially large embeddings
         for i in range(0, len(new_chunks), batch_size):
             batch = new_chunks[i:i + batch_size]
             ids = [str(chunk['id']) for chunk in batch]
@@ -115,19 +114,19 @@ class PolicyVectorDB:
         # Retrieve more results initially to allow for filtering
         results = collection.query(
             query_embeddings=query_embedding,
-            n_results=k * 2,  # Retrieve more to filter by threshold
+            n_results=k * 2, # Retrieve more to filter by threshold
             include=["documents", "metadatas", "distances"]
         )
 
         search_results = []
-        if results and results.get('documents') and results['documents'][0]:
-            for i, doc in enumerate(results['documents'][0]):
+        if results and results.get('documents') and results['documents']:
+            for i, doc in enumerate(results['documents'][0]):  # ← Fixed indexing
                 # The distance for normalized embeddings is often interpreted as 1 - cosine_similarity
-                relevance_score = 1 - results['distances'][0][i]
+                relevance_score = 1 - results['distances'][i]  # ← Fixed indexing
                 if relevance_score >= self.relevance_threshold:
                     search_results.append({
                         'text': doc,
-                        'metadata': results['metadatas'][0][i],
+                        'metadata': results['metadatas'][i],  # ← Fixed indexing
                         'relevance_score': relevance_score
                     })
 
@@ -168,13 +167,13 @@ class PolicyVectorDB:
             )
         
         search_results = []
-        if results and results.get('documents') and results['documents'][0]:
-            for i, doc in enumerate(results['documents']):
-                relevance_score = 1 - results['distances'][i]
+        if results and results.get('documents') and results['documents']:
+            for i, doc in enumerate(results['documents'][0]):  # ← Fixed indexing
+                relevance_score = 1 - results['distances'][i]  # ← Fixed indexing
                 if relevance_score >= self.relevance_threshold:
                     search_results.append({
                         'text': doc,
-                        'metadata': results['metadatas'][0][i],
+                        'metadata': results['metadatas'][i],  # ← Fixed indexing
                         'relevance_score': relevance_score
                     })