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

working embedding encoder

Browse source
This commit is contained in:
Ishaan Jaff 2025-07-24 15:25:37 -07:00
parent 432f89e29a
commit 334b90aad8
2 changed files with 14 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/router_strategy/auto_router/litellm_encoder.py
View file

@ -69,6 +69,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin):
name=model_name,
score_threshold=score_threshold if score_threshold is not None else 0.3,
)
self.model_name = model_name
self.litellm_router_instance = litellm_router_instance
def __call__(self, docs: list[Any], **kwargs) -> list[list[float]]:
@ -90,7 +91,9 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin):
raise ValueError("litellm_router_instance is not set")
try:
embeds = self.litellm_router_instance.embedding(
input=docs, model=f"{self.type}/{self.name}", **kwargs
input=docs,
model=self.model_name,
**kwargs
)
return litellm_to_list(embeds)
except Exception as e:
@ -103,7 +106,9 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin):
raise ValueError("litellm_router_instance is not set")
try:
embeds = self.litellm_router_instance.embedding(
input=docs, model=f"{self.type}/{self.name}", **kwargs
input=docs,
model=self.model_name,
**kwargs
)
return litellm_to_list(embeds)
except Exception as e:
@ -116,7 +121,9 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin):
raise ValueError("litellm_router_instance is not set")
try:
embeds = await self.litellm_router_instance.aembedding(
input=docs, model=f"{self.type}/{self.name}", **kwargs
input=docs,
model=self.model_name,
**kwargs
)
return litellm_to_list(embeds)
except Exception as e:
@ -129,7 +136,9 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin):
raise ValueError("litellm_router_instance is not set")
try:
embeds = await self.litellm_router_instance.aembedding(
input=docs, model=f"{self.type}/{self.name}", **kwargs
input=docs,
model=self.model_name,
**kwargs
)
return litellm_to_list(embeds)
except Exception as e:

2
tests/local_testing/test_router_auto_router.py
View file

@ -56,7 +56,7 @@ router = Router(
"model_name": "auto_router_2",
"litellm_params": {
"model": "auto_router/auto_router_2",
"auto_router_config_path": "auto_router/router_2.json",
"auto_router_config_path": "auto_router/router.json",
"auto_router_default_model": "gpt-4o-mini",
"auto_router_embedding_model": "custom-text-embedding-model-2",
},