phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 19:54:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_completion.py): re-add bedrock + sagemaker testing

Browse source
This commit is contained in:
Krrish Dholakia 2023-10-31 16:49:05 -07:00
parent ed8f1ab5dc
commit b98a58d1b1
2 changed files with 84 additions and 82 deletions
Download patch file Download diff file Expand all files Collapse all files

16
litellm/main.py
View file

@ -906,15 +906,19 @@ def completion(
openai.api_key = get_secret("OPENROUTER_API_KEY") or get_secret( openai.api_key = get_secret("OPENROUTER_API_KEY") or get_secret(
"OR_API_KEY" "OR_API_KEY"
) or litellm.api_key ) or litellm.api_key
data = {
"model": model,
"messages": messages,
**optional_params
}
## LOGGING ## LOGGING
logging.pre_call(input=messages, api_key=openai.api_key) logging.pre_call(input=messages, api_key=openai.api_key, additional_args={"complete_input_dict": data, "headers": headers})
## COMPLETION CALL ## COMPLETION CALL
if litellm.headers: if litellm.headers:
response = openai.ChatCompletion.create( response = openai.ChatCompletion.create(
model=model,
messages=messages,
headers=litellm.headers, headers=litellm.headers,
**optional_params, **data,
) )
else: else:
openrouter_site_url = get_secret("OR_SITE_URL") openrouter_site_url = get_secret("OR_SITE_URL")
@ -926,13 +930,11 @@ def completion(
if openrouter_app_name is None: if openrouter_app_name is None:
openrouter_app_name = "liteLLM" openrouter_app_name = "liteLLM"
response = openai.ChatCompletion.create( response = openai.ChatCompletion.create(
model=model,
messages=messages,
headers={ headers={
"HTTP-Referer": openrouter_site_url, # To identify your site "HTTP-Referer": openrouter_site_url, # To identify your site
"X-Title": openrouter_app_name, # To identify your app "X-Title": openrouter_app_name, # To identify your app
}, },
**optional_params, **data,
) )
## LOGGING ## LOGGING
logging.post_call( logging.post_call(

150
litellm/tests/test_completion.py
View file

@ -812,93 +812,93 @@ def test_completion_together_ai():
# except Exception as e: # except Exception as e:
# pytest.fail(f"Error occurred: {e}") # pytest.fail(f"Error occurred: {e}")
# def test_completion_sagemaker(): def test_completion_sagemaker():
# try: try:
# response = completion( response = completion(
# model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b",
# messages=messages, messages=messages,
# temperature=0.2, temperature=0.2,
# max_tokens=80, max_tokens=80,
# logger_fn=logger_fn logger_fn=logger_fn
# ) )
# # Add any assertions here to check the response # Add any assertions here to check the response
# print(response) print(response)
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_sagemaker() # test_completion_sagemaker()
# def test_completion_bedrock_titan(): def test_completion_bedrock_titan():
# try: try:
# response = completion( response = completion(
# model="bedrock/amazon.titan-tg1-large", model="bedrock/amazon.titan-tg1-large",
# messages=messages, messages=messages,
# temperature=0.2, temperature=0.2,
# max_tokens=200, max_tokens=200,
# top_p=0.8, top_p=0.8,
# logger_fn=logger_fn logger_fn=logger_fn
# ) )
# # Add any assertions here to check the response # Add any assertions here to check the response
# print(response) print(response)
# except RateLimitError: except RateLimitError:
# pass pass
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_bedrock_titan() # test_completion_bedrock_titan()
# def test_completion_bedrock_claude(): def test_completion_bedrock_claude():
# print("calling claude") print("calling claude")
# try: try:
# response = completion( response = completion(
# model="anthropic.claude-instant-v1", model="anthropic.claude-instant-v1",
# messages=messages, messages=messages,
# max_tokens=10, max_tokens=10,
# temperature=0.1, temperature=0.1,
# logger_fn=logger_fn logger_fn=logger_fn
# ) )
# # Add any assertions here to check the response # Add any assertions here to check the response
# print(response) print(response)
# except RateLimitError: except RateLimitError:
# pass pass
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_bedrock_claude() test_completion_bedrock_claude()
# def test_completion_bedrock_claude_completion_auth(): def test_completion_bedrock_claude_completion_auth():
# print("calling bedrock claude completion params auth") print("calling bedrock claude completion params auth")
# import os import os
# aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"]
# aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"]
# aws_region_name = os.environ["AWS_REGION_NAME"] aws_region_name = os.environ["AWS_REGION_NAME"]
# os.environ["AWS_ACCESS_KEY_ID"] = "" os.environ["AWS_ACCESS_KEY_ID"] = ""
# os.environ["AWS_SECRET_ACCESS_KEY"] = "" os.environ["AWS_SECRET_ACCESS_KEY"] = ""
# os.environ["AWS_REGION_NAME"] = "" os.environ["AWS_REGION_NAME"] = ""
# try: try:
# response = completion( response = completion(
# model="bedrock/anthropic.claude-instant-v1", model="bedrock/anthropic.claude-instant-v1",
# messages=messages, messages=messages,
# max_tokens=10, max_tokens=10,
# temperature=0.1, temperature=0.1,
# logger_fn=logger_fn, logger_fn=logger_fn,
# aws_access_key_id=aws_access_key_id, aws_access_key_id=aws_access_key_id,
# aws_secret_access_key=aws_secret_access_key, aws_secret_access_key=aws_secret_access_key,
# aws_region_name=aws_region_name, aws_region_name=aws_region_name,
# ) )
# # Add any assertions here to check the response # Add any assertions here to check the response
# print(response) print(response)
# os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
# os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
# os.environ["AWS_REGION_NAME"] = aws_region_name os.environ["AWS_REGION_NAME"] = aws_region_name
# except RateLimitError: except RateLimitError:
# pass pass
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_bedrock_claude_completion_auth() # test_completion_bedrock_claude_completion_auth()
# def test_completion_bedrock_claude_external_client_auth(): # def test_completion_bedrock_claude_external_client_auth():