|
import openai |
|
import base64 |
|
import os |
|
import io |
|
import time |
|
from PIL import Image |
|
from abc import ABCMeta, abstractmethod |
|
|
|
|
|
def create_vision_chat_completion(vision_model, base64_image, prompt): |
|
try: |
|
response = openai.ChatCompletion.create( |
|
model=vision_model, |
|
messages=[ |
|
{ |
|
"role": "user", |
|
"content": [ |
|
{"type": "text", "text": prompt}, |
|
{ |
|
"type": "image_url", |
|
"image_url": { |
|
"url": f"data:image/jpeg;base64,{base64_image}", |
|
}, |
|
}, |
|
], |
|
} |
|
], |
|
max_tokens=1000, |
|
) |
|
return response.choices[0].message.content |
|
except: |
|
return None |
|
|
|
|
|
def create_image(prompt): |
|
try: |
|
response = openai.Image.create( |
|
model="dall-e-3", |
|
prompt=prompt, |
|
response_format="b64_json" |
|
) |
|
return response.data[0]['b64_json'] |
|
except: |
|
return None |
|
|
|
|
|
def image_to_base64(path): |
|
try: |
|
_, suffix = os.path.splitext(path) |
|
if suffix not in {'.jpg', '.jpeg', '.png', '.webp'}: |
|
img = Image.open(path) |
|
img_png = img.convert('RGB') |
|
img_png.tobytes() |
|
byte_buffer = io.BytesIO() |
|
img_png.save(byte_buffer, 'PNG') |
|
encoded_string = base64.b64encode(byte_buffer.getvalue()).decode('utf-8') |
|
else: |
|
with open(path, "rb") as image_file: |
|
encoded_string = base64.b64encode(image_file.read()).decode('utf-8') |
|
return encoded_string |
|
except: |
|
return None |
|
|
|
|
|
def base64_to_image_bytes(image_base64): |
|
try: |
|
return base64.b64decode(image_base64) |
|
except: |
|
return None |
|
|
|
|
|
def inquire_image(work_dir, vision_model, path, prompt): |
|
image_base64 = image_to_base64(f'{work_dir}/{path}') |
|
hypertext_to_display = None |
|
if image_base64 is None: |
|
return "Error: Image transform error", None |
|
else: |
|
response = create_vision_chat_completion(vision_model, image_base64, prompt) |
|
if response is None: |
|
return "Model response error", None |
|
else: |
|
return response, hypertext_to_display |
|
|
|
|
|
def dalle(unique_id, prompt): |
|
img_base64 = create_image(prompt) |
|
text_to_gpt = "Image has been successfully generated and displayed to user." |
|
|
|
if img_base64 is None: |
|
return "Error: Model response error", None |
|
|
|
img_bytes = base64_to_image_bytes(img_base64) |
|
if img_bytes is None: |
|
return "Error: Image transform error", None |
|
|
|
temp_path = f'cache/temp_{unique_id}' |
|
if not os.path.exists(temp_path): |
|
os.mkdir(temp_path) |
|
path = f'{temp_path}/{hash(time.time())}.png' |
|
|
|
with open(path, 'wb') as f: |
|
f.write(img_bytes) |
|
|
|
hypertext_to_display = f'<img src=\"file={path}\" width="50%" style=\'max-width:none; max-height:none\'>' |
|
return text_to_gpt, hypertext_to_display |
|
|
|
|
|
class Tool(metaclass=ABCMeta): |
|
def __init__(self, config): |
|
self.config = config |
|
|
|
@abstractmethod |
|
def support(self): |
|
pass |
|
|
|
@abstractmethod |
|
def get_tool_data(self): |
|
pass |
|
|
|
|
|
class ImageInquireTool(Tool): |
|
def support(self): |
|
return self.config['model']['GPT-4V']['available'] |
|
|
|
def get_tool_data(self): |
|
return { |
|
"tool_name": "inquire_image", |
|
"tool": inquire_image, |
|
"system_prompt": "If necessary, utilize the 'inquire_image' tool to query an AI model regarding the " |
|
"content of images uploaded by users. Avoid phrases like\"based on the analysis\"; " |
|
"instead, respond as if you viewed the image by yourself. Keep in mind that not every" |
|
"tasks related to images require knowledge of the image content, such as converting " |
|
"an image format or extracting image file attributes, which should use `execute_code` " |
|
"tool instead. Use the tool only when understanding the image content is necessary.", |
|
"tool_description": { |
|
"name": "inquire_image", |
|
"description": "This function enables you to inquire with an AI model about the contents of an image " |
|
"and receive the model's response.", |
|
"parameters": { |
|
"type": "object", |
|
"properties": { |
|
"path": { |
|
"type": "string", |
|
"description": "File path of the image" |
|
}, |
|
"prompt": { |
|
"type": "string", |
|
"description": "The question you want to pose to the AI model about the image" |
|
} |
|
}, |
|
"required": ["path", "prompt"] |
|
} |
|
}, |
|
"additional_parameters": { |
|
"work_dir": lambda bot_backend: bot_backend.jupyter_work_dir, |
|
"vision_model": self.config['model']['GPT-4V']['model_name'] |
|
} |
|
} |
|
|
|
|
|
class DALLETool(Tool): |
|
def support(self): |
|
return True |
|
|
|
def get_tool_data(self): |
|
return { |
|
"tool_name": "dalle", |
|
"tool": dalle, |
|
"system_prompt": "If user ask you to generate an art image, you can translate user's requirements into a " |
|
"prompt and sending it to the `dalle` tool. Please note that this tool is specifically " |
|
"designed for creating art images. For scientific figures, such as plots, please use the " |
|
"Python code execution tool `execute_code` instead.", |
|
"tool_description": { |
|
"name": "dalle", |
|
"description": "This function allows you to access OpenAI's DALL·E-3 model for image generation.", |
|
"parameters": { |
|
"type": "object", |
|
"properties": { |
|
"prompt": { |
|
"type": "string", |
|
"description": "A detailed description of the image you want to generate, should be in " |
|
"English only. " |
|
} |
|
}, |
|
"required": ["prompt"] |
|
} |
|
}, |
|
"additional_parameters": { |
|
"unique_id": lambda bot_backend: bot_backend.unique_id, |
|
} |
|
} |
|
|
|
|
|
def get_available_tools(config): |
|
tools = [ImageInquireTool] |
|
|
|
available_tools = [] |
|
for tool in tools: |
|
tool_instance = tool(config) |
|
if tool_instance.support(): |
|
available_tools.append(tool_instance.get_tool_data()) |
|
return available_tools |
|
|
