scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #4501 from BerriAI/litellm_return_Response_headers

Browse source 
[Feat] Return Response headers for OpenAI / Azure OpenAI when `litellm.return_response_headers=True`
This commit is contained in:
Ishaan Jaff 2024-07-01 21:11:10 -07:00 • committed by GitHub
commit 1c194f0275
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 158 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/__init__.py
View file

@ -125,6 +125,9 @@ llm_guard_mode: Literal["all", "key-specific", "request-specific"] = "all"
##################
### PREVIEW FEATURES ###
enable_preview_features: bool = False
return_response_headers: bool = (
False # get response headers from LLM Api providers - example x-remaining-requests,
)
##################
logging: bool = True
caching: bool = (

54
litellm/llms/azure.py
View file

@ -23,6 +23,7 @@ from typing_extensions import overload
import litellm
from litellm import OpenAIConfig
from litellm.caching import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.utils import (
Choices,
CustomStreamWrapper,
@ -458,6 +459,36 @@ class AzureChatCompletion(BaseLLM):
return azure_client
async def make_azure_openai_chat_completion_request(
self,
azure_client: AsyncAzureOpenAI,
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 = (
await azure_client.chat.completions.with_raw_response.create(
**data, timeout=timeout
)
)
headers = dict(raw_response.headers)
response = raw_response.parse()
return headers, response
else:
response = await azure_client.chat.completions.create(
**data, timeout=timeout
)
return None, response
except Exception as e:
raise e
def completion(
self,
model: str,
@ -470,7 +501,7 @@ class AzureChatCompletion(BaseLLM):
azure_ad_token: str,
print_verbose: Callable,
timeout: Union[float, httpx.Timeout],
logging_obj,
logging_obj: LiteLLMLoggingObj,
optional_params,
litellm_params,
logger_fn,
@ -649,9 +680,9 @@ class AzureChatCompletion(BaseLLM):
data: dict,
timeout: Any,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
azure_ad_token: Optional[str] = None,
client=None, # this is the AsyncAzureOpenAI
logging_obj=None,
):
response = None
try:
@ -701,9 +732,13 @@ class AzureChatCompletion(BaseLLM):
"complete_input_dict": data,
},
)
response = await azure_client.chat.completions.create(
**data, timeout=timeout
headers, response = await self.make_azure_openai_chat_completion_request(
azure_client=azure_client,
data=data,
timeout=timeout,
)
logging_obj.model_call_details["response_headers"] = headers
stringified_response = response.model_dump()
logging_obj.post_call(
@ -812,7 +847,7 @@ class AzureChatCompletion(BaseLLM):
async def async_streaming(
self,
logging_obj,
logging_obj: LiteLLMLoggingObj,
api_base: str,
api_key: str,
api_version: str,
@ -861,9 +896,14 @@ class AzureChatCompletion(BaseLLM):
"complete_input_dict": data,
},
)
response = await azure_client.chat.completions.create(
**data, timeout=timeout
headers, response = await self.make_azure_openai_chat_completion_request(
azure_client=azure_client,
data=data,
timeout=timeout,
)
logging_obj.model_call_details["response_headers"] = headers
# return response
streamwrapper = CustomStreamWrapper(
completion_stream=response,

120
litellm/llms/openai.py
View file

@ -21,6 +21,7 @@ from pydantic import BaseModel
from typing_extensions import overload, override
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import ProviderField
from litellm.utils import (
Choices,
@ -652,6 +653,36 @@ class OpenAIChatCompletion(BaseLLM):
else:
return client
async def make_openai_chat_completion_request(
self,
openai_aclient: AsyncOpenAI,
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 = (
await openai_aclient.chat.completions.with_raw_response.create(
**data, timeout=timeout
)
)
headers = dict(raw_response.headers)
response = raw_response.parse()
return headers, response
else:
response = await openai_aclient.chat.completions.create(
**data, timeout=timeout
)
return None, response
except Exception as e:
raise e
def completion(
self,
model_response: ModelResponse,
@ -836,13 +867,13 @@ class OpenAIChatCompletion(BaseLLM):
self,
data: dict,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
timeout: Union[float, httpx.Timeout],
api_key: Optional[str] = None,
api_base: Optional[str] = None,
organization: Optional[str] = None,
client=None,
max_retries=None,
logging_obj=None,
headers=None,
):
response = None
@ -869,8 +900,8 @@ class OpenAIChatCompletion(BaseLLM):
},
)
response = await openai_aclient.chat.completions.create(
**data, timeout=timeout
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(
@ -879,9 +910,11 @@ class OpenAIChatCompletion(BaseLLM):
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},
)
except Exception as e:
raise e
@ -931,10 +964,10 @@ class OpenAIChatCompletion(BaseLLM):
async def async_streaming(
self,
logging_obj,
timeout: Union[float, httpx.Timeout],
data: dict,
model: str,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
organization: Optional[str] = None,
@ -965,9 +998,10 @@ class OpenAIChatCompletion(BaseLLM):
},
)
response = await openai_aclient.chat.completions.create(
**data, timeout=timeout
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,
@ -992,17 +1026,43 @@ class OpenAIChatCompletion(BaseLLM):
else:
raise OpenAIError(status_code=500, message=f"{str(e)}")
# Embedding
async def make_openai_embedding_request(
self,
openai_aclient: AsyncOpenAI,
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 = await openai_aclient.embeddings.with_raw_response.create(
**data, timeout=timeout
) # type: ignore
headers = dict(raw_response.headers)
response = raw_response.parse()
return headers, response
else:
response = await openai_aclient.embeddings.create(**data, timeout=timeout) # type: ignore
return None, response
except Exception as e:
raise e
async def aembedding(
self,
input: list,
data: dict,
model_response: litellm.utils.EmbeddingResponse,
timeout: float,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
client: Optional[AsyncOpenAI] = None,
max_retries=None,
logging_obj=None,
):
response = None
try:
@ -1014,7 +1074,10 @@ class OpenAIChatCompletion(BaseLLM):
max_retries=max_retries,
client=client,
)
response = await openai_aclient.embeddings.create(**data, timeout=timeout) # type: ignore
headers, response = await self.make_openai_embedding_request(
openai_aclient=openai_aclient, data=data, timeout=timeout
)
logging_obj.model_call_details["response_headers"] = headers
stringified_response = response.model_dump()
## LOGGING
logging_obj.post_call(
@ -1229,6 +1292,34 @@ class OpenAIChatCompletion(BaseLLM):
else:
raise OpenAIError(status_code=500, message=str(e))
# Audio Transcriptions
async def make_openai_audio_transcriptions_request(
self,
openai_aclient: AsyncOpenAI,
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 = (
await openai_aclient.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 = await openai_aclient.audio.transcriptions.create(**data, timeout=timeout) # type: ignore
return None, response
except Exception as e:
raise e
def audio_transcriptions(
self,
model: str,
@ -1286,11 +1377,11 @@ class OpenAIChatCompletion(BaseLLM):
data: dict,
model_response: TranscriptionResponse,
timeout: float,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
client=None,
max_retries=None,
logging_obj=None,
):
try:
openai_aclient = self._get_openai_client(
@ -1302,9 +1393,12 @@ class OpenAIChatCompletion(BaseLLM):
client=client,
)
response = await openai_aclient.audio.transcriptions.create(
**data, timeout=timeout
) # type: ignore
headers, response = await self.make_openai_audio_transcriptions_request(
openai_aclient=openai_aclient,
data=data,
timeout=timeout,
)
logging_obj.model_call_details["response_headers"] = headers
stringified_response = response.model_dump()
## LOGGING
logging_obj.post_call(
@ -1497,9 +1591,9 @@ class OpenAITextCompletion(BaseLLM):
model: str,
messages: list,
timeout: float,
logging_obj: LiteLLMLoggingObj,
print_verbose: Optional[Callable] = None,
api_base: Optional[str] = None,
logging_obj=None,
acompletion: bool = False,
optional_params=None,
litellm_params=None,

1
litellm/proxy/proxy_config.yaml
View file

@ -36,6 +36,7 @@ general_settings:
LANGFUSE_SECRET_KEY: "os.environ/LANGFUSE_DEV_SK_KEY"
litellm_settings:
return_response_headers: true
success_callback: ["prometheus"]
callbacks: ["otel", "hide_secrets"]
failure_callback: ["prometheus"]