101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
|
"""
|
||
|
|
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)}"
|