diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index ed8470c934..7a2b3d8d84 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -6,6 +6,7 @@ import time from typing import Callable, Optional from litellm.utils import ModelResponse import litellm +from .prompt_templates.factory import prompt_factory, custom_prompt class AnthropicConstants(Enum): HUMAN_PROMPT = "\n\nHuman: " @@ -71,6 +72,7 @@ def completion( model: str, messages: list, api_base: str, + custom_prompt_dict: dict, model_response: ModelResponse, print_verbose: Callable, encoding, @@ -81,25 +83,18 @@ def completion( logger_fn=None, ): headers = validate_environment(api_key) - prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}" - for message in messages: - if "role" in message: - if message["role"] == "user": - prompt += ( - f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - ) - elif message["role"] == "system": - prompt += ( - f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - ) - else: - prompt += ( - f"{AnthropicConstants.AI_PROMPT.value}{message['content']}" - ) - else: - prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - prompt += f"{AnthropicConstants.AI_PROMPT.value}" - + if model in custom_prompt_dict: + # check if the model has a registered custom prompt + model_prompt_details = custom_prompt_dict[model] + prompt = custom_prompt( + role_dict=model_prompt_details["roles"], + initial_prompt_value=model_prompt_details["initial_prompt_value"], + final_prompt_value=model_prompt_details["final_prompt_value"], + messages=messages + ) + else: + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic") + ## Load Config config = litellm.AnthropicConfig.get_config() for k, v in config.items(): diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 1fb5e84464..7014ebc42b 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -4,6 +4,7 @@ import time from typing import Callable, Optional import litellm from litellm.utils import ModelResponse, get_secret +from .prompt_templates.factory import prompt_factory, custom_prompt class BedrockError(Exception): def __init__(self, status_code, message): @@ -206,27 +207,20 @@ def init_bedrock_client( return client -def convert_messages_to_prompt(messages, provider): +def convert_messages_to_prompt(model, messages, provider, custom_prompt_dict): # handle anthropic prompts using anthropic constants if provider == "anthropic": - prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}" - for message in messages: - if "role" in message: - if message["role"] == "user": - prompt += ( - f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - ) - elif message["role"] == "system": - prompt += ( - f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - ) - else: - prompt += ( - f"{AnthropicConstants.AI_PROMPT.value}{message['content']}" - ) - else: - prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" - prompt += f"{AnthropicConstants.AI_PROMPT.value}" + if model in custom_prompt_dict: + # check if the model has a registered custom prompt + model_prompt_details = custom_prompt_dict[model] + prompt = custom_prompt( + role_dict=model_prompt_details["roles"], + initial_prompt_value=model_prompt_details["initial_prompt_value"], + final_prompt_value=model_prompt_details["final_prompt_value"], + messages=messages + ) + else: + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic") else: prompt = "" for message in messages: @@ -256,6 +250,7 @@ os.environ['AWS_SECRET_ACCESS_KEY'] = "" def completion( model: str, messages: list, + custom_prompt_dict: dict, model_response: ModelResponse, print_verbose: Callable, encoding, @@ -282,7 +277,7 @@ def completion( model = model provider = model.split(".")[0] - prompt = convert_messages_to_prompt(messages, provider) + prompt = convert_messages_to_prompt(model, messages, provider, custom_prompt_dict) inference_params = copy.deepcopy(optional_params) stream = inference_params.pop("stream", False) if provider == "anthropic": diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index 319e7047f2..959b8759f4 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -1,3 +1,4 @@ +from enum import Enum import requests, traceback import json from jinja2 import Template, exceptions, Environment, meta @@ -201,6 +202,31 @@ def hf_chat_template(model: str, messages: list): except: raise Exception("Error rendering template") +# Anthropic template +def anthropic_pt(messages: list): # format - https://docs.anthropic.com/claude/reference/complete_post + class AnthropicConstants(Enum): + HUMAN_PROMPT = "\n\nHuman: " + AI_PROMPT = "\n\nAssistant: " + + prompt = "" + for idx, message in enumerate(messages): # needs to start with `\n\nHuman: ` and end with `\n\nAssistant: ` + if message["role"] == "user": + prompt += ( + f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" + ) + elif message["role"] == "system": + prompt += ( + f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" + ) + else: + prompt += ( + f"{AnthropicConstants.AI_PROMPT.value}{message['content']}" + ) + if idx == 0 and message["role"] == "assistant": # ensure the prompt always starts with `\n\nHuman: ` + prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}" + prompt + prompt += f"{AnthropicConstants.AI_PROMPT.value}" + return prompt + # Function call template def function_call_prompt(messages: list, functions: list): function_prompt = "The following functions are available to you:" @@ -249,6 +275,8 @@ def prompt_factory(model: str, messages: list, custom_llm_provider: Optional[str if custom_llm_provider == "ollama": return ollama_pt(messages=messages) + elif custom_llm_provider == "anthropic": + return anthropic_pt(messages=messages) try: if "meta-llama/llama-2" in model: diff --git a/litellm/main.py b/litellm/main.py index 37927a3371..b7e54e1ea0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -586,7 +586,7 @@ def completion( return response response = model_response - elif model in litellm.anthropic_models: + elif custom_llm_provider=="anthropic": anthropic_key = ( api_key or litellm.anthropic_key or os.environ.get("ANTHROPIC_API_KEY") or litellm.api_key ) @@ -600,6 +600,7 @@ def completion( model=model, messages=messages, api_base=api_base, + custom_prompt_dict=litellm.custom_prompt_dict, model_response=model_response, print_verbose=print_verbose, optional_params=optional_params, @@ -1036,6 +1037,7 @@ def completion( model_response = bedrock.completion( model=model, messages=messages, + custom_prompt_dict=litellm.custom_prompt_dict, model_response=model_response, print_verbose=print_verbose, optional_params=optional_params,