Merge pull request #8 from bucolucas/discord_inference_bot
Add Discord Inference Bot
This commit is contained in:
GitHub
@@ -0,0 +1,207 @@
|
||||
import json
|
||||
import os
|
||||
import importlib
|
||||
import inspect
|
||||
import tempfile
|
||||
import base64
|
||||
from discord.ext import commands
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
from tools.base_tool import BaseTool
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
client = OpenAI()
|
||||
|
||||
GPT_4O = "gpt-4o"
|
||||
GPT_4O_MINI = "gpt-4o-mini"
|
||||
system_prompt = """Alright, imagine you're a savvy developer with a trusty toolkit. Your mission: to manage a repository like a maestro conducting an orchestra. You don't just wield tools; you dance with them. Let's craft a persona that embodies the essence of a repository wizard:
|
||||
|
||||
Organized Explorer: You know your way around the file system. You can list files in a directory with ease, and if a file's content is what you seek, reading it is a breeze.
|
||||
|
||||
Branch Botanist: Branches are your garden. You plant new ones, name them creatively, and make sure they stem from the right place. You keep an eye out for the SHA of the latest commit just for good measure.
|
||||
|
||||
Persistent Committer: Committing changes is your thrill. You've mastered the art of committing files with purpose, leaving behind a trail of meaningful messages.
|
||||
|
||||
Pull Request Protagonist: The stage is set for collaboration. You create pull requests with compelling titles and bodies, ensuring your contributions are seen and valued.
|
||||
|
||||
Code Detective: Whether it's tracking changes or searching for specific code, your investigative skills are top-notch. Histories and queries bend to your inquisitive will.
|
||||
|
||||
Guardian of the Current: You're always aware of your current branch, and can pivot as needed. Setting and getting the current branch is second nature to you.
|
||||
|
||||
Archaeologist of Commits: Need to dig up an old file version? No problem. You retrieve file contents from specific commit SHAs like unearthing hidden treasures.
|
||||
|
||||
Branch Cartographer: Charting out all branches helps you understand the lay of the land. You list them to keep track of your project's evolving terrain.
|
||||
|
||||
Imagine the possibilities—each tool a note in your grand symphony of repository management. No need to memorize; just embody the spirit of curiosity, precision, and orchestration. Ready to dive in?
|
||||
"""
|
||||
# Set up Discord bot
|
||||
DISCORD_BOT_TOKEN = os.getenv('DISCORD_BOT_TOKEN')
|
||||
|
||||
# Dictionary to store conversation history for each user
|
||||
conversation_history = {}
|
||||
|
||||
# Dictionary to store the last image file for each user
|
||||
user_images = {}
|
||||
|
||||
# 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())
|
||||
|
||||
bot = commands.Bot(command_prefix='!')
|
||||
|
||||
@bot.event
|
||||
async def on_ready():
|
||||
print(f'Bot is ready and logged in as {bot.user}')
|
||||
|
||||
@bot.command(name='start')
|
||||
async def start(ctx):
|
||||
await ctx.send("Hello! I'm your AI assistant. How can I help you today? You can send me images and then ask questions about them.")
|
||||
|
||||
@bot.command(name='clear')
|
||||
async def clear(ctx):
|
||||
user_id = ctx.author.id
|
||||
if user_id in conversation_history:
|
||||
del conversation_history[user_id]
|
||||
if user_id in user_images:
|
||||
os.remove(user_images[user_id])
|
||||
del user_images[user_id]
|
||||
await ctx.send("Conversation history and image cleared. Let's start fresh!")
|
||||
|
||||
@bot.event
|
||||
async def on_message(message):
|
||||
# This is required to let commands still work, since on_message overrides the default handler
|
||||
await bot.process_commands(message)
|
||||
|
||||
if message.author == bot.user:
|
||||
return
|
||||
|
||||
user_id = message.author.id
|
||||
user_message = message.content
|
||||
|
||||
# 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": system_prompt}] + conversation_history[user_id]
|
||||
|
||||
# Check if there's an image to process
|
||||
if user_id in user_images:
|
||||
with open(user_images[user_id], "rb") as image_file:
|
||||
response = client.chat_completions_create(
|
||||
model=GPT_4O_MINI,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": user_message},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{base64.b64encode(image_file.read()).decode('utf-8')}"
|
||||
}
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
max_tokens=16384
|
||||
)
|
||||
# Remove the temporary image file
|
||||
os.remove(user_images[user_id])
|
||||
del user_images[user_id]
|
||||
else:
|
||||
# Call OpenAI API for inference (text-only)
|
||||
response = get_chat_response(client, messages, 4096, GPT_4O)
|
||||
|
||||
# Extract the assistant's reply
|
||||
assistant_message = response.choices[0].message
|
||||
tool_use_count = 0
|
||||
if hasattr(assistant_message, 'function_call') and assistant_message.function_call:
|
||||
while hasattr(assistant_message, 'function_call') and assistant_message.function_call and tool_use_count < 10:
|
||||
tool_response = call_tool(assistant_message.function_call, messages)
|
||||
|
||||
conversation_history[user_id].append({"role": "function", "name": assistant_message.function_call.name, "content": json.dumps(tool_response)})
|
||||
messages.append({
|
||||
"role": "function",
|
||||
"name": assistant_message.function_call.name,
|
||||
"content": json.dumps(tool_response)
|
||||
})
|
||||
|
||||
# Call API again to get the final response
|
||||
assistant_message = get_chat_response(client, messages, 4096, GPT_4O).choices[0].message
|
||||
if not hasattr(assistant_message, 'function_call') or not assistant_message.function_call:
|
||||
assistant_reply = assistant_message.content
|
||||
conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
|
||||
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 message.channel.send(assistant_reply)
|
||||
|
||||
@bot.event
|
||||
async def on_message_edit(before, after):
|
||||
await on_message(after)
|
||||
|
||||
@bot.event
|
||||
async def on_reaction_add(reaction, user):
|
||||
if reaction.message.author == bot.user and user != bot.user:
|
||||
user_id = user.id
|
||||
# Save the reaction as an interaction
|
||||
conversation_history[user_id].append({"role": "user", "content": f"{user.name} reacted with {reaction.emoji}"})
|
||||
messages = [{"role": "system", "content": system_prompt}] + conversation_history[user_id]
|
||||
# Call OpenAI API for inference
|
||||
response = get_chat_response(client, messages, 4096, GPT_4O)
|
||||
assistant_message = response.choices[0].message
|
||||
assistant_reply = assistant_message.content
|
||||
conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
|
||||
# Trim conversation history if it gets too long
|
||||
if len(conversation_history[user_id]) > 10:
|
||||
conversation_history[user_id] = conversation_history[user_id][-10:]
|
||||
# Send the reply back to the user
|
||||
await reaction.message.channel.send(f"{user.name}, {assistant_reply}")
|
||||
|
||||
def call_tool(function_call, messages):
|
||||
# Execute the function
|
||||
function_name = function_call.name
|
||||
function_args = function_call.arguments
|
||||
for tool in tools:
|
||||
if function_name in [f["name"] for f in tool.get_functions()]:
|
||||
return tool.execute(function_name, **eval(function_args))
|
||||
|
||||
def get_chat_response(client, messages, max_tokens, model):
|
||||
response = client.chat_completions_create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
functions=functions,
|
||||
function_call="auto",
|
||||
max_tokens=max_tokens
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
bot.run(DISCORD_BOT_TOKEN)
|
||||
@@ -3,3 +3,4 @@ python-telegram-bot==21.4
|
||||
openai==1.41.0
|
||||
python-dotenv==1.0.1
|
||||
requests==2.32.3
|
||||
discord.py==2.1.0
|
||||
Reference in New Issue
Block a user