Merge pull request #2325 from BerriAI/litellm_claude_3_bedrock

claude 3 bedrock integrations
This commit is contained in:
Krish Dholakia 2024-03-04 16:15:48 -08:00 committed by GitHub
commit 2e6ef6fd11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 349 additions and 44 deletions

View file

@ -2,11 +2,17 @@ import os, types
import json
from enum import Enum
import requests
import time
import time, uuid
from typing import Callable, Optional
from litellm.utils import ModelResponse, Usage
from litellm.utils import ModelResponse, Usage, map_finish_reason
import litellm
from .prompt_templates.factory import prompt_factory, custom_prompt
from .prompt_templates.factory import (
prompt_factory,
custom_prompt,
construct_tool_use_system_prompt,
extract_between_tags,
parse_xml_params,
)
import httpx
@ -41,6 +47,7 @@ class AnthropicConfig:
top_p: Optional[int] = None
top_k: Optional[int] = None
metadata: Optional[dict] = None
system: Optional[str] = None
def __init__(
self,
@ -50,6 +57,7 @@ class AnthropicConfig:
top_p: Optional[int] = None,
top_k: Optional[int] = None,
metadata: Optional[dict] = None,
system: Optional[str] = None,
) -> None:
locals_ = locals()
for key, value in locals_.items():
@ -108,6 +116,7 @@ def completion(
headers={},
):
headers = validate_environment(api_key, headers)
_is_function_call = False
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details = custom_prompt_dict[model]
@ -118,38 +127,19 @@ def completion(
messages=messages,
)
else:
prompt = prompt_factory(
# Separate system prompt from rest of message
system_prompt_idx: Optional[int] = None
for idx, message in enumerate(messages):
if message["role"] == "system":
optional_params["system"] = message["content"]
system_prompt_idx = idx
break
if system_prompt_idx is not None:
messages.pop(system_prompt_idx)
# Format rest of message according to anthropic guidelines
messages = prompt_factory(
model=model, messages=messages, custom_llm_provider="anthropic"
)
"""
format messages for anthropic
1. Anthropic supports roles like "user" and "assistant", (here litellm translates system-> assistant)
2. The first message always needs to be of role "user"
3. Each message must alternate between "user" and "assistant" (this is not addressed as now by litellm)
4. final assistant content cannot end with trailing whitespace (anthropic raises an error otherwise)
"""
# 1. Anthropic only supports roles like "user" and "assistant"
for idx, message in enumerate(messages):
if message["role"] == "system":
message["role"] = "assistant"
# if this is the final assistant message, remove trailing whitespace
# TODO: only do this if it's the final assistant message
if message["role"] == "assistant":
message["content"] = message["content"].strip()
# 2. The first message always needs to be of role "user"
if len(messages) > 0:
if messages[0]["role"] != "user":
# find the index of the first user message
for i, message in enumerate(messages):
if message["role"] == "user":
break
# remove the user message at existing position and add it to the front
messages.pop(i)
# move the first user message to the front
messages = [message] + messages
## Load Config
config = litellm.AnthropicConfig.get_config()
@ -159,6 +149,17 @@ def completion(
): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
## Handle Tool Calling
if "tools" in optional_params:
_is_function_call = True
tool_calling_system_prompt = construct_tool_use_system_prompt(
tools=optional_params["tools"]
)
optional_params["system"] = (
optional_params.get("system", "\n") + tool_calling_system_prompt
) # add the anthropic tool calling prompt to the system prompt
optional_params.pop("tools")
data = {
"model": model,
"messages": messages,
@ -167,7 +168,7 @@ def completion(
## LOGGING
logging_obj.pre_call(
input=prompt,
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
@ -225,8 +226,33 @@ def completion(
)
else:
text_content = completion_response["content"][0].get("text", None)
model_response.choices[0].message.content = text_content # type: ignore
model_response.choices[0].finish_reason = completion_response["stop_reason"]
## TOOL CALLING - OUTPUT PARSE
if _is_function_call == True:
function_name = extract_between_tags("tool_name", text_content)[0]
function_arguments_str = extract_between_tags("invoke", text_content)[
0
].strip()
function_arguments_str = f"<invoke>{function_arguments_str}</invoke>"
function_arguments = parse_xml_params(function_arguments_str)
_message = litellm.Message(
tool_calls=[
{
"id": f"call_{uuid.uuid4()}",
"type": "function",
"function": {
"name": function_name,
"arguments": json.dumps(function_arguments),
},
}
],
content=None,
)
model_response.choices[0].message = _message # type: ignore
else:
model_response.choices[0].message.content = text_content # type: ignore
model_response.choices[0].finish_reason = map_finish_reason(
completion_response["stop_reason"]
)
## CALCULATING USAGE
prompt_tokens = completion_response["usage"]["input_tokens"]

View file

@ -1,8 +1,9 @@
from enum import Enum
import requests, traceback
import json
import json, re, xml.etree.ElementTree as ET
from jinja2 import Template, exceptions, Environment, meta
from typing import Optional, Any
import imghdr, base64
def default_pt(messages):
@ -390,7 +391,7 @@ def format_prompt_togetherai(messages, prompt_format, chat_template):
return prompt
###
### ANTHROPIC ###
def anthropic_pt(
@ -424,6 +425,184 @@ def anthropic_pt(
return prompt
def construct_format_parameters_prompt(parameters: dict):
parameter_str = "<parameter>\n"
for k, v in parameters.items():
parameter_str += f"<{k}>"
parameter_str += f"{v}"
parameter_str += f"</{k}>"
parameter_str += "\n</parameter>"
return parameter_str
def construct_format_tool_for_claude_prompt(name, description, parameters):
constructed_prompt = (
"<tool_description>\n"
f"<tool_name>{name}</tool_name>\n"
"<description>\n"
f"{description}\n"
"</description>\n"
"<parameters>\n"
f"{construct_format_parameters_prompt(parameters)}\n"
"</parameters>\n"
"</tool_description>"
)
return constructed_prompt
def construct_tool_use_system_prompt(
tools,
): # from https://github.com/anthropics/anthropic-cookbook/blob/main/function_calling/function_calling.ipynb
tool_str_list = []
for tool in tools:
tool_str = construct_format_tool_for_claude_prompt(
tool["function"]["name"],
tool["function"].get("description", ""),
tool["function"].get("parameters", {}),
)
tool_str_list.append(tool_str)
tool_use_system_prompt = (
"In this environment you have access to a set of tools you can use to answer the user's question.\n"
"\n"
"You may call them like this:\n"
"<function_calls>\n"
"<invoke>\n"
"<tool_name>$TOOL_NAME</tool_name>\n"
"<parameters>\n"
"<$PARAMETER_NAME>$PARAMETER_VALUE</$PARAMETER_NAME>\n"
"...\n"
"</parameters>\n"
"</invoke>\n"
"</function_calls>\n"
"\n"
"Here are the tools available:\n"
"<tools>\n" + "\n".join([tool_str for tool_str in tool_str_list]) + "\n</tools>"
)
return tool_use_system_prompt
def convert_to_anthropic_image_obj(openai_image_url: str):
"""
Input:
"image_url": "data:image/jpeg;base64,{base64_image}",
Return:
"source": {
"type": "base64",
"media_type": "image/jpeg",
"data": {base64_image},
}
"""
# Extract the base64 image data
base64_data = openai_image_url.split("data:image/")[1].split(";base64,")[1]
# Infer image format from the URL
image_format = openai_image_url.split("data:image/")[1].split(";base64,")[0]
return {
"type": "base64",
"media_type": f"image/{image_format}",
"data": base64_data,
}
def anthropic_messages_pt(messages: list):
"""
format messages for anthropic
1. Anthropic supports roles like "user" and "assistant", (here litellm translates system-> assistant)
2. The first message always needs to be of role "user"
3. Each message must alternate between "user" and "assistant" (this is not addressed as now by litellm)
4. final assistant content cannot end with trailing whitespace (anthropic raises an error otherwise)
5. System messages are a separate param to the Messages API (used for tool calling)
"""
## Ensure final assistant message has no trailing whitespace
last_assistant_message_idx: Optional[int] = None
# reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, add a blank 'user' or 'assistant' message to ensure compatibility
new_messages = []
if len(messages) == 1:
# check if the message is a user message
if messages[0]["role"] == "assistant":
new_messages.append({"role": "user", "content": ""})
# check if content is a list (vision)
if isinstance(messages[0]["content"], list): # vision input
new_content = []
for m in messages[0]["content"]:
if m.get("type", "") == "image_url":
new_content.append(
{
"type": "image",
"source": convert_to_anthropic_image_obj(
m["image_url"]["url"]
),
}
)
elif m.get("type", "") == "text":
new_content.append({"type": "text", "text": m["text"]})
new_messages.append({"role": messages[0]["role"], "content": new_content}) # type: ignore
else:
new_messages.append(messages[0])
return new_messages
for i in range(len(messages) - 1): # type: ignore
if i == 0 and messages[i]["role"] == "assistant":
new_messages.append({"role": "user", "content": ""})
if isinstance(messages[i]["content"], list): # vision input
new_content = []
for m in messages[i]["content"]:
if m.get("type", "") == "image_url":
new_content.append(
{
"type": "image",
"source": convert_to_anthropic_image_obj(
m["image_url"]["url"]
),
}
)
elif m.get("type", "") == "text":
new_content.append({"type": "text", "content": m["text"]})
new_messages.append({"role": messages[i]["role"], "content": new_content}) # type: ignore
else:
new_messages.append(messages[i])
if messages[i]["role"] == messages[i + 1]["role"]:
if messages[i]["role"] == "user":
new_messages.append({"role": "assistant", "content": ""})
else:
new_messages.append({"role": "user", "content": ""})
if messages[i]["role"] == "assistant":
last_assistant_message_idx = i
if last_assistant_message_idx is not None:
new_messages[last_assistant_message_idx]["content"] = new_messages[
last_assistant_message_idx
][
"content"
].strip() # no trailing whitespace for final assistant message
return new_messages
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]
return ext_list
def parse_xml_params(xml_content):
root = ET.fromstring(xml_content)
params = {}
for child in root.findall(".//parameters/*"):
params[child.tag] = child.text
return params
###
def amazon_titan_pt(
messages: list,
): # format - https://github.com/BerriAI/litellm/issues/1896
@ -650,10 +829,9 @@ def prompt_factory(
if custom_llm_provider == "ollama":
return ollama_pt(model=model, messages=messages)
elif custom_llm_provider == "anthropic":
if any(_ in model for _ in ["claude-2.1", "claude-v2:1"]):
return claude_2_1_pt(messages=messages)
else:
if model == "claude-instant-1" or model == "claude-2.1":
return anthropic_pt(messages=messages)
return anthropic_messages_pt(messages=messages)
elif custom_llm_provider == "together_ai":
prompt_format, chat_template = get_model_info(token=api_key, model=model)
return format_prompt_togetherai(

View file

@ -351,7 +351,7 @@ def test_gemini_pro_vision_base64():
load_vertex_ai_credentials()
litellm.set_verbose = True
litellm.num_retries = 3
image_path = "cached_logo.jpg"
image_path = "../proxy/cached_logo.jpg"
# Getting the base64 string
base64_image = encode_image(image_path)
resp = litellm.completion(

View file

@ -99,6 +99,47 @@ def test_completion_claude_3():
pytest.fail(f"Error occurred: {e}")
def test_completion_claude_3_function_call():
litellm.set_verbose = True
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
try:
# test without max tokens
response = completion(
model="anthropic/claude-3-opus-20240229",
messages=messages,
tools=tools,
tool_choice="auto",
)
# Add any assertions, here to check response args
print(response)
assert isinstance(response.choices[0].message.tool_calls[0].function.name, str)
assert isinstance(
response.choices[0].message.tool_calls[0].function.arguments, str
)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_claude_3_stream():
litellm.set_verbose = False
messages = [{"role": "user", "content": "Hello, world"}]
@ -118,6 +159,51 @@ def test_completion_claude_3_stream():
pytest.fail(f"Error occurred: {e}")
def encode_image(image_path):
import base64
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode("utf-8")
@pytest.mark.skip(
reason="we already test claude-3, this is just another way to pass images"
)
def test_completion_claude_3_base64():
try:
litellm.set_verbose = True
litellm.num_retries = 3
image_path = "../proxy/cached_logo.jpg"
# Getting the base64 string
base64_image = encode_image(image_path)
resp = litellm.completion(
model="anthropic/claude-3-opus-20240229",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "data:image/jpeg;base64," + base64_image
},
},
],
}
],
)
print(f"\nResponse: {resp}")
prompt_tokens = resp.usage.prompt_tokens
raise Exception("it worked!")
except Exception as e:
if "500 Internal error encountered.'" in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")
def test_completion_mistral_api():
try:
litellm.set_verbose = True

View file

@ -200,6 +200,10 @@ def map_finish_reason(
return "content_filter"
elif finish_reason == "STOP": # vertex ai
return "stop"
elif finish_reason == "end_turn" or finish_reason == "stop_sequence": # anthropic
return "stop"
elif finish_reason == "max_tokens": # anthropic
return "length"
return finish_reason
@ -4106,6 +4110,7 @@ def get_optional_params(
and custom_llm_provider != "anyscale"
and custom_llm_provider != "together_ai"
and custom_llm_provider != "mistral"
and custom_llm_provider != "anthropic"
):
if custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat":
# ollama actually supports json output
@ -4186,7 +4191,15 @@ def get_optional_params(
## raise exception if provider doesn't support passed in param
if custom_llm_provider == "anthropic":
## check if unsupported param passed in
supported_params = ["stream", "stop", "temperature", "top_p", "max_tokens"]
supported_params = [
"stream",
"stop",
"temperature",
"top_p",
"max_tokens",
"tools",
"tool_choice",
]
_check_valid_arg(supported_params=supported_params)
# handle anthropic params
if stream:
@ -4201,6 +4214,8 @@ def get_optional_params(
optional_params["top_p"] = top_p
if max_tokens is not None:
optional_params["max_tokens"] = max_tokens
if tools is not None:
optional_params["tools"] = tools
elif custom_llm_provider == "cohere":
## check if unsupported param passed in
supported_params = [
@ -9704,4 +9719,4 @@ def _get_base_model_from_metadata(model_call_details=None):
base_model = model_info.get("base_model", None)
if base_model is not None:
return base_model
return None
return None