src/gaia/utils/verify_components.py

"""
Step-by-step verification of GAIA components
This script tests each component individually to identify issues
"""
import os
import sys
import time
import logging
import traceback
from typing import Dict, Any, List

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("verify_components")

class ComponentVerifier:
    """Verifies each component of the GAIA system step by step"""
    
    def __init__(self):
        self.results = {}
        self.env_vars_set = False
        self.supabase_connected = False
        
    def verify_environment(self):
        """Verify environment variables are set"""
        try:
            logger.info("Checking environment variables...")
            required_vars = ["OPENAI_API_KEY", "SUPABASE_URL", "SUPABASE_KEY"]
            missing = []
            
            for var in required_vars:
                if not os.environ.get(var):
                    missing.append(var)
            
            if missing:
                logger.error(f"Missing required environment variables: {', '.join(missing)}")
                return False
            
            logger.info("All required environment variables are set")
            self.env_vars_set = True
            return True
        except Exception as e:
            logger.error(f"Error checking environment variables: {e}")
            return False
    
    def verify_imports(self):
        """Verify that all necessary modules can be imported"""
        try:
            logger.info("Verifying imports...")
            
            # 1. Try importing OpenAI
            try:
                from openai import OpenAI
                logger.info("OpenAI module imported successfully")
                self.results["openai_import"] = True
            except Exception as e:
                logger.error(f"Failed to import OpenAI: {e}")
                self.results["openai_import"] = False
            
            # 2. Try importing Supabase
            try:
                from supabase import create_client
                logger.info("Supabase module imported successfully")
                self.results["supabase_import"] = True
            except Exception as e:
                logger.error(f"Failed to import Supabase: {e}")
                self.results["supabase_import"] = False
            
            # 3. Try importing agent.config
            try:
                from agent.config import get_memory_config
                logger.info("Agent config module imported successfully")
                self.results["agent_config_import"] = True
            except Exception as e:
                logger.error(f"Failed to import agent.config: {e}")
                self.results["agent_config_import"] = False
            
            # 4. Try importing memory module
            try:
                from memory.supabase_memory import (
                    SupabaseMemory,
                    ConversationMemory,
                    ResultCache,
                    WorkingMemory
                )
                logger.info("Memory module imported successfully")
                self.results["memory_import"] = True
            except Exception as e:
                logger.error(f"Failed to import memory module: {e}")
                self.results["memory_import"] = False
            
            # 5. Try importing memory fix
            try:
                from memory.fix_supabase_tables import apply_monkey_patch, ensure_tables_exist
                logger.info("Memory fix module imported successfully")
                self.results["memory_fix_import"] = True
            except Exception as e:
                logger.error(f"Failed to import memory fix module: {e}")
                self.results["memory_fix_import"] = False
            
            # 6. Try importing agent.graph
            try:
                from agent.graph import create_agent_graph, run_agent_graph
                logger.info("Agent graph module imported successfully")
                self.results["agent_graph_import"] = True
            except Exception as e:
                logger.error(f"Failed to import agent.graph: {e}")
                self.results["agent_graph_import"] = False
            
            # 7. Try importing agent
            try:
                from agent.agent import GAIAAgent
                logger.info("GAIAAgent imported successfully")
                self.results["agent_import"] = True
            except Exception as e:
                logger.error(f"Failed to import GAIAAgent: {e}")
                self.results["agent_import"] = False
            
            # 8. Try importing tools
            try:
                from tools.web_tools import duckduckgo_search
                logger.info("Web tools imported successfully")
                self.results["web_tools_import"] = True
            except Exception as e:
                logger.error(f"Failed to import web tools: {e}")
                self.results["web_tools_import"] = False
            
            return True
        except Exception as e:
            logger.error(f"Error verifying imports: {e}")
            return False
    
    def verify_supabase_connection(self):
        """Verify connection to Supabase"""
        if not self.env_vars_set:
            logger.error("Environment variables not verified, skipping Supabase connection test")
            return False
        
        try:
            logger.info("Verifying Supabase connection...")
            from supabase import create_client
            
            supabase_url = os.environ.get("SUPABASE_URL")
            supabase_key = os.environ.get("SUPABASE_KEY")
            
            client = create_client(supabase_url, supabase_key)
            
            # Try a simple query
            response = client.table("_memory").select("count(*)", count="exact").execute()
            count = response.count if hasattr(response, 'count') else 0
            
            logger.info(f"Successfully connected to Supabase, found {count} records in _memory table")
            self.supabase_connected = True
            return True
        except Exception as e:
            logger.error(f"Failed to connect to Supabase: {e}")
            logger.error(traceback.format_exc())
            return False
    
    def apply_supabase_fixes(self):
        """Verify and apply Supabase table fixes"""
        if not self.supabase_connected:
            logger.error("Supabase connection not verified, skipping table fixes")
            return False
        
        try:
            logger.info("Applying Supabase table fixes...")
            from memory.fix_supabase_tables import apply_monkey_patch, ensure_tables_exist
            
            # Ensure tables exist
            ensure_tables_exist()
            logger.info("Successfully ensured all required tables exist")
            
            # Apply monkey patch
            apply_monkey_patch()
            logger.info("Successfully applied monkey patch for Supabase tables")
            
            return True
        except Exception as e:
            logger.error(f"Failed to apply Supabase fixes: {e}")
            logger.error(traceback.format_exc())
            return False
    
    def initialize_memory(self):
        """Initialize and test memory system"""
        try:
            logger.info("Initializing memory system...")
            from memory.supabase_memory import SupabaseMemory
            from agent.config import get_memory_config
            
            # Get memory config
            memory_config = get_memory_config()
            memory_config["enabled"] = True
            
            # Initialize memory
            memory = SupabaseMemory(memory_config)
            
            # Test basic operations
            test_key = f"test_key_{int(time.time())}"
            test_value = {"message": "Test value", "timestamp": time.time()}
            
            # Store
            memory.store(test_key, test_value)
            logger.info(f"Successfully stored test value with key: {test_key}")
            
            # Retrieve
            retrieved = memory.retrieve(test_key)
            
            if retrieved and retrieved.get("message") == "Test value":
                logger.info("Successfully retrieved test value")
                return True
            else:
                logger.error(f"Retrieved value doesn't match. Got: {retrieved}")
                return False
        except Exception as e:
            logger.error(f"Failed to initialize memory: {e}")
            logger.error(traceback.format_exc())
            return False
    
    def test_memory_subtypes(self):
        """Test all memory subtypes"""
        try:
            logger.info("Testing memory subtypes...")
            from memory.supabase_memory import (
                SupabaseMemory, 
                ConversationMemory, 
                ResultCache, 
                WorkingMemory
            )
            from agent.config import get_memory_config
            
            # Initialize base memory
            memory_config = get_memory_config()
            memory_config["enabled"] = True
            memory = SupabaseMemory(memory_config)
            
            # Test session ID
            session_id = f"test_session_{int(time.time())}"
            
            # Test conversation memory
            try:
                conversation = ConversationMemory(memory, session_id)
                conversation.add_message("user", "Test question")
                conversation.add_message("assistant", "Test answer")
                messages = conversation.get_messages()
                
                if len(messages) >= 2:
                    logger.info("ConversationMemory working correctly")
                    self.results["conversation_memory"] = True
                else:
                    logger.error(f"ConversationMemory not working correctly. Got {len(messages)} messages")
                    self.results["conversation_memory"] = False
            except Exception as e:
                logger.error(f"Error testing ConversationMemory: {e}")
                self.results["conversation_memory"] = False
            
            # Test result cache
            try:
                cache = ResultCache(memory)
                cache_key = f"test_cache_{int(time.time())}"
                cache.cache_result(cache_key, "Test cache result")
                cached = cache.get_result(cache_key)
                
                if cached == "Test cache result":
                    logger.info("ResultCache working correctly")
                    self.results["result_cache"] = True
                else:
                    logger.error(f"ResultCache not working correctly. Got: {cached}")
                    self.results["result_cache"] = False
            except Exception as e:
                logger.error(f"Error testing ResultCache: {e}")
                self.results["result_cache"] = False
            
            # Test working memory
            try:
                working = WorkingMemory(memory, session_id)
                working.store_intermediate_result("test_step", "Test intermediate result")
                results = working.list_intermediate_results()
                
                if results and any(r.get("step") == "test_step" for r in results):
                    logger.info("WorkingMemory working correctly")
                    self.results["working_memory"] = True
                else:
                    logger.error(f"WorkingMemory not working correctly. Got: {results}")
                    self.results["working_memory"] = False
            except Exception as e:
                logger.error(f"Error testing WorkingMemory: {e}")
                self.results["working_memory"] = False
            
            # Overall result
            return (self.results.get("conversation_memory", False) and 
                   self.results.get("result_cache", False) and 
                   self.results.get("working_memory", False))
        
        except Exception as e:
            logger.error(f"Failed to test memory subtypes: {e}")
            logger.error(traceback.format_exc())
            return False
    
    def run_all_tests(self):
        """Run all verification tests in sequence"""
        logger.info("=== STARTING COMPONENT VERIFICATION ===")
        
        # Step 1: Verify environment variables
        self.results["environment"] = self.verify_environment()
        
        # Step 2: Verify imports
        self.results["imports"] = self.verify_imports()
        
        # Step 3: Verify Supabase connection
        if self.results["environment"]:
            self.results["supabase_connection"] = self.verify_supabase_connection()
        else:
            logger.warning("Skipping Supabase connection test due to environment variables issue")
            self.results["supabase_connection"] = False
        
        # Step 4: Apply Supabase fixes
        if self.results["supabase_connection"]:
            self.results["supabase_fixes"] = self.apply_supabase_fixes()
        else:
            logger.warning("Skipping Supabase fixes due to connection issue")
            self.results["supabase_fixes"] = False
        
        # Step 5: Initialize memory
        if self.results["supabase_fixes"]:
            self.results["memory_initialize"] = self.initialize_memory()
        else:
            logger.warning("Skipping memory initialization due to Supabase fixes issue")
            self.results["memory_initialize"] = False
        
        # Step 6: Test memory subtypes
        if self.results["memory_initialize"]:
            self.results["memory_subtypes"] = self.test_memory_subtypes()
        else:
            logger.warning("Skipping memory subtypes test due to initialization issue")
            self.results["memory_subtypes"] = False
        
        # Print summary
        logger.info("\n=== COMPONENT VERIFICATION RESULTS ===")
        for component, result in self.results.items():
            status = "SUCCESS" if result else "FAILED"
            logger.info(f"{component}: {status}")
        
        # Return overall result
        return all(self.results.values())

if __name__ == "__main__":
    verifier = ComponentVerifier()
    success = verifier.run_all_tests()
    
    if success:
        logger.info("\n✅ All components verified successfully!")
        sys.exit(0)
    else:
        logger.error("\n❌ Component verification failed!")
        sys.exit(1)