mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
fix(utils.py): return additional kwargs from openai-like response body
Closes https://github.com/BerriAI/litellm/issues/4981
This commit is contained in:
parent
bd68714f51
commit
09ee8c6e2d
2 changed files with 70 additions and 19 deletions
|
@ -1533,9 +1533,53 @@ def test_completion_fireworks_ai_dynamic_params(api_key, api_base):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.skip(reason="this test is flaky")
|
# @pytest.mark.skip(reason="this test is flaky")
|
||||||
def test_completion_perplexity_api():
|
def test_completion_perplexity_api():
|
||||||
try:
|
try:
|
||||||
|
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
|
# litellm.set_verbose= True
|
||||||
messages = [
|
messages = [
|
||||||
{"role": "system", "content": "You're a good bot"},
|
{"role": "system", "content": "You're a good bot"},
|
||||||
|
@ -1552,8 +1596,10 @@ def test_completion_perplexity_api():
|
||||||
model="mistral-7b-instruct",
|
model="mistral-7b-instruct",
|
||||||
messages=messages,
|
messages=messages,
|
||||||
api_base="https://api.perplexity.ai",
|
api_base="https://api.perplexity.ai",
|
||||||
|
client=openai_client,
|
||||||
)
|
)
|
||||||
print(response)
|
print(response)
|
||||||
|
assert hasattr(response, "citations")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
pytest.fail(f"Error occurred: {e}")
|
pytest.fail(f"Error occurred: {e}")
|
||||||
|
|
||||||
|
|
|
@ -5859,6 +5859,11 @@ def convert_to_model_response_object(
|
||||||
if _response_headers is not None:
|
if _response_headers is not None:
|
||||||
model_response_object._response_headers = _response_headers
|
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
|
return model_response_object
|
||||||
elif response_type == "embedding" and (
|
elif response_type == "embedding" and (
|
||||||
model_response_object is None
|
model_response_object is None
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue