diff --git a/README.md b/README.md index 0b7dc63..1602cb3 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,7 @@ # OpenWebUI-Bedrock-Image-Tool +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/bedrock-image-tool.py b/bedrock-image-tool.py new file mode 100644 index 0000000..e88a1d3 --- /dev/null +++ b/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)}"