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

feat(proxy_server.py): support azure batch api endpoints

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-22 15:21:43 -07:00
parent 03a8624379
commit f9ab33cbc2
6 changed files with 83 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/batches/main.py
View file

@ -46,7 +46,7 @@ async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai"] = "openai",
custom_llm_provider: str = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -215,7 +215,7 @@ def create_batch(
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai"] = "openai",
custom_llm_provider: str = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -379,7 +379,7 @@ def retrieve_batch(
async def alist_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
custom_llm_provider: Literal["openai"] = "openai",
custom_llm_provider: str = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,

10
litellm/files/main.py
View file

@ -35,7 +35,7 @@ azure_files_instance = AzureOpenAIFilesAPI()
async def afile_retrieve(
file_id: str,
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: str = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -189,7 +189,7 @@ def file_retrieve(
# Delete file
async def afile_delete(
file_id: str,
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: str = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -339,7 +339,7 @@ def file_delete(
# List files
async def afile_list(
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: str = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -493,7 +493,7 @@ def file_list(
async def acreate_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune"],
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: str = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -654,7 +654,7 @@ def create_file(
async def afile_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: str = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,

9
litellm/proxy/_new_secret_config.yaml
View file

@ -2,12 +2,3 @@ model_list:
- model_name: "*"
litellm_params:
model: "*"
litellm_settings:
success_callback: ["s3"]
cache: true
s3_callback_params:
s3_bucket_name: mytestbucketlitellm # AWS Bucket Name for S3
s3_region_name: us-west-2 # AWS Region Name for S3
s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/<variable name> to pass environment variables. This is AWS Access Key ID for S3
s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3

56
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -66,6 +66,11 @@ def get_files_provider_config(
return None
@router.post(
"/{provider}/v1/files",
dependencies=[Depends(user_api_key_auth)],
tags=["files"],
)
@router.post(
"/v1/files",
dependencies=[Depends(user_api_key_auth)],
@ -80,6 +85,7 @@ async def create_file(
request: Request,
fastapi_response: Response,
purpose: str = Form(...),
provider: Optional[str] = None,
custom_llm_provider: str = Form(default="openai"),
file: UploadFile = File(...),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -110,6 +116,8 @@ async def create_file(
data: Dict = {}
try:
if provider is not None:
custom_llm_provider = provider
# Use orjson to parse JSON data, orjson speeds up requests significantly
# Read the file content
file_content = await file.read()
@ -141,7 +149,9 @@ async def create_file(
_create_file_request.update(llm_provider_config)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
response = await litellm.acreate_file(**_create_file_request)
response = await litellm.acreate_file(
**_create_file_request, custom_llm_provider=custom_llm_provider
)
### ALERTING ###
asyncio.create_task(
@ -195,6 +205,11 @@ async def create_file(
)
@router.get(
"/{provider}/v1/files/{file_id:path}",
dependencies=[Depends(user_api_key_auth)],
tags=["files"],
)
@router.get(
"/v1/files/{file_id:path}",
dependencies=[Depends(user_api_key_auth)],
@ -209,6 +224,7 @@ async def get_file(
request: Request,
fastapi_response: Response,
file_id: str,
provider: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -246,9 +262,10 @@ async def get_file(
proxy_config=proxy_config,
)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None: # default to openai
provider = "openai"
response = await litellm.afile_retrieve(
custom_llm_provider="openai", file_id=file_id, **data
custom_llm_provider=provider, file_id=file_id, **data
)
### ALERTING ###
@ -303,6 +320,11 @@ async def get_file(
)
@router.delete(
"/{provider}/v1/files/{file_id:path}",
dependencies=[Depends(user_api_key_auth)],
tags=["files"],
)
@router.delete(
"/v1/files/{file_id:path}",
dependencies=[Depends(user_api_key_auth)],
@ -317,6 +339,7 @@ async def delete_file(
request: Request,
fastapi_response: Response,
file_id: str,
provider: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -355,9 +378,10 @@ async def delete_file(
proxy_config=proxy_config,
)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None: # default to openai
provider = "openai"
response = await litellm.afile_delete(
custom_llm_provider="openai", file_id=file_id, **data
custom_llm_provider=provider, file_id=file_id, **data
)
### ALERTING ###
@ -412,6 +436,11 @@ async def delete_file(
)
@router.get(
"/{provider}/v1/files",
dependencies=[Depends(user_api_key_auth)],
tags=["files"],
)
@router.get(
"/v1/files",
dependencies=[Depends(user_api_key_auth)],
@ -426,6 +455,7 @@ async def list_files(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
provider: Optional[str] = None,
purpose: Optional[str] = None,
):
"""
@ -463,9 +493,10 @@ async def list_files(
proxy_config=proxy_config,
)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None:
provider = "openai"
response = await litellm.afile_list(
custom_llm_provider="openai", purpose=purpose, **data
custom_llm_provider=provider, purpose=purpose, **data
)
### ALERTING ###
@ -520,6 +551,11 @@ async def list_files(
)
@router.get(
"/{provider}/v1/files/{file_id:path}/content",
dependencies=[Depends(user_api_key_auth)],
tags=["files"],
)
@router.get(
"/v1/files/{file_id:path}/content",
dependencies=[Depends(user_api_key_auth)],
@ -534,6 +570,7 @@ async def get_file_content(
request: Request,
fastapi_response: Response,
file_id: str,
provider: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -571,9 +608,10 @@ async def get_file_content(
proxy_config=proxy_config,
)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None:
provider = "openai"
response = await litellm.afile_content(
custom_llm_provider="openai", file_id=file_id, **data
custom_llm_provider=provider, file_id=file_id, **data
)
### ALERTING ###

2
litellm/proxy/proxy_cli.py
View file

@ -78,7 +78,7 @@ def is_port_in_use(port):
@click.option("--api_base", default=None, help="API base URL.")
@click.option(
"--api_version",
default="2024-02-01",
default="2024-07-01-preview",
help="For azure - pass in the api version.",
)
@click.option(

33
litellm/proxy/proxy_server.py
View file

@ -4877,6 +4877,11 @@ async def run_thread(
######################################################################
@router.get(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
tags=["batch"],
)
@router.post(
"/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -4890,6 +4895,7 @@ async def run_thread(
async def create_batch(
request: Request,
fastapi_response: Response,
provider: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
@ -4936,9 +4942,10 @@ async def create_batch(
_create_batch_data = CreateBatchRequest(**data)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None:
provider = "openai"
response = await litellm.acreate_batch(
custom_llm_provider="openai", **_create_batch_data
custom_llm_provider=provider, **_create_batch_data
)
### ALERTING ###
@ -4994,6 +5001,11 @@ async def create_batch(
)
@router.get(
"/{provider}/v1/batches/{batch_id:path}",
dependencies=[Depends(user_api_key_auth)],
tags=["batch"],
)
@router.get(
"/v1/batches/{batch_id:path}",
dependencies=[Depends(user_api_key_auth)],
@ -5008,6 +5020,7 @@ async def retrieve_batch(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
provider: Optional[str] = None,
batch_id: str = Path(
title="Batch ID to retrieve", description="The ID of the batch to retrieve"
),
@ -5032,9 +5045,10 @@ async def retrieve_batch(
batch_id=batch_id,
)
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None:
provider = "openai"
response = await litellm.aretrieve_batch(
custom_llm_provider="openai", **_retrieve_batch_request
custom_llm_provider=provider, **_retrieve_batch_request
)
### ALERTING ###
@ -5091,6 +5105,11 @@ async def retrieve_batch(
)
@router.get(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
tags=["batch"],
)
@router.get(
"/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -5103,6 +5122,7 @@ async def retrieve_batch(
)
async def list_batches(
fastapi_response: Response,
provider: Optional[str] = None,
limit: Optional[int] = None,
after: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -5123,9 +5143,10 @@ async def list_batches(
global proxy_logging_obj
verbose_proxy_logger.debug("GET /v1/batches after={} limit={}".format(after, limit))
try:
# for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
if provider is None:
provider = "openai"
response = await litellm.alist_batches(
custom_llm_provider="openai",
custom_llm_provider=provider,
after=after,
limit=limit,
)