fix(factory.py): support image url requests for anthropic

This commit is contained in:
Krrish Dholakia 2024-03-06 11:09:50 -08:00
parent 688b903d19
commit 81b92f3434
2 changed files with 97 additions and 15 deletions

View file

@ -2,7 +2,7 @@ from enum import Enum
import requests, traceback
import json, re, xml.etree.ElementTree as ET
from jinja2 import Template, exceptions, Environment, meta
from typing import Optional, Any, List
from typing import Optional, Any
import imghdr, base64
@ -481,6 +481,34 @@ def construct_tool_use_system_prompt(
return tool_use_system_prompt
def convert_url_to_base64(url):
import requests
import base64
response = requests.get(url)
if response.status_code == 200:
image_bytes = response.content
base64_image = base64.b64encode(image_bytes).decode("utf-8")
img_type = url.split(".")[-1].lower()
if img_type == "jpg" or img_type == "jpeg":
img_type = "image/jpeg"
elif img_type == "png":
img_type = "image/png"
elif img_type == "gif":
img_type = "image/gif"
elif img_type == "webp":
img_type = "image/webp"
else:
raise Exception(
f"Error: Unsupported image format. Format={img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']"
)
return f"data:{img_type};base64,{base64_image}"
else:
raise Exception(f"Error: Unable to fetch image from URL. url={url}")
def convert_to_anthropic_image_obj(openai_image_url: str):
"""
Input:
@ -493,6 +521,9 @@ def convert_to_anthropic_image_obj(openai_image_url: str):
"data": {base64_image},
}
"""
try:
if openai_image_url.startswith("http"):
openai_image_url = convert_url_to_base64(url=openai_image_url)
# Extract the base64 image data
base64_data = openai_image_url.split("data:image/")[1].split(";base64,")[1]
@ -504,6 +535,10 @@ def convert_to_anthropic_image_obj(openai_image_url: str):
"media_type": f"image/{image_format}",
"data": base64_data,
}
except Exception as e:
raise Exception(
"""Image url not in expected format. Example Expected input - "image_url": "data:image/jpeg;base64,{base64_image}". Supported formats - ['image/jpeg', 'image/png', 'image/gif', 'image/webp'] """
)
def anthropic_messages_pt(messages: list):
@ -586,7 +621,7 @@ def anthropic_messages_pt(messages: list):
return new_messages
def extract_between_tags(tag: str, string: str, strip: bool = False) -> List[str]:
def extract_between_tags(tag: str, string: str, strip: bool = False) -> list[str]:
ext_list = re.findall(f"<{tag}>(.+?)</{tag}>", string, re.DOTALL)
if strip:
ext_list = [e.strip() for e in ext_list]

View file

@ -151,8 +151,6 @@ def test_completion_claude_3_function_call():
assert isinstance(
response.choices[0].message.tool_calls[0].function.arguments, str
)
except litellm.ServiceUnavailableError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@ -221,6 +219,55 @@ def test_completion_claude_3_base64():
pytest.fail(f"An exception occurred - {str(e)}")
def test_completion_claude_3_function_plus_image():
litellm.set_verbose = True
image_content = [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
},
},
]
image_message = {"role": "user", "content": image_content}
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "text",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
tool_choice = {"type": "function", "function": {"name": "get_current_weather"}}
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
response = completion(
model="claude-3-sonnet-20240229",
messages=[image_message],
tool_choice=tool_choice,
tools=tools,
stream=False,
)
print(response)
def test_completion_mistral_api():
try:
litellm.set_verbose = True
@ -1406,9 +1453,9 @@ def test_completion_replicate_vicuna():
def test_replicate_custom_prompt_dict():
litellm.set_verbose = True
model_name = "replicate/mistralai/mixtral-8x7b-instruct-v0.1"
model_name = "replicate/meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0"
litellm.register_prompt_template(
model="replicate/mistralai/mixtral-8x7b-instruct-v0.1",
model="replicate/meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0",
initial_prompt_value="You are a good assistant", # [OPTIONAL]
roles={
"system": {