From debef7544d94987dfdb7f15e9b27e937f2c10a7d Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 17 Jan 2024 18:43:58 -0800
Subject: [PATCH] (feat) return Azure enahncements used

---
 litellm/utils.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/litellm/utils.py b/litellm/utils.py
index 5f7afc11965a..f7cc5d2a54a7 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -252,7 +252,13 @@ def __setitem__(self, key, value):
 
 class Choices(OpenAIObject):
     def __init__(
-        self, finish_reason=None, index=0, message=None, logprobs=None, **params
+        self,
+        finish_reason=None,
+        index=0,
+        message=None,
+        logprobs=None,
+        enhancements=None,
+        **params,
     ):
         super(Choices, self).__init__(**params)
         self.finish_reason = (
@@ -265,6 +271,8 @@ def __init__(
             self.message = message
         if logprobs is not None:
             self.logprobs = logprobs
+        if enhancements is not None:
+            self.enhancements = enhancements
 
     def __contains__(self, key):
         # Define custom behavior for the 'in' operator
@@ -319,6 +327,7 @@ def __init__(
         index=0,
         delta: Optional[Delta] = None,
         logprobs=None,
+        enhancements=None,
         **params,
     ):
         super(StreamingChoices, self).__init__(**params)
@@ -334,6 +343,8 @@ def __init__(
 
         if logprobs is not None:
             self.logprobs = logprobs
+        if enhancements is not None:
+            self.enhancements = enhancements
 
     def __contains__(self, key):
         # Define custom behavior for the 'in' operator
@@ -5092,8 +5103,13 @@ def convert_to_streaming_response(response_object: Optional[dict] = None):
             # gpt-4 vision can return 'finish_reason' or 'finish_details'
             finish_reason = choice.get("finish_details")
         logprobs = choice.get("logprobs", None)
+        enhancements = choice.get("enhancements", None)
         choice = StreamingChoices(
-            finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs
+            finish_reason=finish_reason,
+            index=idx,
+            delta=delta,
+            logprobs=logprobs,
+            enhancements=enhancements,
         )
 
         choice_list.append(choice)
@@ -5151,11 +5167,13 @@ def convert_to_model_response_object(
                     # gpt-4 vision can return 'finish_reason' or 'finish_details'
                     finish_reason = choice.get("finish_details")
                 logprobs = choice.get("logprobs", None)
+                enhancements = choice.get("enhancements", None)
                 choice = Choices(
                     finish_reason=finish_reason,
                     index=idx,
                     message=message,
                     logprobs=logprobs,
+                    enhancements=enhancements,
                 )
                 choice_list.append(choice)
             model_response_object.choices = choice_list