Merge pull request #3751 from BerriAI/litellm_router_abatch_completion_one_model

[FEAT] - add litellm.Router - `abatch_completion_one_model_multiple_requests`
This commit is contained in:
Ishaan Jaff 2024-05-20 18:42:54 -07:00 committed by GitHub
commit e8a3916d46
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -662,6 +662,10 @@ class Router:
async def abatch_completion( async def abatch_completion(
self, models: List[str], messages: List[Dict[str, str]], **kwargs self, models: List[str], messages: List[Dict[str, str]], **kwargs
): ):
"""
Async Batch Completion - Batch Process 1 request to multiple model_group on litellm.Router
Use this for sending the same request to N models
"""
async def _async_completion_no_exceptions( async def _async_completion_no_exceptions(
model: str, messages: List[Dict[str, str]], **kwargs model: str, messages: List[Dict[str, str]], **kwargs
@ -686,6 +690,51 @@ class Router:
response = await asyncio.gather(*_tasks) response = await asyncio.gather(*_tasks)
return response return response
async def abatch_completion_one_model_multiple_requests(
self, model: str, messages: List[List[Dict[str, str]]], **kwargs
):
"""
Async Batch Completion - Batch Process multiple Messages to one model_group on litellm.Router
Use this for sending multiple requests to 1 model
Args:
model (List[str]): model group
messages (List[List[Dict[str, str]]]): list of messages. Each element in the list is one request
**kwargs: additional kwargs
Usage:
response = await self.abatch_completion_one_model_multiple_requests(
model="gpt-3.5-turbo",
messages=[
[{"role": "user", "content": "hello"}, {"role": "user", "content": "tell me something funny"}],
[{"role": "user", "content": "hello good mornign"}],
]
)
"""
async def _async_completion_no_exceptions(
model: str, messages: List[Dict[str, str]], **kwargs
):
"""
Wrapper around self.async_completion that catches exceptions and returns them as a result
"""
try:
return await self.acompletion(model=model, messages=messages, **kwargs)
except Exception as e:
return e
_tasks = []
for message_request in messages:
# add each task but if the task fails
_tasks.append(
_async_completion_no_exceptions(
model=model, messages=message_request, **kwargs
)
)
response = await asyncio.gather(*_tasks)
return response
def image_generation(self, prompt: str, model: str, **kwargs): def image_generation(self, prompt: str, model: str, **kwargs):
try: try:
kwargs["model"] = model kwargs["model"] = model