test fixes for vertex mistral, this model was deprecated on vertex

This commit is contained in:
Ishaan Jaff 2025-04-16 20:51:45 -07:00
parent c38146e180
commit 198922b26f

View file

@ -4112,98 +4112,6 @@ async def test_async_text_completion_chat_model_stream():
# asyncio.run(test_async_text_completion_chat_model_stream())
@pytest.mark.parametrize("model", ["vertex_ai/mistral-nemo@2407"]) #
@pytest.mark.asyncio
async def test_completion_codestral_fim_api(model):
try:
from test_amazing_vertex_completion import (
load_vertex_ai_credentials,
)
load_vertex_ai_credentials()
litellm.set_verbose = True
import logging
from litellm._logging import verbose_logger
verbose_logger.setLevel(level=logging.DEBUG)
response = await litellm.atext_completion(
model=model,
prompt="def is_odd(n): \n return n % 2 == 1 \ndef test_is_odd():",
suffix="return True",
temperature=0,
top_p=1,
max_tokens=10,
min_tokens=10,
seed=10,
stop=["return"],
)
# Add any assertions here to check the response
print(response)
assert response.choices[0].text is not None
# cost = litellm.completion_cost(completion_response=response)
# print("cost to make mistral completion=", cost)
# assert cost > 0.0
except litellm.ServiceUnavailableError:
print("got ServiceUnavailableError")
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize(
"model",
["vertex_ai/mistral-nemo@2407"],
)
@pytest.mark.asyncio
async def test_completion_codestral_fim_api_stream(model):
try:
from test_amazing_vertex_completion import (
load_vertex_ai_credentials,
)
load_vertex_ai_credentials()
import logging
from litellm._logging import verbose_logger
litellm.set_verbose = False
# verbose_logger.setLevel(level=logging.DEBUG)
response = await litellm.atext_completion(
model=model,
prompt="def is_odd(n): \n return n % 2 == 1 \ndef test_is_odd():",
suffix="return True",
temperature=0,
top_p=1,
stream=True,
seed=10,
stop=["return"],
)
full_response = ""
# Add any assertions here to check the response
async for chunk in response:
print(chunk)
full_response += chunk.get("choices")[0].get("text") or ""
print("full_response", full_response)
# cost = litellm.completion_cost(completion_response=response)
# print("cost to make mistral completion=", cost)
# assert cost > 0.0
except litellm.APIConnectionError as e:
print(e)
pass
except litellm.ServiceUnavailableError as e:
print(e)
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def mock_post(*args, **kwargs):
mock_response = MagicMock()
mock_response.status_code = 200