Refactor knowledge base management functions for improved clarity and error handling

app.py

@@ -185,10 +185,18 @@ def respond(
         
         yield new_history, conversation_id
 
-def build_kb():
-    """Function to create knowledge base"""
+def update_kb():
+    """Function to update existing knowledge base with new documents"""
     try:
-        success, message = create_vector_store()
+        success, message = create_vector_store(mode="update")
+        return message
+    except Exception as e:
+        return f"Error updating knowledge base: {str(e)}"
+
+def rebuild_kb():
+    """Function to create knowledge base from scratch"""
+    try:
+        success, message = create_vector_store(mode="rebuild")
         return message
     except Exception as e:
         return f"Error creating knowledge base: {str(e)}"
@@ -271,7 +279,8 @@ with gr.Blocks() as demo:
                 
                 with gr.Column(scale=1):
                     gr.Markdown("### Knowledge Base Management")
-                    build_kb_btn = gr.Button("Create/Update Knowledge Base", variant="primary")
+                    update_kb_btn = gr.Button("Update Knowledge Base", variant="secondary")
+                    rebuild_kb_btn = gr.Button("Rebuild Knowledge Base", variant="primary")
                     kb_status = gr.Textbox(label="Knowledge Base Status", interactive=False)
 
             submit_btn.click(
@@ -279,7 +288,8 @@ with gr.Blocks() as demo:
                 [msg, chatbot, conversation_id],
                 [chatbot, conversation_id, msg]
             )
-            build_kb_btn.click(build_kb, None, kb_status)
+            update_kb_btn.click(update_kb, None, kb_status)
+            rebuild_kb_btn.click(rebuild_kb, None, kb_status)
             clear_btn.click(lambda: ([], None), None, [chatbot, conversation_id])
 
         with gr.Tab("Model Settings"):

src/knowledge_base/vector_store.py

@@ -15,8 +15,16 @@ def get_embeddings():
         model_kwargs={'device': 'cpu'}
     )
 
-def create_vector_store():
-    """Create vector store and upload to dataset"""
+def create_vector_store(mode: str = "rebuild"):
+    """
+    Create or update vector store and upload to dataset
+    
+    Args:
+        mode: Either "rebuild" (create from scratch) or "update" (add new documents)
+    
+    Returns:
+        (success, message)
+    """
     # Load documents
     documents = load_documents()
     
@@ -33,32 +41,55 @@ def create_vector_store():
     # Initialize embeddings
     embeddings = get_embeddings()
     
-    # Create vector store in temporary directory
-    with tempfile.TemporaryDirectory() as temp_dir:
-        vector_store = FAISS.from_documents(chunks, embeddings)
-        # Save to temporary directory
-        vector_store.save_local(folder_path=temp_dir)
-        
-        # Copy files to VECTOR_STORE_PATH for subsequent loading
-        os.makedirs(VECTOR_STORE_PATH, exist_ok=True)
-        for file in ["index.faiss", "index.pkl"]:
-            shutil.copy2(
-                os.path.join(temp_dir, file),
-                os.path.join(VECTOR_STORE_PATH, file)
-            )
+    try:
+        if mode == "update":
+            # Try to load existing vector store
+            from src.knowledge_base.dataset import DatasetManager
+            dataset = DatasetManager(token=HF_TOKEN)
+            success, result = dataset.download_vector_store()
+            
+            if success:
+                # Add new documents to existing store
+                vector_store = FAISS.load_local(
+                    VECTOR_STORE_PATH,
+                    embeddings,
+                    allow_dangerous_deserialization=True
+                )
+                vector_store.add_documents(chunks)
+            else:
+                return False, "Failed to load existing vector store for update"
+        else:
+            # Create new vector store
+            vector_store = FAISS.from_documents(chunks, embeddings)
         
-        # Upload to dataset with explicit token passing
-        from src.knowledge_base.dataset import DatasetManager
-        dataset = DatasetManager(token=HF_TOKEN)
-        success, message = dataset.upload_vector_store()
+        # Save and upload
+        with tempfile.TemporaryDirectory() as temp_dir:
+            vector_store.save_local(folder_path=temp_dir)
+            
+            # Copy files to VECTOR_STORE_PATH for subsequent loading
+            os.makedirs(VECTOR_STORE_PATH, exist_ok=True)
+            for file in ["index.faiss", "index.pkl"]:
+                shutil.copy2(
+                    os.path.join(temp_dir, file),
+                    os.path.join(VECTOR_STORE_PATH, file)
+                )
+            
+            # Upload to dataset
+            from src.knowledge_base.dataset import DatasetManager
+            dataset = DatasetManager(token=HF_TOKEN)
+            success, message = dataset.upload_vector_store()
+            
+            # Clean up local files
+            shutil.rmtree(VECTOR_STORE_PATH)
+            
+            if not success:
+                return False, f"Error uploading to dataset: {message}"
         
-        # Clean up local files after upload
-        shutil.rmtree(VECTOR_STORE_PATH)
+        action = "updated" if mode == "update" else "created"
+        return True, f"Knowledge base {action} successfully! Processed {len(documents)} documents, {len(chunks)} chunks."
         
-        if not success:
-            return False, f"Error uploading to dataset: {message}"
-    
-    return True, f"Knowledge base created successfully! Loaded {len(documents)} documents, created {len(chunks)} chunks."
+    except Exception as e:
+        return False, f"Error {mode}ing knowledge base: {str(e)}"
 
 def load_vector_store():
     """Load vector store"""