Update db.py

db.py

@@ -1,89 +1,104 @@
-import gradio as gr
-import torch
-import time
-import spaces
 import sqlite3
 import os
-import json
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import time
 
-# --- Database Setup ---
-# Use the persistent storage location in Hugging Face Spaces
-DB_PATH = os.path.join("/data", "model_tests.db")
+# Use the current directory instead of /data
+DB_NAME = "model_tests.db"
+DB_PATH = os.path.join(".", DB_NAME)
 
 def init_db():
     """Initialize the SQLite database with required tables if they don't exist"""
-    with sqlite3.connect(DB_PATH) as conn:
-        cursor = conn.cursor()
-        
-        # Create table for model inputs and outputs
-        cursor.execute('''
-        CREATE TABLE IF NOT EXISTS model_tests (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            analysis_mode TEXT NOT NULL,
-            system_prompt TEXT NOT NULL,
-            input_content TEXT,
-            model_response TEXT NOT NULL,
-            generation_time REAL NOT NULL,
-            tokens_generated INTEGER NOT NULL,
-            temperature REAL NOT NULL,
-            top_p REAL NOT NULL,
-            max_length INTEGER NOT NULL,
-            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
-        )
-        ''')
+    try:
+        print(f"Initializing database at {DB_PATH}")
+        # Ensure the directory exists
+        os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
         
-        conn.commit()
-    print("Test database initialized successfully")
+        with sqlite3.connect(DB_PATH) as conn:
+            cursor = conn.cursor()
+            
+            # Create table for model inputs and outputs
+            cursor.execute('''
+            CREATE TABLE IF NOT EXISTS model_tests (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                analysis_mode TEXT NOT NULL,
+                system_prompt TEXT NOT NULL,
+                input_content TEXT,
+                model_response TEXT NOT NULL,
+                generation_time REAL NOT NULL,
+                tokens_generated INTEGER NOT NULL,
+                temperature REAL NOT NULL,
+                top_p REAL NOT NULL,
+                max_length INTEGER NOT NULL,
+                timestamp TEXT NOT NULL
+            )
+            ''')
+            
+            conn.commit()
+        print("Test database initialized successfully")
+        return True
+    except Exception as e:
+        print(f"Error initializing database: {e}")
+        return False
 
-# Initialize the database when the app starts
-init_db()
-
-# Function to save test results
 def save_test_result(analysis_mode, system_prompt, input_content, model_response, 
                     generation_time, tokens_generated, temperature, top_p, max_length):
-    with sqlite3.connect(DB_PATH) as conn:
-        cursor = conn.cursor()
+    """Save a test result to the database"""
+    try:
+        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
         
-        cursor.execute('''
-        INSERT INTO model_tests 
-        (analysis_mode, system_prompt, input_content, model_response, 
-        generation_time, tokens_generated, temperature, top_p, max_length)
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
-        ''', (analysis_mode, system_prompt, input_content, model_response, 
-             generation_time, tokens_generated, temperature, top_p, max_length))
+        with sqlite3.connect(DB_PATH) as conn:
+            cursor = conn.cursor()
+            
+            cursor.execute('''
+            INSERT INTO model_tests 
+            (analysis_mode, system_prompt, input_content, model_response, 
+            generation_time, tokens_generated, temperature, top_p, max_length, timestamp)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            ''', (analysis_mode, system_prompt, input_content, model_response, 
+                generation_time, tokens_generated, temperature, top_p, max_length, timestamp))
+            
+            test_id = cursor.lastrowid
+            conn.commit()
         
-        test_id = cursor.lastrowid
-        conn.commit()
-    
-    return test_id
+        return test_id
+    except Exception as e:
+        print(f"Error saving test result: {e}")
+        return None
 
-# Function to retrieve test history
 def get_test_history(limit=10):
-    with sqlite3.connect(DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute('''
-        SELECT id, analysis_mode, timestamp, generation_time, tokens_generated 
-        FROM model_tests 
-        ORDER BY id DESC LIMIT ?
-        ''', (limit,))
-        history = cursor.fetchall()
-    
-    return history
+    """Get recent test history"""
+    try:
+        with sqlite3.connect(DB_PATH) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+            SELECT id, analysis_mode, timestamp, generation_time, tokens_generated 
+            FROM model_tests 
+            ORDER BY id DESC LIMIT ?
+            ''', (limit,))
+            history = cursor.fetchall()
+        
+        return history
+    except Exception as e:
+        print(f"Error getting test history: {e}")
+        return []
 
-# Function to get test details by ID
 def get_test_details(test_id):
-    with sqlite3.connect(DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute('''
-        SELECT * FROM model_tests WHERE id = ?
-        ''', (test_id,))
-        test = cursor.fetchone()
-        
-        if test:
-            # Convert to dictionary
-            columns = [col[0] for col in cursor.description]
-            test_dict = {columns[i]: test[i] for i in range(len(columns))}
-            return test_dict
-        else:
-            return None
+    """Get details for a specific test"""
+    try:
+        with sqlite3.connect(DB_PATH) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+            SELECT * FROM model_tests WHERE id = ?
+            ''', (test_id,))
+            test = cursor.fetchone()
+            
+            if test:
+                # Convert to dictionary
+                columns = [col[0] for col in cursor.description]
+                test_dict = {columns[i]: test[i] for i in range(len(columns))}
+                return test_dict
+            else:
+                return None
+    except Exception as e:
+        print(f"Error getting test details: {e}")
+        return None