Coverage for agentos/swarm/task_decomposer.py: 30%

1""" 

2v1.9.4: LLM-driven Task Decomposer. 

3 

4Splits complex tasks into sub-task DAGs with dependencies, 

5assigning each sub-task to appropriate agent roles. 

6""" 

7 

8from __future__ import annotations 

9 

10import json as _json 

11from dataclasses import dataclass, field 

12from typing import Any 

13 

14_DECOMPOSE_PROMPT = """You are a task decomposition expert. Given a complex task, break it into 

15a sequence of sub-tasks that can be executed independently or sequentially. 

16 

17Input task: {task} 

18Available agents: {agents} 

19 

20Output a JSON array of sub-tasks. Each sub-task must have: 

21- "id": unique short string (e.g. "step_1") 

22- "title": human-readable title 

23- "description": what this sub-task should accomplish 

24- "depends_on": list of sub-task IDs that must complete before this one (empty list if none) 

25- "agent_hint": which agent role is best suited (from the available list, or "any") 

26- "expected_output": brief description of expected result 

27 

28Rules: 

291. First sub-tasks should have no dependencies 

302. Each sub-task should be independently executable 

313. Use at most {max_depth} levels of nesting 

324. Sub-tasks should be concrete and actionable 

33 

34Output ONLY the JSON array, no other text. 

35JSON:""" 

36 

37 

38@dataclass 

39class SubTask: 

40 """A single sub-task in the decomposition DAG.""" 

41 

42 id: str 

43 title: str 

44 description: str 

45 depends_on: list[str] = field(default_factory=list) 

46 agent_hint: str = "any" 

47 expected_output: str = "" 

48 status: str = "pending" # pending | running | done | failed 

49 output: Any = None 

50 

51 def to_dict(self) -> dict: 

52 return { 

53 "id": self.id, 

54 "title": self.title, 

55 "description": self.description, 

56 "depends_on": self.depends_on, 

57 "agent_hint": self.agent_hint, 

58 "expected_output": self.expected_output, 

59 "status": self.status, 

60 } 

61 

62 

63@dataclass 

64class Decomposition: 

65 """Result of task decomposition.""" 

66 

67 original_task: str 

68 sub_tasks: list[SubTask] = field(default_factory=list) 

69 total_steps: int = 0 

70 

71 

72class TaskDecomposer: 

73 """LLM-driven task decomposition engine. 

74 

75 Breaks complex tasks into executable sub-task DAGs. 

76 """ 

77 

78 def __init__(self, max_depth: int = 4, llm_model: str = "gpt-4o-mini"): 

79 self.max_depth = max_depth 

80 self._llm_model = llm_model 

81 

82 def decompose( 

83 self, 

84 task: str, 

85 agents: list[str] | None = None, 

86 ) -> Decomposition: 

87 """Decompose a complex task into sub-tasks. 

88 

89 Args: 

90 task: The full task description 

91 agents: List of available agent names for role assignment 

92 

93 Returns: 

94 Decomposition with ordered sub-tasks 

95 """ 

96 agents_list = agents or ["general"] 

97 agent_str = ", ".join(agents_list) 

98 

99 prompt = _DECOMPOSE_PROMPT.format( 

100 task=task, 

101 agents=agent_str, 

102 max_depth=self.max_depth, 

103 ) 

104 

105 # Try LLM-based decomposition first, fall back to rule-based 

106 result = self._llm_decompose(prompt) 

107 if result: 

108 return result 

109 

110 return self._fallback_decompose(task, agents_list) 

111 

112 def _llm_decompose(self, prompt: str) -> Decomposition | None: 

113 """Use LLM to decompose task. Returns None on failure.""" 

114 try: 

115 import os 

116 

117 api_key = os.environ.get("OPENAI_API_KEY", "") 

118 if not api_key: 

119 return None 

120 

121 import requests 

122 

123 resp = requests.post( 

124 "https://api.openai.com/v1/chat/completions", 

125 headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, 

126 json={ 

127 "model": self._llm_model, 

128 "messages": [{"role": "user", "content": prompt}], 

129 "temperature": 0.2, 

130 "max_tokens": 1000, 

131 }, 

132 timeout=30, 

133 ) 

134 if resp.status_code != 200: 

135 return None 

136 

137 text = resp.json()["choices"][0]["message"]["content"] 

138 

139 # Extract JSON from response 

140 start = text.find("[") 

141 end = text.rfind("]") + 1 

142 if start == -1 or end == 0: 

143 return None 

144 

145 data = _json.loads(text[start:end]) 

146 sub_tasks = [] 

147 for item in data: 

148 st = SubTask( 

149 id=item.get("id", f"step_{len(sub_tasks)+1}"), 

150 title=item.get("title", ""), 

151 description=item.get("description", ""), 

152 depends_on=item.get("depends_on", []), 

153 agent_hint=item.get("agent_hint", "any"), 

154 expected_output=item.get("expected_output", ""), 

155 ) 

156 sub_tasks.append(st) 

157 

158 return Decomposition( 

159 original_task=prompt, 

160 sub_tasks=sub_tasks, 

161 total_steps=len(sub_tasks), 

162 ) 

163 except Exception: 

164 return None 

165 

166 def _fallback_decompose(self, task: str, agents: list[str]) -> Decomposition: 

167 """Rule-based fallback when LLM unavailable. 

168 

169 Splits on explicit delimiters ('then', 'after', numbered steps) 

170 or uses keyword-based phase decomposition. 

171 """ 

172 import re 

173 

174 # Try to split on explicit markers 

175 markers = re.split( 

176 r"(?:Step\s*\d+[.:]\s*|\d+\)\s*|(?:then|之后|然后|接着)[,，\s]*|;\s*)", 

177 task, 

178 flags=re.IGNORECASE, 

179 ) 

180 markers = [m.strip() for m in markers if m.strip()] 

181 

182 if len(markers) > 1: 

183 sub_tasks = [] 

184 for i, desc in enumerate(markers): 

185 st = SubTask( 

186 id=f"step_{i+1}", 

187 title=desc[:50], 

188 description=desc, 

189 depends_on=[f"step_{i}"] if i > 0 else [], 

190 agent_hint=agents[0] if agents else "any", 

191 ) 

192 sub_tasks.append(st) 

193 return Decomposition( 

194 original_task=task, 

195 sub_tasks=sub_tasks, 

196 total_steps=len(sub_tasks), 

197 ) 

198 

199 # Single task — keyword-based phase decomposition 

200 phases = [] 

201 task_lower = task.lower() 

202 

203 if any(k in task_lower for k in ("search", "find", "search for", "搜索", "查找")): 

204 phases.append(("search", "Search and gather information")) 

205 if any(k in task_lower for k in ("analyze", "analysis", "分析", "处理")): 

206 phases.append(("analyze", "Analyze collected information")) 

207 if any(k in task_lower for k in ("write", "generate", "create", "写", "生成", "创建")): 

208 phases.append(("generate", "Generate final output")) 

209 if any(k in task_lower for k in ("code", "implement", "build", "代码", "实现", "开发")): 

210 phases.append(("implement", "Implement the solution")) 

211 if any(k in task_lower for k in ("test", "verify", "validate", "测试", "验证")): 

212 phases.append(("verify", "Verify and validate results")) 

213 

214 if not phases: 

215 phases = [("execute", task)] 

216 

217 sub_tasks = [] 

218 for i, (pid, desc) in enumerate(phases): 

219 st = SubTask( 

220 id=f"phase_{i+1}_{pid}", 

221 title=pid.capitalize(), 

222 description=desc, 

223 depends_on=[sub_tasks[-1].id] if sub_tasks else [], 

224 agent_hint=agents[0] if agents else "any", 

225 ) 

226 sub_tasks.append(st) 

227 

228 return Decomposition( 

229 original_task=task, 

230 sub_tasks=sub_tasks, 

231 total_steps=len(sub_tasks), 

232 )