fix(utils.py): return additional kwargs from openai-like response body

Closes https://github.com/BerriAI/litellm/issues/4981

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
+        response_object = {
-        messages = [
+            "id": "a8f37485-026e-45da-81a9-cf0184896840",
-            {"role": "system", "content": "You're a good bot"},
+            "model": "llama-3-sonar-small-32k-online",
-            {
+            "created": 1722186391,
-                "role": "user",
+            "usage": {"prompt_tokens": 17, "completion_tokens": 65, "total_tokens": 82},
-                "content": "Hey",
+            "citations": [
-            },
+                "https://www.sciencedirect.com/science/article/pii/S007961232200156X",
-            {
+                "https://www.britannica.com/event/World-War-II",
-                "role": "user",
+                "https://www.loc.gov/classroom-materials/united-states-history-primary-source-timeline/great-depression-and-world-war-ii-1929-1945/world-war-ii/",
-                "content": "Hey",
+                "https://www.nationalww2museum.org/war/topics/end-world-war-ii-1945",
-            },
+                "https://en.wikipedia.org/wiki/World_War_II",
-        ]
+            ],
-        response = completion(
+            "object": "chat.completion",
-            model="mistral-7b-instruct",
+            "choices": [
-            messages=messages,
+                {
-            api_base="https://api.perplexity.ai",
+                    "index": 0,
-        )
+                    "finish_reason": "stop",
-        print(response)
+                    "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}")
 

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