mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
Feature: Router aembedding
This commit is contained in:
parent
046e1384c3
commit
bf98d48bba
2 changed files with 40 additions and 1 deletions
|
@ -112,7 +112,20 @@ class Router:
|
||||||
data["caching"] = self.cache_responses
|
data["caching"] = self.cache_responses
|
||||||
# call via litellm.embedding()
|
# call via litellm.embedding()
|
||||||
return litellm.embedding(**{**data, **kwargs})
|
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):
|
def set_model_list(self, model_list: list):
|
||||||
self.model_list = model_list
|
self.model_list = model_list
|
||||||
|
|
||||||
|
|
|
@ -181,3 +181,29 @@ def test_acompletion_on_router():
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
pytest.fail(f"Error occurred: {e}")
|
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}")
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue