Coverage for src / mcp_server_langgraph / cli / create_agent.py: 100%

1""" 

2Agent creation command for MCP Server CLI. 

3 

4Generates agent files from templates. 

5""" 

6 

7from pathlib import Path 

8from typing import Literal 

9 

10AgentTemplate = Literal["basic", "research", "customer-support", "code-review", "data-analyst"] 

11 

12 

13AGENT_TEMPLATES = { 

14 "basic": """\"\"\" 

15Basic Agent 

16 

17A simple agent with customizable tools. 

18\"\"\" 

19 

20from langgraph.graph import StateGraph 

21from typing import TypedDict, List 

22 

23 

24class {class_name}State(TypedDict): 

25 \"\"\"State for {agent_name} agent.\"\"\" 

26 query: str 

27 response: str 

28 context: List[str] 

29 

30 

31def process_query(state: {class_name}State) -> {class_name}State: 

32 \"\"\"Process the user query.\"\"\" 

33 # TODO: Implement your agent logic here 

34 state["response"] = f"Processed: {{state['query']}}" 

35 return state 

36 

37 

38# Create agent graph 

39graph = StateGraph({class_name}State) 

40graph.add_node("process", process_query) 

41graph.set_entry_point("process") 

42graph.set_finish_point("process") 

43 

44{agent_name}_agent = graph.compile() 

45""", 

46 "research": """\"\"\" 

47Research Agent 

48 

49Searches and summarizes information from multiple sources. 

50\"\"\" 

51 

52from langgraph.graph import StateGraph 

53from typing import TypedDict, List, Annotated 

54 

55 

56class ResearchState(TypedDict): 

57 \"\"\"State for research agent.\"\"\" 

58 query: str 

59 search_results: List[str] 

60 summary: str 

61 sources: List[str] 

62 

63 

64def search(state: ResearchState) -> ResearchState: 

65 \"\"\"Search for information.\"\"\" 

66 # TODO: Implement search using Tavily, Google, or other search tools 

67 state["search_results"] = ["Result 1", "Result 2", "Result 3"] 

68 state["sources"] = ["https://example.com/1", "https://example.com/2"] 

69 return state 

70 

71 

72def summarize(state: ResearchState) -> ResearchState: 

73 \"\"\"Summarize search results.\"\"\" 

74 # TODO: Use LLM to summarize results 

75 results = " ".join(state["search_results"]) 

76 state["summary"] = f"Summary of: {{results}}" 

77 return state 

78 

79 

80# Create research agent graph 

81graph = StateGraph(ResearchState) 

82graph.add_node("search", search) 

83graph.add_node("summarize", summarize) 

84graph.add_edge("search", "summarize") 

85graph.set_entry_point("search") 

86graph.set_finish_point("summarize") 

87 

88research_agent = graph.compile() 

89""", 

90 "customer-support": """\"\"\" 

91Customer Support Agent 

92 

93Handles customer inquiries with FAQ lookup and escalation. 

94\"\"\" 

95 

96from langgraph.graph import StateGraph 

97from typing import TypedDict, Literal 

98 

99 

100class SupportState(TypedDict): 

101 \"\"\"State for customer support agent.\"\"\" 

102 query: str 

103 intent: Literal["faq", "technical", "billing", "escalate"] 

104 response: str 

105 escalated: bool 

106 

107 

108def classify_intent(state: SupportState) -> SupportState: 

109 \"\"\"Classify the customer query.\"\"\" 

110 # TODO: Use LLM to classify intent 

111 query_lower = state["query"].lower() 

112 if "price" in query_lower or "cost" in query_lower: 

113 state["intent"] = "billing" 

114 elif "error" in query_lower or "bug" in query_lower: 

115 state["intent"] = "technical" 

116 else: 

117 state["intent"] = "faq" 

118 return state 

119 

120 

121def handle_faq(state: SupportState) -> SupportState: 

122 \"\"\"Handle FAQ queries.\"\"\" 

123 # TODO: Lookup in knowledge base 

124 state["response"] = "Here's the answer from our FAQ..." 

125 return state 

126 

127 

128def handle_technical(state: SupportState) -> SupportState: 

129 \"\"\"Handle technical support queries.\"\"\" 

130 # TODO: Check technical documentation or escalate 

131 state["response"] = "Let me help you troubleshoot..." 

132 return state 

133 

134 

135def escalate(state: SupportState) -> SupportState: 

136 \"\"\"Escalate to human agent.\"\"\" 

137 state["escalated"] = True 

138 state["response"] = "I'm connecting you with a specialist..." 

139 return state 

140 

141 

142def route(state: SupportState) -> str: 

143 \"\"\"Route to appropriate handler.\"\"\" 

144 intent_map = {{ 

145 "faq": "handle_faq", 

146 "technical": "handle_technical", 

147 "billing": "escalate", 

148 }} 

149 return intent_map.get(state["intent"], "handle_faq") 

150 

151 

152# Create support agent graph 

153graph = StateGraph(SupportState) 

154graph.add_node("classify", classify_intent) 

155graph.add_node("handle_faq", handle_faq) 

156graph.add_node("handle_technical", handle_technical) 

157graph.add_node("escalate", escalate) 

158 

159graph.set_entry_point("classify") 

160graph.add_conditional_edges("classify", route) 

161 

162support_agent = graph.compile() 

163""", 

164} 

165 

166 

167def generate_agent(name: str, template: AgentTemplate = "basic", tools: list[str] | None = None) -> None: 

168 """ 

169 Generate an agent file from a template. 

170 

171 Args: 

172 name: Name of the agent (e.g., "my_agent") 

173 template: Agent template to use 

174 tools: List of tools to include 

175 

176 Raises: 

177 ValueError: If template is invalid 

178 FileExistsError: If agent file already exists 

179 """ 

180 if template not in AGENT_TEMPLATES: 

181 msg = f"Invalid template: {template}" 

182 raise ValueError(msg) 

183 

184 # Create agents directory if it doesn't exist 

185 agents_dir = Path("src/agents") 

186 agents_dir.mkdir(parents=True, exist_ok=True) 

187 

188 # Generate file name 

189 agent_file = agents_dir / f"{name}_agent.py" 

190 

191 if agent_file.exists(): 

192 msg = f"Agent file already exists: {agent_file}" 

193 raise FileExistsError(msg) 

194 

195 # Get template content 

196 template_content = AGENT_TEMPLATES[template] 

197 

198 # Replace placeholders 

199 class_name = "".join(word.capitalize() for word in name.split("_")) 

200 content = template_content.format( 

201 agent_name=name, 

202 class_name=class_name, 

203 ) 

204 

205 # Write file 

206 agent_file.write_text(content) 

207 

208 # Create test file 

209 tests_dir = Path("tests/agents") 

210 tests_dir.mkdir(parents=True, exist_ok=True) 

211 

212 test_file = tests_dir / f"test_{name}_agent.py" 

213 test_content = f'''""" 

214Tests for {name} agent. 

215""" 

216 

217import pytest 

218from src.agents.{name}_agent import {name}_agent 

219 

220 

221def test_{name}_agent_basic(): 

222 """Test basic agent functionality.""" 

223 # TODO: Implement tests 

224 assert {name}_agent is not None 

225 

226 

227@pytest.mark.asyncio 

228async def test_{name}_agent_async(): 

229 """Test async agent invocation.""" 

230 # TODO: Implement async tests 

231 pass 

232''' 

233 

234 test_file.write_text(test_content)