(test) Added completion and embedding tests for watsonx provider

This commit is contained in:
Simon Sanchez Viloria 2024-04-24 12:52:29 +02:00
parent f9a7456eaa
commit 9fc30e8b31
3 changed files with 73 additions and 0 deletions

View file

@ -2565,6 +2565,41 @@ def test_completion_palm_stream():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_watsonx():
litellm.set_verbose = True
model_name = "watsonx/ibm/granite-13b-chat-v2"
try:
response = completion(
model=model_name,
messages=messages,
stop=["stop"],
max_tokens=20,
)
# Add any assertions here to check the response
print(response)
except litellm.APIError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_acompletion_watsonx():
litellm.set_verbose = True
model_name = "watsonx/deployment/"+os.getenv("WATSONX_DEPLOYMENT_ID")
print("testing watsonx")
try:
response = await litellm.acompletion(
model=model_name,
messages=messages,
temperature=0.2,
max_tokens=80,
space_id=os.getenv("WATSONX_SPACE_ID_TEST"),
)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_palm_stream()