Coverage for agentos/models/routing_strategy.py: 58%

1""" 

2智能路由策略 — 按任务复杂度自动选择模型。 

3""" 

4 

5from __future__ import annotations 

6 

7from enum import StrEnum 

8 

9 

10class Complexity(StrEnum): 

11 """复杂度评估结果。""" 

12 

13 SIMPLE = "simple" 

14 MODERATE = "moderate" 

15 COMPLEX_REASONING = "complex_reasoning" 

16 LARGE_CODEBASE = "large_codebase" 

17 CRITICAL = "critical" 

18 

19 

20class Budget(StrEnum): 

21 """成本预算配置。""" 

22 

23 UNLIMITED = "unlimited" 

24 NORMAL = "normal" 

25 TIGHT = "tight" 

26 

27 

28ROUTING_TABLE: dict[tuple[Complexity, Budget], str] = { 

29 (Complexity.SIMPLE, Budget.TIGHT): "minimax-m2.7", 

30 (Complexity.SIMPLE, Budget.NORMAL): "deepseek-v3.1", 

31 (Complexity.SIMPLE, Budget.UNLIMITED): "deepseek-v3.1", 

32 (Complexity.MODERATE, Budget.TIGHT): "deepseek-v3.1", 

33 (Complexity.MODERATE, Budget.NORMAL): "kimi-k2.6", 

34 (Complexity.MODERATE, Budget.UNLIMITED): "kimi-k2.6", 

35 (Complexity.COMPLEX_REASONING, Budget.TIGHT): "deepseek-r1", 

36 (Complexity.COMPLEX_REASONING, Budget.NORMAL): "deepseek-r1", 

37 (Complexity.COMPLEX_REASONING, Budget.UNLIMITED): "claude-opus-4.8", 

38 (Complexity.LARGE_CODEBASE, Budget.TIGHT): "qwen-3.6-plus", 

39 (Complexity.LARGE_CODEBASE, Budget.NORMAL): "qwen-3.6-plus", 

40 (Complexity.LARGE_CODEBASE, Budget.UNLIMITED): "qwen-3.6-plus", 

41 (Complexity.CRITICAL, Budget.TIGHT): "kimi-k2.6", 

42 (Complexity.CRITICAL, Budget.NORMAL): "kimi-k2.6", 

43 (Complexity.CRITICAL, Budget.UNLIMITED): "claude-opus-4.8", 

44} 

45 

46# 简单启发式的复杂度与预算判定规则 

47SIMPLE_KEYWORDS = ["list", "show", "display", "read", "what is", "how to", "print"] 

48COMPLEX_KEYWORDS = [ 

49 "analyze", 

50 "debug", 

51 "optimize", 

52 "refactor", 

53 "design", 

54 "architecture", 

55 "research", 

56] 

57LARGE_KEYWORDS = ["entire codebase", "all files", "monorepo", "whole project"] 

58CRITICAL_KEYWORDS = ["production", "critical", "must not fail", "urgent"] 

59 

60TIGHT_BUDGET_KEYWORDS = ["cheapest", "save cost", "budget mode", "最便宜", "省钱"] 

61 

62 

63class RoutingStrategy: 

64 """根据任务描述自动决定使用哪个模型。""" 

65 

66 @staticmethod 

67 def assess_complexity(task: str) -> Complexity: 

68 task_lower = task.lower() 

69 if any(kw in task_lower for kw in CRITICAL_KEYWORDS): 

70 return Complexity.CRITICAL 

71 if any(kw in task_lower for kw in LARGE_KEYWORDS): 

72 return Complexity.LARGE_CODEBASE 

73 if any(kw in task_lower for kw in COMPLEX_KEYWORDS): 

74 return Complexity.COMPLEX_REASONING 

75 if any(kw in task_lower for kw in SIMPLE_KEYWORDS): 

76 return Complexity.SIMPLE 

77 return Complexity.MODERATE 

78 

79 @staticmethod 

80 def assess_budget(task: str) -> Budget: 

81 task_lower = task.lower() 

82 if any(kw in task_lower for kw in TIGHT_BUDGET_KEYWORDS): 

83 return Budget.TIGHT 

84 return Budget.NORMAL 

85 

86 @classmethod 

87 def route(cls, task: str, budget: Budget | None = None) -> str: 

88 complexity = cls.assess_complexity(task) 

89 if budget is None: 

90 budget = cls.assess_budget(task) 

91 return ROUTING_TABLE.get((complexity, budget), "kimi-k2.6")