db_ops.py

import sqlite3

from typing import Any, Dict, List

# Database Design 
def init_db():
    """Initialize SQLite database with required tables"""
    conn = sqlite3.connect("chatbot.db")
    c = conn.cursor()

    # Create conversations table 
    c.execute("""
CREATE TABLE IF NOT EXISTS conversations (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id TEXT NOT NULL,  
    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
    summary TEXT        
)
""")
    
    # Create messages table 
    c.execute("""
CREATE TABLE IF NOT EXISTS messages (
    id INTEGER PRIMARY KEY AUTOINCREMENT, 
    conversation_id INTEGER,
    role TEXT NOT NULL, 
    content TEXT NOT NULL, 
    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, 
    FOREIGN KEY (conversation_id) REFERENCES conversations (id)                        
)
""")
    
    conn.commit()
    conn.close()

# Adding Conversations
def store_message(session_id:str, role:str, content:str)->None:
    """Store a message in the database"""
    conn = sqlite3.connect("chatbot.db")
    c = conn.cursor()

    # Get current conversation id or create new one 
    c.execute("SELECT id FROM conversations WHERE session_id=? ORDER BY timestamp DESC LIMIT 1",
               (session_id,))
    result = c.fetchone()

    if result is None:
        # Create a new conversation 
        c.execute("INSERT INTO conversations (session_id) VALUES (?)", (session_id,))
        conversation_id = c.lastrowid
    else:
        conversation_id = result[0] 

    # Store message 
    c.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, ?, ?)", 
              (conversation_id, role, content))
    
    conn.commit()
    conn.close()

# Retrieve Conversations
def get_recent_conversations(session_id:str, limit:int=5)->List[Dict[str, Any]]:
    """Retrieve the most recent conversation for a session"""
    conn = sqlite3.connect("chatbot.db")
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    # Get recent conversation ID 
    c.execute('''
SELECT id, timestamp FROM conversations 
WHERE session_id = ?
ORDER BY timestamp DESC 
LIMIT ?
''', (session_id, limit))
    
    conversations = []

    for conv_row in c.fetchall():
        conv_id = conv_row["id"]

        # Get messages for this conversations 
        c.execute('''
SELECT role, content FROM messages 
WHERE conversation_id = ? 
ORDER BY timestamp ASC
''', (conv_id,)) 
        
        messages = []
        for msg in c.fetchall():
            messages.append({"role": msg["role"], "content": msg["content"]})

        conversations.append({
            "id": conv_id, 
            "timestamp": conv_row["timestamp"], 
            "messages":messages
        })

    conn.close()
    return conversations 

# Storing Conversations 
def store_conversation_summary(session_id:str, conv_id:int, summary:str)->None:
    """Store a generated summary for a conversation"""
    conn = sqlite3.connect("chatbot.db")
    c = conn.cursor()

    # Update the summary for the conversation 
    c.execute('''UPDATE conversations SET summary = ? WHERE id = ?''', (summary, conv_id))
    conn.commit()
    conn.close()