File size: 4,307 Bytes
c922f8b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
"""
Fix Supabase Tables for GAIA Memory

This module applies monkey patches to memory classes to ensure they use the correct tables.
It also ensures that all required tables exist in Supabase.

NOTICE: This is a proxy file that re-exports from the consolidated implementation.
Please use imports from 'memory' module directly in new code instead of importing
from this file.

Usage:
    from memory.fix_supabase_tables import apply_monkey_patch, ensure_tables_exist
"""

import logging
from memory.supabase_memory_consolidated import (
    SupabaseMemory,
    ConversationMemory,
    ResultCache,
    WorkingMemory,
    verify_tables_exist
)
from agent.config import get_memory_config

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

def apply_monkey_patch() -> bool:
    """
    Apply monkey patches to memory classes to ensure they use the correct tables.
    This is now handled by the consolidated implementation.
    
    Returns:
        bool: True if the patches were applied successfully, False otherwise
    """
    logger.info("Using consolidated implementation with proper table handling...")
    
    # The consolidated implementation already ensures correct table usage,
    # so we just need to return True here for backward compatibility
    return True

def ensure_tables_exist() -> bool:
    """
    Ensure all required tables exist in Supabase.
    
    Returns:
        bool: True if all tables exist, False otherwise
    """
    logger.info("Ensuring all required tables exist using consolidated implementation...")
    
    try:
        memory_config = get_memory_config()
        memory_config["enabled"] = True
        
        memory = SupabaseMemory(config=memory_config)
        
        return verify_tables_exist(memory)
    
    except Exception as e:
        logger.error(f"Error ensuring tables exist: {str(e)}")
        return False

def test_tables() -> bool:
    """
    Test all memory tables by exercising their APIs.
    
    Returns:
        bool: True if all tests pass, False otherwise
    """
    logger.info("Testing memory tables...")
    
    try:
        memory_config = get_memory_config()
        memory_config["enabled"] = True
        
        memory = SupabaseMemory(config=memory_config)
        
        # Test conversation memory
        conversation = ConversationMemory(memory, "test-conversation")
        conversation.add_message("user", "test message")
        messages = conversation.get_messages()
        
        if not messages or messages[0].get("content") != "test message":
            logger.error(f"Conversation memory test failed")
            return False
        
        conversation.clear()
        logger.info("✅ Conversation memory test passed")
        
        # Test result cache
        cache = ResultCache(memory)
        cache.cache_result("test_query", "test_answer")
        result = cache.get_result("test_query")
        
        if result != "test_answer":
            logger.error(f"Result cache test failed")
            return False
        
        cache.clear()
        logger.info("✅ Result cache test passed")
        
        # Test working memory
        working = WorkingMemory(memory, "test-session")
        working.store_intermediate_result("test_key", "test_value")
        result = working.get_intermediate_result("test_key")
        
        if result != "test_value":
            logger.error(f"Working memory test failed")
            return False
        
        working.clear()
        logger.info("✅ Working memory test passed")
        
        logger.info("✅ All memory table tests passed")
        return True
    
    except Exception as e:
        logger.error(f"Error testing memory tables: {str(e)}")
        return False

if __name__ == "__main__":
    if apply_monkey_patch() and ensure_tables_exist():
        if test_tables():
            logger.info("All memory tables fixed and tested successfully")
            import sys
            sys.exit(0)
        else:
            logger.error("Memory table tests failed")
            import sys
            sys.exit(1)
    else:
        logger.error("Failed to fix memory tables")
        import sys
        sys.exit(1)