Enhance knowledge base rebuilding functionality: Add method to rebuild knowledge base from selected URLs, improve error handling, and log actions during vector store download process.

app.py

@@ -72,6 +72,51 @@ def update_kb_with_selected(sources_df) -> str:
         logger.error(f"Error updating knowledge base: {str(e)}")
         return f"Error updating knowledge base: {str(e)}"
 
+def rebuild_kb_with_selected(sources_df):
+    """Rebuild knowledge base from scratch using only selected URLs"""
+    try:
+        selected_urls = get_selected_urls(sources_df)
+        
+        if not selected_urls:
+            return "Error: No URLs selected for inclusion"
+        
+        # Временно заменяем URLS на выбранные URL
+        original_urls = constants.URLS.copy()
+        constants.URLS = selected_urls
+        
+        try:
+            # Пересоздаем базу знаний
+            success, message = create_vector_store(mode="rebuild")
+            
+            # Сохраняем метаданные если успешно
+            if success:
+                metadata = {
+                    "last_updated": datetime.datetime.now().isoformat(),
+                    "source_count": len(selected_urls),
+                    "sources": selected_urls
+                }
+                
+                # Сохраняем в датасет
+                json_content = json.dumps(metadata, indent=2).encode('utf-8')
+                api = HfApi(token=HF_TOKEN)
+                
+                api.upload_file(
+                    path_or_fileobj=json_content,
+                    path_in_repo="vector_store/metadata.json",
+                    repo_id=DATASET_ID,
+                    repo_type="dataset"
+                )
+            
+            return message
+            
+        finally:
+            # Восстанавливаем оригинальные URL
+            constants.URLS = original_urls
+            
+    except Exception as e:
+        logger.error(f"Error rebuilding knowledge base: {str(e)}")
+        return f"Error rebuilding knowledge base: {str(e)}"
+
 # Set seed for consistent results
 langdetect.DetectorFactory.seed = 0
 
@@ -1734,17 +1779,15 @@ def rebuild_kb_with_selected(sources_df):
             return "Error: No URLs selected for inclusion"
         
         # Временно заменяем URLS на выбранные URL
-        from config import constants
-        original_urls = constants.URLS
+        original_urls = constants.URLS.copy()
         constants.URLS = selected_urls
         
         try:
             # Пересоздаем базу знаний
             success, message = create_vector_store(mode="rebuild")
             
-            # Сохраняем метаданные с информацией о выбранных URL
+            # Сохраняем метаданные если успешно
             if success:
-                # Создаем метаданные с текущей датой и выбранными URL
                 metadata = {
                     "last_updated": datetime.datetime.now().isoformat(),
                     "source_count": len(selected_urls),
@@ -1763,11 +1806,13 @@ def rebuild_kb_with_selected(sources_df):
                 )
             
             return message
+            
         finally:
             # Восстанавливаем оригинальные URL
             constants.URLS = original_urls
             
     except Exception as e:
+        logger.error(f"Error rebuilding knowledge base: {str(e)}")
         return f"Error rebuilding knowledge base: {str(e)}"
 
 def save_kb_metadata():

src/knowledge_base/dataset.py

@@ -45,17 +45,17 @@ def download_vector_store(self) -> Tuple[bool, Union[FAISS, str]]:
     """Download vector store from dataset"""
     try:
         with tempfile.TemporaryDirectory() as temp_dir:
-            print(f"Downloading to temporary directory: {temp_dir}")
+            logger.debug(f"Downloading to temporary directory: {temp_dir}")
             
-            # Download files to temporary directory
             try:
+                # Download vector store files
                 index_path = self.api.hf_hub_download(
                     repo_id=self.dataset_name,
                     filename="vector_store/index.faiss",
                     repo_type="dataset",
                     local_dir=temp_dir
                 )
-                print(f"Downloaded index.faiss to: {index_path}")
+                logger.debug(f"Downloaded index.faiss to: {index_path}")
                 
                 config_path = self.api.hf_hub_download(
                     repo_id=self.dataset_name,
@@ -63,27 +63,26 @@ def download_vector_store(self) -> Tuple[bool, Union[FAISS, str]]:
                     repo_type="dataset",
                     local_dir=temp_dir
                 )
-                print(f"Downloaded index.pkl to: {config_path}")
+                logger.debug(f"Downloaded index.pkl to: {config_path}")
                 
-                # Verify files exist
-                if not os.path.exists(index_path) or not os.path.exists(config_path):
-                    return False, f"Downloaded files not found at {temp_dir}"
-                
-                # Load vector store from temporary directory
+                # Initialize embeddings
                 embeddings = HuggingFaceEmbeddings(
                     model_name=EMBEDDING_MODEL,
                     model_kwargs={'device': 'cpu'}
                 )
                 
+                # Load vector store
                 vector_store = FAISS.load_local(
-                    os.path.join(temp_dir, "vector_store"),
-                    embeddings
+                    folder_path=os.path.join(temp_dir, "vector_store"),
+                    embeddings=embeddings
                 )
                 
                 return True, vector_store
                 
             except Exception as e:
+                logger.error(f"Error downloading vector store: {str(e)}")
                 return False, f"Error downloading vector store: {str(e)}"
+                
     except Exception as e:
         logger.error(f"Error in download_vector_store: {str(e)}")
         return False, str(e)