Update app.py

app.py

@@ -152,7 +152,7 @@ def main(query: str, client: QdrantClient, collection_name: str, tokenizer: Auto
     else:
         return f'Internal Knowledge :\n\n{answer}' if 'knowledge_topic' in text else f'Documents Based :\n\n{answer_with_metadatas}'
 
-def collect_files(directory, pattern):
+def collect_files(conn, cursor, directory, pattern):
     array = []
 
     for filename in os.listdir(directory):
@@ -177,12 +177,12 @@ def collect_files(directory, pattern):
                 array.extend(sparse_embeddings)
             elif (filename.endswith('.npy')):
                 ids_list = np.load(os.path.join(directory, filename), allow_pickle=True).tolist()
-                insert_data(os.path.splitext(filename)[0], ids_list)
+                insert_data(conn, cursor, os.path.splitext(filename)[0], ids_list)
                 array.extend(ids_list)
                 
     return array
 
-def insert_data(name, ids_array):
+def insert_data(conn, cursor, name, ids_array):
     cursor.execute('INSERT INTO table_names (id) VALUES (?)', (name,))
     for ids in ids_array:
         cursor.execute('INSERT INTO table_ids (name, ids_value) VALUES (?, ?)', (name, ids))
@@ -343,7 +343,7 @@ def load_models_and_documents():
             FOREIGN KEY(name) REFERENCES table_names(doc_name)
             )
             ''')
-            insert_data(name, unique_ids)
+            insert_data(conn, cursor, name, unique_ids)
             
             np.save(ids_path, np.array(unique_ids), allow_pickle=True)
         else:
@@ -368,7 +368,7 @@ def load_models_and_documents():
             ''')
 
             unique_ids, payload_docs, dense_embeddings, sparse_embeddings = [
-                collect_files(embeddings_path, pattern) for pattern in patterns
+                collect_files(conn, cursor, embeddings_path, pattern) for pattern in patterns
             ]
 
         st.write('Ingesting saved documents on disk into our Qdrant Vector Database...')
@@ -398,7 +398,7 @@ def load_models_and_documents():
     time.sleep(5)
     container.empty()
     
-    return client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp
+    return client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp, conn, cursor
 
 def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding):
     text_splitter = SemanticChunker(
@@ -458,7 +458,7 @@ def on_change_documents_only():
 if __name__ == '__main__':
     st.set_page_config(page_title="Multipurpose AI Agent",layout="wide", initial_sidebar_state='collapsed')
 
-    client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp = load_models_and_documents()
+    client, collection_name, tokenizer, model, llm, dense_model, sparse_model, nlp, conn, cursor = load_models_and_documents()
 
     if 'menu_id' not in st.session_state:
         st.session_state.menu_id = 'ChatBot'
@@ -847,7 +847,7 @@ if __name__ == '__main__':
                     combined_sparse_matrix = vstack(sparse_matrices)
                     save_npz(sparse_path, combined_sparse_matrix)
 
-                    insert_data(base_name, ids)
+                    insert_data(conn, cursor, base_name, ids)
                     np.save(ids_path, np.array(ids), allow_pickle=True)
 
                     st.toast('Document(s) Ingested !', icon='🎉')