Update src/vector_store.py

src/vector_store.py

@@ -1,20 +1,22 @@
 """
-Vector Store and Embeddings Module using ChromaDB with CLIP embeddings
+Vector Store and Embeddings Module using ChromaDB with sentence-transformers
+UPDATED for ChromaDB v0.4.22+ (auto-persist, no manual persist needed)
 """
 import os
 import json
 from typing import List, Dict
 import chromadb
-from chromadb.config import Settings
 from sentence_transformers import SentenceTransformer
 import numpy as np
 from config import CHROMA_DB_PATH, EMBEDDING_MODEL, EMBEDDING_DIM
 
 
 class CLIPEmbedder:
-    """Custom embedder using CLIP model for multimodal content"""
+    """Custom embedder using sentence-transformers for multimodal content"""
     def __init__(self, model_name: str = EMBEDDING_MODEL):
+        print(f"🔄 Loading embedding model: {model_name}")
         self.model = SentenceTransformer(model_name)
+        print(f"✅ Model loaded successfully")
     
     def embed(self, text: str) -> List[float]:
         """Generate embedding for text"""
@@ -36,21 +38,39 @@ class CLIPEmbedder:
 
 
 class VectorStore:
-    """Vector store manager using ChromaDB"""
+    """Vector store manager using ChromaDB (v0.4.22+ with auto-persist)"""
     def __init__(self):
         self.persist_directory = CHROMA_DB_PATH
         self.embedder = CLIPEmbedder()
         
-        # Initialize ChromaDB with persistent storage
-        self.client = chromadb.PersistentClient(
-            path=self.persist_directory
-        )
+        print(f"\n🔄 Initializing ChromaDB at: {self.persist_directory}")
+        
+        # NEW ChromaDB v0.4.22+ - PersistentClient auto-persists
+        try:
+            self.client = chromadb.PersistentClient(
+                path=self.persist_directory
+            )
+            print(f"✅ ChromaDB PersistentClient initialized")
+        except Exception as e:
+            print(f"❌ Error initializing ChromaDB: {e}")
+            print(f"Trying fallback initialization...")
+            self.client = chromadb.PersistentClient(
+                path=self.persist_directory
+            )
         
         # Get or create collection
-        self.collection = self.client.get_or_create_collection(
-            name="multimodal_rag",
-            metadata={"hnsw:space": "cosine"}
-        )
+        try:
+            self.collection = self.client.get_or_create_collection(
+                name="multimodal_rag",
+                metadata={"hnsw:space": "cosine"}
+            )
+            count = self.collection.count()
+            print(f"✅ Collection loaded: {count} items in store")
+        except Exception as e:
+            print(f"Error with collection: {e}")
+            self.collection = self.client.get_or_create_collection(
+                name="multimodal_rag"
+            )
 
     def add_documents(self, documents: List[Dict], doc_id: str):
         """Add documents to vector store"""
@@ -58,6 +78,8 @@ class VectorStore:
         metadatas = []
         ids = []
         
+        print(f"\n📚 Adding documents for: {doc_id}")
+        
         # Add text chunks
         if 'text' in documents and documents['text']:
             chunks = self._chunk_text(documents['text'], chunk_size=1000, overlap=200)
@@ -69,9 +91,11 @@ class VectorStore:
                     'chunk_idx': str(idx)
                 })
                 ids.append(f"{doc_id}_text_{idx}")
+            print(f"  ✅ Text: {len(chunks)} chunks")
         
         # Add image descriptions and OCR text
         if 'images' in documents:
+            image_count = 0
             for idx, image_data in enumerate(documents['images']):
                 if image_data.get('ocr_text'):
                     texts.append(f"Image {idx}: {image_data['ocr_text']}")
@@ -82,31 +106,44 @@ class VectorStore:
                         'image_path': image_data.get('path', '')
                     })
                     ids.append(f"{doc_id}_image_{idx}")
+                    image_count += 1
+            if image_count > 0:
+                print(f"  ✅ Images: {image_count} with OCR text")
         
         # Add table content
         if 'tables' in documents:
+            table_count = 0
             for idx, table_data in enumerate(documents['tables']):
-                texts.append(f"Table {idx}: {table_data.get('content', '')}")
-                metadatas.append({
-                    'doc_id': doc_id,
-                    'type': 'table',
-                    'table_idx': str(idx)
-                })
-                ids.append(f"{doc_id}_table_{idx}")
+                if table_data.get('content'):
+                    texts.append(f"Table {idx}: {table_data.get('content', '')}")
+                    metadatas.append({
+                        'doc_id': doc_id,
+                        'type': 'table',
+                        'table_idx': str(idx)
+                    })
+                    ids.append(f"{doc_id}_table_{idx}")
+                    table_count += 1
+            if table_count > 0:
+                print(f"  ✅ Tables: {table_count}")
         
         if texts:
             # Generate embeddings
+            print(f"  🔄 Generating {len(texts)} embeddings...")
             embeddings = self.embedder.embed_batch(texts)
             
             # Add to collection
-            self.collection.add(
-                ids=ids,
-                documents=texts,
-                embeddings=embeddings,
-                metadatas=metadatas
-            )
-            
-            print(f"Added {len(texts)} documents to vector store for {doc_id}")
+            try:
+                self.collection.add(
+                    ids=ids,
+                    documents=texts,
+                    embeddings=embeddings,
+                    metadatas=metadatas
+                )
+                print(f"✅ Successfully added {len(texts)} items to vector store")
+                # Auto-persist happens here
+                print(f"✅ Data persisted automatically to: {self.persist_directory}")
+            except Exception as e:
+                print(f"❌ Error adding to collection: {e}")
 
     def search(self, query: str, n_results: int = 5) -> List[Dict]:
         """Search vector store for similar documents"""
@@ -154,7 +191,8 @@ class VectorStore:
             return {
                 'name': 'multimodal_rag',
                 'count': count,
-                'status': 'active'
+                'status': 'active',
+                'persist_path': self.persist_directory
             }
         except Exception as e:
             print(f"Error getting collection info: {e}")
@@ -167,14 +205,29 @@ class VectorStore:
             results = self.collection.get(where={'doc_id': doc_id})
             if results['ids']:
                 self.collection.delete(ids=results['ids'])
-                print(f"Deleted {len(results['ids'])} documents for {doc_id}")
+                print(f"✅ Deleted {len(results['ids'])} documents for {doc_id}")
+                # Auto-persist on delete
+                print(f"✅ Changes persisted automatically")
         except Exception as e:
             print(f"Error deleting documents: {e}")
 
     def persist(self):
-        """Persist vector store to disk"""
+        """
+        No-op for compatibility with older code.
+        ChromaDB v0.4.22+ uses PersistentClient which auto-persists.
+        This method kept for backward compatibility.
+        """
+        print("✅ Vector store is using auto-persist (no manual persist needed)")
+
+    def clear_all(self):
+        """Clear all documents from collection"""
         try:
-            self.client.persist()
-            print("Vector store persisted to disk")
+            # Delete collection and recreate
+            self.client.delete_collection(name="multimodal_rag")
+            self.collection = self.client.get_or_create_collection(
+                name="multimodal_rag",
+                metadata={"hnsw:space": "cosine"}
+            )
+            print("✅ Collection cleared and reset")
         except Exception as e:
-            print(f"Error persisting vector store: {e}")
+            print(f"Error clearing collection: {e}")