From c23bf54e2767f6e5207f57145973ce6e61a7deb3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 24 Aug 2023 18:20:39 -0700 Subject: [PATCH] adding coverage for additional baseten output formats --- litellm/llms/baseten.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/llms/baseten.py b/litellm/llms/baseten.py index b8016a32e47..218efa68320 100644 --- a/litellm/llms/baseten.py +++ b/litellm/llms/baseten.py @@ -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)