phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-24 18:24:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix

Browse source
This commit is contained in:
adriensas 2023-08-23 10:59:36 +02:00
parent d8ed7cd179
commit a0cb56efd5
1 changed files with 11 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

36
litellm/llms/anthropic.py
View file

@ -15,15 +15,11 @@ class AnthropicError(Exception):
def __init__(self, status_code, message):
self.status_code = status_code
self.message = message
super().__init__(
self.message
) # Call the base class constructor with the parameters it needs
super().__init__(self.message) # Call the base class constructor with the parameters it needs
class AnthropicLLM:
def __init__(
self, encoding, default_max_tokens_to_sample, logging_obj, api_key=None
):
def __init__(self, encoding, default_max_tokens_to_sample, logging_obj, api_key=None):
self.encoding = encoding
self.default_max_tokens_to_sample = default_max_tokens_to_sample
self.completion_url = "https://api.anthropic.com/v1/complete"
@ -31,9 +27,7 @@ class AnthropicLLM:
self.logging_obj = logging_obj
self.validate_environment(api_key=api_key)
def validate_environment(
self, api_key
): # set up the environment required to run the model
def validate_environment(self, api_key): # set up the environment required to run the model
# set the api key
if self.api_key == None:
raise ValueError(
@ -62,19 +56,13 @@ class AnthropicLLM:
for message in messages:
if "role" in message:
if message["role"] == "user":
prompt += (
f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
)
prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
else:
prompt += (
f"{AnthropicConstants.AI_PROMPT.value}{message['content']}"
)
prompt += f"{AnthropicConstants.AI_PROMPT.value}{message['content']}"
else:
prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
prompt += f"{AnthropicConstants.AI_PROMPT.value}"
if "max_tokens" in optional_params and optional_params["max_tokens"] != float(
"inf"
):
if "max_tokens" in optional_params and optional_params["max_tokens"] != float("inf"):
max_tokens = optional_params["max_tokens"]
else:
max_tokens = self.default_max_tokens_to_sample
@ -93,9 +81,11 @@ class AnthropicLLM:
)
## COMPLETION CALL
response = requests.post(
self.completion_url, headers=self.headers, data=json.dumps(data)
self.completion_url, headers=self.headers, data=json.dumps(data), stream=optional_params["stream"]
)
print(optional_params)
if "stream" in optional_params and optional_params["stream"] == True:
print("IS STREAMING")
return response.iter_lines()
else:
## LOGGING
@ -114,14 +104,10 @@ class AnthropicLLM:
status_code=response.status_code,
)
else:
model_response["choices"][0]["message"][
"content"
] = completion_response["completion"]
model_response["choices"][0]["message"]["content"] = completion_response["completion"]
## CALCULATING USAGE
prompt_tokens = len(
self.encoding.encode(prompt)
) ##[TODO] use the anthropic tokenizer here
prompt_tokens = len(self.encoding.encode(prompt)) ##[TODO] use the anthropic tokenizer here
completion_tokens = len(
self.encoding.encode(model_response["choices"][0]["message"]["content"])
) ##[TODO] use the anthropic tokenizer here