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

adding coverage for additional baseten output formats

Browse source
This commit is contained in:
Krrish Dholakia 2023-08-24 18:20:39 -07:00
parent 964f146937
commit c23bf54e27
1 changed files with 8 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/llms/baseten.py
View file

@ -100,12 +100,15 @@ class BasetenLLM:
)
else:
if "model_output" in completion_response:
if isinstance(completion_response["model_output"], str):
model_response["choices"][0]["message"]["content"] = completion_response["model_output"]
elif isinstance(completion_response["model_output"], dict) and "data" in completion_response["model_output"] and isinstance(completion_response["model_output"]["data"], list):
if isinstance(completion_response["model_output"], dict) and "data" in completion_response["model_output"] and isinstance(completion_response["model_output"]["data"], list):
model_response["choices"][0]["message"]["content"] = completion_response["model_output"]["data"][0]
else:
raise ValueError(f"Unable to parse response. Original response: {response.text}")
elif isinstance(completion_response["model_output"], str):
model_response["choices"][0]["message"]["content"] = completion_response["model_output"]
elif "completion" in completion_response and isinstance(completion_response["completion"], str):
model_response["choices"][0]["message"]["content"] = completion_response["completion"]
else:
raise ValueError(f"Unable to parse response. Original response: {response.text}")
## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
prompt_tokens = len(
self.encoding.encode(prompt)