From 334b90aad8d11d752b3a4be4daba7bfa9e9c492d Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 24 Jul 2025 15:25:37 -0700
Subject: [PATCH] working embedding encoder

---
 .../auto_router/litellm_encoder.py              | 17 +++++++++++++----
 tests/local_testing/test_router_auto_router.py  |  2 +-
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py
index 57093514a43..e0fd7c3625c 100644
--- a/litellm/router_strategy/auto_router/litellm_encoder.py
+++ b/litellm/router_strategy/auto_router/litellm_encoder.py
@@ -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:
diff --git a/tests/local_testing/test_router_auto_router.py b/tests/local_testing/test_router_auto_router.py
index 944e4df92ee..88ced3010ba 100644
--- a/tests/local_testing/test_router_auto_router.py
+++ b/tests/local_testing/test_router_auto_router.py
@@ -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",
                 },