diff --git a/litellm/__init__.py b/litellm/__init__.py index 2428dad156..ef9374aa1c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -649,4 +649,4 @@ from .exceptions import ( ) from .budget_manager import BudgetManager from .proxy.proxy_cli import run_server -from .router import Router, LiteLLM_Params +from .router import Router diff --git a/litellm/router.py b/litellm/router.py index 39053e05b2..3ea2de18eb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -29,124 +29,7 @@ from litellm.utils import ModelResponse, CustomStreamWrapper import copy from litellm._logging import verbose_router_logger import logging -from pydantic import BaseModel, validator - - -class ModelInfo(BaseModel): - id: Optional[ - str - ] # Allow id to be optional on input, but it will always be present as a str in the model instance - - def __init__(self, id: Optional[Union[str, int]] = None, **params): - if id is None: - id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided - elif isinstance(id, int): - id = str(id) - super().__init__(id=id, **params) - - class Config: - extra = "allow" - - def __contains__(self, key): - # Define custom behavior for the 'in' operator - return hasattr(self, key) - - def get(self, key, default=None): - # Custom .get() method to access attributes with a default value if the attribute doesn't exist - return getattr(self, key, default) - - def __getitem__(self, key): - # Allow dictionary-style access to attributes - return getattr(self, key) - - def __setitem__(self, key, value): - # Allow dictionary-style assignment of attributes - setattr(self, key, value) - - -class LiteLLM_Params(BaseModel): - model: str - tpm: Optional[int] = None - rpm: Optional[int] = None - api_key: Optional[str] = None - api_base: Optional[str] = None - api_version: Optional[str] = None - timeout: Optional[Union[float, str]] = None # if str, pass in as os.environ/ - stream_timeout: Optional[Union[float, str]] = ( - None # timeout when making stream=True calls, if str, pass in as os.environ/ - ) - max_retries: int = 2 # follows openai default of 2 - organization: Optional[str] = None # for openai orgs - ## VERTEX AI ## - vertex_project: Optional[str] = None - vertex_location: Optional[str] = None - ## AWS BEDROCK / SAGEMAKER ## - aws_access_key_id: Optional[str] = None - aws_secret_access_key: Optional[str] = None - aws_region_name: Optional[str] = None - - def __init__(self, max_retries: Optional[Union[int, str]] = None, **params): - if max_retries is None: - max_retries = 2 - elif isinstance(max_retries, str): - max_retries = int(max_retries) # cast to int - super().__init__(max_retries=max_retries, **params) - - class Config: - extra = "allow" - - def __contains__(self, key): - # Define custom behavior for the 'in' operator - return hasattr(self, key) - - def get(self, key, default=None): - # Custom .get() method to access attributes with a default value if the attribute doesn't exist - return getattr(self, key, default) - - def __getitem__(self, key): - # Allow dictionary-style access to attributes - return getattr(self, key) - - def __setitem__(self, key, value): - # Allow dictionary-style assignment of attributes - setattr(self, key, value) - - -class Deployment(BaseModel): - model_name: str - litellm_params: LiteLLM_Params - model_info: ModelInfo - - def __init__(self, model_info: Optional[ModelInfo] = None, **params): - if model_info is None: - model_info = ModelInfo() - super().__init__(model_info=model_info, **params) - - def to_json(self, **kwargs): - try: - return self.model_dump(**kwargs) # noqa - except Exception as e: - # if using pydantic v1 - return self.dict(**kwargs) - - class Config: - extra = "allow" - - def __contains__(self, key): - # Define custom behavior for the 'in' operator - return hasattr(self, key) - - def get(self, key, default=None): - # Custom .get() method to access attributes with a default value if the attribute doesn't exist - return getattr(self, key, default) - - def __getitem__(self, key): - # Allow dictionary-style access to attributes - return getattr(self, key) - - def __setitem__(self, key, value): - # Allow dictionary-style assignment of attributes - setattr(self, key, value) +from litellm.types.router import Deployment, ModelInfo, LiteLLM_Params class Router: diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f0ab444e4..dc29bb9491 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -3,6 +3,7 @@ from typing import List, Optional, Union, Dict, Tuple, Literal from pydantic import BaseModel, validator from .completion import CompletionRequest from .embedding import EmbeddingRequest +import uuid class ModelConfig(BaseModel): @@ -39,3 +40,120 @@ class RouterConfig(BaseModel): "usage-based-routing", "latency-based-routing", ] = "simple-shuffle" + + +class ModelInfo(BaseModel): + id: Optional[ + str + ] # Allow id to be optional on input, but it will always be present as a str in the model instance + + def __init__(self, id: Optional[Union[str, int]] = None, **params): + if id is None: + id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided + elif isinstance(id, int): + id = str(id) + super().__init__(id=id, **params) + + class Config: + extra = "allow" + + def __contains__(self, key): + # Define custom behavior for the 'in' operator + return hasattr(self, key) + + def get(self, key, default=None): + # Custom .get() method to access attributes with a default value if the attribute doesn't exist + return getattr(self, key, default) + + def __getitem__(self, key): + # Allow dictionary-style access to attributes + return getattr(self, key) + + def __setitem__(self, key, value): + # Allow dictionary-style assignment of attributes + setattr(self, key, value) + + +class LiteLLM_Params(BaseModel): + model: str + tpm: Optional[int] = None + rpm: Optional[int] = None + api_key: Optional[str] = None + api_base: Optional[str] = None + api_version: Optional[str] = None + timeout: Optional[Union[float, str]] = None # if str, pass in as os.environ/ + stream_timeout: Optional[Union[float, str]] = ( + None # timeout when making stream=True calls, if str, pass in as os.environ/ + ) + max_retries: int = 2 # follows openai default of 2 + organization: Optional[str] = None # for openai orgs + ## VERTEX AI ## + vertex_project: Optional[str] = None + vertex_location: Optional[str] = None + ## AWS BEDROCK / SAGEMAKER ## + aws_access_key_id: Optional[str] = None + aws_secret_access_key: Optional[str] = None + aws_region_name: Optional[str] = None + + def __init__(self, max_retries: Optional[Union[int, str]] = None, **params): + if max_retries is None: + max_retries = 2 + elif isinstance(max_retries, str): + max_retries = int(max_retries) # cast to int + super().__init__(max_retries=max_retries, **params) + + class Config: + extra = "allow" + + def __contains__(self, key): + # Define custom behavior for the 'in' operator + return hasattr(self, key) + + def get(self, key, default=None): + # Custom .get() method to access attributes with a default value if the attribute doesn't exist + return getattr(self, key, default) + + def __getitem__(self, key): + # Allow dictionary-style access to attributes + return getattr(self, key) + + def __setitem__(self, key, value): + # Allow dictionary-style assignment of attributes + setattr(self, key, value) + + +class Deployment(BaseModel): + model_name: str + litellm_params: LiteLLM_Params + model_info: ModelInfo + + def __init__(self, model_info: Optional[ModelInfo] = None, **params): + if model_info is None: + model_info = ModelInfo() + super().__init__(model_info=model_info, **params) + + def to_json(self, **kwargs): + try: + return self.model_dump(**kwargs) # noqa + except Exception as e: + # if using pydantic v1 + return self.dict(**kwargs) + + class Config: + extra = "allow" + + def __contains__(self, key): + # Define custom behavior for the 'in' operator + return hasattr(self, key) + + def get(self, key, default=None): + # Custom .get() method to access attributes with a default value if the attribute doesn't exist + return getattr(self, key, default) + + def __getitem__(self, key): + # Allow dictionary-style access to attributes + return getattr(self, key) + + def __setitem__(self, key, value): + # Allow dictionary-style assignment of attributes + setattr(self, key, value) diff --git a/litellm/utils.py b/litellm/utils.py index ecc2957303..1f6258c537 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -13,7 +13,6 @@ import dotenv, json, traceback, threading, base64, ast import subprocess, os from os.path import abspath, join, dirname import litellm, openai -from litellm import LiteLLM_Params import itertools import random, uuid, requests from functools import wraps @@ -54,6 +53,7 @@ os.environ["TIKTOKEN_CACHE_DIR"] = ( encoding = tiktoken.get_encoding("cl100k_base") import importlib.metadata from ._logging import verbose_logger +from .types.router import LiteLLM_Params from .integrations.traceloop import TraceloopLogger from .integrations.athina import AthinaLogger from .integrations.helicone import HeliconeLogger