Merge pull request BerriAI#4807 from BerriAI/litellm_return-response_headers

[Feat] Return response headers on `litellm.completion` , `litellm.embedding`

Author: ishaan-jaff

docs/my-website/docs/providers/openai.md

@@ -238,6 +238,104 @@ response = completion(
 
 ## Advanced
 
+### Getting OpenAI API Response Headers 
+
+Set `litellm.return_response_headers = True` to get raw response headers from OpenAI
+
+You can expect to always get the `_response_headers` field from `litellm.completion()`, `litellm.embedding()` functions
+
+<Tabs>
+<TabItem value="litellm.completion" label="litellm.completion">
+
+```python
+litellm.return_response_headers = True
+
+# /chat/completion
+response = completion(
+    model="gpt-4o-mini",
+    messages=[
+        {
+            "role": "user",
+            "content": "hi",
+        }
+    ],
+)
+print(f"response: {response}")
+print("_response_headers=", response._response_headers)
+```
+</TabItem>
+
+<TabItem value="litellm.completion - streaming" label="litellm.completion + stream">
+
+```python
+litellm.return_response_headers = True
+
+# /chat/completion
+response = completion(
+    model="gpt-4o-mini",
+    stream=True,
+    messages=[
+        {
+            "role": "user",
+            "content": "hi",
+        }
+    ],
+)
+print(f"response: {response}")
+print("response_headers=", response._response_headers)
+for chunk in response:
+    print(chunk)
+```
+</TabItem>
+
+<TabItem value="litellm.embedding" label="litellm.embedding">
+
+```python
+litellm.return_response_headers = True
+
+# embedding
+embedding_response = litellm.embedding(
+    model="text-embedding-ada-002",
+    input="hello",
+)
+
+embedding_response_headers = embedding_response._response_headers
+print("embedding_response_headers=", embedding_response_headers)
+```
+
+</TabItem>
+</Tabs>
+Expected Response Headers from OpenAI
+
+```json
+{
+  "date": "Sat, 20 Jul 2024 22:05:23 GMT",
+  "content-type": "application/json",
+  "transfer-encoding": "chunked",
+  "connection": "keep-alive",
+  "access-control-allow-origin": "*",
+  "openai-model": "text-embedding-ada-002",
+  "openai-organization": "*****",
+  "openai-processing-ms": "20",
+  "openai-version": "2020-10-01",
+  "strict-transport-security": "max-age=15552000; includeSubDomains; preload",
+  "x-ratelimit-limit-requests": "5000",
+  "x-ratelimit-limit-tokens": "5000000",
+  "x-ratelimit-remaining-requests": "4999",
+  "x-ratelimit-remaining-tokens": "4999999",
+  "x-ratelimit-reset-requests": "12ms",
+  "x-ratelimit-reset-tokens": "0s",
+  "x-request-id": "req_cc37487bfd336358231a17034bcfb4d9",
+  "cf-cache-status": "DYNAMIC",
+  "set-cookie": "__cf_bm=E_FJY8fdAIMBzBE2RZI2.OkMIO3lf8Hz.ydBQJ9m3q8-1721513123-1.0.1.1-6OK0zXvtd5s9Jgqfz66cU9gzQYpcuh_RLaUZ9dOgxR9Qeq4oJlu.04C09hOTCFn7Hg.k.2tiKLOX24szUE2shw; path=/; expires=Sat, 20-Jul-24 22:35:23 GMT; domain=.api.openai.com; HttpOnly; Secure; SameSite=None, *cfuvid=SDndIImxiO3U0aBcVtoy1TBQqYeQtVDo1L6*Nlpp7EU-1721513123215-0.0.1.1-604800000; path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None",
+  "x-content-type-options": "nosniff",
+  "server": "cloudflare",
+  "cf-ray": "8a66409b4f8acee9-SJC",
+  "content-encoding": "br",
+  "alt-svc": "h3=\":443\"; ma=86400"
+}
+```
+
 ### Parallel Function calling
 See a detailed walthrough of parallel function calling with litellm [here](https://docs.litellm.ai/docs/completion/function_call)
 ```python

litellm/llms/openai.py

@@ -784,6 +784,34 @@ async def make_openai_chat_completion_request(
         except Exception as e:
             raise e
 
+    def make_sync_openai_chat_completion_request(
+        self,
+        openai_client: OpenAI,
+        data: dict,
+        timeout: Union[float, httpx.Timeout],
+    ):
+        """
+        Helper to:
+        - call chat.completions.create.with_raw_response when litellm.return_response_headers is True
+        - call chat.completions.create by default
+        """
+        try:
+            if litellm.return_response_headers is True:
+                raw_response = openai_client.chat.completions.with_raw_response.create(
+                    **data, timeout=timeout
+                )
+
+                headers = dict(raw_response.headers)
+                response = raw_response.parse()
+                return headers, response
+            else:
+                response = openai_client.chat.completions.create(
+                    **data, timeout=timeout
+                )
+                return None, response
+        except Exception as e:
+            raise e
+
     def completion(
         self,
         model_response: ModelResponse,
@@ -916,7 +944,15 @@ def completion(
                             },
                         )
 
-                        response = openai_client.chat.completions.create(**data, timeout=timeout)  # type: ignore
+                        headers, response = (
+                            self.make_sync_openai_chat_completion_request(
+                                openai_client=openai_client,
+                                data=data,
+                                timeout=timeout,
+                            )
+                        )
+
+                        logging_obj.model_call_details["response_headers"] = headers
                         stringified_response = response.model_dump()
                         logging_obj.post_call(
                             input=messages,
@@ -927,6 +963,7 @@ def completion(
                         return convert_to_model_response_object(
                             response_object=stringified_response,
                             model_response_object=model_response,
+                            _response_headers=headers,
                         )
                 except openai.UnprocessableEntityError as e:
                     ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800
@@ -1043,6 +1080,25 @@ async def acompletion(
                     },
                 )
 
+            headers, response = await self.make_openai_chat_completion_request(
+                openai_aclient=openai_aclient, data=data, timeout=timeout
+            )
+            stringified_response = response.model_dump()
+            logging_obj.post_call(
+                input=data["messages"],
+                api_key=api_key,
+                original_response=stringified_response,
+                additional_args={"complete_input_dict": data},
+            )
+            logging_obj.model_call_details["response_headers"] = headers
+            return convert_to_model_response_object(
+                response_object=stringified_response,
+                model_response_object=model_response,
+                hidden_params={"headers": headers},
+                _response_headers=headers,
+            )
+        except Exception as e:
+            raise e
                 headers, response = await self.make_openai_chat_completion_request(
                     openai_aclient=openai_aclient, data=data, timeout=timeout
                 )
@@ -1122,13 +1178,20 @@ def streaming(
                 "complete_input_dict": data,
             },
         )
-        response = openai_client.chat.completions.create(**data, timeout=timeout)
+        headers, response = self.make_sync_openai_chat_completion_request(
+            openai_client=openai_client,
+            data=data,
+            timeout=timeout,
+        )
+
+        logging_obj.model_call_details["response_headers"] = headers
         streamwrapper = CustomStreamWrapper(
             completion_stream=response,
             model=model,
             custom_llm_provider="openai",
             logging_obj=logging_obj,
             stream_options=data.get("stream_options", None),
+            _response_headers=headers,
         )
         return streamwrapper
 
@@ -1170,8 +1233,30 @@ async def async_streaming(
                     },
                 )
 
+            headers, response = await self.make_openai_chat_completion_request(
+                openai_aclient=openai_aclient, data=data, timeout=timeout
+            )
+            logging_obj.model_call_details["response_headers"] = headers
+            streamwrapper = CustomStreamWrapper(
+                completion_stream=response,
+                model=model,
+                custom_llm_provider="openai",
+                logging_obj=logging_obj,
+                stream_options=data.get("stream_options", None),
+                _response_headers=headers,
+            )
+            return streamwrapper
+        except (
+            Exception
+        ) as e:  # need to exception handle here. async exceptions don't get caught in sync functions.
+            if response is not None and hasattr(response, "text"):
+                raise OpenAIError(
+                    status_code=500,
+                    message=f"{str(e)}\n\nOriginal Response: {response.text}",
+
                 headers, response = await self.make_openai_chat_completion_request(
                     openai_aclient=openai_aclient, data=data, timeout=timeout
+
                 )
                 logging_obj.model_call_details["response_headers"] = headers
                 streamwrapper = CustomStreamWrapper(
@@ -1252,6 +1337,32 @@ async def make_openai_embedding_request(
         except Exception as e:
             raise e
 
+    def make_sync_openai_embedding_request(
+        self,
+        openai_client: OpenAI,
+        data: dict,
+        timeout: Union[float, httpx.Timeout],
+    ):
+        """
+        Helper to:
+        - call embeddings.create.with_raw_response when litellm.return_response_headers is True
+        - call embeddings.create by default
+        """
+        try:
+            if litellm.return_response_headers is True:
+                raw_response = openai_client.embeddings.with_raw_response.create(
+                    **data, timeout=timeout
+                )  # type: ignore
+
+                headers = dict(raw_response.headers)
+                response = raw_response.parse()
+                return headers, response
+            else:
+                response = openai_client.embeddings.create(**data, timeout=timeout)  # type: ignore
+                return None, response
+        except Exception as e:
+            raise e
+
     async def aembedding(
         self,
         input: list,
@@ -1286,7 +1397,12 @@ async def aembedding(
                 additional_args={"complete_input_dict": data},
                 original_response=stringified_response,
             )
-            return convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, response_type="embedding")  # type: ignore
+            return convert_to_model_response_object(
+                response_object=stringified_response,
+                model_response_object=model_response,
+                response_type="embedding",
+                _response_headers=headers,
+            )  # type: ignore
         except Exception as e:
             ## LOGGING
             logging_obj.post_call(
@@ -1347,17 +1463,26 @@ def embedding(
                 client=client,
             )
 
-            ## COMPLETION CALL
-            response = openai_client.embeddings.create(**data, timeout=timeout)  # type: ignore
+            ## embedding CALL
+            headers: Optional[Dict] = None
+            headers, sync_embedding_response = self.make_sync_openai_embedding_request(
+                openai_client=openai_client, data=data, timeout=timeout
+            )  # type: ignore
+
             ## LOGGING
+            logging_obj.model_call_details["response_headers"] = headers
             logging_obj.post_call(
                 input=input,
                 api_key=api_key,
                 additional_args={"complete_input_dict": data},
-                original_response=response,
+                original_response=sync_embedding_response,
             )
-
-            return convert_to_model_response_object(response_object=response.model_dump(), model_response_object=model_response, response_type="embedding")  # type: ignore
+            return convert_to_model_response_object(
+                response_object=sync_embedding_response.model_dump(),
+                model_response_object=model_response,
+                _response_headers=headers,
+                response_type="embedding",
+            )  # type: ignore
         except OpenAIError as e:
             exception_mapping_worked = True
             raise e
@@ -1520,6 +1645,33 @@ async def make_openai_audio_transcriptions_request(
         except Exception as e:
             raise e
 
+    def make_sync_openai_audio_transcriptions_request(
+        self,
+        openai_client: OpenAI,
+        data: dict,
+        timeout: Union[float, httpx.Timeout],
+    ):
+        """
+        Helper to:
+        - call openai_aclient.audio.transcriptions.with_raw_response when litellm.return_response_headers is True
+        - call openai_aclient.audio.transcriptions.create by default
+        """
+        try:
+            if litellm.return_response_headers is True:
+                raw_response = (
+                    openai_client.audio.transcriptions.with_raw_response.create(
+                        **data, timeout=timeout
+                    )
+                )  # type: ignore
+                headers = dict(raw_response.headers)
+                response = raw_response.parse()
+                return headers, response
+            else:
+                response = openai_client.audio.transcriptions.create(**data, timeout=timeout)  # type: ignore
+                return None, response
+        except Exception as e:
+            raise e
+
     def audio_transcriptions(
         self,
         model: str,
@@ -1555,8 +1707,10 @@ def audio_transcriptions(
             timeout=timeout,
             max_retries=max_retries,
         )
-        response = openai_client.audio.transcriptions.create(
-            **data, timeout=timeout  # type: ignore
+        response = self.make_sync_openai_audio_transcriptions_request(
+            openai_client=openai_client,
+            data=data,
+            timeout=timeout,
         )
 
         if isinstance(response, BaseModel):