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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225
| """ LiteLLM 增强版复杂度路由 Hook 支持工具感知路由,根据任务类型和复杂度自动选择模型
路由策略: - 简单任务 → 本地模型 (oMLX Qwen3.5-4B) - 工具操作 → 根据工具类型选择 - 复杂推理 → 云端推理模型 (DeepSeek Reasoner) """
import re import json from typing import Any, Dict, Literal, Optional, List from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ModelResponse import litellm
class EnhancedComplexityRouter(CustomLogger): """ 增强版复杂度路由器:支持工具感知的智能路由 """
def __init__(self): super().__init__()
self.text_threshold = 2
self.tool_model_mapping = { "simple": { "tools": [ "weather", "time", "date", "calculator", "simple_search", "read_file", "list_directory", "get_time", "currency_convert" ], "model": "qwen-local" }, "medium": { "tools": [ "write_file", "edit_file", "web_search", "web_scrape", "image_generation", "document_parser", "api_call", "send_message", "create_reminder" ], "model": "astron-code-latest" }, "complex": { "tools": [ "debug", "code_review", "architecture_design", "security_audit", "multi_agent", "workflow_orchestration", "data_analysis", "report_generation", "research_synthesis" ], "model": "astron-code-latest" } }
self.complex_keywords = [ "分析", "比较", "设计", "架构", "推理", "证明", "评估", "总结", "深度", "详细", "全面", "系统", "研究", "解释原因", "为什么", "如何实现", "如何解决", "论述", "阐述", "判断", "analyze", "compare", "design", "architect", "reason", "prove", "evaluate", "synthesize", "comprehensive", "debug", "optimize" ]
self.simple_keywords = [ "你好", "hi", "hello", "嗨", "hey", "谢谢", "thanks", "查询", "翻译", "天气", "时间", "问候", "打招呼", "是什么", "what is", "define" ]
def _detect_tool_intent(self, text: str, messages: List[Dict]) -> Optional[str]: """ 检测是否有工具调用意图 返回工具类型:simple, medium, complex, 或 None """ text_lower = text.lower()
tool_indicators = { "simple": [ "查一下天气", "现在几点了", "今天多少号", "帮我算一下", "weather", "time", "date", "what time is it" ], "medium": [ "帮我写一个文件", "搜索一下", "查找网页", "发消息", "write", "search", "scrape", "send", "create file", "帮我发送", "生成图片" ], "complex": [ "帮我debug", "代码审查", "架构设计", "安全审计", "帮我分析这段代码", "多步骤完成", "自动执行", "debug", "review", "architect", "audit", "analyze code", "帮我规划", "research" ] }
for level, indicators in tool_indicators.items(): for indicator in indicators: if indicator.lower() in text_lower: return level
for msg in messages[-3:]: if isinstance(msg, dict): content = msg.get("content", "") tool_calls = msg.get("tool_calls", []) if tool_calls: for call in tool_calls: func_name = call.get("function", {}).get("name", "").lower() for level, config in self.tool_model_mapping.items(): if func_name in config["tools"]: return level
return None
def _calculate_complexity_score(self, text: str) -> int: """ 计算纯文本复杂度分数 """ score = 0 text_lower = text.lower()
tokens = len(text.split()) if tokens > 1000: score += 4 elif tokens > 500: score += 3 elif tokens > 200: score += 2 elif tokens > 50: score += 1
for keyword in self.complex_keywords: if keyword.lower() in text_lower: score += 1
for keyword in self.simple_keywords: if keyword.lower() in text_lower: score -= 2
question_count = text.count('?') + text.count('?') if question_count >= 3: score += 2 elif question_count >= 1: score += 1
greeting_patterns = [ r'^你好[吗呀啊]?', r'^hi\s*$', r'^hello\s*$', r'^hey\s*$' ] for pattern in greeting_patterns: if re.match(pattern, text.strip(), re.IGNORECASE): score = -10
return score
def _select_model_for_tool(self, tool_level: str) -> str: """ 根据工具类型选择模型 """ if tool_level in self.tool_model_mapping: return self.tool_model_mapping[tool_level]["model"] return "astron-code-latest"
async def async_pre_call_hook( self, user_api_key_dict: Any, cache: Any, data: Dict, call_type: Literal["completion", "text_completion", "embeddings", "image_generation", "moderation", "audio_transcription"] ) -> Dict: """ LiteLLM Pre-Call Hook 支持工具感知的智能路由 """ if call_type != "completion": return data
messages = data.get("messages", []) text_content = ""
for message in messages: if isinstance(message, dict): content = message.get("content", "") if isinstance(content, str): text_content += content + " " elif isinstance(message, str): text_content += message + " "
tool_level = self._detect_tool_intent(text_content, messages)
if tool_level: target_model = self._select_model_for_tool(tool_level) route_reason = f"tool-{tool_level}" print(f"[EnhancedRouter] {route_reason} → {target_model}") else: complexity_score = self._calculate_complexity_score(text_content)
if complexity_score >= self.text_threshold: target_model = "astron-code-latest" route_reason = f"complex-text (score={complexity_score})" else: target_model = "qwen-local" route_reason = f"simple-text (score={complexity_score})"
print(f"[EnhancedRouter] {route_reason} → {target_model}")
data["model"] = target_model return data
enhanced_router_instance = EnhancedComplexityRouter()
|