From dd1048cb3519ccba4804d53a2805efbb44922981 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Jul 2024 12:11:50 -0700 Subject: [PATCH] fix(main.py): fix linting errors --- litellm/main.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 949466642ca..bc40194b435 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -522,7 +522,7 @@ def mock_completion( ) return response if n is None: - model_response["choices"][0]["message"]["content"] = mock_response + model_response.choices[0].message.content = mock_response # type: ignore else: _all_choices = [] for i in range(n): @@ -533,12 +533,12 @@ def mock_completion( ), ) _all_choices.append(_choice) - model_response["choices"] = _all_choices - model_response["created"] = int(time.time()) - model_response["model"] = model + model_response.choices = _all_choices # type: ignore + model_response.created = int(time.time()) + model_response.model = model if mock_tool_calls: - model_response["choices"][0]["message"]["tool_calls"] = [ + model_response.choices[0].message.tool_calls = [ # type: ignore ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls ] @@ -2672,9 +2672,9 @@ def completion( """ string_response = response_json["data"][0]["output"][0] ## RESPONSE OBJECT - model_response["choices"][0]["message"]["content"] = string_response - model_response["created"] = int(time.time()) - model_response["model"] = model + model_response.choices[0].message.content = string_response # type: ignore + model_response.created = int(time.time()) + model_response.model = model response = model_response else: raise ValueError(