@@ -784,73 +784,20 @@ def test_unittest_discovery_with_pytest_class_fixture():
784784
785785import hashlib
786786import json
787- from typing import Any, Dict, List, Optional, Required, TypedDict, Union # noqa: UP035
788-
789-
790- class LiteLLMParamsTypedDict(TypedDict, total=False):
791- model: str
792- custom_llm_provider: Optional[str]
793- tpm: Optional[int]
794- rpm: Optional[int]
795- order: Optional[int]
796- weight: Optional[int]
797- max_parallel_requests: Optional[int]
798- api_key: Optional[str]
799- api_base: Optional[str]
800- api_version: Optional[str]
801- stream_timeout: Optional[Union[float, str]]
802- max_retries: Optional[int]
803- organization: Optional[Union[list, str]] # for openai orgs
804- ## DROP PARAMS ##
805- drop_params: Optional[bool]
806- ## UNIFIED PROJECT/REGION ##
807- region_name: Optional[str]
808- ## VERTEX AI ##
809- vertex_project: Optional[str]
810- vertex_location: Optional[str]
811- ## AWS BEDROCK / SAGEMAKER ##
812- aws_access_key_id: Optional[str]
813- aws_secret_access_key: Optional[str]
814- aws_region_name: Optional[str]
815- ## IBM WATSONX ##
816- watsonx_region_name: Optional[str]
817- ## CUSTOM PRICING ##
818- input_cost_per_token: Optional[float]
819- output_cost_per_token: Optional[float]
820- input_cost_per_second: Optional[float]
821- output_cost_per_second: Optional[float]
822- num_retries: Optional[int]
823- ## MOCK RESPONSES ##
824-
825- # routing params
826- # use this for tag-based routing
827- tags: Optional[list[str]]
828-
829- # deployment budgets
830- max_budget: Optional[float]
831- budget_duration: Optional[str]
832-
833- class DeploymentTypedDict(TypedDict, total=False):
834- model_name: Required[str]
835- litellm_params: Required[LiteLLMParamsTypedDict]
836- model_info: dict
837787
838788class Router:
839- model_names: set = set() # noqa: RUF012
840- cache_responses: Optional[bool] = False
841- default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour
789+ model_names: list
790+ cache_responses = False
842791 tenacity = None
843792
844793 def __init__( # noqa: PLR0915
845794 self,
846- model_list: Optional[
847- Union[list[DeploymentTypedDict], list[dict[str, Any]]]
848- ] = None,
795+ model_list = None,
849796 ) -> None:
850- self.model_list = model_list # noqa: ARG002
851- self.model_id_to_deployment_index_map: dict[str, int] = {}
852- self.model_name_to_deployment_indices: dict[str, list[int]] = {}
853- def _generate_model_id(self, model_group: str , litellm_params: dict): # noqa: ANN202
797+ self.model_list = model_list
798+ self.model_id_to_deployment_index_map = {}
799+ self.model_name_to_deployment_indices = {}
800+ def _generate_model_id(self, model_group, litellm_params):
854801 # Optimized: Use list and join instead of string concatenation in loop
855802 # This avoids creating many temporary string objects (O(n) vs O(n²) complexity)
856803 parts = [model_group]
@@ -874,7 +821,7 @@ def _generate_model_id(self, model_group: str, litellm_params: dict): # noqa: A
874821
875822 return hash_object.hexdigest()
876823 def _add_model_to_list_and_index_map(
877- self, model: dict , model_id: Optional[str] = None
824+ self, model, model_id = None
878825 ) -> None:
879826 idx = len(self.model_list)
880827 self.model_list.append(model)
@@ -892,7 +839,7 @@ def _add_model_to_list_and_index_map(
892839 self.model_name_to_deployment_indices[model_name] = []
893840 self.model_name_to_deployment_indices[model_name].append(idx)
894841
895- def _build_model_id_to_deployment_index_map(self, model_list: list) -> None :
842+ def _build_model_id_to_deployment_index_map(self, model_list) :
896843 # First populate the model_list
897844 self.model_list = []
898845 for _, model in enumerate(model_list):
@@ -911,6 +858,7 @@ def _build_model_id_to_deployment_index_map(self, model_list: list) -> None:
911858 model["model_info"]["id"] = model_id
912859
913860 self._add_model_to_list_and_index_map(model=model, model_id=model_id)
861+
914862"""
915863 code_file_path .write_text (code_file_content )
916864
@@ -924,9 +872,9 @@ def _build_model_id_to_deployment_index_map(self, model_list: list) -> None:
924872
925873class TestRouterIndexManagement:
926874 @pytest.fixture
927- def router(self): # noqa: ANN201
875+ def router(self):
928876 return Router(model_list=[])
929- def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001
877+ def test_build_model_id_to_deployment_index_map(self, router):
930878 model_list = [
931879 {
932880 "model_name": "gpt-3.5-turbo",
@@ -941,7 +889,7 @@ def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa:
941889 ]
942890
943891 # Test: Build index from model list
944- router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001
892+ router._build_model_id_to_deployment_index_map(model_list)
945893
946894 # Verify: model_list is populated
947895 assert len(router.model_list) == 2
@@ -1630,9 +1578,9 @@ def test_analyze_imports_class_fixture():
16301578
16311579class TestRouterIndexManagement:
16321580 @pytest.fixture
1633- def router(self): # noqa: ANN201
1581+ def router(self):
16341582 return Router(model_list=[])
1635- def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001
1583+ def test_build_model_id_to_deployment_index_map(self, router):
16361584 model_list = [
16371585 {
16381586 "model_name": "gpt-3.5-turbo",
@@ -1647,7 +1595,7 @@ def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa:
16471595 ]
16481596
16491597 # Test: Build index from model list
1650- router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001
1598+ router._build_model_id_to_deployment_index_map(model_list)
16511599
16521600 # Verify: model_list is populated
16531601 assert len(router.model_list) == 2
0 commit comments