File size: 6,555 Bytes
728ba08 4f82c19 728ba08 42075bf ce233ea 728ba08 ce233ea f45fdfa 728ba08 4f82c19 728ba08 4f82c19 728ba08 4f82c19 728ba08 f45fdfa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import os
import yaml
import logging
from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field
logger = logging.getLogger("proxy.config")
class ModelEndpointConfig(BaseModel):
"""
Object-oriented definition of a specific physical LLM endpoint.
Includes rate limits (TPM/RPM), context limits (TPR), and API configurations.
"""
model_name: str = Field(..., description="Logical virtual model name (e.g. oss-chat-fast)")
model: str = Field(..., description="Physical model path (e.g. groq/llama3-8b-8192)")
api_key: Optional[str] = Field(None, description="API Key for the provider")
api_base: Optional[str] = Field(None, description="Optional custom base URL for the API (e.g. for Ollama)")
rpm: int = Field(default=1000, description="Requests Per Minute rate limit")
tpm: int = Field(default=100000, description="Tokens Per Minute rate limit")
tpr: int = Field(default=8192, description="Tokens Per Request (Context Window) limit")
cost_per_million: float = Field(default=0.1, description="Average cost in USD per million tokens")
max_tokens: Optional[int] = Field(None, description="Max response tokens allowed")
complexity_tier: Optional[str] = Field(None, description="Complexity tier for this endpoint: low, medium, or high")
class ProxyConfig:
"""
Loads, parses, and validates the LiteLLM Proxy configuration from YAML and Environment variables.
"""
def __init__(self, config_path: str = "config.yaml"):
self.config_path = config_path
self.endpoints: List[ModelEndpointConfig] = []
self.routing_strategy: str = "simple-shuffle"
self.fallback_policy: str = "retry_next_suitable"
self.num_retries: int = 3
self.timeout: int = 10
self.context_window_fallbacks: List[Dict[str, List[str]]] = []
self.general_fallbacks: List[Dict[str, List[str]]] = []
self.load_config()
def load_config(self):
"""Reads config.yaml, resolves env vars, and populates OOP configurations."""
if not os.path.exists(self.config_path):
raise FileNotFoundError(f"Configuration file not found at: {self.config_path}")
try:
with open(self.config_path, "r") as f:
raw_data = yaml.safe_load(f) or {}
# Parse general router settings
router_settings = raw_data.get("router_settings", {})
self.routing_strategy = router_settings.get("routing_strategy", "simple-shuffle")
self.fallback_policy = router_settings.get("fallback_policy", "retry_next_suitable")
self.num_retries = int(router_settings.get("num_retries", 3))
self.timeout = int(router_settings.get("timeout", 10))
# Parse fallbacks
self.context_window_fallbacks = raw_data.get("context_window_fallbacks", [])
self.general_fallbacks = raw_data.get("general_fallbacks", [])
# Parse model list
model_list = raw_data.get("model_list", [])
self.endpoints = []
for m in model_list:
model_name = m.get("model_name")
litellm_params = m.get("litellm_params", {})
# Resolve env vars in model params
raw_model = litellm_params.get("model", "")
raw_key = litellm_params.get("api_key", "")
raw_base = litellm_params.get("api_base", "")
api_key = None
if raw_key:
if raw_key.startswith("os.environ/"):
env_var_name = raw_key.replace("os.environ/", "")
api_key = os.environ.get(env_var_name)
else:
api_key = raw_key
api_base = None
if raw_base:
if raw_base.startswith("os.environ/"):
env_var_name = raw_base.replace("os.environ/", "")
api_base = os.environ.get(env_var_name)
else:
api_base = raw_base
cost = litellm_params.get("cost_per_million", 0.1)
tpr = litellm_params.get("tpr", 8192)
complexity_tier = litellm_params.get("complexity_tier")
# Intelligent tier fallback if not explicitly provided
if not complexity_tier:
if cost >= 0.50:
complexity_tier = "high"
elif cost >= 0.04:
complexity_tier = "medium"
else:
complexity_tier = "low"
endpoint = ModelEndpointConfig(
model_name=model_name,
model=raw_model,
api_key=api_key,
api_base=api_base,
rpm=litellm_params.get("rpm", 1000),
tpm=litellm_params.get("tpm", 100000),
tpr=tpr,
cost_per_million=cost,
max_tokens=litellm_params.get("max_tokens"),
complexity_tier=complexity_tier
)
self.endpoints.append(endpoint)
logger.info(f"Loaded {len(self.endpoints)} endpoints from {self.config_path}")
except Exception as e:
logger.error(f"Error loading configuration: {e}")
raise e
def get_endpoints_for_model(self, virtual_model_name: str) -> List[ModelEndpointConfig]:
"""Returns physical endpoints matching the requested virtual model name."""
return [e for e in self.endpoints if e.model_name == virtual_model_name]
def to_litellm_model_list(self) -> List[Dict[str, Any]]:
"""Converts OOP configuration list to the dictionary format expected by LiteLLM Router."""
litellm_list = []
for e in self.endpoints:
params = {
"model": e.model,
"rpm": e.rpm,
"tpm": e.tpm
}
if e.api_key:
params["api_key"] = e.api_key
if e.api_base:
params["api_base"] = e.api_base
if e.max_tokens:
params["max_tokens"] = e.max_tokens
litellm_list.append({
"model_name": e.model_name,
"litellm_params": params
})
return litellm_list |