Updating chroma db to be singleton class

src/main.py

@@ -46,6 +46,8 @@ import random
 from typing import List
 from src.utils.logger import logger
 from config.config import settings
+from src.vectorstores.chroma_manager import ChromaManager
+from src.utils.llm_utils import cleanup_vectorstore
 
 os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'
 # os.environ["OAUTHLIB_RELAX_TOKEN_SCOPE"] = "1"
@@ -678,6 +680,33 @@ async def debug_config():
     return debug_info
 
 
+@app.post("/admin/reset-chroma")
+async def reset_chroma_db(api_key: str = Depends(verify_api_key)):
+    """
+    Reset ChromaDB completely - use with caution
+
+    This endpoint cleans up all resources and recreates ChromaDB
+    """
+    try:
+        # Reset all vector store resources
+        await cleanup_vectorstore()
+
+        return {
+            "status": "success",
+            "message": "ChromaDB reset complete. You may need to restart the application for changes to take effect.",
+            "details": {
+                "chroma_path": settings.CHROMA_PATH
+            }
+        }
+
+    except Exception as e:
+        logger.error(f"Error resetting ChromaDB: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to reset ChromaDB: {str(e)}"
+        )
+
+
 @app.post("/admin/cleanup")
 async def cleanup_databases(
     include_files: bool = True,

src/utils/llm_utils.py

@@ -2,6 +2,9 @@
 from fastapi import HTTPException
 from typing import Tuple
 
+import asyncio
+import logging
+
 from src.llms.openai_llm import OpenAILanguageModel
 from src.llms.ollama_llm import OllamaLanguageModel
 from src.llms.bert_llm import BERTLanguageModel
@@ -9,20 +12,26 @@ from src.llms.falcon_llm import FalconLanguageModel
 from src.llms.llama_llm import LlamaLanguageModel
 from src.embeddings.huggingface_embedding import HuggingFaceEmbedding
 from src.vectorstores.chroma_vectorstore import ChromaVectorStore
-from src.vectorstores.optimized_vectorstore import get_optimized_vector_store
+from src.vectorstores.chroma_manager import ChromaManager
 from src.utils.logger import logger
 from config.config import settings
 
+# Global vector store instance for reuse
+_vector_store = None
+_embedding_model = None
+_vs_lock = asyncio.Lock()
+
+
 def get_llm_instance(provider: str):
     """
     Get LLM instance based on provider
-    
+
     Args:
         provider (str): Name of the LLM provider
-        
+
     Returns:
         BaseLLM: Instance of the LLM
-        
+
     Raises:
         ValueError: If provider is not supported
     """
@@ -33,29 +42,84 @@ def get_llm_instance(provider: str):
         'falcon': lambda: FalconLanguageModel(),
         'llama': lambda: LlamaLanguageModel(),
     }
-    
+
     if provider not in llm_map:
         raise ValueError(f"Unsupported LLM provider: {provider}")
     return llm_map[provider]()
 
+
 async def get_vector_store() -> Tuple[ChromaVectorStore, HuggingFaceEmbedding]:
     """
-    Get vector store and embedding model instances
-    Uses optimized implementation while maintaining backward compatibility
-    
+    Get vector store and embedding model instances with proper initialization
+
     Returns:
         Tuple[ChromaVectorStore, HuggingFaceEmbedding]: 
             Vector store and embedding model instances
     """
-    try:
-        return await get_optimized_vector_store()
-    except Exception as e:
-        logger.error(f"Error getting optimized vector store: {str(e)}")
-        # Fallback to original implementation if optimization fails
-        logger.warning("Falling back to standard vector store implementation")
-        embedding = HuggingFaceEmbedding(model_name=settings.EMBEDDING_MODEL)
-        vector_store = ChromaVectorStore(
-            embedding_function=embedding.embed_documents,
-            persist_directory=settings.CHROMA_PATH
-        )
-        return vector_store, embedding
+    global _vector_store, _embedding_model, _vs_lock
+
+    async with _vs_lock:
+        if _vector_store is not None and _embedding_model is not None:
+            return _vector_store, _embedding_model
+
+        try:
+            # Load embedding model
+            _embedding_model = HuggingFaceEmbedding(
+                model_name=settings.EMBEDDING_MODEL)
+            logger.info(f"Loaded embedding model: {settings.EMBEDDING_MODEL}")
+
+            # Get ChromaDB client through the manager
+            try:
+                client = await ChromaManager.get_client(
+                    persist_directory=settings.CHROMA_PATH,
+                    reset_if_needed=True
+                )
+                logger.info("Successfully initialized ChromaDB client")
+            except Exception as e:
+                logger.error(f"Error getting ChromaDB client: {str(e)}")
+
+                # Try to reset ChromaDB completely
+                await ChromaManager.reset_chroma(settings.CHROMA_PATH)
+                client = await ChromaManager.get_client(
+                    persist_directory=settings.CHROMA_PATH
+                )
+                logger.info("Recreated ChromaDB client after reset")
+
+            # Create and initialize vector store
+            _vector_store = ChromaVectorStore(
+                embedding_function=_embedding_model.embed_documents,
+                persist_directory=settings.CHROMA_PATH,
+                collection_name="documents",
+                client=client
+            )
+
+            # Initialize the vector store
+            await _vector_store.initialize()
+            logger.info("Vector store successfully initialized")
+
+            return _vector_store, _embedding_model
+
+        except Exception as e:
+            logger.error(f"Error initializing vector store: {str(e)}")
+            raise HTTPException(
+                status_code=500,
+                detail=f"Failed to initialize vector store: {str(e)}"
+            )
+
+
+async def cleanup_vectorstore():
+    """
+    Cleanup and reset vector store resources
+    """
+    global _vector_store, _embedding_model, _vs_lock
+
+    async with _vs_lock:
+        _vector_store = None
+        _embedding_model = None
+
+        # Force garbage collection
+        import gc
+        gc.collect()
+
+        # Reset ChromaDB completely
+        await ChromaManager.reset_chroma(settings.CHROMA_PATH)

src/vectorstores/chroma_manager.py

@@ -0,0 +1,165 @@
+# src/vectorstores/chroma_manager.py
+"""
+ChromaDB connection manager to ensure consistent settings and connection handling
+"""
+
+import os
+import shutil
+import asyncio
+import logging
+import chromadb
+from chromadb.config import Settings
+from typing import Optional, Dict, Any
+
+# Global connection lock and instance
+_instance_lock = asyncio.Lock()
+_chroma_instance = None
+_collection_lock = asyncio.Lock()
+_collection_instances = {}
+
+
+class ChromaManager:
+    """Singleton manager for ChromaDB connections"""
+
+    @staticmethod
+    async def get_client(
+        persist_directory: str,
+        reset_if_needed: bool = False
+    ) -> chromadb.PersistentClient:
+        """
+        Get a shared ChromaDB client with consistent settings
+
+        Args:
+            persist_directory (str): Directory to persist ChromaDB
+            reset_if_needed (bool): Whether to reset the database if connection fails
+
+        Returns:
+            chromadb.PersistentClient: Shared client instance
+        """
+        global _chroma_instance, _instance_lock
+
+        async with _instance_lock:
+            if _chroma_instance is not None:
+                return _chroma_instance
+
+            # Try to create a client
+            try:
+                settings = Settings(
+                    allow_reset=True,
+                    anonymized_telemetry=False,
+                    is_persistent=True
+                )
+
+                _chroma_instance = chromadb.PersistentClient(
+                    path=persist_directory,
+                    settings=settings
+                )
+                logging.info(
+                    f"Successfully created ChromaDB client at {persist_directory}")
+                return _chroma_instance
+
+            except ValueError as e:
+                if "already exists" in str(e) and reset_if_needed:
+                    logging.warning(
+                        f"ChromaDB instance exists with different settings. Attempting reset: {str(e)}")
+                    await ChromaManager.reset_chroma(persist_directory)
+
+                    # Try again after reset
+                    _chroma_instance = chromadb.PersistentClient(
+                        path=persist_directory,
+                        settings=settings
+                    )
+                    return _chroma_instance
+                raise
+
+    @staticmethod
+    async def get_or_create_collection(
+        client: chromadb.PersistentClient,
+        collection_name: str,
+        embedding_dimension: int = 1024
+    ):
+        """
+        Get or create a collection with proper error handling
+
+        Args:
+            client (chromadb.PersistentClient): ChromaDB client
+            collection_name (str): Name of the collection
+            embedding_dimension (int): Dimension of embeddings
+
+        Returns:
+            Collection: ChromaDB collection
+        """
+        global _collection_lock, _collection_instances
+
+        # Use just the collection name as key
+        collection_key = collection_name
+
+        async with _collection_lock:
+            if collection_key in _collection_instances:
+                return _collection_instances[collection_key]
+
+            try:
+                # Try to get existing collection
+                collection = client.get_collection(
+                    name=collection_name,
+                    embedding_function=None
+                )
+                logging.info(f"Found existing collection: {collection_name}")
+                _collection_instances[collection_key] = collection
+                return collection
+
+            except Exception as e:
+                logging.info(
+                    f"Collection {collection_name} does not exist, creating new one: {str(e)}")
+
+                # Create new collection with minimal metadata
+                # Removed the problematic "hnsw:dim" parameter
+                try:
+                    collection = client.create_collection(
+                        name=collection_name,
+                        metadata={"hnsw:space": "cosine"}
+                    )
+                except Exception as create_error:
+                    # If that fails too, try with no metadata
+                    logging.warning(
+                        f"Error creating collection with metadata: {str(create_error)}")
+                    collection = client.create_collection(
+                        name=collection_name
+                    )
+
+                _collection_instances[collection_key] = collection
+                return collection
+
+    @staticmethod
+    async def reset_chroma(persist_directory: str):
+        """
+        Reset ChromaDB completely by removing the directory
+
+        Args:
+            persist_directory (str): Directory to remove
+        """
+        global _chroma_instance, _collection_instances
+
+        # Clear global instances first
+        _chroma_instance = None
+        _collection_instances = {}
+
+        try:
+            # Force garbage collection to release file handles
+            import gc
+            gc.collect()
+
+            # Remove the entire directory
+            if os.path.exists(persist_directory):
+                shutil.rmtree(persist_directory)
+                logging.info(
+                    f"Removed ChromaDB directory: {persist_directory}")
+
+            # Recreate empty directory
+            os.makedirs(persist_directory, exist_ok=True)
+            logging.info(
+                f"Created fresh ChromaDB directory: {persist_directory}")
+
+        except Exception as e:
+            logging.error(f"Error resetting ChromaDB: {str(e)}")
+            raise

src/vectorstores/chroma_settings.py

@@ -0,0 +1,15 @@
+# src/vectorstores/chroma_settings.py
+"""
+Shared settings for ChromaDB to ensure consistency across the application
+"""
+
+import chromadb
+
+
+def get_chroma_settings():
+    """Get consistent ChromaDB settings"""
+    return chromadb.Settings(
+        allow_reset=True,
+        is_persistent=True,
+        anonymized_telemetry=False
+    )

src/vectorstores/chroma_vectorstore.py

@@ -1,10 +1,12 @@
 # src/vectorstores/chroma_vectorstore.py
+from pathlib import Path
 import chromadb
 from typing import List, Callable, Any, Dict, Optional
-from chromadb.config import Settings
 import logging
+import asyncio
 
 from .base_vectorstore import BaseVectorStore
+from .chroma_manager import ChromaManager
 
 
 class ChromaVectorStore(BaseVectorStore):
@@ -13,7 +15,8 @@ class ChromaVectorStore(BaseVectorStore):
         embedding_function: Callable[[List[str]], List[List[float]]],
         persist_directory: str = './chroma_db',
         collection_name: str = "documents",
-        client_settings: Optional[Dict[str, Any]] = None
+        client_settings: Optional[Dict[str, Any]] = None,
+        client=None  # Allow passing an existing client
     ):
         """
         Initialize Chroma Vector Store
@@ -23,25 +26,78 @@ class ChromaVectorStore(BaseVectorStore):
             persist_directory (str): Directory to persist the vector store
             collection_name (str): Name of the collection to use
             client_settings (Optional[Dict[str, Any]]): Additional settings for ChromaDB client
+            client: Optional existing ChromaDB client to use
         """
+        self.embedding_function = embedding_function
+        self.persist_directory = persist_directory
+        self.collection_name = collection_name
+        self.client = client  # Store client for later initialization
+
+        # Will be populated during async initialization
+        self.collection = None
+        self.initialized = False
+
+    async def initialize(self):
+        """Asynchronously initialize the vector store with enhanced error handling"""
+        if self.initialized:
+            return
+
         try:
-            settings = Settings(
-                persist_directory=persist_directory,
-                **(client_settings or {})
-            )
+            # Get client via manager if not provided
+            if self.client is None:
+                self.client = await ChromaManager.get_client(self.persist_directory)
+
+            # Validate client
+            if not self.client:
+                raise ValueError("Failed to obtain ChromaDB client")
+
+            # Get or create collection with more robust handling
+            try:
+                self.collection = await ChromaManager.get_or_create_collection(
+                    client=self.client,
+                    collection_name=self.collection_name,
+                    embedding_dimension=1024  # Default for most models
+                )
+            except Exception as collection_error:
+                logging.error(
+                    f"Error creating collection: {str(collection_error)}")
+
+                # Try to reset and recreate
+                try:
+                    # Attempt to delete existing collection
+                    self.client.delete_collection(self.collection_name)
+                except:
+                    pass
+
+                # Recreate collection
+                self.collection = self.client.create_collection(
+                    name=self.collection_name,
+                    metadata={"hnsw:space": "cosine"}
+                )
+
+            # Additional validation
+            if not self.collection:
+                raise ValueError(
+                    "Failed to create or obtain ChromaDB collection")
+
+            self.initialized = True
+            logging.info(
+                f"ChromaVectorStore initialized with collection: {self.collection_name}")
 
-            self.client = chromadb.PersistentClient(settings=settings)
-            self.collection = self.client.get_or_create_collection(
-                name=collection_name,
-                # Using cosine similarity by default
-                metadata={"hnsw:space": "cosine"}
-            )
-            self.embedding_function = embedding_function
         except Exception as e:
-            logging.error(f"Error initializing ChromaDB: {str(e)}")
+            logging.error(
+                f"Critical error initializing ChromaVectorStore: {str(e)}")
+            # Reset initialization state
+            self.initialized = False
+            self.collection = None
             raise
 
-    def add_documents(
+    async def _ensure_initialized(self):
+        """Make sure the vector store is initialized before use"""
+        if not self.initialized:
+            await self.initialize()
+
+    async def add_documents_async(
         self,
         documents: List[str],
         embeddings: Optional[List[List[float]]] = None,
@@ -49,7 +105,7 @@ class ChromaVectorStore(BaseVectorStore):
         ids: Optional[List[str]] = None
     ) -> None:
         """
-        Add documents to the vector store
+        Add documents asynchronously with enhanced error handling
 
         Args:
             documents (List[str]): List of document texts
@@ -57,48 +113,141 @@ class ChromaVectorStore(BaseVectorStore):
             metadatas (Optional[List[Dict[str, Any]]]): Metadata for each document
             ids (Optional[List[str]]): Custom IDs for the documents
         """
-        try:
-            if not documents:
-                logging.warning("No documents provided to add_documents")
-                return
+        await self._ensure_initialized()
+
+        if not documents:
+            logging.warning("No documents provided to add_documents")
+            return
+
+        # Validate input lists
+        if embeddings and len(documents) != len(embeddings):
+            raise ValueError("Number of documents and embeddings must match")
 
-            if not embeddings:
+        if metadatas and len(documents) != len(metadatas):
+            raise ValueError("Number of documents and metadatas must match")
+
+        # Generate embeddings if not provided
+        if not embeddings:
+            try:
                 embeddings = self.embedding_function(documents)
+            except Exception as e:
+                logging.error(f"Error generating embeddings: {str(e)}")
+                raise
+
+        # Use provided IDs or generate them
+        if not ids:
+            ids = [f"doc_{i}" for i in range(len(documents))]
+
+        # Ensure collection exists and is usable
+        if not self.collection:
+            logging.error("ChromaDB collection is not initialized")
+            await self.initialize()
+
+        # Prepare add parameters
+        add_params = {
+            "documents": documents,
+            "embeddings": embeddings,
+            "ids": ids
+        }
+
+        # Add metadatas if provided
+        if metadatas is not None:
+            add_params["metadatas"] = metadatas
+
+        try:
+            # Add documents to collection with retry mechanism
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    # Clear any cached state
+                    import gc
+                    gc.collect()
+
+                    # Attempt to add documents
+                    self.collection.add(**add_params)
+                    logging.info(
+                        f"Successfully added {len(documents)} documents")
+                    break
+                except (StopIteration, RuntimeError) as retry_error:
+                    if attempt < max_retries - 1:
+                        logging.warning(
+                            f"Retry attempt {attempt + 1}: {str(retry_error)}")
+                        # Optional: Add a small delay between retries
+                        await asyncio.sleep(0.5)
+                    else:
+                        logging.error(
+                            f"Failed to add documents after {max_retries} attempts")
+                        raise
+
+                except (StopIteration, RuntimeError) as retry_error:
+                    if attempt < max_retries - 1:
+                        logging.warning(
+                            f"Retry attempt {attempt + 1}: {str(retry_error)}")
+                        # Optional: Add a small delay between retries
+                        await asyncio.sleep(0.5)
+                    else:
+                        logging.error(
+                            f"Failed to add documents after {max_retries} attempts")
+                        raise
 
-            if len(documents) != len(embeddings):
-                raise ValueError(
-                    "Number of documents and embeddings must match")
-
-            # Use provided IDs or generate them
-            doc_ids = ids if ids is not None else [
-                f"doc_{i}" for i in range(len(documents))]
-
-            # Prepare add parameters
-            add_params = {
-                "documents": documents,
-                "embeddings": embeddings,
-                "ids": doc_ids
-            }
-
-            # Only include metadatas if provided
-            if metadatas is not None:
-                if len(metadatas) != len(documents):
-                    raise ValueError(
-                        "Number of documents and metadatas must match")
-                add_params["metadatas"] = metadatas
-
-            self.collection.add(**add_params)
         except Exception as e:
-            logging.error(f"Error adding documents to ChromaDB: {str(e)}")
-            raise
+            logging.error(
+                f"Unexpected error adding documents to ChromaDB: {str(e)}")
+
+            # Additional debugging information
+            try:
+                logging.info(f"Collection status: {self.collection}")
+                logging.info(f"Documents count: {len(documents)}")
+                logging.info(
+                    f"Embeddings count: {len(add_params.get('embeddings', []))}")
+                logging.info(
+                    f"Metadatas count: {len(add_params.get('metadatas', []))}")
+                logging.info(f"IDs count: {len(add_params.get('ids', []))}")
+            except Exception as debug_error:
+                logging.error(f"Error during debugging: {str(debug_error)}")
 
-    def similarity_search(
+    def add_documents(
+        self,
+        documents: List[str],
+        embeddings: Optional[List[List[float]]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        ids: Optional[List[str]] = None
+    ) -> None:
+        """
+        Synchronous wrapper for add_documents_async
+        """
+        # Create and run a new event loop if needed
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Create a future that can be run in the existing loop
+                asyncio.create_task(self.add_documents_async(
+                    documents, embeddings, metadatas, ids
+                ))
+            else:
+                # Run in a new event loop
+                loop.run_until_complete(self.add_documents_async(
+                    documents, embeddings, metadatas, ids
+                ))
+        except RuntimeError:
+            # No event loop, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            loop.run_until_complete(self.add_documents_async(
+                documents, embeddings, metadatas, ids
+            ))
+
+    async def similarity_search_async(
         self,
         query_embedding: List[float],
         top_k: int = 3,
         **kwargs
     ) -> List[Dict[str, Any]]:
-        """Perform similarity search with improved chunk handling"""
+        """
+        Perform similarity search asynchronously
+        """
+        await self._ensure_initialized()
+
         try:
             # Get more initial results to account for sequential chunks
             results = self.collection.query(
@@ -177,19 +326,122 @@ class ChromaVectorStore(BaseVectorStore):
             logging.error(f"Error in similarity search: {str(e)}")
             raise
 
-    def get_all_documents(
+    def similarity_search(
         self,
-        include_embeddings: bool = False
+        query_embedding: List[float],
+        top_k: int = 3,
+        **kwargs
     ) -> List[Dict[str, Any]]:
         """
-        Retrieve all documents from the vector store
+        Synchronous wrapper for similarity_search_async
+        """
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # We're in an async context, but need to process directly
+                try:
+                    # Get more initial results to account for sequential chunks
+                    results = self.collection.query(
+                        query_embeddings=[query_embedding],
+                        n_results=max(top_k * 2, 10),
+                        include=['documents', 'metadatas', 'distances']
+                    )
+
+                    if not results or 'documents' not in results:
+                        return []
+
+                    formatted_results = []
+                    documents = results['documents'][0]
+                    metadatas = results['metadatas'][0]
+                    distances = results['distances'][0]
+
+                    # Group chunks by document_id
+                    doc_chunks = {}
+                    for doc, meta, dist in zip(documents, metadatas, distances):
+                        doc_id = meta.get('document_id')
+                        chunk_index = meta.get('chunk_index', 0)
+
+                        if doc_id not in doc_chunks:
+                            doc_chunks[doc_id] = []
+
+                        doc_chunks[doc_id].append({
+                            'text': doc,
+                            'metadata': meta,
+                            'score': 1.0 - dist,
+                            'chunk_index': chunk_index
+                        })
+
+                    # Process each document's chunks
+                    for doc_id, chunks in doc_chunks.items():
+                        # Sort chunks by index
+                        chunks.sort(key=lambda x: x['chunk_index'])
+
+                        # Find sequences of chunks with good scores
+                        good_sequences = []
+                        current_sequence = []
 
-        Args:
-            include_embeddings (bool): Whether to include embeddings in the response
+                        for chunk in chunks:
+                            if chunk['score'] > 0.3:  # Adjust threshold as needed
+                                if not current_sequence or \
+                                   chunk['chunk_index'] == current_sequence[-1]['chunk_index'] + 1:
+                                    current_sequence.append(chunk)
+                                else:
+                                    if current_sequence:
+                                        good_sequences.append(current_sequence)
+                                    current_sequence = [chunk]
+                            else:
+                                if current_sequence:
+                                    good_sequences.append(current_sequence)
+                                current_sequence = []
+
+                        if current_sequence:
+                            good_sequences.append(current_sequence)
+
+                        # Add best sequences to results
+                        for sequence in good_sequences:
+                            avg_score = sum(c['score']
+                                            for c in sequence) / len(sequence)
+                            combined_text = ' '.join(
+                                c['text'] for c in sequence)
+
+                            formatted_results.append({
+                                'text': combined_text,
+                                'metadata': sequence[0]['metadata'],
+                                'score': avg_score
+                            })
+
+                    # Sort by score and return top_k
+                    formatted_results.sort(
+                        key=lambda x: x['score'], reverse=True)
+                    return formatted_results[:top_k]
+
+                except Exception as e:
+                    logging.error(
+                        f"Error in direct similarity search: {str(e)}")
+                    return []
+            else:
+                # Run in existing loop
+                return loop.run_until_complete(
+                    self.similarity_search_async(
+                        query_embedding, top_k, **kwargs)
+                )
+        except RuntimeError:
+            # No event loop, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop.run_until_complete(
+                self.similarity_search_async(query_embedding, top_k, **kwargs)
+            )
 
-        Returns:
-            List[Dict[str, Any]]: List of documents with their IDs and optionally embeddings
+    async def get_all_documents_async(
+        self,
+        include_embeddings: bool = False
+    ) -> List[Dict[str, Any]]:
         """
+        Retrieve all documents asynchronously
+        """
+        await self._ensure_initialized()
+
         try:
             include = ["documents", "metadatas"]
             if include_embeddings:
@@ -227,16 +479,70 @@ class ChromaVectorStore(BaseVectorStore):
                 f"Error retrieving documents from ChromaDB: {str(e)}")
             raise
 
-    def get_document_chunks(self, document_id: str) -> List[Dict[str, Any]]:
+    def get_all_documents(
+        self,
+        include_embeddings: bool = False
+    ) -> List[Dict[str, Any]]:
         """
-        Retrieve all chunks for a specific document
-
-        Args:
-            document_id (str): ID of the document to retrieve chunks for
+        Synchronous wrapper for get_all_documents_async
+        """
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # We're in an async context, but need to return synchronously
+                # Process the results just like in the async version
+                try:
+                    include = ["documents", "metadatas"]
+                    if include_embeddings:
+                        include.append("embeddings")
+
+                    results = self.collection.get(
+                        include=include
+                    )
+
+                    if not results or 'documents' not in results:
+                        return []
+
+                    documents = []
+                    for i in range(len(results['documents'])):
+                        doc = {
+                            'id': str(i),  # Generate sequential IDs
+                            'text': results['documents'][i],
+                        }
+
+                        if include_embeddings and 'embeddings' in results:
+                            doc['embedding'] = results['embeddings'][i]
+
+                        if 'metadatas' in results and results['metadatas'][i]:
+                            doc['metadata'] = results['metadatas'][i]
+
+                            # Use document_id from metadata if available
+                            if 'document_id' in results['metadatas'][i]:
+                                doc['id'] = results['metadatas'][i]['document_id']
+
+                        documents.append(doc)
+
+                    return documents
+                except:
+                    return []
+            else:
+                return loop.run_until_complete(
+                    self.get_all_documents_async(include_embeddings)
+                )
+        except RuntimeError:
+            # No event loop, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop.run_until_complete(
+                self.get_all_documents_async(include_embeddings)
+            )
 
-        Returns:
-            List[Dict[str, Any]]: List of document chunks with their metadata
+    async def get_document_chunks_async(self, document_id: str) -> List[Dict[str, Any]]:
         """
+        Retrieve all chunks for a specific document asynchronously
+        """
+        await self._ensure_initialized()
+
         try:
             results = self.collection.get(
                 where={"document_id": document_id},
@@ -263,18 +569,53 @@ class ChromaVectorStore(BaseVectorStore):
             logging.error(f"Error retrieving document chunks: {str(e)}")
             raise
 
-    def delete_document(self, document_id: str) -> None:
+    def get_document_chunks(self, document_id: str) -> List[Dict[str, Any]]:
+        """
+        Synchronous wrapper for get_document_chunks_async
         """
-        Delete all chunks associated with a document_id
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Fall back to direct query which may fail
+                try:
+                    results = self.collection.get(
+                        where={"document_id": document_id},
+                        include=["documents", "metadatas"]
+                    )
+
+                    chunks = []
+                    for i in range(len(results['documents'])):
+                        chunk = {
+                            'text': results['documents'][i],
+                            'metadata': results['metadatas'][i] if results.get('metadatas') else None
+                        }
+                        chunks.append(chunk)
+                    return chunks
+                except:
+                    return []
+            else:
+                return loop.run_until_complete(
+                    self.get_document_chunks_async(document_id)
+                )
+        except RuntimeError:
+            # No event loop, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop.run_until_complete(
+                self.get_document_chunks_async(document_id)
+            )
 
-        Args:
-            document_id (str): ID of the document to delete
+    async def delete_document_async(self, document_id: str) -> None:
         """
+        Delete all chunks associated with a document_id asynchronously
+        """
+        await self._ensure_initialized()
+
         try:
             # Get all chunks with the given document_id
             results = self.collection.get(
                 where={"document_id": document_id},
-                include=["metadatas"]
+                include=["ids"]
             )
 
             if not results or 'ids' not in results:
@@ -282,11 +623,28 @@ class ChromaVectorStore(BaseVectorStore):
                 return
 
             # Delete all chunks associated with the document
-            chunk_ids = [
-                f"{document_id}-chunk-{i}" for i in range(len(results['metadatas']))]
-            self.collection.delete(ids=chunk_ids)
+            self.collection.delete(ids=results['ids'])
 
         except Exception as e:
             logging.error(
                 f"Error deleting document {document_id} from ChromaDB: {str(e)}")
             raise
+
+    def delete_document(self, document_id: str) -> None:
+        """
+        Synchronous wrapper for delete_document_async
+        """
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Create a future that can be run in the existing loop
+                asyncio.create_task(self.delete_document_async(document_id))
+            else:
+                # Run in a new event loop
+                loop.run_until_complete(
+                    self.delete_document_async(document_id))
+        except RuntimeError:
+            # No event loop, create a new one
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            loop.run_until_complete(self.delete_document_async(document_id))

src/vectorstores/optimized_vectorstore.py

@@ -13,6 +13,7 @@ from .chroma_vectorstore import ChromaVectorStore
 from src.embeddings.huggingface_embedding import HuggingFaceEmbedding
 from src.utils.logger import logger
 from config.config import settings
+from src.vectorstores.chroma_settings import get_chroma_settings
 
 
 class OptimizedVectorStore(ChromaVectorStore):
@@ -56,30 +57,21 @@ class OptimizedVectorStore(ChromaVectorStore):
             # Clean up existing database if dimensions don't match
             await self._cleanup_if_needed(required_dim)
 
-            # Create ChromaDB client with fresh settings
+            # Create ChromaDB client with consistent settings
             client = chromadb.PersistentClient(
                 path=self._persist_directory,
-                settings=Settings(
-                    allow_reset=True,
-                    is_persistent=True,
-                    anonymized_telemetry=False
-                )
+                settings=get_chroma_settings()  # Use shared settings function
             )
 
             # Create new collection with correct dimensions
-            collection = client.create_collection(
-                name=self._collection_name,
-                metadata={
-                    "hnsw:space": "cosine",
-                    "hnsw:dim": required_dim
-                }
-            )
+            # collection = self._setup_collection(client)
 
             # Initialize parent class
             super().__init__(
                 embedding_function=self._embedding_model.embed_documents,
                 persist_directory=self._persist_directory,
-                collection_name=self._collection_name
+                collection_name=self._collection_name,
+                client=client  # Pass the existing client
             )
 
             self._initialized = True