scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge pull request #2405 from BerriAI/litellm_load_balancing_transcription_endpoints

Browse source 
load balancing transcription endpoints
This commit is contained in:
Krish Dholakia 2024-03-08 23:08:57 -08:00 • committed by GitHub
commit 15446ee6aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 345 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.circleci/config.yml
View file

@ -45,6 +45,7 @@ jobs:
pip install "asyncio==3.4.3"
pip install "apscheduler==3.10.4"
pip install "PyGithub==1.59.1"
pip install python-multipart
- save_cache:
paths:
- ./venv

26
litellm/llms/azure.py
View file

@ -794,9 +794,8 @@ class AzureChatCompletion(BaseLLM):
api_version: Optional[str] = None,
client=None,
azure_ad_token: Optional[str] = None,
max_retries=None,
logging_obj=None,
atranscriptions: bool = False,
atranscription: bool = False,
):
data = {"model": model, "file": audio_file, **optional_params}
@ -805,9 +804,11 @@ class AzureChatCompletion(BaseLLM):
"api_version": api_version,
"azure_endpoint": api_base,
"azure_deployment": model,
"max_retries": max_retries,
"timeout": timeout,
}
max_retries = optional_params.pop("max_retries", None)
azure_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
@ -816,7 +817,10 @@ class AzureChatCompletion(BaseLLM):
elif azure_ad_token is not None:
azure_client_params["azure_ad_token"] = azure_ad_token
if atranscriptions == True:
if max_retries is not None:
azure_client_params["max_retries"] = max_retries
if atranscription == True:
return self.async_audio_transcriptions(
audio_file=audio_file,
data=data,
@ -900,15 +904,25 @@ class AzureChatCompletion(BaseLLM):
response = await async_azure_client.audio.transcriptions.create(
**data, timeout=timeout
) # type: ignore
stringified_response = response.model_dump()
## LOGGING
logging_obj.post_call(
input=audio_file.name,
api_key=api_key,
additional_args={"complete_input_dict": data},
additional_args={
"headers": {
"Authorization": f"Bearer {async_azure_client.api_key}"
},
"api_base": async_azure_client._base_url._uri_reference,
"atranscription": True,
"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="image_generation") # type: ignore
response = convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, response_type="audio_transcription") # type: ignore
return response
except Exception as e:
## LOGGING
logging_obj.post_call(

4
litellm/llms/openai.py
View file

@ -787,10 +787,10 @@ class OpenAIChatCompletion(BaseLLM):
client=None,
max_retries=None,
logging_obj=None,
atranscriptions: bool = False,
atranscription: bool = False,
):
data = {"model": model, "file": audio_file, **optional_params}
if atranscriptions == True:
if atranscription == True:
return self.async_audio_transcriptions(
audio_file=audio_file,
data=data,

8
litellm/main.py
View file

@ -3318,6 +3318,7 @@ def image_generation(
##### Transcription #######################
@client
async def atranscription(*args, **kwargs):
"""
Calls openai + azure whisper endpoints.
@ -3389,7 +3390,7 @@ def transcription(
Allows router to load balance between them
"""
atranscriptions = kwargs.get("atranscriptions", False)
atranscription = kwargs.get("atranscription", False)
litellm_call_id = kwargs.get("litellm_call_id", None)
logger_fn = kwargs.get("logger_fn", None)
proxy_server_request = kwargs.get("proxy_server_request", None)
@ -3425,12 +3426,13 @@ def transcription(
or litellm.azure_key
or get_secret("AZURE_API_KEY")
)
response = azure_chat_completions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscriptions=atranscriptions,
atranscription=atranscription,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
@ -3444,7 +3446,7 @@ def transcription(
audio_file=file,
optional_params=optional_params,
model_response=model_response,
atranscriptions=atranscriptions,
atranscription=atranscription,
timeout=timeout,
logging_obj=litellm_logging_obj,
)

168
litellm/proxy/proxy_server.py
View file

@ -120,6 +120,8 @@ from fastapi import (
Header,
Response,
Form,
UploadFile,
File,
)
from fastapi.routing import APIRouter
from fastapi.security import OAuth2PasswordBearer
@ -3071,13 +3073,13 @@ async def embeddings(
"/v1/images/generations",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse,
tags=["image generation"],
tags=["images"],
)
@router.post(
"/images/generations",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse,
tags=["image generation"],
tags=["images"],
)
async def image_generation(
request: Request,
@ -3218,6 +3220,168 @@ async def image_generation(
)
@router.post(
"/v1/audio/transcriptions",
dependencies=[Depends(user_api_key_auth)],
tags=["audio"],
)
@router.post(
"/audio/transcriptions",
dependencies=[Depends(user_api_key_auth)],
tags=["audio"],
)
async def audio_transcriptions(
request: Request,
file: UploadFile = File(...),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Same params as:
https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl
"""
global proxy_logging_obj
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
form_data = await request.form()
data: Dict = {key: value for key, value in form_data.items() if key != "file"}
# Include original request and headers in the data
data["proxy_server_request"] = { # type: ignore
"url": str(request.url),
"method": request.method,
"headers": dict(request.headers),
"body": copy.copy(data), # use copy instead of deepcopy
}
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
data["user"] = user_api_key_dict.user_id
data["model"] = (
general_settings.get("moderation_model", None) # server default
or user_model # model name passed via cli args
or data["model"] # default passed in http request
)
if user_model:
data["model"] = user_model
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
_headers = dict(request.headers)
_headers.pop(
"authorization", None
) # do not store the original `sk-..` api key in the db
data["metadata"]["headers"] = _headers
data["metadata"]["user_api_key_alias"] = getattr(
user_api_key_dict, "key_alias", None
)
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
data["metadata"]["user_api_key_team_id"] = getattr(
user_api_key_dict, "team_id", None
)
data["metadata"]["endpoint"] = str(request.url)
### TEAM-SPECIFIC PARAMS ###
if user_api_key_dict.team_id is not None:
team_config = await proxy_config.load_team_config(
team_id=user_api_key_dict.team_id
)
if len(team_config) == 0:
pass
else:
team_id = team_config.pop("team_id", None)
data["metadata"]["team_id"] = team_id
data = {
**team_config,
**data,
} # add the team-specific configs to the completion call
router_model_names = (
[m["model_name"] for m in llm_model_list]
if llm_model_list is not None
else []
)
assert (
file.filename is not None
) # make sure filename passed in (needed for type)
with open(file.filename, "wb+") as f:
f.write(await file.read())
try:
data["file"] = open(file.filename, "rb")
### CALL HOOKS ### - modify incoming data / reject request before calling the model
data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type="moderation",
)
## ROUTE TO CORRECT ENDPOINT ##
# skip router if user passed their key
if "api_key" in data:
response = await litellm.atranscription(**data)
elif (
llm_router is not None and data["model"] in router_model_names
): # model in router model list
response = await llm_router.atranscription(**data)
elif (
llm_router is not None
and data["model"] in llm_router.deployment_names
): # model in router deployments, calling a specific deployment on the router
response = await llm_router.atranscription(
**data, specific_deployment=True
)
elif (
llm_router is not None
and llm_router.model_group_alias is not None
and data["model"] in llm_router.model_group_alias
): # model set in model_group_alias
response = await llm_router.atranscription(
**data
) # ensure this goes the llm_router, router will do the correct alias mapping
elif user_model is not None: # `litellm --model <your-model-name>`
response = await litellm.atranscription(**data)
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "Invalid model name passed in"},
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
os.remove(file.filename) # Delete the saved file
### ALERTING ###
data["litellm_status"] = "success" # used for alerting
return response
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e
)
traceback.print_exc()
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e.detail)),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
else:
error_traceback = traceback.format_exc()
error_msg = f"{str(e)}\n\n{error_traceback}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", 500),
)
@router.post(
"/v1/moderations",
dependencies=[Depends(user_api_key_auth)],

102
litellm/router.py
View file

@ -9,7 +9,7 @@
import copy, httpx
from datetime import datetime
from typing import Dict, List, Optional, Union, Literal, Any
from typing import Dict, List, Optional, Union, Literal, Any, BinaryIO
import random, threading, time, traceback, uuid
import litellm, openai
from litellm.caching import RedisCache, InMemoryCache, DualCache
@ -633,6 +633,106 @@ class Router:
self.fail_calls[model_name] += 1
raise e
async def atranscription(self, file: BinaryIO, model: str, **kwargs):
"""
Example Usage:
```
from litellm import Router
client = Router(model_list = [
{
"model_name": "whisper",
"litellm_params": {
"model": "whisper-1",
},
},
])
audio_file = open("speech.mp3", "rb")
transcript = await client.atranscription(
model="whisper",
file=audio_file
)
```
"""
try:
kwargs["model"] = model
kwargs["file"] = file
kwargs["original_function"] = self._atranscription
kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
timeout = kwargs.get("request_timeout", self.timeout)
kwargs.setdefault("metadata", {}).update({"model_group": model})
response = await self.async_function_with_fallbacks(**kwargs)
return response
except Exception as e:
raise e
async def _atranscription(self, file: BinaryIO, model: str, **kwargs):
try:
verbose_router_logger.debug(
f"Inside _atranscription()- model: {model}; kwargs: {kwargs}"
)
deployment = self.get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
)
kwargs.setdefault("metadata", {}).update(
{
"deployment": deployment["litellm_params"]["model"],
"model_info": deployment.get("model_info", {}),
}
)
kwargs["model_info"] = deployment.get("model_info", {})
data = deployment["litellm_params"].copy()
model_name = data["model"]
for k, v in self.default_litellm_params.items():
if (
k not in kwargs
): # prioritize model-specific params > default router params
kwargs[k] = v
elif k == "metadata":
kwargs[k].update(v)
potential_model_client = self._get_client(
deployment=deployment, kwargs=kwargs, client_type="async"
)
# check if provided keys == client keys #
dynamic_api_key = kwargs.get("api_key", None)
if (
dynamic_api_key is not None
and potential_model_client is not None
and dynamic_api_key != potential_model_client.api_key
):
model_client = None
else:
model_client = potential_model_client
self.total_calls[model_name] += 1
response = await litellm.atranscription(
**{
**data,
"file": file,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
self.success_calls[model_name] += 1
verbose_router_logger.info(
f"litellm.atranscription(model={model_name})\033[32m 200 OK\033[0m"
)
return response
except Exception as e:
verbose_router_logger.info(
f"litellm.atranscription(model={model_name})\033[31m Exception {str(e)}\033[0m"
)
if model_name is not None:
self.fail_calls[model_name] += 1
raise e
async def amoderation(self, model: str, input: str, **kwargs):
try:
kwargs["model"] = model

14
litellm/utils.py
View file

@ -2332,7 +2332,7 @@ def client(original_function):
or call_type == CallTypes.transcription.value
):
_file_name: BinaryIO = args[1] if len(args) > 1 else kwargs["file"]
messages = _file_name.name
messages = "audio_file"
stream = True if "stream" in kwargs and kwargs["stream"] == True else False
logging_obj = Logging(
model=model,
@ -2630,6 +2630,8 @@ def client(original_function):
return result
elif "aimg_generation" in kwargs and kwargs["aimg_generation"] == True:
return result
elif "atranscription" in kwargs and kwargs["atranscription"] == True:
return result
### POST-CALL RULES ###
post_call_processing(original_response=result, model=model or None)
@ -7964,7 +7966,9 @@ def exception_type(
message=f"AzureException - {original_exception.message}",
llm_provider="azure",
model=model,
request=original_exception.request,
request=httpx.Request(
method="POST", url="https://openai.com/"
),
)
else:
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
@ -7972,7 +7976,11 @@ def exception_type(
__cause__=original_exception.__cause__,
llm_provider="azure",
model=model,
request=original_exception.request,
request=getattr(
original_exception,
"request",
httpx.Request(method="POST", url="https://openai.com/"),
),
)
if (
"BadRequestError.__init__() missing 1 required positional argument: 'param'"

41
tests/test_whisper.py
View file

@ -1,8 +1,9 @@
# What is this?
## Tests `litellm.transcription` endpoint
## Tests `litellm.transcription` endpoint. Outside litellm module b/c of audio file used in testing (it's ~700kb).
import pytest
import asyncio, time
import aiohttp
import aiohttp, traceback
from openai import AsyncOpenAI
import sys, os, dotenv
from typing import Optional
@ -13,6 +14,7 @@ pwd = os.path.dirname(os.path.realpath(__file__))
print(pwd)
file_path = os.path.join(pwd, "gettysburg.wav")
audio_file = open(file_path, "rb")
load_dotenv()
@ -21,6 +23,7 @@ sys.path.insert(
0, os.path.abspath("../")
) # Adds the parent directory to the system path
import litellm
from litellm import Router
def test_transcription():
@ -77,3 +80,37 @@ async def test_transcription_async_openai():
assert transcript.text is not None
assert isinstance(transcript.text, str)
@pytest.mark.asyncio
async def test_transcription_on_router():
litellm.set_verbose = True
print("\n Testing async transcription on router\n")
try:
model_list = [
{
"model_name": "whisper",
"litellm_params": {
"model": "whisper-1",
},
},
{
"model_name": "whisper",
"litellm_params": {
"model": "azure/azure-whisper",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com/",
"api_key": os.getenv("AZURE_EUROPE_API_KEY"),
"api_version": "2024-02-15-preview",
},
},
]
router = Router(model_list=model_list)
response = await router.atranscription(
model="whisper",
file=audio_file,
)
print(response)
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")