diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 99542a9a20..f38c58b664 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1533,27 +1533,73 @@ def test_completion_fireworks_ai_dynamic_params(api_key, api_base): pass -@pytest.mark.skip(reason="this test is flaky") +# @pytest.mark.skip(reason="this test is flaky") def test_completion_perplexity_api(): try: - # litellm.set_verbose= True - messages = [ - {"role": "system", "content": "You're a good bot"}, - { - "role": "user", - "content": "Hey", - }, - { - "role": "user", - "content": "Hey", - }, - ] - response = completion( - model="mistral-7b-instruct", - messages=messages, - api_base="https://api.perplexity.ai", - ) - print(response) + response_object = { + "id": "a8f37485-026e-45da-81a9-cf0184896840", + "model": "llama-3-sonar-small-32k-online", + "created": 1722186391, + "usage": {"prompt_tokens": 17, "completion_tokens": 65, "total_tokens": 82}, + "citations": [ + "https://www.sciencedirect.com/science/article/pii/S007961232200156X", + "https://www.britannica.com/event/World-War-II", + "https://www.loc.gov/classroom-materials/united-states-history-primary-source-timeline/great-depression-and-world-war-ii-1929-1945/world-war-ii/", + "https://www.nationalww2museum.org/war/topics/end-world-war-ii-1945", + "https://en.wikipedia.org/wiki/World_War_II", + ], + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "World War II was won by the Allied powers, which included the United States, the Soviet Union, Great Britain, France, China, and other countries. The war concluded with the surrender of Germany on May 8, 1945, and Japan on September 2, 1945[2][3][4].", + }, + "delta": {"role": "assistant", "content": ""}, + } + ], + } + + from openai import OpenAI + from openai.types.chat.chat_completion import ChatCompletion + + pydantic_obj = ChatCompletion(**response_object) + + def _return_pydantic_obj(*args, **kwargs): + return pydantic_obj + + print(f"pydantic_obj: {pydantic_obj}") + + openai_client = OpenAI() + + openai_client.chat.completions.create = MagicMock() + + with patch.object( + openai_client.chat.completions, "create", side_effect=_return_pydantic_obj + ) as mock_client: + pass + # litellm.set_verbose= True + messages = [ + {"role": "system", "content": "You're a good bot"}, + { + "role": "user", + "content": "Hey", + }, + { + "role": "user", + "content": "Hey", + }, + ] + response = completion( + model="mistral-7b-instruct", + messages=messages, + api_base="https://api.perplexity.ai", + client=openai_client, + ) + print(response) + assert hasattr(response, "citations") except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index 02ec7415a9..70e000a1a0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5859,6 +5859,11 @@ def convert_to_model_response_object( if _response_headers is not None: model_response_object._response_headers = _response_headers + special_keys = litellm.ModelResponse.model_fields.keys() + for k, v in response_object.items(): + if k not in special_keys: + setattr(model_response_object, k, v) + return model_response_object elif response_type == "embedding" and ( model_response_object is None