Update vector_store.py

vector_store.py

@@ -3,12 +3,11 @@
 Vector store integration for legal document embeddings using InLegalBERT and Pinecone
 """
 import os
-import pinecone
-from langchain_pinecone import PineconeVectorStore  # ⭐ Modern import
-from langchain.embeddings.base import Embeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 import numpy as np
 from typing import List, Dict, Any
+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"""
@@ -61,31 +60,30 @@ class LegalDocumentVectorStore:
         
         self._initialized = True
     
-    def save_document_embeddings(self, document_text: str, document_id: str, 
-                               analysis_results: Dict[str, Any], clause_tagger) -> bool:
-        """Save document embeddings using InLegalBERT model"""
+    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 - NO RE-EMBEDDING"""
         try:
             self._initialize_pinecone()
             
-            # Use the clause tagger's InLegalBERT model
-            legal_embeddings = InLegalBERTEmbeddings(clause_tagger.embedding_model)
-            
-            # Split document into chunks
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=1000,
-                chunk_overlap=200,
-                separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
-            )
+            # 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]
             
-            chunks = text_splitter.split_text(document_text)
+            if not embeddings:
+                print("⚠️ No embeddings found in chunk_data")
+                return False
             
-            # Prepare metadata with analysis results
+            # Prepare metadata
             metadatas = []
-            for i, chunk in enumerate(chunks):
+            for i, chunk_info in enumerate(chunk_data):
+                if chunk_info["embedding"] is None:
+                    continue
+                    
                 metadata = {
                     'document_id': document_id,
                     'chunk_index': i,
-                    'total_chunks': len(chunks),
+                    'total_chunks': len(chunk_data),
                     'source': 'legal_document',
                     'has_key_clauses': len(analysis_results.get('key_clauses', [])) > 0,
                     'risk_count': len(analysis_results.get('risky_terms', [])),
@@ -94,26 +92,25 @@ class LegalDocumentVectorStore:
                 }
                 metadatas.append(metadata)
             
-            # Create vector store using modern API
+            # Add to Pinecone using pre-computed embeddings
             index = self.pc.Index(self.index_name)
-            vectorstore = PineconeVectorStore(
-                index=index,
-                embedding=legal_embeddings,
-                text_key="text"
-            )
             
-            # Add documents to Pinecone
-            vectorstore.add_texts(
-                texts=chunks,
-                metadatas=metadatas,
-                ids=[f"{document_id}_chunk_{i}" for i in range(len(chunks))]
-            )
+            vectors = [
+                {
+                    "id": f"{document_id}_chunk_{i}",
+                    "values": embedding,
+                    "metadata": metadata
+                }
+                for i, (embedding, metadata) in enumerate(zip(embeddings, metadatas))
+            ]
             
-            print(f"✅ Saved {len(chunks)} chunks using InLegalBERT embeddings for document {document_id}")
+            index.upsert(vectors=vectors)
+            
+            print(f"✅ Saved {len(vectors)} pre-computed embeddings to Pinecone")
             return True
             
         except Exception as e:
-            print(f"❌ Error saving to Pinecone: {e}")
+            print(f"❌ Error saving pre-computed embeddings: {e}")
             return False
     
     def get_retriever(self, clause_tagger, document_id: str = None):
@@ -143,3 +140,4 @@ class LegalDocumentVectorStore:
 
 # Global instance
 vector_store = LegalDocumentVectorStore()
+