#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ RAG Manager - RAG 检索与管理模式 功能： 1. 检索 - 查询项目/系统/插件各层 RAG 2. 添加 - 将学习到的模式存入指定层 3. 删除 - 从指定层删除模式 4. 列表 - 列出各层的所有模式 5. 同步 - 将项目学习成果同步到系统共享层 用法： python rag_manager.py --project-root . list --layer system python rag_manager.py --project-root . search "打脸爽点" python rag_manager.py --project-root . add --pattern-id xxx --layer project python rag_manager.py --project-root . delete --pattern-id xxx --layer system python rag_manager.py --project-root . sync --from project --to system """ import argparse import asyncio import json import sys from dataclasses import asdict from datetime import datetime from pathlib import Path from typing import Optional try: from runtime_compat import enable_windows_utf8_stdio except ImportError: enable_windows_utf8_stdio = lambda: None # 添加 scripts 目录到路径 sys.path.insert(0, str(Path(__file__).parent.parent)) from data_modules.cross_project_rag import CrossProjectRAG, RAGLayer, LearnedPattern from data_modules.chapter_analyzer import ChapterAnalyzer, ChapterAnalysisResult class RAGManager: """ RAG 管理器 提供： - 检索：跨三层 RAG 检索 - 添加：存储模式到指定层 - 删除：删除指定模式 - 列表：列出各层模式 - 同步：将模式从项目层同步到系统层 """ def __init__(self, project_root: Path): self.project_root = project_root self.cross_rag = CrossProjectRAG(project_root) self.analyzer = ChapterAnalyzer(project_root) # ==================== 检索 ==================== async def search( self, query: str, top_k: int = 5, layers: Optional[list[str]] = None ) -> list: """检索 RAG""" if layers is None: layer_list = [RAGLayer.PROJECT, RAGLayer.SYSTEM, RAGLayer.PLUGIN] else: layer_map = { "project": RAGLayer.PROJECT, "system": RAGLayer.SYSTEM, "plugin": RAGLayer.PLUGIN } layer_list = [layer_map[l] for l in layers if l in layer_map] results = await self.cross_rag.search(query, top_k, layer_list) return [ { "chunk_id": r.chunk_id, "title": r.content.split('\n')[0][:50] if r.content else r.chunk_id, "content": r.content[:200] + "..." if len(r.content) > 200 else r.content, "score": r.score, "layer": r.source_layer.value, "project": r.source_project, "chapter": r.chapter, "type": r.chunk_type } for r in results ] # ==================== 添加 ==================== def add_pattern( self, pattern_id: str, pattern_type: str, title: str, description: str, catharsis_model: str, source_chapter: int, layer: str = "project", tension_curve: Optional[list] = None, structure: Optional[dict] = None, style_tags: Optional[list] = None ) -> bool: """添加模式""" pattern = LearnedPattern( pattern_id=pattern_id, pattern_type=pattern_type, title=title, description=description, tension_curve=tension_curve or [], catharsis_model=catharsis_model, structure=structure or {}, hot_spots=[], style_tags=style_tags or [], source_project=self.project_root.name, source_chapter=source_chapter, learned_at=datetime.now().isoformat() ) rag_layer = RAGLayer.SYSTEM if layer == "system" else RAGLayer.PROJECT return self.cross_rag.store_learned_pattern(pattern, rag_layer) def learn_and_add( self, chapter_file: Path, layer: str = "project" ) -> Optional[str]: """从章节学习并添加模式""" # 分析章节 result = self.analyzer.analyze_chapter(chapter_file) # 生成模式 pattern = self.analyzer.learn_pattern(result, layer) # 存储 rag_layer = RAGLayer.SYSTEM if layer == "system" else RAGLayer.PROJECT success = self.cross_rag.store_learned_pattern(pattern, rag_layer) if success: return pattern.pattern_id return None # ==================== 删除 ==================== def delete_pattern(self, pattern_id: str, layer: str) -> bool: """删除模式""" if layer == "project": return self._delete_from_project_db(pattern_id) elif layer == "system": return self._delete_from_system(pattern_id) return False def _delete_from_project_db(self, pattern_id: str) -> bool: """从项目数据库删除""" db_path = self.project_root / ".noma" / "rag" / "learned.db" if not db_path.exists(): return False try: import sqlite3 conn = sqlite3.connect(str(db_path)) cursor = conn.cursor() cursor.execute("DELETE FROM learned_patterns WHERE pattern_id = ?", (pattern_id,)) affected = cursor.rowcount conn.commit() conn.close() return affected > 0 except Exception: return False def _delete_from_system(self, pattern_id: str) -> bool: """从系统目录删除""" pattern_file = self.cross_rag.system_learned_dir / f"{pattern_id}.json" if pattern_file.exists(): pattern_file.unlink() return True return False # ==================== 列表 ==================== def list_patterns(self, layer: str) -> list: """列出指定层的模式""" if layer == "project": return self._list_project_patterns() elif layer == "system": return self._list_system_patterns() elif layer == "plugin": return self._list_plugin_patterns() return [] def _list_project_patterns(self) -> list: """列出项目层模式""" db_path = self.project_root / ".noma" / "rag" / "learned.db" if not db_path.exists(): return [] try: import sqlite3 conn = sqlite3.connect(str(db_path)) cursor = conn.cursor() cursor.execute(""" SELECT pattern_id, pattern_type, title, description, catharsis_model, source_chapter, learned_at, usage_count FROM learned_patterns ORDER BY learned_at DESC """) rows = cursor.fetchall() conn.close() return [ { "pattern_id": r[0], "type": r[1], "title": r[2], "description": r[3][:100] + "..." if r[3] and len(r[3]) > 100 else r[3], "catharsis_model": r[4], "source_chapter": r[5], "learned_at": r[6], "usage_count": r[7] } for r in rows ] except Exception: return [] def _list_system_patterns(self) -> list: """列出系统层模式""" if not self.cross_rag.system_learned_dir.exists(): return [] patterns = [] for f in self.cross_rag.system_learned_dir.glob("*.json"): try: data = json.loads(f.read_text(encoding="utf-8")) patterns.append({ "pattern_id": data.get("pattern_id", f.stem), "type": data.get("pattern_type", "unknown"), "title": data.get("title", f.stem), "description": data.get("description", "")[:100], "catharsis_model": data.get("catharsis_model", "unknown"), "source_project": data.get("source_project", "unknown"), "source_chapter": data.get("source_chapter", 0), "learned_at": data.get("learned_at", "") }) except Exception: continue return sorted(patterns, key=lambda x: x.get("learned_at", ""), reverse=True) def _list_plugin_patterns(self) -> list: """列出插件层模式""" patterns = [] plugin_dir = self.cross_rag.plugin_matrices_dir # catharsis models catharsis_dir = plugin_dir / "catharsis_models" if catharsis_dir.exists(): for f in catharsis_dir.glob("*.md"): patterns.append({ "pattern_id": f"plugin:{f.stem}", "type": "catharsis_model", "title": f.stem, "description": "内置爽感模型", "catharsis_model": f.stem, "source": "noma_plugin" }) return patterns # ==================== 同步 ==================== def sync_to_system(self, pattern_id: str) -> bool: """将项目模式同步到系统层""" # 从项目数据库读取 db_path = self.project_root / ".noma" / "rag" / "learned.db" if not db_path.exists(): return False try: import sqlite3 conn = sqlite3.connect(str(db_path)) cursor = conn.cursor() cursor.execute(""" SELECT pattern_id, pattern_type, title, description, tension_curve, catharsis_model, structure, hot_spots, style_tags, source_project, source_chapter, learned_at, usage_count FROM learned_patterns WHERE pattern_id = ? """, (pattern_id,)) row = cursor.fetchone() conn.close() if not row: return False pattern = LearnedPattern( pattern_id=row[0], pattern_type=row[1], title=row[2], description=row[3] or "", tension_curve=json.loads(row[4]) if row[4] else [], catharsis_model=row[5] or "unknown", structure=json.loads(row[6]) if row[6] else {}, hot_spots=json.loads(row[7]) if row[7] else [], style_tags=json.loads(row[8]) if row[8] else [], source_project=row[9] or self.project_root.name, source_chapter=row[10] or 0, learned_at=row[11] or datetime.now().isoformat(), usage_count=row[12] or 0 ) return self.cross_rag.store_learned_pattern(pattern, RAGLayer.SYSTEM) except Exception: return False def sync_all_to_system(self) -> dict: """同步所有项目模式到系统层""" project_patterns = self._list_project_patterns() synced = 0 failed = 0 for p in project_patterns: if self.sync_to_system(p["pattern_id"]): synced += 1 else: failed += 1 return {"synced": synced, "failed": failed, "total": len(project_patterns)} # ==================== CLI ==================== def main(): if sys.platform == "win32": enable_windows_utf8_stdio() parser = argparse.ArgumentParser(description="RAG Manager - RAG 检索与管理") parser.add_argument("--project-root", type=str, default=".", help="项目根目录") subparsers = parser.add_subparsers(dest="command") # 搜索 search_parser = subparsers.add_parser("search", help="检索 RAG") search_parser.add_argument("query", help="检索 query") search_parser.add_argument("--top-k", type=int, default=5) search_parser.add_argument("--layers", type=str, default="project,system", help="检索层级，逗号分隔") # 列表 list_parser = subparsers.add_parser("list", help="列出模式") list_parser.add_argument("--layer", choices=["project", "system", "plugin"], default="project", help="RAG 层") # 添加 add_parser = subparsers.add_parser("add", help="添加模式") add_parser.add_argument("--pattern-id", required=True) add_parser.add_argument("--pattern-type", required=True) add_parser.add_argument("--title", required=True) add_parser.add_argument("--description", required=True) add_parser.add_argument("--catharsis-model", required=True) add_parser.add_argument("--source-chapter", type=int, required=True) add_parser.add_argument("--layer", choices=["project", "system"], default="project") add_parser.add_argument("--learn", help="从章节文件学习") add_parser.add_argument("--chapter-file", help="章节文件路径") # 删除 delete_parser = subparsers.add_parser("delete", help="删除模式") delete_parser.add_argument("--pattern-id", required=True) delete_parser.add_argument("--layer", choices=["project", "system"], required=True) # 同步 sync_parser = subparsers.add_parser("sync", help="同步到系统层") sync_parser.add_argument("--pattern-id", help="同步单个模式（可选）") sync_parser.add_argument("--all", action="store_true", help="同步所有") args = parser.parse_args() if not args.project_root: print("Error: --project-root is required") sys.exit(1) project_root = Path(args.project_root).resolve() manager = RAGManager(project_root) if args.command == "search": layers = [l.strip() for l in args.layers.split(",")] results = asyncio.run(manager.search(args.query, args.top_k, layers)) print(f"\n=== Search Results ({len(results)}) ===") for r in results: print(f"\n[{r['layer']}] {r['title']}") print(f" Type: {r['type']}, Chapter: {r['chapter']}") print(f" Score: {r['score']:.2f}") print(f" Content: {r['content']}") elif args.command == "list": patterns = manager.list_patterns(args.layer) print(f"\n=== {args.layer.upper()} Patterns ({len(patterns)}) ===") for p in patterns: print(f"\n[{p['pattern_id']}] {p['title']}") print(f" Type: {p['type']}, Catharsis: {p.get('catharsis_model', 'N/A')}") print(f" Source: {p.get('source_project', 'N/A')} Ch.{p.get('source_chapter', 0)}") if p.get("description"): print(f" Desc: {p['description'][:100]}") elif args.command == "add": if hasattr(args, 'learn') and args.learn: # 从章节学习 chapter_file = project_root / args.chapter_file pattern_id = manager.learn_and_add(chapter_file, args.layer) if pattern_id: print(f"✓ Pattern learned and added: {pattern_id}") else: print("✗ Failed to learn pattern") sys.exit(1) else: success = manager.add_pattern( args.pattern_id, args.pattern_type, args.title, args.description, args.catharsis_model, args.source_chapter, args.layer ) if success: print(f"✓ Pattern added to {args.layer}") else: print("✗ Failed to add pattern") sys.exit(1) elif args.command == "delete": success = manager.delete_pattern(args.pattern_id, args.layer) if success: print(f"✓ Pattern deleted from {args.layer}") else: print("✗ Failed to delete pattern") sys.exit(1) elif args.command == "sync": if args.all: result = manager.sync_all_to_system() print(f"✓ Synced {result['synced']}/{result['total']} patterns") if result['failed'] > 0: print(f" Failed: {result['failed']}") elif args.pattern_id: success = manager.sync_to_system(args.pattern_id) if success: print(f"✓ Pattern synced to system") else: print("✗ Failed to sync pattern") sys.exit(1) else: print("Specify --pattern-id or --all") sys.exit(1) else: parser.print_help() if __name__ == "__main__": main()