From 2f6fed30ad2eee2e1850b0d978c7fec8d5c890a2 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 31 Oct 2023 22:26:41 -0700 Subject: [PATCH] (test) bedrock.cohere testing --- litellm/tests/test_completion.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index db9e254206..a8ef59d64a 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -408,7 +408,7 @@ def test_completion_openai_prompt_array(): response_str = response["choices"][0]["text"] except Exception as e: pytest.fail(f"Error occurred: {e}") -test_completion_openai_prompt_array() +# test_completion_openai_prompt_array() def test_completion_hf_prompt_array(): try: @@ -423,7 +423,7 @@ def test_completion_hf_prompt_array(): response_str = response["choices"][0]["text"] except Exception as e: pytest.fail(f"Error occurred: {e}") -test_completion_hf_prompt_array() +# test_completion_hf_prompt_array() def test_completion_text_openai(): try: @@ -879,6 +879,26 @@ def test_completion_bedrock_claude(): pytest.fail(f"Error occurred: {e}") # test_completion_bedrock_claude() +def test_completion_bedrock_cohere(): + print("calling bedrock cohere") + try: + response = completion( + model="bedrock/cohere.command-text-v14", + messages=[{"role": "user", "content": "hi"}], + temperature=0.1, + max_tokens=10, + stream=True + ) + # Add any assertions here to check the response + print(response) + for chunk in response: + print(chunk) + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") +# test_completion_bedrock_cohere() + def test_completion_bedrock_claude_completion_auth(): print("calling bedrock claude completion params auth")