mai_dx_wrapper.py

#!/usr/bin/env python3
"""
Обгортка для MAI-DX що виправляє проблеми з Rich форматуванням
"""
import os
import sys
import warnings

# Налаштування середовища
os.environ.update({
    "SWARMS_VERBOSITY": "ERROR",
    "RICH_TRACEBACK": "0", 
    "SWARMS_SHOW_PANEL": "true",
    "SWARMS_AUTO_PRINT": "true",
    "PYTHONWARNINGS": "ignore"
})

def patch_rich_and_swarms():
    """Патчимо Rich та Swarms для уникнення помилок виводу"""
    try:
        # Патчимо Rich Console
        import rich.console
        
        original_print = rich.console.Console.print
        
        def safe_print(self, *args, **kwargs):
            """Безпечний print що ігнорує None значення"""
            try:
                # Фільтруємо None значення
                filtered_args = [arg for arg in args if arg is not None]
                if filtered_args:
                    return original_print(self, *filtered_args, **kwargs)
            except Exception:
                # Ігноруємо всі помилки виводу
                pass
                
        rich.console.Console.print = safe_print
        
        # Патчимо Swarms formatter
        try:
            import swarms.utils.formatter as formatter
            
            # # Заміняємо _print_panel на заглушку
            # def dummy_print_panel(*args, **kwargs):
            #     pass
                
            # if hasattr(formatter, 'Formatter'):
            #     formatter.Formatter._print_panel = dummy_print_panel
                
            # Патчимо print_panel_token_usage
            def safe_print_panel_token_usage(self, *args, **kwargs):
                try:
                    # Просто виводимо базову інформацію без форматування
                    if args and hasattr(args[0], '__dict__'):
                        data = args[0].__dict__ if hasattr(args[0], '__dict__') else str(args[0])
                        print(f"[Token Usage] {data}")
                except:
                    pass
                    
            if hasattr(formatter, 'print_panel_token_usage'):
                formatter.print_panel_token_usage = safe_print_panel_token_usage
                
        except ImportError:
            pass
            
        # Патчимо threading для уникнення помилок в потоках
        import threading
        original_thread_init = threading.Thread.__init__
        
        def safe_thread_init(self, *args, **kwargs):
            """Обгортаємо target функцію в try/except"""
            original_target = kwargs.get('target')
            
            if original_target:
                def safe_target(*args, **kwargs):
                    try:
                        return original_target(*args, **kwargs)
                    except Exception:
                        # Ігноруємо помилки в потоках
                        pass
                        
                kwargs['target'] = safe_target
                
            return original_thread_init(self, *args, **kwargs)
            
        threading.Thread.__init__ = safe_thread_init
        
    except Exception as e:
        print(f"Warning: Could not patch Rich/Swarms: {e}")

# Застосовуємо патчі перед імпортом MAI-DX
patch_rich_and_swarms()

# Тепер можна безпечно імпортувати MAI-DX
try:
    from mai_dx import MaiDxOrchestrator
    MAI_DX_AVAILABLE = True
except ImportError as e:
    MAI_DX_AVAILABLE = False
    print(f"MAI-DX not available: {e}")

class SafeMaiDxOrchestrator:
    """Безпечна обгортка для MaiDxOrchestrator"""
    
    def __init__(self, *args, **kwargs):
        if not MAI_DX_AVAILABLE:
            raise ImportError("MAI-DX is not available")
            
        # Додаткові патчі безпосередньо перед створенням
        patch_rich_and_swarms()
        
        # Створюємо оригінальний orchestrator
        self.orchestrator = MaiDxOrchestrator(*args, **kwargs)
        
    def run(self, *args, **kwargs):
        """Виконати діагностику з додатковим захистом"""
        # Переконуємося що патчі активні
        patch_rich_and_swarms()
        
        # Тимчасово змінюємо stderr щоб приховати помилки Rich
        import io
        original_stderr = sys.stderr
        sys.stderr = io.StringIO()
        
        try:
            # Виконуємо оригінальний run
            result = self.orchestrator.run(*args, **kwargs)
            return result
        finally:
            # Відновлюємо stderr
            sys.stderr = original_stderr
            
    def __getattr__(self, name):
        """Проксі для всіх інших методів"""
        return getattr(self.orchestrator, name)

# Експортуємо класи
__all__ = ['SafeMaiDxOrchestrator', 'MAI_DX_AVAILABLE', 'patch_rich_and_swarms']