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

1""" 

2Database Tools with Letta Visual Standard 

3 

4AI-friendly database tools using the Letta formatting style: 

5- Action-result trees with ● and ⎿ 

6- Summarized outputs (no walls of text) 

7- Thinking blocks with ✻ 

8""" 

9 

10from typing import Dict, Any, Optional 

11from pathlib import Path 

12 

13from .database_manager import db_manager, DatabaseType 

14from .letta_formatter import ( 

15 letta_formatter, 

16 action_result, 

17 thinking, 

18 summarize_query, 

19 summarize_schema, 

20 LettaGlyphs 

21) 

22from .logger import get_logger 

23from .errors import DatabaseError 

24 

25logger = get_logger() 

26 

27 

28# Tool definitions for AI (OpenAI function calling format) 

29DATABASE_TOOLS_LETTA = [ 

30 { 

31 "type": "function", 

32 "function": { 

33 "name": "db_connect", 

34 "description": "Connect to a database (PostgreSQL, MySQL, SQLite, or MongoDB). Returns connection ID for subsequent operations.", 

35 "parameters": { 

36 "type": "object", 

37 "properties": { 

38 "connection_id": { 

39 "type": "string", 

40 "description": "Unique identifier for this connection" 

41 }, 

42 "db_type": { 

43 "type": "string", 

44 "enum": ["postgresql", "mysql", "sqlite", "mongodb"], 

45 "description": "Database type" 

46 }, 

47 "host": { 

48 "type": "string", 

49 "description": "Database host (not needed for SQLite)" 

50 }, 

51 "port": { 

52 "type": "integer", 

53 "description": "Database port (optional, uses defaults)" 

54 }, 

55 "database": { 

56 "type": "string", 

57 "description": "Database name" 

58 }, 

59 "username": { 

60 "type": "string", 

61 "description": "Database username" 

62 }, 

63 "password": { 

64 "type": "string", 

65 "description": "Database password" 

66 }, 

67 "file_path": { 

68 "type": "string", 

69 "description": "File path for SQLite database" 

70 } 

71 }, 

72 "required": ["connection_id", "db_type"] 

73 } 

74 } 

75 }, 

76 { 

77 "type": "function", 

78 "function": { 

79 "name": "db_query", 

80 "description": "Execute SQL query on connected database. Safe mode only allows SELECT queries. Returns results as JSON.", 

81 "parameters": { 

82 "type": "object", 

83 "properties": { 

84 "connection_id": { 

85 "type": "string", 

86 "description": "Connection ID from db_connect" 

87 }, 

88 "query": { 

89 "type": "string", 

90 "description": "SQL query to execute" 

91 }, 

92 "safe_mode": { 

93 "type": "boolean", 

94 "description": "If true, only SELECT queries allowed (default: true)" 

95 } 

96 }, 

97 "required": ["connection_id", "query"] 

98 } 

99 } 

100 }, 

101 { 

102 "type": "function", 

103 "function": { 

104 "name": "db_schema", 

105 "description": "Get database schema (tables, columns, types). Useful for understanding database structure before querying.", 

106 "parameters": { 

107 "type": "object", 

108 "properties": { 

109 "connection_id": { 

110 "type": "string", 

111 "description": "Connection ID from db_connect" 

112 } 

113 }, 

114 "required": ["connection_id"] 

115 } 

116 } 

117 }, 

118 { 

119 "type": "function", 

120 "function": { 

121 "name": "db_disconnect", 

122 "description": "Disconnect from database and clean up resources.", 

123 "parameters": { 

124 "type": "object", 

125 "properties": { 

126 "connection_id": { 

127 "type": "string", 

128 "description": "Connection ID to disconnect" 

129 } 

130 }, 

131 "required": ["connection_id"] 

132 } 

133 } 

134 } 

135] 

136 

137 

138# Tool implementations with Letta formatting 

139def tool_db_connect_letta( 

140 connection_id: str, 

141 db_type: str, 

142 **kwargs 

143) -> Dict[str, Any]: 

144 """Connect to database with Letta visual output""" 

145 

146 # Prepare args for display (hide password) 

147 display_args = { 

148 "connection_id": connection_id, 

149 "db_type": db_type 

150 } 

151 if "host" in kwargs: 

152 display_args["host"] = kwargs["host"] 

153 if "database" in kwargs: 

154 display_args["database"] = kwargs["database"] 

155 if "username" in kwargs: 

156 display_args["username"] = kwargs["username"] 

157 if "password" in kwargs: 

158 display_args["password"] = "***" 

159 

160 try: 

161 # Attempt connection 

162 success = db_manager.connect(connection_id, db_type, **kwargs) 

163 

164 if success: 

165 # Format success message 

166 summary = f"Successfully connected to {db_type} database" 

167 if kwargs.get('database'): 

168 summary += f" ({kwargs['database']})" 

169 

170 action_result( 

171 "db_connect", 

172 display_args, 

173 summary, 

174 success=True 

175 ) 

176 

177 return { 

178 "success": True, 

179 "connection_id": connection_id, 

180 "message": summary 

181 } 

182 else: 

183 raise DatabaseError(f"Failed to connect to {db_type}") 

184 

185 except Exception as e: 

186 logger.error(f"Database connection failed: {e}") 

187 

188 error_msg = str(e) 

189 action_result( 

190 "db_connect", 

191 display_args, 

192 f"Connection failed: {error_msg}", 

193 success=False 

194 ) 

195 

196 return { 

197 "success": False, 

198 "error": error_msg 

199 } 

200 

201 

202def tool_db_query_letta( 

203 connection_id: str, 

204 query: str, 

205 safe_mode: bool = True 

206) -> Dict[str, Any]: 

207 """Execute database query with Letta visual output""" 

208 

209 # Truncate query for display if too long 

210 display_query = query if len(query) < 60 else query[:57] + "..." 

211 

212 display_args = { 

213 "connection_id": connection_id, 

214 "query": display_query 

215 } 

216 

217 try: 

218 # Execute query 

219 result = db_manager.execute_query(connection_id, query, safe_mode=safe_mode) 

220 

221 # Summarize results 

222 summary = summarize_query(result) 

223 

224 action_result( 

225 "db_query", 

226 display_args, 

227 summary, 

228 success=True 

229 ) 

230 

231 # Return full result for AI to process 

232 return result 

233 

234 except Exception as e: 

235 logger.error(f"Query failed: {e}") 

236 

237 error_msg = str(e) 

238 action_result( 

239 "db_query", 

240 display_args, 

241 f"Query failed: {error_msg}", 

242 success=False 

243 ) 

244 

245 return { 

246 "success": False, 

247 "error": error_msg 

248 } 

249 

250 

251def tool_db_schema_letta(connection_id: str) -> Dict[str, Any]: 

252 """Get database schema with Letta visual output""" 

253 

254 display_args = {"connection_id": connection_id} 

255 

256 try: 

257 schema = db_manager.get_schema(connection_id) 

258 

259 # Summarize schema 

260 summary = summarize_schema(schema) 

261 

262 action_result( 

263 "db_schema", 

264 display_args, 

265 summary, 

266 success=True 

267 ) 

268 

269 return schema 

270 

271 except Exception as e: 

272 logger.error(f"Schema fetch failed: {e}") 

273 

274 error_msg = str(e) 

275 action_result( 

276 "db_schema", 

277 display_args, 

278 f"Schema fetch failed: {error_msg}", 

279 success=False 

280 ) 

281 

282 return { 

283 "success": False, 

284 "error": error_msg 

285 } 

286 

287 

288def tool_db_disconnect_letta(connection_id: str) -> Dict[str, Any]: 

289 """Disconnect from database with Letta visual output""" 

290 

291 display_args = {"connection_id": connection_id} 

292 

293 try: 

294 success = db_manager.disconnect(connection_id) 

295 

296 if success: 

297 action_result( 

298 "db_disconnect", 

299 display_args, 

300 f"Disconnected from {connection_id}", 

301 success=True 

302 ) 

303 

304 return { 

305 "success": True, 

306 "message": f"Disconnected from {connection_id}" 

307 } 

308 else: 

309 action_result( 

310 "db_disconnect", 

311 display_args, 

312 f"Connection {connection_id} not found", 

313 success=False 

314 ) 

315 

316 return { 

317 "success": False, 

318 "error": f"Connection {connection_id} not found" 

319 } 

320 

321 except Exception as e: 

322 error_msg = str(e) 

323 action_result( 

324 "db_disconnect", 

325 display_args, 

326 f"Disconnect failed: {error_msg}", 

327 success=False 

328 ) 

329 

330 return { 

331 "success": False, 

332 "error": error_msg 

333 } 

334 

335 

336# Tool registry for Letta-style tools 

337TOOL_HANDLERS_LETTA = { 

338 "db_connect": tool_db_connect_letta, 

339 "db_query": tool_db_query_letta, 

340 "db_schema": tool_db_schema_letta, 

341 "db_disconnect": tool_db_disconnect_letta 

342} 

343 

344 

345# Example workflow demonstrating the Letta visual style 

346def example_database_workflow(): 

347 """ 

348 Example database workflow showing the Letta visual style. 

349  

350 Output will look like: 

351 ✻ Thinking... 

352 I need to connect to the PostgreSQL database first, then query the users table. 

353  

354 ● db_connect(connection_id="mydb", db_type="postgresql", host="localhost") 

355 ⎿ Successfully connected to postgresql database (users_db) ✓ 

356  

357 ● db_schema(connection_id="mydb") 

358 ⎿ Schema: 12 tables, 85 columns total ✓ 

359  

360 ● db_query(connection_id="mydb", query="SELECT COUNT(*) FROM users") 

361 ⎿ Found 150 rows with 1 columns ✓ 

362  

363 ● db_disconnect(connection_id="mydb") 

364 ⎿ Disconnected from mydb ✓ 

365 """ 

366 

367 # Show reasoning 

368 thinking("I need to connect to the PostgreSQL database first, then query the users table to get the count.") 

369 

370 # Connect 

371 tool_db_connect_letta( 

372 connection_id="mydb", 

373 db_type="postgresql", 

374 host="localhost", 

375 port=5432, 

376 database="users_db", 

377 username="admin", 

378 password="secret123" 

379 ) 

380 

381 # Get schema 

382 tool_db_schema_letta("mydb") 

383 

384 # Query 

385 tool_db_query_letta( 

386 connection_id="mydb", 

387 query="SELECT COUNT(*) FROM users" 

388 ) 

389 

390 # Disconnect 

391 tool_db_disconnect_letta("mydb") 

392 

393 

394if __name__ == "__main__": 

395 # Run example 

396 print("\n" + "="*60) 

397 print("LETTA VISUAL STANDARD - DATABASE OPERATIONS") 

398 print("="*60 + "\n") 

399 

400 example_database_workflow() 

401 

402 print("\n" + "="*60) 

403 print("Status Bar:") 

404 print(letta_formatter.format_status_bar()) 

405 print("="*60 + "\n")