File size: 3,004 Bytes
66dbebd |
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 |
"""
Database initialization and management
"""
import sqlite3
import logging
import os
from pathlib import Path
logger = logging.getLogger(__name__)
class DatabaseManager:
def __init__(self, db_path: str = "sessions.db"):
self.db_path = db_path
self.connection = None
self._init_db()
def _init_db(self):
"""Initialize database with required tables"""
try:
# Create database directory if needed
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
self.connection = sqlite3.connect(self.db_path, check_same_thread=False)
self.connection.row_factory = sqlite3.Row
# Create tables
self._create_tables()
logger.info(f"Database initialized at {self.db_path}")
except Exception as e:
logger.error(f"Database initialization failed: {e}")
# Fallback to in-memory database
self.connection = sqlite3.connect(":memory:", check_same_thread=False)
self._create_tables()
logger.info("Using in-memory database as fallback")
def _create_tables(self):
"""Create required database tables"""
cursor = self.connection.cursor()
# Sessions table
cursor.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
last_activity TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
context_data TEXT,
user_metadata TEXT
)
""")
# Interactions table
cursor.execute("""
CREATE TABLE IF NOT EXISTS interactions (
interaction_id TEXT PRIMARY KEY,
session_id TEXT REFERENCES sessions(session_id),
user_input TEXT NOT NULL,
agent_trace TEXT,
final_response TEXT,
processing_time INTEGER,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
self.connection.commit()
logger.info("Database tables created successfully")
def get_connection(self):
"""Get database connection"""
return self.connection
def close(self):
"""Close database connection"""
if self.connection:
self.connection.close()
logger.info("Database connection closed")
# Global database instance
db_manager = None
def init_database(db_path: str = "sessions.db"):
"""Initialize global database instance"""
global db_manager
if db_manager is None:
db_manager = DatabaseManager(db_path)
return db_manager
def get_db():
"""Get database connection"""
global db_manager
if db_manager is None:
init_database()
return db_manager.get_connection()
# Initialize database on import
init_database()
|