Update telegram_inference_bot.py to use SQLite database

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