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

Merge pull request #2474 from BerriAI/litellm_support_command_r

Browse source 
[New-Model] Cohere/command-r
This commit is contained in:
Ishaan Jaff 2024-03-12 11:11:56 -07:00 • committed by GitHub
commit 5172fb1de9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 386 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

36
litellm/utils.py
View file

@ -7411,7 +7411,9 @@ def exception_type(
model=model,
response=original_exception.response,
)
elif custom_llm_provider == "cohere": # Cohere
elif (
custom_llm_provider == "cohere" or custom_llm_provider == "cohere_chat"
): # Cohere
if (
"invalid api token" in error_str
or "No API key provided." in error_str
@ -8544,6 +8546,29 @@ class CustomStreamWrapper:
except:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_cohere_chat_chunk(self, chunk):
chunk = chunk.decode("utf-8")
data_json = json.loads(chunk)
print_verbose(f"chunk: {chunk}")
try:
text = ""
is_finished = False
finish_reason = ""
if "text" in data_json:
text = data_json["text"]
elif "is_finished" in data_json and data_json["is_finished"] == True:
is_finished = data_json["is_finished"]
finish_reason = data_json["finish_reason"]
else:
return
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
except:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_azure_chunk(self, chunk):
is_finished = False
finish_reason = ""
@ -9073,6 +9098,15 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = response_obj[
"finish_reason"
]
elif self.custom_llm_provider == "cohere_chat":
response_obj = self.handle_cohere_chat_chunk(chunk)
if response_obj is None:
return
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj[
"finish_reason"
]
elif self.custom_llm_provider == "bedrock":
if self.sent_last_chunk:
raise StopIteration