added the file stats info to the UI

app.py

@@ -20,7 +20,6 @@ def restart_system():
         log_message("Начало перезапуска системы...")
         log_message("Очистка кэша HuggingFace...")
         
-        # Clear HuggingFace cache to force fresh download
         import shutil
         cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
         if os.path.exists(cache_dir):
@@ -41,15 +40,20 @@ def restart_system():
         )
         
         if query_engine:
+            # Get updated stats
+            stats = get_repository_stats(HF_REPO_ID, HF_TOKEN, JSON_FILES_DIR, 
+                                        TABLE_DATA_DIR, IMAGE_DATA_DIR)
+            stats_display = format_stats_display(stats)
+            
             log_message("Система успешно перезапущена")
-            return "✅ Система успешно перезапущена! Новые документы загружены."
+            return "✅ Система успешно перезапущена! Новые документы загружены.", stats_display
         else:
-            return "❌ Ошибка при перезапуске системы"
+            return "❌ Ошибка при перезапуске системы", "Статистика недоступна"
             
     except Exception as e:
         error_msg = f"Ошибка перезапуска: {str(e)}"
         log_message(error_msg)
-        return f"❌ {error_msg}"
+        return f"❌ {error_msg}", "Статистика недоступна"
 
 def initialize_system(repo_id, hf_token, download_dir, chunks_filename=None, 
                      json_files_dir=None, table_data_dir=None, image_data_dir=None, 
@@ -458,6 +462,17 @@ Rerank Top K: {retrieval_params['rerank_top_k']}"""
             Выберите тип документа и загрузите файл. Система автоматически обработает и добавит его в базу знаний.
             """)
             
+            # Add stats display at the top
+            stats_display = gr.Markdown(
+                value=format_stats_display(
+                    get_repository_stats(HF_REPO_ID, HF_TOKEN, JSON_FILES_DIR, 
+                                    TABLE_DATA_DIR, IMAGE_DATA_DIR)
+                ),
+                label=""
+            )
+            
+            gr.Markdown("---")  # Separator
+            
             with gr.Row():
                 with gr.Column(scale=2):
                     file_type_radio = gr.Radio(
@@ -518,17 +533,18 @@ Rerank Top K: {retrieval_params['rerank_top_k']}"""
                     4. Нажмите "Перезапустить систему"
                     """)
 
-        upload_btn.click(
-            fn=process_uploaded_file,
-            inputs=[file_upload, file_type_radio],
-            outputs=[upload_status]
-        )
+            upload_btn.click(
+                fn=process_uploaded_file,
+                inputs=[file_upload, file_type_radio],
+                outputs=[upload_status]
+            )
+
+            restart_btn.click(
+                fn=restart_system,
+                inputs=[],
+                outputs=[restart_status, stats_display]  
+            )
 
-        restart_btn.click(
-            fn=restart_system,
-            inputs=[],
-            outputs=[restart_status]
-        )
         switch_btn.click(
             fn=switch_model_func,
             inputs=[model_dropdown],

main_utils.py

@@ -209,7 +209,48 @@ def enhance_query_with_keywords(query):
         return enhanced
     return f"{query}"
 
+def get_repository_stats(repo_id, hf_token, json_dir, table_dir, image_dir):
+    """Get statistics about documents in the repository"""
+    try:
+        from huggingface_hub import list_repo_files
+        
+        files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
+        
+        # Count JSON text files
+        json_files = [f for f in files if f.startswith(json_dir) and f.endswith('.json')]
+        zip_files = [f for f in files if f.startswith(json_dir) and f.endswith('.zip')]
+        
+        # Count table files
+        table_files = [f for f in files if f.startswith(table_dir) and 
+                      (f.endswith('.json') or f.endswith('.xlsx') or f.endswith('.xls'))]
+        
+        # Count image files
+        image_files = [f for f in files if f.startswith(image_dir) and 
+                      (f.endswith('.csv') or f.endswith('.xlsx') or f.endswith('.xls'))]
+        
+        stats = {
+            'text_files': len(json_files) + len(zip_files),
+            'table_files': len(table_files),
+            'image_files': len(image_files),
+            'total_files': len(json_files) + len(zip_files) + len(table_files) + len(image_files)
+        }
+        
+        log_message(f"Repository stats: {stats}")
+        return stats
+    except Exception as e:
+        log_message(f"Error getting repository stats: {e}")
+        return {'text_files': 0, 'table_files': 0, 'image_files': 0, 'total_files': 0}
+
+def format_stats_display(stats):
+    """Format statistics for display"""
+    return f"""📊 **Статистика базы данных:**
 
+📝 Текстовые документы (JSON): **{stats['text_files']}**
+📊 Табличные данные: **{stats['table_files']}**
+🖼️ Изображения: **{stats['image_files']}**
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+📦 Всего файлов: **{stats['total_files']}**
+"""
 
 def merge_table_chunks(chunk_info):
     merged = {}