Files
cyclop/telegram_inference_bot.py
T
2024-08-16 12:43:59 -05:00

123 lines
4.6 KiB
Python

import os
import importlib
import inspect
from telegram import Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
from openai import OpenAI
from dotenv import load_dotenv
from tools.base_tool import BaseTool
# Load environment variables
load_dotenv()
client = OpenAI()
# Set up Telegram bot
TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
# Dictionary to store conversation history for each user
conversation_history = {}
# Load tools
tools = []
tools_dir = os.path.join(os.path.dirname(__file__), 'tools')
for filename in os.listdir(tools_dir):
if filename.endswith('.py') and filename != '__init__.py' and filename != 'base_tool.py':
module_name = f'tools.{filename[:-3]}'
module = importlib.import_module(module_name)
for name, obj in inspect.getmembers(module):
if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool:
tools.append(obj())
# Collect all function definitions
functions = []
for tool in tools:
functions.extend(tool.get_functions())
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await update.message.reply_text("Hello! I'm your AI assistant. How can I help you today?")
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]
await update.message.reply_text("Conversation history cleared. Let's start fresh!")
async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
try:
user_id = update.effective_user.id
user_message = update.message.text
# Initialize conversation history for new users
if user_id not in conversation_history:
conversation_history[user_id] = []
# Add user message to conversation history
conversation_history[user_id].append({"role": "user", "content": user_message})
# Prepare messages for OpenAI API
messages = [{"role": "system", "content": "You are a helpful assistant."}] + conversation_history[user_id]
# Call OpenAI API for inference
response = client.chat.completions.create(
model="gpt-4",
messages=messages,
functions=functions,
function_call="auto"
)
# Extract the assistant's reply
assistant_message = response.choices[0].message
if assistant_message.function_call:
# Execute the function
function_name = assistant_message.function_call.name
function_args = assistant_message.function_call.arguments
for tool in tools:
if function_name in [f["name"] for f in tool.get_functions()]:
function_response = tool.execute(function_name, **eval(function_args))
messages.append({
"role": "function",
"name": function_name,
"content": function_response
})
# Call API again to get the final response
response = client.chat.completions.create(
model="gpt-4",
messages=messages
)
assistant_reply = response.choices[0].message.content
break
else:
assistant_reply = assistant_message.content
# Add assistant's reply to conversation history
conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
# Trim conversation history if it gets too long (e.g., keep last 10 messages)
if len(conversation_history[user_id]) > 10:
conversation_history[user_id] = conversation_history[user_id][-10:]
# Send the reply back to the user
await update.message.reply_text(assistant_reply)
except Exception as e:
print(f"An error occurred: {str(e)}")
await update.message.reply_text("Sorry, an error occurred while processing your request.")
def main() -> None:
# Create the Application and pass it your bot's token
application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
# Add handlers
application.add_handler(CommandHandler("start", start))
application.add_handler(CommandHandler("clear", clear))
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
# Start the Bot
print("Bot is running...")
application.run_polling()
if __name__ == '__main__':
main()