Update retrival.py

retrival.py

@@ -258,13 +258,14 @@ def split_text(documents: list[Document]):
 ####---------------------------------------------------- Creating and Storeing Data in Vector DB  --------------------------------------------------####
 ########################################################################################################################################################
 
-def save_to_chroma(chunks: list[Document], name: str, tables: list[Document]):
-    CHROMA_PATH = f"./VectorDB/chroma_{name}"
-    TABLE_PATH = f"./TableDB/chroma_{name}"
+#def save_to_chroma(chunks: list[Document], name: str, tables: list[Document]):
+def save_to_chroma(chunks: list[Document], name: str):
+CHROMA_PATH = f"./VectorDB/chroma_{name}"
+    #TABLE_PATH = f"./TableDB/chroma_{name}"
     if os.path.exists(CHROMA_PATH):
         shutil.rmtree(CHROMA_PATH)
-    if os.path.exists(TABLE_PATH):
-        shutil.rmtree(TABLE_PATH)
+    # if os.path.exists(TABLE_PATH):
+    #     shutil.rmtree(TABLE_PATH)
 
     try:
         # Load the embedding model
@@ -279,19 +280,21 @@ def save_to_chroma(chunks: list[Document], name: str, tables: list[Document]):
         )
         print("Document database successfully saved.")
 
-        # Create Chroma DB for tables if available [NOTE: Some of the data is converted to string because int and float show null if added]
-        if tables:
-            print("Creating table vector database...")
-            tdb = Chroma.from_documents(
-                documents=tables,
-                embedding=embedding_function,
-                persist_directory=TABLE_PATH,
-            )
-            print("Table database successfully saved.")
-        else:
-            tdb = None
-
-        return db, tdb
+        # # Create Chroma DB for tables if available [NOTE: Some of the data is converted to string because int and float show null if added]
+        # if tables:
+        #     print("Creating table vector database...")
+        #     tdb = Chroma.from_documents(
+        #         documents=tables,
+        #         embedding=embedding_function,
+        #         persist_directory=TABLE_PATH,
+        #     )
+        #     print("Table database successfully saved.")
+        # else:
+        #     tdb = None
+
+        #return db, tdb
+        return db
+    
     except Exception as e:
         print("Error while saving to Chroma:", e)
         return None
@@ -306,54 +309,54 @@ def save_to_chroma(chunks: list[Document], name: str, tables: list[Document]):
 ####----------------------------------------------------------- Updating Existing Data in Vector DB  -----------------------------------------------####
 ########################################################################################################################################################
 
-def add_document_to_existing_db(new_documents: list[Document], db_name: str):
-    CHROMA_PATH = f"./VectorDB/chroma_{db_name}"
-
-    if not os.path.exists(CHROMA_PATH):
-        print(f"Database '{db_name}' does not exist. Please create it first.")
-        return
-
-    try:
-        embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-        #embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
-        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
-
-        print("Adding new documents to the existing database...")
-        chunks = split_text(new_documents)
-        db.add_documents(chunks)
-        db.persist()
-        print("New documents added and database updated successfully.")
-    except Exception as e:
-        print("Error while adding documents to existing database:", e)
-
-def delete_chunks_by_source(chroma_path, source_to_delete):
-    if not os.path.exists(chroma_path):
-        print(f"Database at path '{chroma_path}' does not exist.")
-        return
-
-    try:
-        #embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-        embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
-        db = Chroma(persist_directory=chroma_path, embedding_function=embedding_function)
+# def add_document_to_existing_db(new_documents: list[Document], db_name: str):
+#     CHROMA_PATH = f"./VectorDB/chroma_{db_name}"
 
-        print(f"Retrieving all metadata to identify chunks with source '{source_to_delete}'...")
-        metadata_list = db.get()["metadatas"]
+#     if not os.path.exists(CHROMA_PATH):
+#         print(f"Database '{db_name}' does not exist. Please create it first.")
+#         return
 
-        # Identify indices of chunks to delete
-        indices_to_delete = [
-            idx for idx, metadata in enumerate(metadata_list) if metadata.get("source") == source_to_delete
-        ]
+#     try:
+#         embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+#         #embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
+#         db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+
+#         print("Adding new documents to the existing database...")
+#         chunks = split_text(new_documents)
+#         db.add_documents(chunks)
+#         db.persist()
+#         print("New documents added and database updated successfully.")
+#     except Exception as e:
+#         print("Error while adding documents to existing database:", e)
 
-        if not indices_to_delete:
-            print(f"No chunks found with source '{source_to_delete}'.")
-            return
+# def delete_chunks_by_source(chroma_path, source_to_delete):
+#     if not os.path.exists(chroma_path):
+#         print(f"Database at path '{chroma_path}' does not exist.")
+#         return
 
-        print(f"Deleting {len(indices_to_delete)} chunks with source '{source_to_delete}'...")
-        db.delete(indices=indices_to_delete)
-        db.persist()
-        print("Chunks deleted and database updated successfully.")
-    except Exception as e:
-        print(f"Error while deleting chunks by source: {e}")
+#     try:
+#         #embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+#         embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
+#         db = Chroma(persist_directory=chroma_path, embedding_function=embedding_function)
+
+#         print(f"Retrieving all metadata to identify chunks with source '{source_to_delete}'...")
+#         metadata_list = db.get()["metadatas"]
+
+#         # Identify indices of chunks to delete
+#         indices_to_delete = [
+#             idx for idx, metadata in enumerate(metadata_list) if metadata.get("source") == source_to_delete
+#         ]
+
+#         if not indices_to_delete:
+#             print(f"No chunks found with source '{source_to_delete}'.")
+#             return
+
+#         print(f"Deleting {len(indices_to_delete)} chunks with source '{source_to_delete}'...")
+#         db.delete(indices=indices_to_delete)
+#         db.persist()
+#         print("Chunks deleted and database updated successfully.")
+#     except Exception as e:
+#         print(f"Error while deleting chunks by source: {e}")
         
 # # update a data store        
 # def update_data_store(file_path, db_name):
@@ -390,21 +393,22 @@ def generate_data_store(file_path, db_name):
     print(f"Filepath ===> {file_path}  DB Name ====> {db_name}")
 
     try:
-        documents,processed_documents,table_document = load_document(file_path)
+        documents,grouped_documents = load_document(file_path)
         print("Documents loaded successfully.")
     except Exception as e:
         print(f"Error loading documents: {e}")
         return
 
     try:
-        chunks = split_text(documents)
+        chunks = split_text(grouped_documents)
         print(f"Text split into {len(chunks)} chunks.")
     except Exception as e:
         print(f"Error splitting text: {e}")
         return
 
     try:
-        asyncio.run(save_to_chroma(save_to_chroma(chunks, db_name, table_document)))
+        #asyncio.run(save_to_chroma(save_to_chroma(chunks, db_name, table_document)))
+        asyncio.run(save_to_chroma(save_to_chroma(chunks, db_name)))
         print(f"Data saved to Chroma for database {db_name}.")
     except Exception as e:
         print(f"Error saving to Chroma: {e}")