Update vector_store.py

vector_store.py

@@ -9,6 +9,7 @@ from langchain_pinecone import PineconeVectorStore
 from langchain.embeddings.base import Embeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
+
 class InLegalBERTEmbeddings(Embeddings):
     """Custom LangChain embeddings wrapper for InLegalBERT"""
     
@@ -23,6 +24,7 @@ class InLegalBERTEmbeddings(Embeddings):
         """Embed a single query"""
         return self.model.encode([text])[0].tolist()
 
+
 class LegalDocumentVectorStore:
     """Manages vector storage for legal documents"""
     
@@ -59,27 +61,44 @@ class LegalDocumentVectorStore:
             print(f"✅ Created Pinecone index: {self.index_name}")
         
         self._initialized = True
-    
-    def save_document_embeddings_optimized(self, chunk_data: List[Dict], document_id: str, 
-                                     analysis_results: Dict[str, Any]) -> bool:
+
+    def _normalize_embedding(self, embedding):
+        """Ensure embedding is always a list of floats"""
+        if embedding is None:
+            return None
+        if isinstance(embedding, np.ndarray):
+            return embedding.tolist()
+        if isinstance(embedding, list):
+            # Already a Python list
+            return embedding
+        # Fallback: try converting if it's a torch tensor or similar
+        try:
+            return embedding.tolist()
+        except Exception:
+            return list(embedding)
+
+    def save_document_embeddings_optimized(
+        self,
+        chunk_data: List[Dict],
+        document_id: str,
+        analysis_results: Dict[str, Any]
+    ) -> bool:
         """Save embeddings using pre-computed vectors with proper text storage"""
         try:
             self._initialize_pinecone()
-            
-            # Use pre-computed embeddings instead of re-generating
-            texts = [chunk["text"] for chunk in chunk_data]
-            embeddings = [chunk["embedding"].tolist() for chunk in chunk_data if chunk["embedding"] is not None]
-            
-            if not embeddings:
+
+            # Normalize embeddings safely
+            valid_chunks = [c for c in chunk_data if c.get("embedding") is not None]
+            if not valid_chunks:
                 print("⚠️ No embeddings found in chunk_data")
                 return False
-            
-            # Prepare vectors with BOTH metadata AND text content
+
             vectors = []
-            for i, chunk_info in enumerate(chunk_data):
-                if chunk_info["embedding"] is None:
+            for i, chunk_info in enumerate(valid_chunks):
+                normalized_embedding = self._normalize_embedding(chunk_info["embedding"])
+                if normalized_embedding is None:
                     continue
-                
+
                 metadata = {
                     'document_id': document_id,
                     'chunk_index': i,
@@ -91,20 +110,20 @@ class LegalDocumentVectorStore:
                     'timestamp': str(np.datetime64('now')),
                     'text': chunk_info["text"]  # Store text in metadata for retrieval
                 }
-                
+
                 vectors.append({
                     "id": f"{document_id}_chunk_{i}",
-                    "values": chunk_info["embedding"].tolist(),
+                    "values": normalized_embedding,
                     "metadata": metadata
                 })
-            
-            # Add to Pinecone using pre-computed embeddings
+
+            # Add to Pinecone
             index = self.pc.Index(self.index_name)
             index.upsert(vectors=vectors)
-            
+
             print(f"✅ Saved {len(vectors)} pre-computed embeddings with text to Pinecone")
             return True
-            
+
         except Exception as e:
             print(f"❌ Error saving pre-computed embeddings: {e}")
             return False
@@ -113,33 +132,30 @@ class LegalDocumentVectorStore:
         """Get retriever for chat functionality with improved settings"""
         try:
             self._initialize_pinecone()
-            
+
             legal_embeddings = InLegalBERTEmbeddings(clause_tagger.embedding_model)
             index = self.pc.Index(self.index_name)
-            
+
             vectorstore = PineconeVectorStore(
                 index=index,
                 embedding=legal_embeddings,
                 text_key="text"  # Use text stored in metadata
             )
-            
+
             # Configure search parameters
-            search_kwargs = {
-                'k': 10  # Increased from default 5 for better context
-            }
-            
+            search_kwargs = {'k': 10}
             if document_id:
                 search_kwargs['filter'] = {'document_id': document_id}
-            
-            # Use similarity search without threshold initially
+
             return vectorstore.as_retriever(
                 search_type="similarity",
                 search_kwargs=search_kwargs
             )
-            
+
         except Exception as e:
             print(f"❌ Error creating retriever: {e}")
             return None
 
+
 # Global instance
 vector_store = LegalDocumentVectorStore()