This commit is contained in:
Krrish Dholakia 2023-08-30 16:05:42 -07:00
parent 0ea59702fd
commit fcaf514546
2 changed files with 12 additions and 10 deletions

View file

@ -48,6 +48,8 @@ class HuggingfaceRestAPILLM:
completion_url: str = ""
if "https" in model:
completion_url = model
elif custom_api_base:
completion_url = custom_api_base
elif "HF_API_BASE" in os.environ:
completion_url = os.getenv("HF_API_BASE", "")
else:

View file

@ -83,16 +83,16 @@ def test_completion_claude_stream():
# pytest.fail(f"Error occurred: {e}")
def test_completion_hf_deployed_api():
try:
user_message = "There's a llama in my garden 😱 What should I do?"
messages = [{ "content": user_message,"role": "user"}]
response = completion(model="huggingface/https://ji16r2iys9a8rjk2.us-east-1.aws.endpoints.huggingface.cloud", messages=messages, logger_fn=logger_fn)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_completion_hf_deployed_api()
# def test_completion_hf_deployed_api():
# try:
# user_message = "There's a llama in my garden 😱 What should I do?"
# messages = [{ "content": user_message,"role": "user"}]
# response = completion(model="huggingface/https://ji16r2iys9a8rjk2.us-east-1.aws.endpoints.huggingface.cloud", messages=messages, logger_fn=logger_fn)
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# def test_completion_cohere(): # commenting for now as the cohere endpoint is being flaky
# try: