From b6066d1eced2738c2b33ccb013e753063d191795 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 20 Jun 2024 13:49:44 -0700
Subject: [PATCH] feat - set custom routing strategy

---
 ...odel_prices_and_context_window_backup.json | 11 +++++++++-
 litellm/router.py                             | 13 +++++++++++
 litellm/types/router.py                       | 22 +++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index d1d221b4550..1441d92a213 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -865,7 +865,7 @@
     },
     "deepseek-coder": {
         "max_tokens": 4096,
-        "max_input_tokens": 16000,
+        "max_input_tokens": 32000,
         "max_output_tokens": 4096,
         "input_cost_per_token": 0.00000014,
         "output_cost_per_token": 0.00000028,
@@ -1984,6 +1984,15 @@
         "litellm_provider": "replicate",
         "mode": "chat"
     },
+    "openrouter/deepseek/deepseek-coder": {
+        "max_tokens": 4096,
+        "max_input_tokens": 32000,
+        "max_output_tokens": 4096,
+        "input_cost_per_token": 0.00000014,
+        "output_cost_per_token": 0.00000028,
+        "litellm_provider": "openrouter",
+        "mode": "chat"
+    },
     "openrouter/microsoft/wizardlm-2-8x22b:nitro": {
         "max_tokens": 65536,
         "input_cost_per_token": 0.000001,
diff --git a/litellm/router.py b/litellm/router.py
index 9200089d5b6..08efbc4147c 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -69,6 +69,7 @@ from litellm.types.router import (
     AlertingConfig,
     AllowedFailsPolicy,
     AssistantsTypedDict,
+    CustomRoutingStrategy,
     Deployment,
     DeploymentTypedDict,
     LiteLLM_Params,
@@ -4814,6 +4815,18 @@ class Router:
         except Exception as e:
             pass
 
+    def set_custom_routing_strategy(self, CustomRoutingStrategy: CustomRoutingStrategy):
+        setattr(
+            self,
+            "get_available_deployment",
+            CustomRoutingStrategy.get_available_deployment,
+        )
+        setattr(
+            self,
+            "async_get_available_deployment",
+            CustomRoutingStrategy.async_get_available_deployment,
+        )
+
     def flush_cache(self):
         litellm.cache = None
         self.cache.flush_cache()
diff --git a/litellm/types/router.py b/litellm/types/router.py
index da3c999dc88..25b1b5c9c00 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -451,3 +451,25 @@ class ModelGroupInfo(BaseModel):
 class AssistantsTypedDict(TypedDict):
     custom_llm_provider: Literal["azure", "openai"]
     litellm_params: LiteLLMParamsTypedDict
+
+
+class CustomRoutingStrategy:
+    async def async_get_available_deployment(
+        self,
+        model: str,
+        messages: Optional[List[Dict[str, str]]] = None,
+        input: Optional[Union[str, List]] = None,
+        specific_deployment: Optional[bool] = False,
+        request_kwargs: Optional[Dict] = None,
+    ):
+        pass
+
+    def get_available_deployment(
+        self,
+        model: str,
+        messages: Optional[List[Dict[str, str]]] = None,
+        input: Optional[Union[str, List]] = None,
+        specific_deployment: Optional[bool] = False,
+        request_kwargs: Optional[Dict] = None,
+    ):
+        pass