Spaces:
Build error
Build error
File size: 3,248 Bytes
3132f43 |
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 99 100 101 102 103 104 105 106 107 108 109 110 111 |
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()
|