db.py

import sqlite3
import os
import time

# Use a non-hidden filename in the current directory
DB_NAME = "model_tests.db"  # No leading dot
DB_PATH = os.path.join(os.getcwd(), DB_NAME)  # Use absolute path with getcwd()

def init_db():
    """Initialize the SQLite database with required tables if they don't exist"""
    try:
        print(f"Initializing database at {DB_PATH}")
        # Ensure the directory exists
        os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
        
        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(f"Test database initialized successfully at {DB_PATH}")
        return True
    except Exception as e:
        print(f"Error initializing database: {e}")
        return False

def save_test_result(analysis_mode, system_prompt, input_content, model_response, 
                    generation_time, tokens_generated, temperature, top_p, max_length):
    """Save a test result to the database"""
    try:
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        
        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()
        
        return test_id
    except Exception as e:
        print(f"Error saving test result: {e}")
        return None

def get_test_history(limit=10):
    """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 []

def get_test_details(test_id):
    """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