test volcengine

This commit is contained in:
Ishaan Jaff 2024-06-26 17:09:30 -07:00
parent fcdda417bb
commit 0f489b68eb

View file

@ -1222,44 +1222,6 @@ def test_completion_fireworks_ai():
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
def test_fireworks_ai_tool_calling():
litellm.set_verbose = True
model_name = "fireworks_ai/accounts/fireworks/models/firefunction-v2"
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
response = completion(
model=model_name,
messages=messages,
tools=tools,
tool_choice="required",
)
print(response)
@pytest.mark.skip(reason="this test is flaky") @pytest.mark.skip(reason="this test is flaky")
def test_completion_perplexity_api(): def test_completion_perplexity_api():
try: try:
@ -3508,6 +3470,30 @@ def test_completion_deep_infra_mistral():
# test_completion_deep_infra_mistral() # test_completion_deep_infra_mistral()
@pytest.mark.skip(reason="Local test - don't have a volcengine account as yet")
def test_completion_volcengine():
litellm.set_verbose = True
model_name = "volcengine/<OUR_ENDPOINT_ID>"
try:
response = completion(
model=model_name,
messages=[
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
],
api_key="<OUR_API_KEY>",
)
# Add any assertions here to check the response
print(response)
except litellm.exceptions.Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_nvidia_nim(): def test_completion_nvidia_nim():
model_name = "nvidia_nim/databricks/dbrx-instruct" model_name = "nvidia_nim/databricks/dbrx-instruct"
try: try: