Merge pull request #3561 from simonsanvil/feature/watsonx-integration

(fix) Fixed linting and other bugs with watsonx provider
This commit is contained in:
Krish Dholakia 2024-05-11 09:56:02 -07:00 committed by GitHub
commit d33e49411d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 310 additions and 101 deletions

View file

@ -10430,7 +10430,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(
@ -10445,8 +10445,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: