From 91bbe303baed193f3497e037b9c30fa3eb74b7cb Mon Sep 17 00:00:00 2001 From: Josh Knapp Date: Thu, 2 Jan 2025 19:38:42 -0800 Subject: [PATCH] pushing updates for bot --- .gitignore | 3 +- open-webui-tool/bedrock-image-tool.py | 100 ++++++++++++++++++++++++++ open-webui-tool/readme.md | 5 ++ v2/bot.py | 67 +++++++++++++---- 4 files changed, 160 insertions(+), 15 deletions(-) create mode 100644 open-webui-tool/bedrock-image-tool.py create mode 100644 open-webui-tool/readme.md diff --git a/.gitignore b/.gitignore index 3cf45bf..1ee9cad 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -scripts/.env \ No newline at end of file +scripts/.env +v2/.env \ No newline at end of file diff --git a/open-webui-tool/bedrock-image-tool.py b/open-webui-tool/bedrock-image-tool.py new file mode 100644 index 0000000..e88a1d3 --- /dev/null +++ b/open-webui-tool/bedrock-image-tool.py @@ -0,0 +1,100 @@ +""" +title: Bedrock Image Description +author: Josh Knapp +version: 0.1.0 +description="Provide Direct Bedrock call for image generation" +""" + +import subprocess +import json +from pydantic import BaseModel, Field + + +# Try to import boto3, install if not present +try: + import boto3 +except ImportError: + print("boto3 package not found. Attempting to install...") + try: + subprocess.check_call([sys.executable, "-m", "pip", "install", "boto3"]) + import boto3 + + print("boto3 package installed successfully") + except subprocess.CalledProcessError as e: + print(f"Failed to install boto3 package: {str(e)}") + + +class Tools: + class Valves(BaseModel): + AWS_ACCESS_KEY: str = Field( + default="", + description="AWS Access Key", + ) + AWS_SECRET_KEY: str = Field( + default="", + description="AWS Secret Key", + ) + AWS_BEDROCK_MODEL: str = Field( + default="", + description="AWS Bedrock Model to use" + ) + + def __init__(self): + self.valves = self.Valves() + pass + + def analyze_image(self, base64_image: str) -> str: + """ + Analyze an image using AWS Bedrock's vision model + Args: + base64_image (str): Base64 encoded image string + Returns: + str: Description of the image + """ + try: + # Initialize Bedrock runtime client + bedrock = boto3.client( + service_name="bedrock-runtime", + aws_access_key_id=self.valves.AWS_ACCESS_KEY, + aws_secret_access_key=self.valves.AWS_SECRET_KEY, + region_name="us-east-1" # or your preferred region + ) + + # Prepare the request body + request_body = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 1000, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": base64_image + } + }, + { + "type": "text", + "text": "Please describe this image in detail." + } + ] + } + ] + } + + # Invoke the model + response = bedrock.invoke_model( + modelId=self.valves.AWS_BEDROCK_MODEL, + body=json.dumps(request_body) + ) + + # Parse and return the response + response_body = json.loads(response['body'].read()) + return response_body['messages'][0]['content'][0]['text'] + + except Exception as e: + print(f"Error analyzing image: {str(e)}") + return f"Error analyzing image: {str(e)}" diff --git a/open-webui-tool/readme.md b/open-webui-tool/readme.md new file mode 100644 index 0000000..240d292 --- /dev/null +++ b/open-webui-tool/readme.md @@ -0,0 +1,5 @@ +For any model to use this tool, you must set the valve values, and add something to the model to let it know to use the tool. + +``` +You have access to a tool that allows you to get descriptions of images called "Bedrock Image Description". Any image handling should be sent through this tool. +``` \ No newline at end of file diff --git a/v2/bot.py b/v2/bot.py index 9edf967..cf24c9e 100644 --- a/v2/bot.py +++ b/v2/bot.py @@ -1,9 +1,10 @@ import os import discord from discord.ext import commands -import openai from openai import OpenAI - +import base64 +import requests +from io import BytesIO from collections import deque from dotenv import load_dotenv @@ -35,28 +36,55 @@ bot = commands.Bot(command_prefix='!', intents=intents) # Message history cache channel_history = {} +async def download_image(url): + response = requests.get(url) + if response.status_code == 200: + image_data = BytesIO(response.content) + base64_image = base64.b64encode(image_data.read()).decode('utf-8') + return base64_image + return None + async def get_chat_history(channel, limit=100): messages = [] async for message in channel.history(limit=limit): - messages.append(f"{message.author.name}: {message.content}") + content = f"{message.author.name}: {message.content}" + + # Handle attachments (images) + for attachment in message.attachments: + if any(attachment.filename.lower().endswith(ext) for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']): + content += f" [Image: {attachment.url}]" + + messages.append(content) return "\n".join(reversed(messages)) -async def get_ai_response(context, user_message): - formatted_prompt = f"##CONTEXT##\n{context}\n##ENDCONTEXT##\n\n{user_message}" +async def get_ai_response(context, user_message, image_urls=None): + messages = [{"role": "user", "content": []}] + + # Add text content + text_content = f"##CONTEXT##\n{context}\n##ENDCONTEXT##\n\n{user_message}" + messages[0]["content"].append({"type": "text", "text": text_content}) + + # Add image content if present + if image_urls: + for url in image_urls: + base64_image = await download_image(url) + if base64_image: + messages[0]["content"].append({ + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{base64_image}" + } + }) try: - response = client.chat.completions.create(model=MODEL_NAME, - messages=[ - {"role": "user", "content": formatted_prompt} - ]) + response = client.chat.completions.create( + model=MODEL_NAME, + messages=messages + ) return response.choices[0].message.content except Exception as e: return f"Error: {str(e)}" -@bot.event -async def on_ready(): - print(f'{bot.user} has connected to Discord!') - @bot.event async def on_message(message): # Ignore messages from the bot itself @@ -81,14 +109,25 @@ async def on_message(message): # Remove bot mention from the message user_message = message.content.replace(f'<@{bot.user.id}>', '').strip() + # Collect image URLs from the message + image_urls = [] + for attachment in message.attachments: + if any(attachment.filename.lower().endswith(ext) for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']): + image_urls.append(attachment.url) + # Get AI response - response = await get_ai_response(history, user_message) + response = await get_ai_response(history, user_message, image_urls) # Send response await message.reply(response) await bot.process_commands(message) +@bot.event +async def on_ready(): + print(f'{bot.user} has connected to Discord!') + + def main(): if not all([DISCORD_TOKEN, OPENAI_API_KEY, OPENWEBUI_API_BASE, MODEL_NAME]): print("Error: Missing required environment variables")