phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-24 18:24:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor(azure.py): working azure completion calls with openai v1 sdk

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-11 16:44:39 -08:00
parent d0bd932b3c
commit 39c2597c33
9 changed files with 70 additions and 58 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/__init__.py
View file

@ -2,6 +2,7 @@
import threading, requests
from typing import Callable, List, Optional, Dict, Union
from litellm.caching import Cache
import httpx
input_callback: List[Union[str, Callable]] = []
success_callback: List[Union[str, Callable]] = []
@ -44,7 +45,7 @@ max_budget: float = 0.0 # set the max budget across all providers
_current_cost = 0 # private variable, used if max budget is set
error_logs: Dict = {}
add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt
client_session: Optional[requests.Session] = None
client_session: Optional[httpx.Client] = None
model_fallbacks: Optional[List] = None
model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
num_retries: Optional[int] = None

90
litellm/llms/azure.py
View file

@ -4,12 +4,14 @@ from .base import BaseLLM
from litellm.utils import ModelResponse, Choices, Message, CustomStreamWrapper, convert_to_model_response_object
from typing import Callable, Optional
from litellm import OpenAIConfig
import aiohttp
import httpx
class AzureOpenAIError(Exception):
def __init__(self, status_code, message):
def __init__(self, status_code, message, request: httpx.Request, response: httpx.Response):
self.status_code = status_code
self.message = message
self.request = request
self.response = response
super().__init__(
self.message
) # Call the base class constructor with the parameters it needs
@ -64,7 +66,7 @@ class AzureOpenAIConfig(OpenAIConfig):
top_p)
class AzureChatCompletion(BaseLLM):
_client_session: requests.Session
_client_session: httpx.Client
def __init__(self) -> None:
super().__init__()
@ -126,17 +128,7 @@ class AzureChatCompletion(BaseLLM):
else:
return self.acompletion(api_base=api_base, data=data, headers=headers, model_response=model_response)
elif "stream" in optional_params and optional_params["stream"] == True:
response = self._client_session.post(
url=api_base,
json=data,
headers=headers,
stream=optional_params["stream"]
)
if response.status_code != 200:
raise AzureOpenAIError(status_code=response.status_code, message=response.text)
## RESPONSE OBJECT
return response.iter_lines()
return self.streaming(logging_obj=logging_obj, api_base=api_base, data=data, headers=headers, model_response=model_response, model=model)
else:
response = self._client_session.post(
url=api_base,
@ -144,7 +136,7 @@ class AzureChatCompletion(BaseLLM):
headers=headers,
)
if response.status_code != 200:
raise AzureOpenAIError(status_code=response.status_code, message=response.text)
raise AzureOpenAIError(status_code=response.status_code, message=response.text, request=response.request, response=response)
## RESPONSE OBJECT
return convert_to_model_response_object(response_object=response.json(), model_response_object=model_response)
@ -152,39 +144,61 @@ class AzureChatCompletion(BaseLLM):
exception_mapping_worked = True
raise e
except Exception as e:
if exception_mapping_worked:
raise e
else:
import traceback
raise AzureOpenAIError(status_code=500, message=traceback.format_exc())
raise e
async def acompletion(self, api_base: str, data: dict, headers: dict, model_response: ModelResponse):
async with aiohttp.ClientSession() as session:
async with session.post(api_base, json=data, headers=headers, ssl=False) as response:
response_json = await response.json()
if response.status != 200:
raise AzureOpenAIError(status_code=response.status, message=response.text)
async with httpx.AsyncClient(timeout=600) as client:
response = await client.post(api_base, json=data, headers=headers)
response_json = response.json()
if response.status_code != 200:
raise AzureOpenAIError(status_code=response.status_code, message=response.text)
## RESPONSE OBJECT
return convert_to_model_response_object(response_object=response_json, model_response_object=model_response)
## RESPONSE OBJECT
return convert_to_model_response_object(response_object=response_json, model_response_object=model_response)
def streaming(self,
logging_obj,
api_base: str,
data: dict,
headers: dict,
model_response: ModelResponse,
model: str
):
with self._client_session.stream(
url=f"{api_base}",
json=data,
headers=headers,
method="POST"
) as response:
if response.status_code != 200:
raise AzureOpenAIError(status_code=response.status_code, message=response.text(), request=self._client_session.request, response=response)
completion_stream = response.iter_lines()
streamwrapper = CustomStreamWrapper(completion_stream=completion_stream, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
for transformed_chunk in streamwrapper:
yield transformed_chunk
async def async_streaming(self,
logging_obj,
api_base: str,
data: dict, headers: dict,
data: dict,
headers: dict,
model_response: ModelResponse,
model: str):
async with aiohttp.ClientSession() as session:
async with session.post(api_base, json=data, headers=headers, ssl=False) as response:
# Check if the request was successful (status code 200)
if response.status != 200:
raise AzureOpenAIError(status_code=response.status, message=await response.text())
# Handle the streamed response
stream_wrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="azure",logging_obj=logging_obj)
async for transformed_chunk in stream_wrapper:
yield transformed_chunk
client = httpx.AsyncClient()
async with client.stream(
url=f"{api_base}",
json=data,
headers=headers,
method="POST"
) as response:
if response.status_code != 200:
raise AzureOpenAIError(status_code=response.status_code, message=response.text(), request=self._client_session.request, response=response)
streamwrapper = CustomStreamWrapper(completion_stream=response.aiter_lines(), model=model, custom_llm_provider="azure",logging_obj=logging_obj)
async for transformed_chunk in streamwrapper:
yield transformed_chunk
def embedding(self,
model: str,

2
litellm/llms/openai.py
View file

@ -249,7 +249,7 @@ class OpenAIChatCompletion(BaseLLM):
response = await client.post(api_base, json=data, headers=headers)
response_json = response.json()
if response.status_code != 200:
raise OpenAIError(status_code=response.status_code, message=response.text)
raise OpenAIError(status_code=response.status_code, message=response.text, request=response.request, response=response)
## RESPONSE OBJECT

4
litellm/main.py
View file

@ -447,9 +447,7 @@ def completion(
logging_obj=logging,
acompletion=acompletion
)
if optional_params.get("stream", False) and acompletion is False:
response = CustomStreamWrapper(response, model, custom_llm_provider=custom_llm_provider, logging_obj=logging)
return response
## LOGGING
logging.post_call(
input=messages,

8
litellm/tests/test_async_fn.py
View file

@ -28,13 +28,13 @@ def test_async_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="gpt-3.5-turbo", messages=messages)
response = await acompletion(model="azure/chatgpt-v-2", messages=messages)
print(f"response: {response}")
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
asyncio.run(test_get_response())
test_async_response()
# test_async_response()
def test_get_response_streaming():
import asyncio
@ -42,7 +42,7 @@ def test_get_response_streaming():
user_message = "write a short poem in one sentence"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="gpt-3.5-turbo", messages=messages, stream=True)
response = await acompletion(model="azure/chatgpt-v-2", messages=messages, stream=True)
print(type(response))
import inspect
@ -65,7 +65,7 @@ def test_get_response_streaming():
asyncio.run(test_async_call())
# test_get_response_streaming()
test_get_response_streaming()
def test_get_response_non_openai_streaming():
import asyncio

4
litellm/tests/test_completion.py
View file

@ -421,7 +421,7 @@ def test_completion_openai():
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_completion_openai()
# test_completion_openai()
def test_completion_text_openai():
try:
@ -634,7 +634,7 @@ def test_completion_azure():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_azure()
test_completion_azure()
def test_completion_azure2():
# test if we can pass api_base, api_version and api_key in compleition()
try:

10
litellm/tests/test_exceptions.py
View file

@ -59,7 +59,7 @@ def test_context_window(model):
@pytest.mark.parametrize("model", models)
def test_context_window_with_fallbacks(model):
ctx_window_fallback_dict = {"command-nightly": "claude-2", "gpt-3.5-turbo-instruct": "gpt-3.5-turbo-16k"}
ctx_window_fallback_dict = {"command-nightly": "claude-2", "gpt-3.5-turbo-instruct": "gpt-3.5-turbo-16k", "azure/chatgpt-v-2": "gpt-3.5-turbo-16k"}
sample_text = "how does a court case get to the Supreme Court?" * 1000
messages = [{"content": sample_text, "role": "user"}]
@ -67,8 +67,8 @@ def test_context_window_with_fallbacks(model):
# for model in litellm.models_by_provider["bedrock"]:
# test_context_window(model=model)
test_context_window(model="azure/chatgpt-v-2")
# test_context_window_with_fallbacks(model="gpt-3.5-turbo")
# test_context_window(model="azure/chatgpt-v-2")
# test_context_window_with_fallbacks(model="azure/chatgpt-v-2")
# Test 2: InvalidAuth Errors
@pytest.mark.parametrize("model", models)
def invalid_auth(model): # set the model key to an invalid key, depending on the model
@ -163,7 +163,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
# for model in litellm.models_by_provider["bedrock"]:
# invalid_auth(model=model)
# invalid_auth(model="gpt-3.5-turbo")
# invalid_auth(model="azure/chatgpt-v-2")
# Test 3: Invalid Request Error
@pytest.mark.parametrize("model", models)
@ -173,7 +173,7 @@ def test_invalid_request_error(model):
with pytest.raises(BadRequestError):
completion(model=model, messages=messages, max_tokens="hello world")
# test_invalid_request_error(model="gpt-3.5-turbo")
# test_invalid_request_error(model="azure/chatgpt-v-2")
# Test 3: Rate Limit Errors
# def test_model_call(model):
# try:

2
litellm/tests/test_streaming.py
View file

@ -372,7 +372,7 @@ def test_completion_azure_stream():
print(f"completion_response: {complete_response}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_azure_stream()
test_completion_azure_stream()
def test_completion_claude_stream():
try:

5
litellm/utils.py
View file

@ -3845,7 +3845,8 @@ def exception_type(
raise AuthenticationError(
message=f"AzureException - {original_exception.message}",
llm_provider="azure",
model=model
model=model,
response=original_exception.response
)
elif original_exception.status_code == 408:
exception_mapping_worked = True
@ -4225,7 +4226,6 @@ class CustomStreamWrapper:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_azure_chunk(self, chunk):
chunk = chunk.decode("utf-8")
is_finished = False
finish_reason = ""
text = ""
@ -4299,7 +4299,6 @@ class CustomStreamWrapper:
def handle_openai_text_completion_chunk(self, chunk):
try:
# str_line = chunk.decode("utf-8") # Convert bytes to string
str_line = chunk
text = ""
is_finished = False