Update embedding model to Google Generative AI and enhance vector store functionality

- Changed the embedding model from OpenAI to Google Generative AI in the EmbeddingManager class.
- Updated the configuration to reflect the new embedding model path.
- Modified validation checks to ensure the presence of the Google API key for RAG embeddings.
- Added a new method to reset the vector store, allowing for a complete clear and recreation of the collection.
- Enhanced logging to provide clearer feedback on embedding model initialization and vector store operations.

src/core/config.py

@@ -86,7 +86,7 @@ class RAGConfig:
     chat_history_path: str = "./data/chat_history"
     
     # Embedding settings
-    embedding_model: str = "text-embedding-3-small"
+    embedding_model: str = "models/text-embedding-004"
     embedding_chunk_size: int = 1000
     
     # Chunking settings
@@ -182,8 +182,8 @@ class Config:
             validation_results["warnings"].append("Mistral API key not found - Mistral parser will be unavailable")
         
         # Check RAG dependencies
-        if not self.api.openai_api_key:
-            validation_results["warnings"].append("OpenAI API key not found - RAG embeddings will be unavailable")
+        if not self.api.google_api_key:
+            validation_results["warnings"].append("Google API key not found - RAG embeddings will be unavailable")
         
         if not self.api.google_api_key:
             validation_results["warnings"].append("Google API key not found - RAG chat will be unavailable")

src/rag/embeddings.py

@@ -2,7 +2,7 @@
 
 import os
 from typing import Optional
-from langchain_openai import OpenAIEmbeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from src.core.config import config
 from src.core.logging_config import get_logger
 
@@ -12,30 +12,28 @@ class EmbeddingManager:
     """Manages embedding models for document vectorization."""
     
     def __init__(self):
-        self._embedding_model: Optional[OpenAIEmbeddings] = None
+        self._embedding_model: Optional[GoogleGenerativeAIEmbeddings] = None
         
-    def get_embedding_model(self) -> OpenAIEmbeddings:
-        """Get or create the OpenAI embedding model."""
+    def get_embedding_model(self) -> GoogleGenerativeAIEmbeddings:
+        """Get or create the Gemini embedding model."""
         if self._embedding_model is None:
             try:
-                # Get OpenAI API key from config/environment
-                openai_api_key = config.api.openai_api_key or os.getenv("OPENAI_API_KEY")
+                # Get Google API key from config/environment
+                google_api_key = config.api.google_api_key or os.getenv("GOOGLE_API_KEY")
                 
-                if not openai_api_key:
-                    raise ValueError("OpenAI API key not found. Please set OPENAI_API_KEY in environment variables.")
+                if not google_api_key:
+                    raise ValueError("Google API key not found. Please set GOOGLE_API_KEY in environment variables.")
                 
-                self._embedding_model = OpenAIEmbeddings(
-                    model="text-embedding-3-small",
-                    openai_api_key=openai_api_key,
-                    chunk_size=1000,  # Process documents in chunks
-                    max_retries=3,
-                    timeout=30
+                self._embedding_model = GoogleGenerativeAIEmbeddings(
+                    model=config.rag.embedding_model,
+                    google_api_key=google_api_key,
+                    task_type="RETRIEVAL_DOCUMENT"
                 )
                 
-                logger.info("OpenAI embedding model initialized successfully")
+                logger.info(f"Gemini embedding model ({config.rag.embedding_model}) initialized successfully")
                 
             except Exception as e:
-                logger.error(f"Failed to initialize OpenAI embedding model: {e}")
+                logger.error(f"Failed to initialize Gemini embedding model: {e}")
                 raise
                 
         return self._embedding_model
@@ -50,14 +48,14 @@ class EmbeddingManager:
             
             # Check if we got a valid embedding (list of floats)
             if isinstance(embedding, list) and len(embedding) > 0 and isinstance(embedding[0], float):
-                logger.info("Embedding model test successful")
+                logger.info("Gemini embedding model test successful")
                 return True
             else:
-                logger.error("Embedding model test failed: Invalid embedding format")
+                logger.error("Gemini embedding model test failed: Invalid embedding format")
                 return False
                 
         except Exception as e:
-            logger.error(f"Embedding model test failed: {e}")
+            logger.error(f"Gemini embedding model test failed: {e}")
             return False
 
 # Global embedding manager instance

src/rag/vector_store.py

@@ -70,6 +70,7 @@ class VectorStoreManager:
         
         logger.info(f"VectorStoreManager initialized with persist_directory={self.persist_directory}")
     
+
     def get_vector_store(self) -> Chroma:
         """Get or create the Chroma vector store."""
         if self._vector_store is None:
@@ -314,7 +315,7 @@ class VectorStoreManager:
                 "collection_name": self.collection_name,
                 "persist_directory": self.persist_directory,
                 "document_count": count,
-                "embedding_model": "text-embedding-3-small"
+                "embedding_model": config.rag.embedding_model
             }
             
             logger.info(f"Collection info: {info}")
@@ -371,6 +372,36 @@ class VectorStoreManager:
             logger.error(f"Error searching with metadata filter: {e}")
             return []
     
+    def reset_vector_store(self) -> bool:
+        """
+        Reset the vector store completely.
+        This will clear all documents and recreate the collection.
+        
+        Returns:
+            True if successful, False otherwise
+        """
+        try:
+            logger.info("Resetting vector store...")
+            
+            # Clear all documents and reset the vector store
+            success = self.clear_all_documents()
+            
+            if success:
+                # Also delete the collection to ensure clean state
+                if self._vector_store is not None:
+                    self._vector_store.delete_collection()
+                    self._vector_store = None
+                
+                logger.info("Vector store reset successfully")
+                return True
+            else:
+                logger.error("Failed to reset vector store")
+                return False
+                
+        except Exception as e:
+            logger.error(f"Error resetting vector store: {e}")
+            return False
+    
     def clear_all_documents(self) -> bool:
         """
         Clear all documents from the vector store collection.