Merge pull request #56 from bucolucas/add-sqlite-storage

Add SQLite storage for conversation persistence

database.py

@@ -0,0 +1,42 @@
+import sqlite3
+import json
+from typing import List, Dict, Any
+
+class Database:
+    def __init__(self, db_file='telegram_bot.db'):
+        self.db_file = db_file
+        self.conn = sqlite3.connect(self.db_file)
+        self.create_tables()
+
+    def create_tables(self):
+        with self.conn:
+            self.conn.execute('''
+                CREATE TABLE IF NOT EXISTS conversations (
+                    user_id INTEGER PRIMARY KEY,
+                    history TEXT
+                )
+            ''')
+
+    def save_conversation(self, user_id: int, history: List[Dict[str, Any]]):
+        with self.conn:
+            self.conn.execute('''
+                INSERT OR REPLACE INTO conversations (user_id, history)
+                VALUES (?, ?)
+            ''', (user_id, json.dumps(history)))
+
+    def get_conversation(self, user_id: int) -> List[Dict[str, Any]]:
+        cursor = self.conn.execute('SELECT history FROM conversations WHERE user_id = ?', (user_id,))
+        result = cursor.fetchone()
+        if result:
+            return json.loads(result[0])
+        return []
+
+    def clear_conversation(self, user_id: int):
+        with self.conn:
+            self.conn.execute('DELETE FROM conversations WHERE user_id = ?', (user_id,))
+
+    def close(self):
+        self.conn.close()
+
+# Create a global instance of the Database class
+db = Database()

telegram_inference_bot.py

@@ -9,6 +9,7 @@ from telegram.ext import Application, CommandHandler, MessageHandler, filters, C
 from dotenv import load_dotenv
 from tools.base_tool import BaseTool
 from anthropic import Anthropic
+from database import db  # Import the database handler
 
 # Load environment variables
 load_dotenv()
@@ -31,9 +32,6 @@ TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
 with open("prompts/developer_prompt.txt", "r") as file:
     system_prompt = file.read().strip()
 
-# Dictionary to store conversation history for each user
-conversation_history = {}
-
 # Dictionary to store processing status for each user
 processing_status = {}
 
@@ -61,8 +59,7 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
 
 async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     user_id = update.effective_user.id
-    if user_id in conversation_history:
-        del conversation_history[user_id]
+    db.clear_conversation(user_id)
     for tool in tools:
         tool.clear()
 
@@ -88,16 +85,14 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
 
         logging.info(f"Message from user {user_id}: {user_message}")
 
-        if user_id not in conversation_history:
-            conversation_history[user_id] = []
-
-        conversation_history[user_id].append({"role": "user", "content": user_message})
+        conversation_history = db.get_conversation(user_id)
+        conversation_history.append({"role": "user", "content": user_message})
 
         # Send initial status message
         status_message = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]]))
         processing_status[user_id] = {"processing": True, "message_id": status_message.message_id}
 
-        messages = conversation_history[user_id]
+        messages = conversation_history
 
         response = get_chat_response(messages)
         tool_calls = []
@@ -140,10 +135,13 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
 
         if (toolUseCount == 0):
             assistant_reply = response.content
-            conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
+            messages.append({"role": "assistant", "content": assistant_reply})
         
-        if len(conversation_history[user_id]) > 20:
-            conversation_history[user_id] = conversation_history[user_id][-20:]
+        if len(messages) > 20:
+            messages = messages[-20:]
+
+        # Save the updated conversation history
+        db.save_conversation(user_id, messages)
 
         # Remove the status message
         await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message.message_id)