test_clarifai_completion.py

importsys, os
importtraceback
fromdotenvimportload_dotenv
importasyncio, logging

load_dotenv()
importos, io

sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
importpytest
importlitellm
fromlitellmimport (
embedding,
completion,
acompletion,
acreate,
completion_cost,
Timeout,
ModelResponse,
)
fromlitellmimportRateLimitError

# litellm.num_retries = 3
litellm.cache=None
litellm.success_callback= []
user_message="Write a short poem about the sky"
messages= [{"content": user_message, "role": "user"}]


@pytest.fixture(autouse=True)
defreset_callbacks():
print("\npytest fixture - resetting callbacks")
litellm.success_callback= []
litellm._async_success_callback= []
litellm.failure_callback= []
litellm.callbacks= []


@pytest.mark.skip(reason="Account rate limited.")
deftest_completion_clarifai_claude_2_1():
print("calling clarifai claude completion")
importos

clarifai_pat=os.environ["CLARIFAI_API_KEY"]

try:
response=completion(
model="clarifai/anthropic.completion.claude-2_1",
num_retries=3,
messages=messages,
max_tokens=10,
temperature=0.1,
 )
print(response)

exceptRateLimitError:
pass

exceptExceptionase:
pytest.fail(f"Error occured: {e}")


@pytest.mark.skip(reason="Account rate limited")
deftest_completion_clarifai_mistral_large():
try:
litellm.set_verbose=True
response: ModelResponse=completion(
model="clarifai/mistralai.completion.mistral-small",
messages=messages,
num_retries=3,
max_tokens=10,
temperature=0.78,
 )
# Add any assertions here to check the response
assertlen(response.choices) >0
assertlen(response.choices[0].message.content) >0
exceptRateLimitError:
pass
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


@pytest.mark.skip(reason="Account rate limited")
@pytest.mark.asyncio
deftest_async_completion_clarifai():
importasyncio

litellm.set_verbose=True

asyncdeftest_get_response():
user_message="Hello, how are you?"
messages= [{"content": user_message, "role": "user"}]
try:
response=awaitacompletion(
model="clarifai/openai.chat-completion.GPT-4",
messages=messages,
num_retries=3,
timeout=10,
api_key=os.getenv("CLARIFAI_API_KEY"),
 )
print(f"response: {response}")
exceptlitellm.Timeoutase:
pass
exceptExceptionase:
pytest.fail(f"An exception occurred: {e}")

asyncio.run(test_get_response())