Feature: Router aembedding

This commit is contained in:
mc-marcocheng 2023-10-25 11:19:45 +08:00
parent 046e1384c3
commit bf98d48bba
2 changed files with 40 additions and 1 deletions

View file

@ -113,6 +113,19 @@ class Router:
# call via litellm.embedding()
return litellm.embedding(**{**data, **kwargs})
async def aembedding(self,
model: str,
input: Union[str, List],
is_async: Optional[bool] = True,
**kwargs) -> Union[List[float], None]:
# pick the one that is available (lowest TPM/RPM)
deployment = self.get_available_deployment(model=model, input=input)
data = deployment["litellm_params"]
data["input"] = input
data["caching"] = self.cache_responses
return await litellm.aembedding(**{**data, **kwargs})
def set_model_list(self, model_list: list):
self.model_list = model_list

View file

@ -181,3 +181,29 @@ def test_acompletion_on_router():
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
def test_aembedding_on_router():
try:
model_list = [
{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "text-embedding-ada-002",
},
"tpm": 100000,
"rpm": 10000,
},
]
async def embedding_call():
router = Router(model_list=model_list)
response = await router.aembedding(
model="text-embedding-ada-002",
input=["good morning from litellm", "this is another item"],
)
print(response)
asyncio.run(embedding_call())
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")