scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

Merge pull request #2290 from ti3x/bedrock_mistral

Browse source 
Add support for Bedrock Mistral models
This commit is contained in:
Ishaan Jaff 2024-03-04 08:42:47 -08:00 • committed by GitHub
commit 84415ef7b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 151 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
docs/my-website/docs/providers/bedrock.md
View file

@ -286,18 +286,20 @@ response = litellm.embedding(
## Supported AWS Bedrock Models
Here's an example of using a bedrock model with LiteLLM
| Model Name | Command |
|--------------------------|------------------------------------------------------------------|
| Model Name | Command |
|----------------------------|------------------------------------------------------------------|
| Anthropic Claude-V2.1 | `completion(model='bedrock/anthropic.claude-v2:1', messages=messages)` | `os.environ['ANTHROPIC_ACCESS_KEY_ID']`, `os.environ['ANTHROPIC_SECRET_ACCESS_KEY']` |
| Anthropic Claude-V2 | `completion(model='bedrock/anthropic.claude-v2', messages=messages)` | `os.environ['ANTHROPIC_ACCESS_KEY_ID']`, `os.environ['ANTHROPIC_SECRET_ACCESS_KEY']` |
| Anthropic Claude-V2 | `completion(model='bedrock/anthropic.claude-v2', messages=messages)` | `os.environ['ANTHROPIC_ACCESS_KEY_ID']`, `os.environ['ANTHROPIC_SECRET_ACCESS_KEY']` |
| Anthropic Claude-Instant V1 | `completion(model='bedrock/anthropic.claude-instant-v1', messages=messages)` | `os.environ['ANTHROPIC_ACCESS_KEY_ID']`, `os.environ['ANTHROPIC_SECRET_ACCESS_KEY']` |
| Amazon Titan Lite | `completion(model='bedrock/amazon.titan-text-lite-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Amazon Titan Express | `completion(model='bedrock/amazon.titan-text-express-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Cohere Command | `completion(model='bedrock/cohere.command-text-v14', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| AI21 J2-Mid | `completion(model='bedrock/ai21.j2-mid-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Amazon Titan Lite | `completion(model='bedrock/amazon.titan-text-lite-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Amazon Titan Express | `completion(model='bedrock/amazon.titan-text-express-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Cohere Command | `completion(model='bedrock/cohere.command-text-v14', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| AI21 J2-Mid | `completion(model='bedrock/ai21.j2-mid-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| AI21 J2-Ultra | `completion(model='bedrock/ai21.j2-ultra-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Meta Llama 2 Chat 13b | `completion(model='bedrock/meta.llama2-13b-chat-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Meta Llama 2 Chat 70b | `completion(model='bedrock/meta.llama2-70b-chat-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Meta Llama 2 Chat 13b | `completion(model='bedrock/meta.llama2-13b-chat-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Meta Llama 2 Chat 70b | `completion(model='bedrock/meta.llama2-70b-chat-v1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Mistral 7B Instruct | `completion(model='bedrock/mistral.mistral-7b-instruct-v0:2', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
## Bedrock Embedding

1
litellm/__init__.py
View file

@ -594,6 +594,7 @@ from .llms.bedrock import (
AmazonCohereConfig,
AmazonLlamaConfig,
AmazonStabilityConfig,
AmazonMistralConfig
)
from .llms.openai import OpenAIConfig, OpenAITextCompletionConfig
from .llms.azure import AzureOpenAIConfig, AzureOpenAIError

65
litellm/llms/bedrock.py
View file

@ -282,6 +282,55 @@ class AmazonLlamaConfig:
}
class AmazonMistralConfig:
"""
Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html
Supported Params for the Amazon / Mistral models:
- `max_tokens` (integer) max tokens,
- `temperature` (float) temperature for model,
- `top_p` (float) top p for model
- `stop` [string] A list of stop sequences that if generated by the model, stops the model from generating further output.
- `top_k` (float) top k for model
"""
max_tokens: Optional[int] = None
temperature: Optional[float] = None
top_p: Optional[float] = None
top_k: Optional[float] = None
stop: Optional[list[str]] = None
def __init__(
self,
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
top_p: Optional[int] = None,
top_k: Optional[float] = None,
stop: Optional[list[str]] = None,
) -> None:
locals_ = locals()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return {
k: v
for k, v in cls.__dict__.items()
if not k.startswith("__")
and not isinstance(
v,
(
types.FunctionType,
types.BuiltinFunctionType,
classmethod,
staticmethod,
),
)
and v is not None
}
class AmazonStabilityConfig:
"""
Reference: https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/providers?model=stability.stable-diffusion-xl-v0
@ -492,6 +541,8 @@ def convert_messages_to_prompt(model, messages, provider, custom_prompt_dict):
prompt = prompt_factory(
model=model, messages=messages, custom_llm_provider="bedrock"
)
elif provider == "mistral":
prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock")
else:
prompt = ""
for message in messages:
@ -623,7 +674,16 @@ def completion(
"textGenerationConfig": inference_params,
}
)
elif provider == "mistral":
## LOAD CONFIG
config = litellm.AmazonMistralConfig.get_config()
for k, v in config.items():
if (
k not in inference_params
): # completion(top_k=3) > amazon_config(top_k=3) <- allows for dynamic variables to be passed in
inference_params[k] = v
data = json.dumps({"prompt": prompt, **inference_params})
else:
data = json.dumps({})
@ -729,6 +789,9 @@ def completion(
outputText = response_body["generations"][0]["text"]
elif provider == "meta":
outputText = response_body["generation"]
elif provider == "mistral":
outputText = response_body["outputs"][0]["text"]
model_response["finish_reason"] = response_body["outputs"][0]["stop_reason"]
else: # amazon titan
outputText = response_body.get("results")[0].get("outputText")
@ -1055,4 +1118,4 @@ def image_generation(
image_dict = {"url": artifact["base64"]}
model_response.data = image_dict
return model_response
return model_response

6
litellm/llms/prompt_templates/factory.py
View file

@ -110,9 +110,9 @@ def mistral_instruct_pt(messages):
"post_message": " [/INST]\n",
},
"user": {"pre_message": "[INST] ", "post_message": " [/INST]\n"},
"assistant": {"pre_message": " ", "post_message": " "},
"assistant": {"pre_message": " ", "post_message": "</s> "},
},
final_prompt_value="</s>",
final_prompt_value="",
messages=messages,
)
return prompt
@ -674,6 +674,8 @@ def prompt_factory(
return claude_2_1_pt(messages=messages)
else:
return anthropic_pt(messages=messages)
elif "mistral." in model:
return mistral_instruct_pt(messages=messages)
try:
if "meta-llama/llama-2" in model and "chat" in model:
return llama_2_chat_pt(messages=messages)

46
litellm/tests/test_bedrock_completion.py
View file

@ -2,25 +2,25 @@
# import sys, os
# import traceback
# from dotenv import load_dotenv
#
# load_dotenv()
# import os, io
#
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import embedding, completion, completion_cost, Timeout
# from litellm import embedding, completion, completion_cost, Timeout, ModelResponse
# from litellm import RateLimitError
#
# # litellm.num_retries = 3
# litellm.cache = None
# litellm.success_callback = []
# user_message = "Write a short poem about the sky"
# messages = [{"content": user_message, "role": "user"}]
#
#
# @pytest.fixture(autouse=True)
# def reset_callbacks():
# print("\npytest fixture - resetting callbacks")
@ -257,3 +257,37 @@
# # test_provisioned_throughput()
# def test_completion_bedrock_mistral_completion_auth():
# print("calling bedrock mistral completion params auth")
# import os
#
# # aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"]
# # aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"]
# # aws_region_name = os.environ["AWS_REGION_NAME"]
#
# # os.environ.pop("AWS_ACCESS_KEY_ID", None)
# # os.environ.pop("AWS_SECRET_ACCESS_KEY", None)
# # os.environ.pop("AWS_REGION_NAME", None)
# try:
# response:ModelResponse = completion(
# model="bedrock/mistral.mistral-7b-instruct-v0:2",
# messages=messages,
# max_tokens=10,
# temperature=0.1,
# )
# # Add any assertions here to check the response
# assert len(response.choices) > 0
# assert len(response.choices[0].message.content) > 0
#
# # os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
# # os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
# # os.environ["AWS_REGION_NAME"] = aws_region_name
# except RateLimitError:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
#
#
# test_completion_bedrock_mistral_completion_auth()

17
litellm/utils.py
View file

@ -4551,6 +4551,21 @@ def get_optional_params(
optional_params["temperature"] = temperature
if max_tokens is not None:
optional_params["max_tokens"] = max_tokens
elif "mistral" in model:
supported_params = ["max_tokens", "temperature", "stop", "top_p", "stream"]
_check_valid_arg(supported_params=supported_params)
# mistral params on bedrock
# \"max_tokens\":400,\"temperature\":0.7,\"top_p\":0.7,\"stop\":[\"\\\\n\\\\nHuman:\"]}"
if max_tokens is not None:
optional_params["max_tokens"] = max_tokens
if temperature is not None:
optional_params["temperature"] = temperature
if top_p is not None:
optional_params["top_p"] = top_p
if stop is not None:
optional_params["stop"] = stop
if stream is not None:
optional_params["stream"] = stream
elif custom_llm_provider == "aleph_alpha":
supported_params = [
"max_tokens",
@ -9677,4 +9692,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

16
model_prices_and_context_window.json
View file

@ -1236,6 +1236,20 @@
"litellm_provider": "bedrock",
"mode": "embedding"
},
"bedrock/us-west-2/mistral.mixtral-8x7b-instruct": {
"max_tokens": 32000,
"input_cost_per_token": 0.00000045,
"output_cost_per_token": 0.0000007,
"litellm_provider": "bedrock",
"mode": "completion"
},
"bedrock/us-west-2/mistral.mistral-7b-instruct": {
"max_tokens": 32000,
"input_cost_per_token": 0.00000015,
"output_cost_per_token": 0.0000002,
"litellm_provider": "bedrock",
"mode": "completion"
},
"anthropic.claude-v1": {
"max_tokens": 100000,
"max_output_tokens": 8191,
@ -2220,4 +2234,4 @@
"mode": "embedding"
}
}
}