From 9afdc8b4eef2a6d24dd96f5b50043c503b28f0b4 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 2 Jan 2024 13:24:43 +0530
Subject: [PATCH] (feat) add Router init Pydantic Type

---
 litellm/types/router.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 litellm/types/router.py

diff --git a/litellm/types/router.py b/litellm/types/router.py
new file mode 100644
index 000000000..5ffdc71e3
--- /dev/null
+++ b/litellm/types/router.py
@@ -0,0 +1,41 @@
+from typing import List, Optional, Union, Dict, Tuple, Literal
+
+from pydantic import BaseModel, validator
+from completion import CompletionRequest
+from embedding import EmbeddingRequest
+
+
+class ModelConfig(BaseModel):
+    model_name: str
+    litellm_params: Union[CompletionRequest, EmbeddingRequest]
+    tpm: int
+    rpm: int
+
+
+class RouterConfig(BaseModel):
+    model_list: List[ModelConfig]
+
+    redis_url: Optional[str]
+    redis_host: Optional[str]
+    redis_port: Optional[int]
+    redis_password: Optional[str]
+
+    cache_responses: Optional[bool] = False
+    cache_kwargs: Optional[Dict] = {}
+    caching_groups: Optional[List[Tuple[str, List[str]]]] = None
+    client_ttl: Optional[int] = 3600
+    num_retries: Optional[int] = 0
+    timeout: Optional[float] = None
+    default_litellm_params: Optional[Dict[str, str]] = {}
+    set_verbose: Optional[bool] = False
+    fallbacks: Optional[List[str]] = []
+    allowed_fails: Optional[int] = None
+    context_window_fallbacks: Optional[List[str]] = []
+    model_group_alias: Optional[Dict[str, List[str]]] = {}
+    retry_after: Optional[int] = 0
+    routing_strategy: Literal[
+        "simple-shuffle",
+        "least-busy",
+        "usage-based-routing",
+        "latency-based-routing",
+    ] = "simple-shuffle"