1#!/usr/bin/env python3

2"""Memory and reflection management MCP tools.

3

4 This module provides tools for storing, searching, and managing reflections and conversation memories.

5

6 Refactored to use utility modules for reduced code duplication.

7"""

8

9 from __future__ import annotations

10

11 import asyncio

12 import operator

13 import typing as t

14 from datetime import datetime

15 from typing import TYPE_CHECKING, Any

16

17 from session_buddy.utils.database_helpers import require_reflection_database

18 from session_buddy.utils.error_handlers import (

19 DatabaseUnavailableError,

20 ValidationError,

21 _get_logger,

22 validate_required,

23)

24 from session_buddy.utils.messages import ToolMessages

25 from session_buddy.utils.tool_wrapper import format_reflection_result

26

27 if TYPE_CHECKING:

28 from session_buddy.adapters.reflection_adapter import ReflectionDatabaseAdapter

29

30

31_reflection_tools_available: bool | None = None

32_reflection_db: ReflectionDatabaseAdapter | None = None

33

34

35 def _check_reflection_tools_available() -> bool:

36 """Check if reflection tools are available, cached for reuse."""

37 global _reflection_tools_available

38 if _reflection_tools_available is not None: 38 ↛ 40 line 38 didn't jump to line 40 because the condition on line 38 was always true

39 return _reflection_tools_available

40 try:

41 import duckdb

42

43 _reflection_tools_available = True

44 except ImportError:

45 _reflection_tools_available = False

46 return _reflection_tools_available

47

48

49 async def _get_reflection_database() -> ReflectionDatabaseAdapter:

50 """Get reflection database instance (patchable for tests)."""

51 global _reflection_db

52 if _reflection_db is not None:

53 return _reflection_db

54 _reflection_db = await require_reflection_database()

55 return _reflection_db

56

57

58 async def _execute_database_tool(

59 operation: t.Callable[[ReflectionDatabaseAdapter], t.Awaitable[t.Any]],

60 formatter: t.Callable[[t.Any], str],

61 operation_name: str,

62 validator: t.Callable[[], None] | None = None,

63) -> str:

64 try:

65 if validator:

66 validator()

67

68 db = await _get_reflection_database()

69 result = await operation(db)

70 return formatter(result)

71 except ValidationError as e:

72 return ToolMessages.validation_error(operation_name, str(e))

73 except DatabaseUnavailableError as e:

74 return ToolMessages.not_available(operation_name, str(e))

75 except Exception as e:

76 _get_logger().exception(f"Error in {operation_name}: {e}")

77 return ToolMessages.operation_failed(operation_name, e)

78

79

80 async def _execute_simple_database_tool(

81 operation: t.Callable[[ReflectionDatabaseAdapter], t.Awaitable[str]],

82 operation_name: str,

83) -> str:

84 try:

85 db = await _get_reflection_database()

86 return await operation(db)

87 except DatabaseUnavailableError as e:

88 return ToolMessages.not_available(operation_name, str(e))

89 except Exception as e:

90 _get_logger().exception(f"Error in {operation_name}: {e}")

91 return ToolMessages.operation_failed(operation_name, e)

92

93

94 def _format_score(score: float) -> str:

95 """Format a score as a percentage or relevance indicator."""

96 return f"{score:.2f}"

97

98

99# ============================================================================

100# Store Reflection Tool

101# ============================================================================

102

103

104 async def _store_reflection_operation(

105 db: ReflectionDatabaseAdapter, content: str, tags: list[str]

106) -> dict[str, Any]:

107 """Execute reflection storage operation."""

108 success = await db.store_reflection(content, tags=tags)

109 return {

110 "success": success,

111 "content": content,

112 "tags": tags,

113 "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),

114 }

115

116

117 def _format_store_reflection_result(result: dict[str, Any]) -> str:

118 """Format reflection storage result."""

119 return format_reflection_result(

120 result["success"],

121 result["content"],

122 result.get("tags"),

123 result.get("timestamp"),

124 )

125

126

127 async def _store_reflection_impl(content: str, tags: list[str] | None = None) -> str:

128 """Implementation for store_reflection tool."""

129 if not _check_reflection_tools_available():

130 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

131

132 try:

133 validate_required(content, "content")

134 db = await _get_reflection_database()

135 result = await _store_reflection_operation(db, content, tags or [])

136 return _format_store_reflection_result(result)

137 except ValidationError as e:

138 return ToolMessages.validation_error("Store reflection", str(e))

139 except DatabaseUnavailableError as e:

140 return ToolMessages.not_available("Store reflection", str(e))

141 except Exception as e:

142 _get_logger().exception(f"Error storing reflection: {e}")

143 return f"Error storing reflection: {e}"

144

145

146# ============================================================================

147# Quick Search Tool

148# ============================================================================

149

150

151 async def _quick_search_operation(

152 db: ReflectionDatabaseAdapter,

153 query: str,

154 project: str | None,

155 min_score: float,

156) -> str:

157 """Execute quick search operation and format results."""

158 results = await db.search_conversations(

159 query=query,

160 project=project,

161 limit=1,

162 min_score=min_score,

163 )

164

165 lines = [f"🔍 Quick search for: '{query}'"]

166

167 if results:

168 result = results[0]

169 lines.extend(

170 (

171 "📊 Found results (showing top 1)",

172 f"📝 {ToolMessages.truncate_text(result['content'], 150)}",

173 )

174 )

175 if result.get("project"): 175 ↛ 177 line 175 didn't jump to line 177 because the condition on line 175 was always true

176 lines.append(f"📁 Project: {result['project']}")

177 if result.get("score") is not None: 177 ↛ 179 line 177 didn't jump to line 179 because the condition on line 177 was always true

178 lines.append(f"⭐ Relevance: {_format_score(result['score'])}")

179 lines.append(f"📅 Date: {result.get('timestamp', 'Unknown')}")

180 else:

181 lines.extend(

182 (

183 "🔍 No results found",

184 "💡 Try adjusting your search terms or lowering min_score",

185 )

186 )

187

188 return "\n".join(lines)

189

190

191 async def _quick_search_impl(

192 query: str,

193 min_score: float = 0.7,

194 project: str | None = None,

195) -> str:

196 """Implementation for quick_search tool."""

197 if not _check_reflection_tools_available():

198 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

199

200 async def operation(db: ReflectionDatabaseAdapter) -> str:

201 return await _quick_search_operation(db, query, project, min_score)

202

203 return await _execute_simple_database_tool(operation, "Quick search")

204

205

206# ============================================================================

207# Search Summary Tool

208# ============================================================================

209

210

211 async def _analyze_project_distribution(

212 results: list[dict[str, Any]],

213) -> dict[str, int]:

214 """Analyze project distribution of search results."""

215 projects: dict[str, int] = {}

216 for result in results:

217 proj = result.get("project", "Unknown")

218 projects[proj] = projects.get(proj, 0) + 1

219 return projects

220

221

222 async def _analyze_relevance_scores(

223 results: list[dict[str, Any]],

224) -> tuple[float, list[float]]:

225 """Analyze relevance scores of search results."""

226 scores = [r.get("score", 0.0) for r in results if r.get("score") is not None]

227 avg_score = sum(scores) / len(scores) if scores else 0.0

228 return avg_score, scores

229

230

231 async def _extract_common_themes(

232 results: list[dict[str, Any]],

233) -> list[tuple[str, int]]:

234 """Extract common themes from search results."""

235 all_content = " ".join([r["content"] for r in results])

236 words = all_content.lower().split()

237 word_freq: dict[str, int] = {}

238

239 for word in words:

240 if len(word) > 4: # Skip short words

241 word_freq[word] = word_freq.get(word, 0) + 1

242

243 if word_freq: 243 ↛ 245 line 243 didn't jump to line 245 because the condition on line 243 was always true

244 return sorted(word_freq.items(), key=operator.itemgetter(1), reverse=True)[:5]

245 return []

246

247

248 async def _format_search_summary(

249 query: str,

250 results: list[dict[str, Any]],

251) -> str:

252 """Format complete search summary."""

253 lines = [

254 f"📊 Search Summary for: '{query}'",

255 "=" * 50,

256 ]

257

258 if not results:

259 lines.extend(

260 [

261 "🔍 No results found",

262 "💡 Try different search terms or lower the min_score threshold",

263 ]

264 )

265 return "\n".join(lines)

266

267 # Basic stats

268 lines.append(f"📈 Total results: {len(results)}")

269

270 # Project distribution

271 projects = await _analyze_project_distribution(results)

272 if len(projects) > 1: 272 ↛ 280 line 272 didn't jump to line 280 because the condition on line 272 was always true

273 lines.append("📁 Project distribution:")

274 for proj, count in sorted(

275 projects.items(), key=operator.itemgetter(1), reverse=True

276 ):

277 lines.append(f" • {proj}: {count} results")

278

279 # Time distribution

280 timestamps = [r.get("timestamp") for r in results if r.get("timestamp")]

281 if timestamps: 281 ↛ 285 line 281 didn't jump to line 285 because the condition on line 281 was always true

282 lines.append(f"📅 Time range: {len(timestamps)} results with dates")

283

284 # Relevance scores

285 avg_score, scores = await _analyze_relevance_scores(results)

286 if scores: 286 ↛ 290 line 286 didn't jump to line 290 because the condition on line 286 was always true

287 lines.append(f"⭐ Average relevance: {_format_score(avg_score)}")

288

289 # Common themes

290 top_words = await _extract_common_themes(results)

291 if top_words: 291 ↛ 296 line 291 didn't jump to line 296 because the condition on line 291 was always true

292 lines.append("🔤 Common themes:")

293 for word, freq in top_words:

294 lines.append(f" • {word}: {freq} mentions")

295

296 return "\n".join(lines)

297

298

299 async def _search_summary_operation(

300 db: ReflectionDatabaseAdapter,

301 query: str,

302 project: str | None,

303 min_score: float,

304) -> str:

305 """Execute search summary operation."""

306 results = await db.search_conversations(

307 query=query,

308 project=project,

309 limit=20,

310 min_score=min_score,

311 )

312 return await _format_search_summary(query, results)

313

314

315 async def _search_summary_impl(

316 query: str,

317 min_score: float = 0.7,

318 project: str | None = None,

319) -> str:

320 """Implementation for search_summary tool."""

321 if not _check_reflection_tools_available():

322 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

323

324 try:

325 db = await _get_reflection_database()

326 return await _search_summary_operation(db, query, project, min_score)

327 except DatabaseUnavailableError as e:

328 return ToolMessages.not_available("Search summary", str(e))

329 except Exception as e:

330 _get_logger().exception(f"Search summary error: {e}")

331 return f"Search summary error: {e}"

332

333

334# ============================================================================

335# Search by File Tool

336# ============================================================================

337

338

339 async def _format_file_search_results(

340 file_path: str,

341 results: list[dict[str, Any]],

342) -> str:

343 """Format file search results."""

344 lines = [

345 f"📁 Searching conversations about: {file_path}",

346 "=" * 50,

347 ]

348

349 if not results:

350 lines.extend(

351 [

352 "🔍 No conversations found about this file",

353 "💡 The file might not have been discussed in previous sessions",

354 ]

355 )

356 return "\n".join(lines)

357

358 lines.append(f"📈 Found {len(results)} relevant conversations:")

359

360 for i, result in enumerate(results, 1):

361 lines.append(

362 f"\n{i}. 📝 {ToolMessages.truncate_text(result['content'], 200)}",

363 )

364 if result.get("project"): 364 ↛ 366 line 364 didn't jump to line 366 because the condition on line 364 was always true

365 lines.append(f" 📁 Project: {result['project']}")

366 if result.get("score") is not None: 366 ↛ 368 line 366 didn't jump to line 368 because the condition on line 366 was always true

367 lines.append(f" ⭐ Relevance: {_format_score(result['score'])}")

368 if result.get("timestamp"): 368 ↛ 360 line 368 didn't jump to line 360 because the condition on line 368 was always true

369 lines.append(f" 📅 Date: {result['timestamp']}")

370

371 return "\n".join(lines)

372

373

374 async def _search_by_file_operation(

375 db: ReflectionDatabaseAdapter,

376 file_path: str,

377 limit: int,

378 project: str | None,

379) -> str:

380 """Execute file search operation."""

381 results = await db.search_conversations(

382 query=file_path,

383 project=project,

384 limit=limit,

385 )

386 return await _format_file_search_results(file_path, results)

387

388

389 async def _search_by_file_impl(

390 file_path: str,

391 limit: int = 10,

392 project: str | None = None,

393) -> str:

394 """Implementation for search_by_file tool."""

395 if not _check_reflection_tools_available():

396 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

397

398 try:

399 db = await _get_reflection_database()

400 return await _search_by_file_operation(db, file_path, limit, project)

401 except DatabaseUnavailableError as e:

402 return ToolMessages.not_available("Search by file", str(e))

403 except Exception as e:

404 _get_logger().exception(f"File search error: {e}")

405 return f"File search error: {e}"

406

407

408# ============================================================================

409# Search by Concept Tool

410# ============================================================================

411

412

413 async def _format_concept_search_results(

414 concept: str,

415 results: list[dict[str, Any]],

416 include_files: bool,

417) -> str:

418 """Format concept search results."""

419 lines = [

420 f"🧠 Searching for concept: '{concept}'",

421 "=" * 50,

422 ]

423

424 if not results:

425 lines.extend(

426 [

427 "🔍 No conversations found about this concept",

428 "💡 Try related terms or broader concepts",

429 ]

430 )

431 return "\n".join(lines)

432

433 lines.append(f"📈 Found {len(results)} related conversations:")

434

435 for i, result in enumerate(results, 1):

436 lines.append(

437 f"\n{i}. 📝 {ToolMessages.truncate_text(result['content'], 250)}",

438 )

439 if result.get("project"): 439 ↛ 441 line 439 didn't jump to line 441 because the condition on line 439 was always true

440 lines.append(f" 📁 Project: {result['project']}")

441 if result.get("score") is not None: 441 ↛ 443 line 441 didn't jump to line 443 because the condition on line 441 was always true

442 lines.append(f" ⭐ Relevance: {_format_score(result['score'])}")

443 if result.get("timestamp"): 443 ↛ 446 line 443 didn't jump to line 446 because the condition on line 443 was always true

444 lines.append(f" 📅 Date: {result['timestamp']}")

445

446 if include_files and result.get("files"): 446 ↛ 447 line 446 didn't jump to line 447 because the condition on line 446 was never true

447 files = result["files"][:3]

448 if files:

449 lines.append(f" 📄 Files: {', '.join(files)}")

450

451 return "\n".join(lines)

452

453

454 async def _search_by_concept_operation(

455 db: ReflectionDatabaseAdapter,

456 concept: str,

457 include_files: bool,

458 limit: int,

459 project: str | None,

460) -> str:

461 """Execute concept search operation."""

462 results = await db.search_conversations(

463 query=concept,

464 project=project,

465 limit=limit,

466 )

467 return await _format_concept_search_results(concept, results, include_files)

468

469

470 async def _search_by_concept_impl(

471 concept: str,

472 include_files: bool = True,

473 limit: int = 10,

474 project: str | None = None,

475) -> str:

476 """Implementation for search_by_concept tool."""

477 if not _check_reflection_tools_available():

478 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

479

480 try:

481 db = await _get_reflection_database()

482 return await _search_by_concept_operation(

483 db, concept, include_files, limit, project

484 )

485 except DatabaseUnavailableError as e:

486 return ToolMessages.not_available("Search by concept", str(e))

487 except Exception as e:

488 _get_logger().exception(f"Concept search error: {e}")

489 return f"Concept search error: {e}"

490

491

492# ============================================================================

493# Reflection Stats Tool

494# ============================================================================

495

496

497 def _format_stats_new(stats: dict[str, t.Any]) -> list[str]:

498 """Format statistics in new format (conversations_count, reflections_count)."""

499 conv_count = stats.get("conversations_count", 0)

500 refl_count = stats.get("reflections_count", 0)

501 provider = stats.get("embedding_provider", "unknown")

502

503 return [

504 f"📈 Total conversations: {conv_count}",

505 f"💭 Total reflections: {refl_count}",

506 f"🔧 Embedding provider: {provider}",

507 f"\n🏥 Database health: {'✅ Healthy' if (conv_count + refl_count) > 0 else '⚠️ Empty'}",

508 ]

509

510

511 def _format_new_stats(stats: dict[str, t.Any]) -> list[str]:

512 """Backward-compatible alias for _format_stats_new."""

513 return _format_stats_new(stats)

514

515

516 def _format_stats_old(stats: dict[str, t.Any]) -> list[str]:

517 """Format statistics in old/test format (total_reflections, projects, date_range)."""

518 output = [

519 f"📈 Total reflections: {stats.get('total_reflections', 0)}",

520 f"📁 Projects: {stats.get('projects', 0)}",

521 ]

522

523 # Add date range if present

524 date_range = stats.get("date_range")

525 if isinstance(date_range, dict):

526 output.append(

527 f"📅 Date range: {date_range.get('start')} to {date_range.get('end')}"

528 )

529

530 # Add recent activity if present

531 recent_activity = stats.get("recent_activity", [])

532 if recent_activity:

533 output.append("\n🕐 Recent activity:")

534 output.extend([f" • {activity}" for activity in recent_activity[:5]])

535

536 # Database health

537 is_healthy = stats.get("total_reflections", 0) > 0

538 output.append(f"\n🏥 Database health: {'✅ Healthy' if is_healthy else '⚠️ Empty'}")

539

540 return output

541

542

543 def _format_old_stats(stats: dict[str, t.Any]) -> list[str]:

544 """Backward-compatible alias for _format_stats_old."""

545 return _format_stats_old(stats)

546

547

548 async def _reflection_stats_operation(db: ReflectionDatabaseAdapter) -> str:

549 """Execute reflection stats operation."""

550 stats = await db.get_stats()

551

552 lines = ["📊 Reflection Database Statistics", "=" * 40]

553

554 if stats and "error" not in stats: 554 ↛ 561 line 554 didn't jump to line 561 because the condition on line 554 was always true

555 # Format based on stat structure

556 if "conversations_count" in stats: 556 ↛ 559 line 556 didn't jump to line 559 because the condition on line 556 was always true

557 lines.extend(_format_stats_new(stats))

558 else:

559 lines.extend(_format_stats_old(stats))

560 else:

561 lines.extend(

562 [

563 "📊 No statistics available",

564 "💡 Database may be empty or inaccessible",

565 ]

566 )

567

568 return "\n".join(lines)

569

570

571 async def _reflection_stats_impl() -> str:

572 """Implementation for reflection_stats tool."""

573 if not _check_reflection_tools_available():

574 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

575

576 async def operation(db: ReflectionDatabaseAdapter) -> str:

577 return await _reflection_stats_operation(db)

578

579 return await _execute_simple_database_tool(operation, "Reflection stats")

580

581

582# ============================================================================

583# Reset Database Tool

584# ============================================================================

585

586

587 async def _close_db_connection(conn: t.Any) -> None:

588 """Close database connection, handling both async and sync cases."""

589 close_method = getattr(conn, "close", None)

590 if not callable(close_method): 590 ↛ 591 line 590 didn't jump to line 591 because the condition on line 590 was never true

591 return

592

593 result = close_method()

594 if asyncio.iscoroutine(result): 594 ↛ exit line 594 didn't return from function '_close_db_connection' because the condition on line 594 was always true

595 await result

596

597

598 async def _close_db_object(db_obj: t.Any) -> None:

599 """Close database object using async or sync close method."""

600 # Try async close first

601 aclose_method = getattr(db_obj, "aclose", None)

602 if callable(aclose_method): 602 ↛ 609 line 602 didn't jump to line 609 because the condition on line 602 was always true

603 result = aclose_method()

604 if asyncio.iscoroutine(result): 604 ↛ 606 line 604 didn't jump to line 606 because the condition on line 604 was always true

605 await result

606 return

607

608 # Fallback to sync close

609 close_method = getattr(db_obj, "close", None)

610 if callable(close_method):

611 close_method()

612

613

614 async def _close_reflection_db_safely(db_obj: t.Any) -> None:

615 """Safely close reflection database and its connection.

616

617 Handles both legacy and adapter-style DB objects.

618 """

619 # Close connection if it exists (legacy style)

620 conn = getattr(db_obj, "conn", None)

621 if conn: 621 ↛ 625 line 621 didn't jump to line 625 because the condition on line 621 was always true

622 await _close_db_connection(conn)

623

624 # Close the database object itself

625 await _close_db_object(db_obj)

626

627

628 async def _reset_reflection_database_impl() -> str:

629 """Implementation for reset_reflection_database tool."""

630 if not _check_reflection_tools_available():

631 return "Reflection tools not available. Install dependencies: uv sync --extra embeddings"

632

633 global _reflection_db

634 try:

635 if _reflection_db:

636 await _close_reflection_db_safely(_reflection_db)

637

638 _reflection_db = None

639 await _get_reflection_database()

640

641 lines = [

642 "🔄 Reflection database connection reset",

643 "✅ New connection established successfully",

644 "💡 Database locks should be resolved",

645 ]

646 return "\n".join(lines)

647

648 except Exception as e:

649 return ToolMessages.operation_failed("Reset database", e)

650

651

652# ============================================================================

653# MCP Tool Registration

654# ============================================================================

655

656

657 def register_memory_tools(mcp_server: Any) -> None:

658 """Register all memory management tools with the MCP server."""

659

660 @mcp_server.tool() # type: ignore[misc]

661 async def store_reflection(content: str, tags: list[str] | None = None) -> str:

662 """Store an important insight or reflection for future reference."""

663 return await _store_reflection_impl(content, tags)

664

665 @mcp_server.tool() # type: ignore[misc]

666 async def quick_search(

667 query: str,

668 min_score: float = 0.7,

669 project: str | None = None,

670 ) -> str:

671 """Quick search that returns only the count and top result for fast overview."""

672 return await _quick_search_impl(query, min_score, project)

673

674 @mcp_server.tool() # type: ignore[misc]

675 async def search_summary(

676 query: str,

677 limit: int = 10,

678 project: str | None = None,

679 min_score: float = 0.7,

680 ) -> str:

681 """Get aggregated insights from search results without individual result details."""

682 return await _search_summary_impl(query, min_score, project)

683

684 @mcp_server.tool() # type: ignore[misc]

685 async def search_by_file(

686 file_path: str,

687 limit: int = 10,

688 project: str | None = None,

689 min_score: float = 0.7,

690 ) -> str:

691 """Search for conversations that analyzed a specific file."""

692 return await _search_by_file_impl(file_path, limit, project)

693

694 @mcp_server.tool() # type: ignore[misc]

695 async def search_by_concept(

696 concept: str,

697 include_files: bool = True,

698 limit: int = 10,

699 project: str | None = None,

700 min_score: float = 0.7,

701 ) -> str:

702 """Search for conversations about a specific development concept."""

703 return await _search_by_concept_impl(concept, include_files, limit, project)

704

705 @mcp_server.tool() # type: ignore[misc]

706 async def reflection_stats(project: str | None = None) -> str:

707 """Get statistics about the reflection database."""

708 return await _reflection_stats_impl()

709

710 @mcp_server.tool() # type: ignore[misc]

711 async def reset_reflection_database() -> str:

712 """Reset the reflection database connection to fix lock issues."""

713 return await _reset_reflection_database_impl()