Latest version of Agent

Merging latest status of LLM agent to main branch

llm implementations/db_manager.py

@@ -6,13 +6,42 @@ class DatabaseHandler:
     def __init__(self, db_name: str = "code_optimizer.db"):
         self.db_name = db_name
         self.init_database()
+        self._connection = None
+        self._cursor = None
+
+    property
+    def connection(self):
+        if self._connection is None:
+            self._connection = sqlite3.connect(self.db_name)
+            self._connection.row_factory = sqlite3.Row
+        return self._connection
+
+    @property
+    def cursor(self):
+        if self._cursor is None:
+            self._cursor = self.connection.cursor()
+        return self._cursor
+
+    def close(self):
+        if self._cursor:
+            self._cursor.close()
+            self._cursor = None
+        if self._connection:
+            self._connection.close()
+            self._connection = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
 
     def init_database(self):
         """Initialize the database with required tables"""
         with sqlite3.connect(self.db_name) as conn:
             cursor = conn.cursor()
             
-            # Create conversations table
+            # Add indexes for better query performance
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS conversations (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -21,7 +50,6 @@ class DatabaseHandler:
                 )
             """)
             
-            # Create entries table for individual code optimizations
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS optimization_entries (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -29,11 +57,10 @@ class DatabaseHandler:
                     original_code TEXT NOT NULL,
                     num_techniques INTEGER NOT NULL,
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                    FOREIGN KEY (conversation_id) REFERENCES conversations (id)
+                    FOREIGN KEY (conversation_id) REFERENCES conversations (id) ON DELETE CASCADE
                 )
             """)
             
-            # Create results table for optimization results
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS optimization_results (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -41,12 +68,46 @@ class DatabaseHandler:
                     improved_code TEXT NOT NULL,
                     output TEXT,
                     execution_time FLOAT,
-                    techniques_applied TEXT,
-                    FOREIGN KEY (entry_id) REFERENCES optimization_entries (id)
+                    memory_usage FLOAT,
+                    techniques TEXT,
+                    FOREIGN KEY (entry_id) REFERENCES optimization_entries (id) ON DELETE CASCADE
                 )
             """)
+            
+            # Add indexes
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_conversation_created 
+                ON conversations(created_at)
+            """)
+            
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_entries_conversation 
+                ON optimization_entries(conversation_id)
+            """)
+            
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_results_entry 
+                ON optimization_results(entry_id)
+            """)
+            
             conn.commit()
 
+    def execute_with_retry(self, query, params=None, max_retries=3):
+        """Execute SQL with retry logic for better reliability"""
+        for attempt in range(max_retries):
+            try:
+                if params is None:
+                    self.cursor.execute(query)
+                else:
+                    self.cursor.execute(query, params)
+                self.connection.commit()
+                return self.cursor
+            except sqlite3.OperationalError as e:
+                if attempt == max_retries - 1:
+                    raise
+                time.sleep(0.1 * (attempt + 1))
+                continue
+
     def create_conversation(self, title: str) -> int:
         """Create a new conversation and return its ID"""
         with sqlite3.connect(self.db_name) as conn:
@@ -77,13 +138,14 @@ class DatabaseHandler:
             for result in results:
                 cursor.execute("""
                     INSERT INTO optimization_results 
-                    (entry_id, improved_code, output, execution_time, techniques_applied)
-                    VALUES (?, ?, ?, ?, ?)
+                    (entry_id, improved_code, output, execution_time, memory_usage, techniques)
+                    VALUES (?, ?, ?, ?, ?, ?)
                 """, (
                     entry_id,
                     result['code'],
                     result['output'],
                     result['execution_time'],
+                    result['memory_usage'],  # Insert memory_usage
                     result['techniques']
                 ))
             
@@ -121,7 +183,7 @@ class DatabaseHandler:
             conn.row_factory = sqlite3.Row
             cursor = conn.cursor()
             cursor.execute("""
-                SELECT improved_code, output, execution_time, techniques_applied
+                SELECT improved_code, output, execution_time, memory_usage, techniques
                 FROM optimization_results
                 WHERE entry_id = ?
             """, (entry_id,))
@@ -164,7 +226,7 @@ class DatabaseHandler:
         except Exception as e:
             print(f"Error deleting conversation: {e}")
             raise
-    
+
     def clear_all_conversations(self) -> None:
         """
         Delete all conversations and their associated entries and results
@@ -184,4 +246,4 @@ class DatabaseHandler:
                 conn.commit()
         except Exception as e:
             print(f"Error clearing all conversations: {e}")
-            raise
\ No newline at end of file
+            raise