Update utils/database.py

utils/database.py

@@ -1,19 +1,30 @@
-import streamlit as st  # Add this import at the top
-from langchain.vectorstores import FAISS
+import chromadb
+from chromadb.config import Settings
+from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 import os
 
-def initialize_vector_db():
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={'device': 'cpu'}
+def initialize_vector_db(persist_dir: str):
+    client = chromadb.PersistentClient(
+        path=persist_dir,
+        settings=Settings(anonymized_telemetry=False)
+    
+    embedding_func = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
-    return FAISS.from_texts(["Initial text"], embeddings)
+    
+    return Chroma(
+        client=client,
+        collection_name="document_embeddings",
+        embedding_function=embedding_func,
+        persist_directory=persist_dir
+    )
+
+def add_to_collection(chunks: list):
+    db = get_vector_db()
+    db.add_texts(texts=chunks)
+    db.persist()
+    return True
 
-def add_to_collection(chunks):
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    if 'vector_db' in st.session_state:  # Now st is properly defined
-        st.session_state.vector_db.add_texts(chunks)
-    else:
-        st.session_state.vector_db = FAISS.from_texts(chunks, embeddings)
-    return True
+def get_vector_db():
+    return st.session_state.vector_db