phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) watsonx.py: Fixed linting errors and make sure stream chunk always return usage

Browse source
This commit is contained in:
Simon Sanchez Viloria 2024-05-10 11:53:33 +02:00
parent 66a1b581e5
commit 170fd11c82
2 changed files with 279 additions and 92 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/utils.py
View file

@ -10285,7 +10285,7 @@ class CustomStreamWrapper:
response = chunk.replace("data: ", "").strip()
parsed_response = json.loads(response)
else:
return {"text": "", "is_finished": False}
return {"text": "", "is_finished": False, "prompt_tokens": 0, "completion_tokens": 0}
else:
print_verbose(f"chunk: {chunk} (Type: {type(chunk)})")
raise ValueError(
@ -10300,8 +10300,8 @@ class CustomStreamWrapper:
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
"prompt_tokens": results[0].get("input_token_count", None),
"completion_tokens": results[0].get("generated_token_count", None),
"prompt_tokens": results[0].get("input_token_count", 0),
"completion_tokens": results[0].get("generated_token_count", 0),
}
return {"text": "", "is_finished": False}
except Exception as e: