src/gaia/utils/fix_graph_imports.py

#!/usr/bin/env python
"""
Fix for the ModuleNotFoundError: No module named 'langchain.tools.render' error.

This script patches the agent/graph.py file to remove the unused import that's
causing the runtime error.
"""

import os
import re
import sys
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def patch_graph_file():
    """Patch the agent/graph.py file to remove the problematic import."""
    graph_file = "agent/graph.py"
    
    if not os.path.exists(graph_file):
        logger.error(f"Could not find {graph_file}")
        return False
    
    try:
        with open(graph_file, "r") as f:
            content = f.read()
        
        # Create a backup
        backup_file = f"{graph_file}.bak"
        with open(backup_file, "w") as f:
            f.write(content)
        logger.info(f"Created backup at {backup_file}")
        
        # Remove the problematic import
        pattern = r"from langgraph\.prebuilt import ToolNode\n"
        updated_content = re.sub(pattern, "# Removed problematic import: from langgraph.prebuilt import ToolNode\n", content)
        
        with open(graph_file, "w") as f:
            f.write(updated_content)
        
        logger.info(f"Successfully patched {graph_file}")
        return True
    
    except Exception as e:
        logger.error(f"Error patching file: {str(e)}")
        return False

def main():
    """Main function to run the patch."""
    logger.info("Starting graph.py import patch...")
    
    if patch_graph_file():
        logger.info("Patch applied successfully!")
        logger.info("This fixes the 'ModuleNotFoundError: No module named 'langchain.tools.render'' error")
        logger.info("by removing an unused import from agent/graph.py")
    else:
        logger.error("Failed to apply patch")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())