470 lines
16 KiB
Python
470 lines
16 KiB
Python
#!/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()
|