services/vibe-bot
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

human cleanup

Browse Source
This commit is contained in:
Reese Wells
2026-03-09 22:36:04 -04:00
parent 3defce1365
commit 488912a991
8 changed files with 377 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files

248
vibe_bot/main.py
View File

@@ -5,7 +5,19 @@ import base64
from io import BytesIO
from openai import OpenAI
import logging
from database import get_database, CustomBotManager
from database import get_database, CustomBotManager # type: ignore
from config import ( # type: ignore
CHAT_ENDPOINT_KEY,
DISCORD_TOKEN,
CHAT_ENDPOINT,
CHAT_MODEL,
IMAGE_EDIT_ENDPOINT_KEY,
IMAGE_GEN_ENDPOINT,
IMAGE_EDIT_ENDPOINT,
MAX_COMPLETION_TOKENS,
)
import llama_wrapper # type: ignore
import requests
# Configure logging
logging.basicConfig(
@@ -13,31 +25,11 @@ logging.basicConfig(
)
logger = logging.getLogger(__name__)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN", "placeholder")
OPENAI_API_ENDPOINT = os.getenv("OPENAI_API_ENDPOINT")
IMAGE_GEN_ENDPOINT = os.getenv("IMAGE_GEN_ENDPOINT")
IMAGE_EDIT_ENDPOINT = os.getenv("IMAGE_EDIT_ENDPOINT")
MAX_COMPLETION_TOKENS = int(os.getenv("MAX_COMPLETION_TOKENS", "1000"))
if not OPENAI_API_ENDPOINT:
raise Exception("OPENAI_API_ENDPOINT required.")
if not IMAGE_GEN_ENDPOINT:
raise Exception("IMAGE_GEN_ENDPOINT required.")
# Set your OpenAI API key as an environment variable
# You can also pass it directly but environment variables are safer
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "placeholder")
# Initialize the bot
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="!", intents=intents)
# OpenAI Completions API endpoint
OPENAI_COMPLETIONS_URL = f"{OPENAI_API_ENDPOINT}/chat/completions"
@bot.event
async def on_ready():
@@ -46,7 +38,7 @@ async def on_ready():
logger.info(f"Bot logged in as {bot.user}")
@bot.command(name="custom-bot")
@bot.command(name="custom-bot") # type: ignore
async def custom_bot(ctx, bot_name: str, *, personality: str):
"""Create a custom bot with a name and personality
@@ -129,14 +121,14 @@ async def list_custom_bots(ctx):
f"Found {len(bots)} custom bots, displaying top 10 for {ctx.author.name}"
)
bot_list = "🤖 **Available Custom Bots**:\n\n"
for name, prompt, creator in bots[:10]: # Limit to 10 bots
bot_list += f"• **{name}** (created by {creator})\n"
for name, prompt, creator in bots:
bot_list += f"• **{name}**\n"
logger.info(f"Sending bot list response to {ctx.author.name}")
await ctx.send(bot_list)
@bot.command(name="delete-custom-bot")
@bot.command(name="delete-custom-bot") # type: ignore
async def delete_custom_bot(ctx, bot_name: str):
"""Delete a custom bot (only the creator can delete)
@@ -194,16 +186,16 @@ async def on_message(message):
if message.author == bot.user:
return
message_author = message.author.name
message_content = message.content.lower()
logger.debug(
f"Processing message from {message.author.name}: '{message.content[:50]}...'"
f"Processing message from {message_author}: '{message_content[:50]}...'"
)
ctx = await bot.get_context(message)
# Check if the message starts with a custom bot command
content = message.content.lower()
logger.info(f"Initializing CustomBotManager to check for custom bot commands")
logger.info("Initializing CustomBotManager to check for custom bot commands")
custom_bot_manager = CustomBotManager()
logger.info("Fetching list of custom bots to check for matching commands")
@@ -212,7 +204,7 @@ async def on_message(message):
logger.info(f"Checking {len(custom_bots)} custom bots for command match")
for bot_name, system_prompt, _ in custom_bots:
# Check if message starts with the custom bot name followed by a space
if content.startswith(f"!{bot_name} "):
if message_content.startswith(f"!{bot_name} "):
logger.info(
f"Custom bot command detected: '{bot_name}' triggered by {message.author.name}"
)
@@ -224,25 +216,14 @@ async def on_message(message):
)
# Prepare the payload with custom personality
payload = {
"model": "qwen3-vl-30b-a3b-instruct",
"messages": [
{
"role": "system",
"content": system_prompt,
},
{"role": "user", "content": user_message},
],
"max_completion_tokens": MAX_COMPLETION_TOKENS,
}
response_prefix = f"**{bot_name} response**"
logger.info(f"Sending request to OpenAI API for bot '{bot_name}'")
await handle_chat(
ctx=ctx,
bot_name=bot_name,
message=user_message,
payload=payload,
system_prompt=system_prompt,
response_prefix=response_prefix,
)
return
@@ -258,24 +239,22 @@ async def doodlebob(ctx, *, message: str):
logger.info(f"Doodlebob command triggered by {ctx.author.name}: {message[:100]}")
await ctx.send(f"**Doodlebob erasing {message[:100]}...**")
image_prompt_payload = {
"model": "qwen3-vl-30b-a3b-instruct",
"messages": [
{
"role": "system",
"content": (
"Given the following message, convert it to a detailed image generation prompt that will be passed directly into an image generation model."
"If told to generate an image of yourself, generate a picture of a rat. If told to generate a picture of 'me', 'myself', or some other self"
" reference, generate a picture of a rat. Only respond with a valid image generation prompt, do not affirm the user or respond to the user's"
" questions."
),
},
{"role": "user", "content": message},
],
}
system_prompt = (
"Given the following message, convert it to a detailed image generation prompt that will be passed directly into an image generation model."
"If told to generate an image of yourself, generate a picture of a rat. If told to generate a picture of 'me', 'myself', or some other self"
" reference, generate a picture of a rat. Only respond with a valid image generation prompt, do not affirm the user or respond to the user's"
" questions."
)
# Wait for the generated image prompt
image_prompt = await call_llm(ctx, image_prompt_payload)
image_prompt = llama_wrapper.chat_completion_instruct(
system_prompt=system_prompt,
user_prompt=message,
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
# If the string is empty we had an error
if image_prompt == "":
@@ -285,32 +264,16 @@ async def doodlebob(ctx, *, message: str):
# Alert the user we're generating the image
await ctx.send(f"**Doodlebob calling drone strike on {image_prompt[:100]}...**")
# Create the image prompt payload
image_payload = {
"model": "default",
"prompt": image_prompt,
"n": 1,
"size": "1024x1024",
}
# Call the image generation endpoint
response = requests.post(
f"{IMAGE_GEN_ENDPOINT}/images/generations",
json=image_payload,
timeout=120,
image_b64 = llama_wrapper.image_generation(
prompt=message,
openai_url=IMAGE_EDIT_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
)
if response.status_code == 200:
result = response.json()
# Send image
image_data = BytesIO(base64.b64decode(result["data"][0]["b64_json"]))
send_img = discord.File(image_data, filename="image.png")
await ctx.send(file=send_img)
else:
print(f"❌ Error: {response.status_code}")
print(response.text)
return None
# Save the image to a file
edited_image_data = BytesIO(base64.b64decode(image_b64))
send_img = discord.File(edited_image_data, filename="image.png")
await ctx.send(file=send_img)
@bot.command(name="retcon")
@@ -321,31 +284,23 @@ async def retcon(ctx, *, message: str):
await ctx.send(f"**Rewriting history to match {message[:100]}...**")
client = OpenAI(base_url=IMAGE_EDIT_ENDPOINT, api_key=OPENAI_API_KEY)
result = client.images.edit(
model="placeholder",
image=[image_bytestream],
image_b64 = llama_wrapper.image_edit(
image=image_bytestream,
prompt=message,
size="1024x1024",
openai_url=IMAGE_EDIT_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
)
image_base64 = result.data[0].b64_json
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
edited_image_data = BytesIO(image_bytes)
edited_image_data = BytesIO(base64.b64decode(image_b64))
send_img = discord.File(edited_image_data, filename="image.png")
await ctx.send(file=send_img)
async def handle_chat(ctx, *, message: str, payload: dict, response_prefix: str):
# Check if API key is set
if not OPENAI_API_KEY:
await ctx.send(
"Error: OpenAI API key is not configured. Please set the OPENAI_API_KEY environment variable."
)
return
async def handle_chat(
ctx, *, bot_name: str, message: str, system_prompt: str, response_prefix: str
):
await ctx.send(f"{bot_name} is searching its databanks for {message[:50]}...")
# Get database instance
db = get_database()
@@ -356,32 +311,27 @@ async def handle_chat(ctx, *, message: str, payload: dict, response_prefix: str)
)
if context:
payload["messages"][0][
"content"
] += f"\n\nRelevant conversation history:\n{context}"
user_message = f"\n\nRelevant conversation history:\n{context}\n\n{message}"
else:
user_message = message
payload["messages"][1]["content"] = message
logger.info(user_message)
print(payload)
system_prompt_edit = (
"Keep your responses somewhat short, limited to 500 words or less. "
f"{system_prompt}"
)
try:
# Initialize OpenAI client
client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_API_ENDPOINT)
# Call OpenAI API
response = client.chat.completions.create(
model=payload["model"],
messages=payload["messages"],
max_completion_tokens=MAX_COMPLETION_TOKENS,
frequency_penalty=1.5,
presence_penalty=1.5,
temperature=1,
seed=-1,
bot_response = llama_wrapper.chat_completion_instruct(
system_prompt=system_prompt_edit,
user_prompt=user_message,
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
# Extract the generated text
generated_text = response.choices[0].message.content.strip()
# Store both user message and bot response in the database
db.add_message(
message_id=f"{ctx.message.id}",
@@ -394,68 +344,24 @@ async def handle_chat(ctx, *, message: str, payload: dict, response_prefix: str)
db.add_message(
message_id=f"{ctx.message.id}_response",
user_id=str(bot.user.id),
username=bot.user.name,
content=f"Bot: {generated_text}",
user_id=str(bot.user.id), # type: ignore
username=bot.user.name, # type: ignore
content=f"Bot: {bot_response}",
channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None,
)
# Send the response back to the chat
await ctx.send(response_prefix)
while generated_text:
send_chunk = generated_text[:1000]
generated_text = generated_text[1000:]
while bot_response:
send_chunk = bot_response[:1000]
bot_response = bot_response[1000:]
await ctx.send(send_chunk)
except requests.exceptions.HTTPError as e:
await ctx.send(f"Error: OpenAI API error - {e}")
except requests.exceptions.Timeout:
await ctx.send("Error: Request timed out. Please try again.")
except Exception as e:
await ctx.send(f"Error: {str(e)}")
async def call_llm(ctx, payload: dict) -> str:
# Check if API key is set
if not OPENAI_API_KEY:
await ctx.send(
"Error: OpenAI API key is not configured. Please set the OPENAI_API_KEY environment variable."
)
return ""
# Set headers
headers = {
"Authorization": f"Bearer {OPENAI_API_KEY}",
"Content-Type": "application/json",
}
try:
# Initialize OpenAI client
client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_API_ENDPOINT)
# Call OpenAI API
response = client.chat.completions.create(
model=payload["model"],
messages=payload["messages"],
max_tokens=MAX_COMPLETION_TOKENS,
)
# Extract the generated text
generated_text = response.choices[0].message.content.strip()
print(generated_text)
return generated_text
except requests.exceptions.HTTPError as e:
await ctx.send(f"Error: OpenAI API error - {e}")
except requests.exceptions.Timeout:
await ctx.send("Error: Request timed out. Please try again.")
except Exception as e:
await ctx.send(f"Error: {str(e)}")
return ""
# Run the bot
if __name__ == "__main__":
bot.run(DISCORD_TOKEN)