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()
|