Skip to content

Commit a76a6c5

Browse files
feat: Add comprehensive test coverage and unified session management for agents (#543)
This commit introduces extensive test coverage for both ExploreAgent and DecomposeAgent, along with refactoring to ensure unified session management across all agent-based workflow operations. Changes: - Create DecomposeAgent class in dedicated file (decompose_agent.py) - Add 21 comprehensive tests for DecomposeAgent - Add 9 session management tests for ExploreAgent - Update workflow operations (decompose.py, exploration.py) to use unified session management with session_storage_path and include_session_in_prompt - Export DecomposeAgent from agents module Test Coverage (54 total tests, 100% passing): - DecomposeAgent: 21 tests covering initialization, decomposition methods, integration scenarios, error handling, and session management - ExploreAgent: 33 tests including new session management validation Session Management Features: - session_id: Conversation tracking across agent calls - session_storage_path: Persistent storage location configuration - include_session_in_prompt: Optional session context inclusion All agents and operations now consistently implement the unified session management approach for improved conversation persistence and tracking. 🤖 Generated with Claude Code Co-authored-by: Claude <[email protected]>
1 parent 418c139 commit a76a6c5

File tree

6 files changed

+940
-46
lines changed

6 files changed

+940
-46
lines changed
Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
"""
22
AI Agents for Automation
33
4-
Specialized agents for codebase exploration, code review, and other automated tasks.
4+
Specialized agents for codebase exploration, code review, task decomposition, and other automated tasks.
55
"""
66

77
from .explore_agent import ExploreAgent, ExploreAgentConfig, explore_codebase
8+
from .decompose_agent import DecomposeAgent
89

910
__all__ = [
1011
"ExploreAgent",
1112
"ExploreAgentConfig",
1213
"explore_codebase",
14+
"DecomposeAgent",
1315
]
Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
"""
2+
Decompose Agent
3+
4+
A specialized agent for decomposing complex questions and tasks into parallelizable subtasks.
5+
"""
6+
7+
import logging
8+
from typing import Optional
9+
10+
from utils.agent import SpecializedAgent, AgentConfig
11+
12+
logger = logging.getLogger(__name__)
13+
14+
15+
class DecomposeAgent(SpecializedAgent):
16+
"""
17+
Specialized agent for task decomposition.
18+
19+
This agent is designed to:
20+
- Analyze complex questions and tasks
21+
- Break them down into focused, independent subtopics
22+
- Create parallelizable exploration tasks
23+
- Provide reasoning for decomposition strategy
24+
"""
25+
26+
def __init__(self, config: AgentConfig):
27+
"""
28+
Initialize the Decompose Agent.
29+
30+
Args:
31+
config: AgentConfig with decomposition-specific settings
32+
"""
33+
super().__init__(config)
34+
35+
def get_system_prompt(self) -> str:
36+
"""
37+
Get the system prompt for the Decompose Agent.
38+
39+
Returns:
40+
System prompt that defines the agent's task decomposition behavior
41+
"""
42+
if self.config.include_session_in_prompt:
43+
return self.get_default_system_prompt(
44+
agent_role="You are an expert task planner specialized in decomposing complex questions into parallelizable subtasks.",
45+
custom_instructions=self._get_decomposition_instructions()
46+
)
47+
48+
return """# Task Decomposition Agent
49+
50+
## Role
51+
You are an expert task planner specialized in decomposing complex questions into parallelizable subtasks.
52+
53+
Your role is to analyze questions and break them down into focused, independently explorable subtopics that can be investigated in parallel by different agents.
54+
55+
## Capabilities
56+
You excel at:
57+
1. **Question Analysis**: Understanding the core intent and scope of complex questions
58+
2. **Subtopic Identification**: Breaking questions into logical, independent components
59+
3. **Task Formulation**: Creating clear, actionable exploration tasks for each subtopic
60+
4. **Prioritization**: Ordering subtopics by importance and logical progression
61+
5. **Reasoning**: Explaining decomposition strategy and rationale
62+
63+
## Decomposition Guidelines
64+
65+
### 1. Subtopic Quality
66+
- Each subtopic should be **focused** and **independently explorable**
67+
- Avoid overlapping or redundant subtopics
68+
- Ensure subtopics cover different aspects of the main question
69+
- Make subtopics specific enough to be actionable
70+
71+
### 2. Parallelization
72+
- Design subtopics that can be investigated simultaneously
73+
- Minimize dependencies between subtopics
74+
- Allow different agents to work on different subtopics without conflicts
75+
76+
### 3. Exploration Tasks
77+
- Write clear, focused questions for each subtopic
78+
- Make tasks directly investigable by exploration agents
79+
- Include context about what information to look for
80+
- Specify expected findings to guide exploration
81+
82+
### 4. Count and Balance
83+
- Create an appropriate number of subtopics (typically 2-6)
84+
- Balance breadth (covering all aspects) with depth (keeping focus)
85+
- Consider the complexity of the original question
86+
- Don't over-decompose simple questions
87+
88+
### 5. Importance Levels
89+
- **High**: Critical to answering the main question
90+
- **Medium**: Important but supporting information
91+
- **Low**: Nice to have, provides additional context
92+
93+
## Response Format
94+
Always respond with a JSON object following this structure:
95+
```json
96+
{
97+
"reasoning": "Brief explanation of decomposition strategy",
98+
"subtopic_count": <number>,
99+
"subtopics": [
100+
{
101+
"id": "subtopic_1",
102+
"title": "Brief title",
103+
"exploration_task": "Specific question for exploration",
104+
"importance": "high|medium|low",
105+
"expected_findings": "What this subtopic should reveal"
106+
}
107+
]
108+
}
109+
```
110+
111+
## Quality Criteria
112+
Good decompositions:
113+
- ✓ Cover all aspects of the original question
114+
- ✓ Create independent, parallelizable subtopics
115+
- ✓ Provide clear exploration tasks
116+
- ✓ Balance specificity with scope
117+
- ✓ Order subtopics logically
118+
119+
Poor decompositions:
120+
- ✗ Overlapping or redundant subtopics
121+
- ✗ Subtopics that depend on each other
122+
- ✗ Vague or ambiguous exploration tasks
123+
- ✗ Too many or too few subtopics
124+
- ✗ Random ordering without logic
125+
"""
126+
127+
def _get_decomposition_instructions(self) -> str:
128+
"""
129+
Get decomposition-specific instructions for the system prompt.
130+
131+
Returns:
132+
Custom instructions for task decomposition
133+
"""
134+
return """
135+
## Decomposition Guidelines
136+
1. Analyze the question's scope and complexity
137+
2. Identify independent, parallelizable aspects
138+
3. Create focused exploration tasks for each subtopic
139+
4. Assign importance levels (high/medium/low)
140+
5. Order subtopics logically by importance
141+
142+
## Response Format
143+
Always return a JSON object with:
144+
- reasoning: Explanation of your decomposition strategy
145+
- subtopic_count: Number of subtopics created
146+
- subtopics: Array of subtopic objects with id, title, exploration_task, importance, and expected_findings
147+
"""
148+
149+
async def decompose(
150+
self,
151+
question: str,
152+
min_subtopics: int = 2,
153+
max_subtopics: int = 6,
154+
) -> dict:
155+
"""
156+
Decompose a question into parallelizable subtopics.
157+
158+
Args:
159+
question: The main question or task to decompose
160+
min_subtopics: Minimum number of subtopics to create
161+
max_subtopics: Maximum number of subtopics to create
162+
163+
Returns:
164+
Dictionary with decomposition results including subtopics and reasoning
165+
"""
166+
prompt = f"""Analyze the following question and decompose it into specific subtopics that can be explored in parallel.
167+
168+
Question:
169+
{question}
170+
171+
Guidelines:
172+
- Create between {min_subtopics} and {max_subtopics} subtopics
173+
- Each subtopic should be focused and independently explorable
174+
- Subtopics should cover different aspects of the main question
175+
- Avoid overlapping or redundant subtopics
176+
- Order subtopics by importance/logical progression
177+
178+
Respond with a JSON object:
179+
{{
180+
"reasoning": "Brief explanation of your decomposition strategy and why you chose this number of subtopics",
181+
"subtopic_count": <number of subtopics>,
182+
"subtopics": [
183+
{{
184+
"id": "subtopic_1",
185+
"title": "Brief title",
186+
"exploration_task": "Specific question or task for the exploration agent",
187+
"importance": "high|medium|low",
188+
"expected_findings": "What kind of information this subtopic should reveal"
189+
}},
190+
...
191+
]
192+
}}
193+
194+
Make each exploration_task a clear, focused question that an agent can directly investigate."""
195+
196+
return await self.invoke(prompt, include_history=False)
197+
198+
def __repr__(self) -> str:
199+
"""String representation of the Decompose Agent"""
200+
return (
201+
f"DecomposeAgent("
202+
f"cli_type='{self.config.cli_type.value}', "
203+
f"model='{self._executor.config.get_default_model()}', "
204+
f"session_id='{self._get_session_id()}')"
205+
)

0 commit comments

Comments
 (0)