cura/vector_store.py

"""
Vector Store for Mindify Chat
"""

import chromadb

def set_up_chromadb(collection_name: str):
    """
    Set up a ChromaDB collection for storing vectors.
    
    Args:
    collection_name: str
        The name of the collection to create or retrieve.
    
    Returns:
    ChromaDB Collection
        The ChromaDB collection object.
    """
    chroma_client = chromadb.Client()
    
    try: 
        # Check if the collection already exists
        collection = chroma_client.get_collection(name=collection_name)
        return collection
    except:
        # Create a new collection
        collection = chroma_client.create_collection(name=collection_name)
        return collection


def index_vector_store(collection_name:str, files: list):
    """
    Index the files in the ChromaDB collection.
    
    Args:
    collection: ChromaDB Collection
        The collection to store the vectors in.
    files: list
        A list of strings containing the contents of the files.
        
    Returns:
    bool
        True if the data is stored successfully, False otherwise.
    """
    # Set up collection 
    try:
        collection = chromadb.Client().get_collection(name=collection_name)
    except:
        collection = chromadb.Client().create_collection(name=collection_name)
    
    print("Indexing files...")
    ids = []
    for i in range(len(files[0])):
        ids.append(str(i))
        
    print("Storing GitHub data in ChromaDB...")
    try:
        collection.add(ids=ids, documents=files[0])
        print("Data stored successfully!")
        
        return True
    except:
        print("Error storing data in ChromaDB")
        return False
    
def query_vector_store(collection_name: str, query: str):
    """
    Query the ChromaDB collection for similar vectors to the query vector.
    """
    print("Querying ChromaDB...")
    try:
        list_collection = chromadb.Client().list_collections()
        print(list_collection)
        collection = chromadb.Client().get_collection(name=collection_name)
        return collection.query(query_texts=query, n_results=5)
    except:
        print("Error querying ChromaDB")
        return None