Coverage for src / osiris_cli / integrated_commands.py: 0%

1""" 

2Integrated Commands for Osiris CLI 

3 

4New commands that integrate all the advanced features: 

5- Cost tracking and statistics 

6- Plan mode for complex tasks 

7- Checkpoint system 

8- Database operations  

9- Cloud operations 

10""" 

11 

12from rich.console import Console 

13from rich.table import Table 

14from rich.panel import Panel 

15from rich.prompt import Prompt, Confirm 

16 

17console = Console() 

18 

19 

20async def show_cost_stats(): 

21 """Show cost tracking statistics""" 

22 from .cost_tracker import cost_tracker 

23 from .letta_formatter import Icons 

24 

25 stats = cost_tracker.get_stats() 

26 

27 # Create cost table 

28 table = Table(title=f"{Icons.COST} Cost Tracking Statistics", show_header=True, header_style="bold cyan") 

29 table.add_column("Metric", style="cyan") 

30 table.add_column("Value", style="yellow") 

31 

32 table.add_row("Total Calls", str(stats["total_calls"])) 

33 table.add_row("Total Tokens", f"{stats['total_tokens']:,}") 

34 table.add_row("Total Cost", f"${stats['total_cost']:.4f}") 

35 table.add_row("Average Cost/Call", f"${stats['avg_cost_per_call']:.4f}") 

36 

37 if stats["total_calls"] > 0: 

38 table.add_row("Input Tokens", f"{stats['total_input_tokens']:,}") 

39 table.add_row("Output Tokens", f"{stats['total_output_tokens']:,}") 

40 

41 console.print(table) 

42 

43 # Provider breakdown 

44 if stats["by_provider"]: 

45 console.print("\n[bold cyan]By Provider:[/bold cyan]") 

46 for provider, prov_stats in stats["by_provider"].items(): 

47 console.print(f" {Icons.API} {provider}: {prov_stats['calls']} calls, ${prov_stats['cost']:.4f}") 

48 

49 

50async def plan_mode(): 

51 """Interactive plan mode""" 

52 from .plan_mode import PlanMode 

53 from .letta_formatter import Icons 

54 

55 console.print(Panel( 

56 f"{Icons.TOOL} [bold cyan]Plan Mode[/bold cyan]\n" 

57 "Break down complex tasks into systematic steps", 

58 border_style="cyan" 

59 )) 

60 

61 plan_mgr = PlanMode() 

62 

63 # Show existing plans 

64 plans = plan_mgr.list_plans() 

65 if plans: 

66 console.print(f"\n[cyan]Existing Plans:[/cyan] {len(plans)}") 

67 for i, plan in enumerate(plans[:5], 1): 

68 status = f"{len([s for s in plan['steps'] if s['status'] == 'completed'])}/{len(plan['steps'])} steps" 

69 console.print(f" {i}. {plan['title']} ({status})") 

70 

71 # Options 

72 console.print("\n[cyan]Options:[/cyan]") 

73 console.print(" 1. Create new plan") 

74 console.print(" 2. Continue existing plan") 

75 console.print(" 3. View plan details") 

76 console.print(" 4. Back") 

77 

78 choice = Prompt.ask("Choose an option", choices=["1", "2", "3", "4"], default="4") 

79 

80 if choice == "1": 

81 title = Prompt.ask("Plan title") 

82 description = Prompt.ask("Description") 

83 steps_input = Prompt.ask("Steps (comma-separated)") 

84 steps = [s.strip() for s in steps_input.split(",")] 

85 

86 plan = plan_mgr.create_plan(title, description, steps) 

87 console.print(f"\n{Icons.SUCCESS} Created plan with {len(steps)} steps") 

88 

89 elif choice == "2" and plans: 

90 # Continue existing plan 

91 next_step = plan_mgr.get_next_step() 

92 if next_step: 

93 console.print(f"\n{Icons.ARROW} Next step: {next_step.description}") 

94 

95 

96async def checkpoint_menu(): 

97 """Interactive checkpoint menu""" 

98 from .checkpoint import CheckpointManager 

99 from .session import session 

100 from .letta_formatter import Icons 

101 

102 console.print(Panel( 

103 f"{Icons.FILE} [bold cyan]Checkpoint System[/bold cyan]\n" 

104 "Save and restore conversation states", 

105 border_style="cyan" 

106 )) 

107 

108 mgr = CheckpointManager(session.id) 

109 

110 # Show existing checkpoints 

111 checkpoints = mgr.list_checkpoints() 

112 if checkpoints: 

113 console.print(f"\n[cyan]Checkpoints:[/cyan] {len(checkpoints)}") 

114 for i, cp in enumerate(checkpoints[:5], 1): 

115 console.print(f" {i}. {cp['description']} ({cp['timestamp']})") 

116 

117 # Options 

118 console.print("\n[cyan]Options:[/cyan]") 

119 console.print(" 1. Create checkpoint") 

120 console.print(" 2. Restore checkpoint") 

121 console.print(" 3. View checkpoint details") 

122 console.print(" 4. Back") 

123 

124 choice = Prompt.ask("Choose an option", choices=["1", "2", "3", "4"], default="4") 

125 

126 if choice == "1": 

127 description = Prompt.ask("Checkpoint description", default="Manual checkpoint") 

128 cp = mgr.create_checkpoint(description=description, messages=session.messages) 

129 console.print(f"\n{Icons.SUCCESS} Created checkpoint: {cp.id}") 

130 

131 

132async def database_menu(): 

133 """Interactive database operations menu""" 

134 from .database_manager import db_manager 

135 from .letta_formatter import Icons, thinking, action_result 

136 

137 console.print(Panel( 

138 f"{Icons.DATABASE} [bold cyan]Database Operations[/bold cyan]\n" 

139 "Connect to and query databases", 

140 border_style="cyan" 

141 )) 

142 

143 # Show active connections 

144 connections = db_manager.list_connections() 

145 if connections: 

146 console.print(f"\n[cyan]Active Connections:[/cyan]") 

147 for conn in connections: 

148 console.print(f" {Icons.DATABASE} {conn['id']} ({conn['type']})") 

149 

150 # Options 

151 console.print("\n[cyan]Options:[/cyan]") 

152 console.print(" 1. Connect to database") 

153 console.print(" 2. Query database") 

154 console.print(" 3. View schema") 

155 console.print(" 4. Disconnect") 

156 console.print(" 5. Back") 

157 

158 choice = Prompt.ask("Choose an option", choices=["1", "2", "3", "4", "5"], default="5") 

159 

160 if choice == "1": 

161 # Connect 

162 console.print("\n[cyan]Database Types:[/cyan] postgresql, mysql, sqlite, mongodb") 

163 db_type = Prompt.ask("Database type", choices=["postgresql", "mysql", "sqlite", "mongodb"]) 

164 connection_id = Prompt.ask("Connection ID", default="mydb") 

165 

166 if db_type == "sqlite": 

167 file_path = Prompt.ask("SQLite file path", default="test.db") 

168 

169 thinking("Connecting to SQLite database...") 

170 

171 success = db_manager.connect(connection_id, db_type, file_path=file_path) 

172 

173 if success: 

174 action_result("db_connect", {"connection_id": connection_id, "db_type": db_type}, 

175 f"Successfully connected to {db_type} database", success=True) 

176 else: 

177 action_result("db_connect", {"connection_id": connection_id, "db_type": db_type}, 

178 "Connection failed", success=False) 

179 else: 

180 host = Prompt.ask("Host", default="localhost") 

181 database = Prompt.ask("Database name") 

182 username = Prompt.ask("Username") 

183 password = Prompt.ask("Password", password=True) 

184 

185 thinking(f"Connecting to {db_type} database at {host}...") 

186 

187 success = db_manager.connect( 

188 connection_id, db_type, 

189 host=host, database=database, 

190 username=username, password=password 

191 ) 

192 

193 if success: 

194 action_result("db_connect", {"connection_id": connection_id, "db_type": db_type}, 

195 f"Successfully connected to {db_type} database ({database})", success=True) 

196 else: 

197 action_result("db_connect", {"connection_id": connection_id, "db_type": db_type}, 

198 "Connection failed", success=False) 

199 

200 elif choice == "2" and connections: 

201 # Query 

202 conn_id = Prompt.ask("Connection ID", default=connections[0]['id']) 

203 query = Prompt.ask("SQL Query") 

204 safe_mode = Confirm.ask("Safe mode (SELECT only)?", default=True) 

205 

206 thinking(f"Executing query on {conn_id}...") 

207 

208 try: 

209 result = db_manager.execute_query(conn_id, query, safe_mode=safe_mode) 

210 

211 if "rows" in result: 

212 from .letta_formatter import summarize_query 

213 summary = summarize_query(result) 

214 action_result("db_query", {"connection_id": conn_id, "query": query[:50]}, 

215 summary, success=True) 

216 

217 # Show sample data 

218 if result["rows"]: 

219 console.print("\n[cyan]Sample Results:[/cyan]") 

220 for row in result["rows"][:5]: 

221 console.print(f" {row}") 

222 else: 

223 action_result("db_query", {"connection_id": conn_id, "query": query[:50]}, 

224 result.get("message", "Query executed"), success=True) 

225 except Exception as e: 

226 action_result("db_query", {"connection_id": conn_id, "query": query[:50]}, 

227 f"Query failed: {str(e)}", success=False) 

228 

229 elif choice == "3" and connections: 

230 # Schema 

231 conn_id = Prompt.ask("Connection ID", default=connections[0]['id']) 

232 

233 thinking(f"Fetching schema for {conn_id}...") 

234 

235 try: 

236 schema = db_manager.get_schema(conn_id) 

237 from .letta_formatter import summarize_schema 

238 summary = summarize_schema(schema) 

239 action_result("db_schema", {"connection_id": conn_id}, summary, success=True) 

240 

241 # Show tables 

242 if "tables" in schema: 

243 console.print("\n[cyan]Tables:[/cyan]") 

244 for table_name, columns in list(schema["tables"].items())[:10]: 

245 console.print(f"\n {Icons.DATABASE} {table_name}") 

246 for col in columns[:5]: 

247 console.print(f" - {col['column']} ({col['type']})") 

248 except Exception as e: 

249 action_result("db_schema", {"connection_id": conn_id}, 

250 f"Schema fetch failed: {str(e)}", success=False) 

251 

252 

253async def cloud_menu(): 

254 """Interactive cloud operations menu""" 

255 from .cloud_manager import cloud_manager, CloudProvider 

256 from .letta_formatter import Icons, thinking, action_result 

257 

258 console.print(Panel( 

259 f"{Icons.CLOUD} [bold cyan]Cloud Operations[/bold cyan]\n" 

260 "Manage AWS, GCP, and Azure resources", 

261 border_style="cyan" 

262 )) 

263 

264 # Show provider status 

265 status = cloud_manager.get_provider_status() 

266 console.print("\n[cyan]Provider Status:[/cyan]") 

267 for provider, prov_status in status.items(): 

268 cli = f"{Icons.SUCCESS} CLI" if prov_status["cli_available"] else f"{Icons.ERROR} CLI" 

269 auth = f"{Icons.SUCCESS} Auth" if prov_status["authenticated"] else f"{Icons.ERROR} Auth" 

270 console.print(f" {Icons.CLOUD} {provider}: {cli}, {auth}") 

271 

272 # Options 

273 console.print("\n[cyan]Options:[/cyan]") 

274 console.print(" 1. List AWS EC2 instances") 

275 console.print(" 2. List AWS S3 buckets") 

276 console.print(" 3. List AWS Lambda functions") 

277 console.print(" 4. List GCP Compute instances") 

278 console.print(" 5. List Azure VMs") 

279 console.print(" 6. Back") 

280 

281 choice = Prompt.ask("Choose an option", choices=["1", "2", "3", "4", "5", "6"], default="6") 

282 

283 if choice == "1": 

284 region = Prompt.ask("AWS Region", default="us-east-1") 

285 thinking(f"Listing EC2 instances in {region}...") 

286 

287 try: 

288 result = cloud_manager.aws_list_ec2(region) 

289 summary = f"Found {result['count']} instances" 

290 if result['instances']: 

291 states = {} 

292 for inst in result['instances']: 

293 state = inst.get('state', 'unknown') 

294 states[state] = states.get(state, 0) + 1 

295 summary += f": {', '.join(f'{count} {state}' for state, count in states.items())}" 

296 

297 action_result("aws_list_ec2", {"region": region}, summary, success=True) 

298 

299 # Show instances 

300 if result['instances']: 

301 console.print("\n[cyan]Instances:[/cyan]") 

302 for inst in result['instances'][:10]: 

303 console.print(f" {Icons.CLOUD} {inst['id']} - {inst['type']} ({inst['state']})") 

304 except Exception as e: 

305 action_result("aws_list_ec2", {"region": region}, 

306 f"Failed: {str(e)}", success=False) 

307 

308 elif choice == "2": 

309 thinking("Listing S3 buckets...") 

310 

311 try: 

312 result = cloud_manager.aws_list_s3() 

313 action_result("aws_list_s3", {}, f"Found {result['count']} S3 buckets", success=True) 

314 

315 if result['buckets']: 

316 console.print("\n[cyan]Buckets:[/cyan]") 

317 for bucket in result['buckets'][:10]: 

318 console.print(f" {Icons.FILE} {bucket['name']}") 

319 except Exception as e: 

320 action_result("aws_list_s3", {}, f"Failed: {str(e)}", success=False)