From f8901d7eb97d5a03a01aae4e0593e7f61104c23c Mon Sep 17 00:00:00 2001 From: bucolucas Date: Sun, 18 Aug 2024 18:16:19 -0500 Subject: [PATCH 1/2] Add SQLite database handler --- database.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 database.py diff --git a/database.py b/database.py new file mode 100644 index 0000000..5429bee --- /dev/null +++ b/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() \ No newline at end of file From 589d4644c70748021a64d5aaa72ae091754ad3c2 Mon Sep 17 00:00:00 2001 From: bucolucas Date: Sun, 18 Aug 2024 18:17:07 -0500 Subject: [PATCH 2/2] Update telegram_inference_bot.py to use SQLite database --- telegram_inference_bot.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py index 527d05b..f07437c 100644 --- a/telegram_inference_bot.py +++ b/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)