diff --git a/litellm/main.py b/litellm/main.py index 6ed5275052..59261fc863 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -277,13 +277,21 @@ def completion( **optional_params ) else: + openrouter_site_url = get_secret("OR_SITE_URL") + openrouter_app_name = get_secret("OR_APP_NAME") + # if openrouter_site_url is None, set it to https://litellm.ai + if openrouter_site_url is None: + openrouter_site_url = "https://litellm.ai" + # if openrouter_app_name is None, set it to liteLLM + if openrouter_app_name is None: + openrouter_app_name = "liteLLM" response = openai.ChatCompletion.create( model=model, messages = messages, headers = { - "HTTP-Referer": os.environ.get("OR_SITE_URL"), # To identify your app - "X-Title": os.environ.get("OR_APP_NAME") + "HTTP-Referer": openrouter_site_url, # To identify your site + "X-Title": openrouter_app_name # To identify your app }, **optional_params ) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 43c416e75e..b4f3866b82 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -87,14 +87,13 @@ def test_completion_openai_with_optional_params(): except Exception as e: pytest.fail(f"Error occurred: {e}") -# TODO: Fix this test -# def test_completion_openrouter(): -# try: -# response = completion(model="google/palm-2-chat-bison", messages=messages, temperature=0.5, top_p=0.1, user="ishaan_dev@berri.ai") -# # Add any assertions here to check the response -# print(response) -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") +def test_completion_openrouter(): + try: + response = completion(model="google/palm-2-chat-bison", messages=messages, temperature=0.5, top_p=0.1, user="ishaan_dev@berri.ai") + # Add any assertions here to check the response + print(response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") def test_completion_openai_with_more_optional_params(): try: