#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Cross-Project RAG - 跨层级 RAG 检索模块 实现三层 RAG 架构（向下继承）： 1. 小说私有层 - novel/.noma/rag/ 2. 工作空间共享层 - workspaces/{ws}/.noma/rag/ 3. 工程共享层 - project_root/.noma/rag/ 4. 插件内置层 - plugin/matrices/ 读取时：小说 → 工作空间 → 工程 → 插件（向下继承） 写入时：默认写入小说层，可选择向上沉淀 """ import json import os import sqlite3 import asyncio from pathlib import Path from dataclasses import dataclass, field from typing import List, Dict, Any, Optional from enum import Enum from datetime import datetime try: from runtime_compat import enable_windows_utf8_stdio except ImportError: enable_windows_utf8_stdio = lambda: None class RAGLayer(Enum): """RAG 层级 - 四层架构""" NOVEL = "novel" # 小说私有（最高权重） WORKSPACE = "workspace" # 工作空间共享 PROJECT = "project" # 工程目录共享 PLUGIN = "plugin" # 插件内置 @dataclass class LearnedPattern: """学习到的爽点模式""" pattern_id: str pattern_type: str title: str description: str tension_curve: List[List[float]] catharsis_model: str structure: Dict[str, int] hot_spots: List[List[Any]] style_tags: List[str] source_project: str source_chapter: int learned_at: str usage_count: int = 0 metadata: Dict[str, Any] = field(default_factory=dict) @dataclass class CrossProjectSearchResult: """检索结果""" chunk_id: str content: str score: float source_layer: RAGLayer source_project: Optional[str] chapter: Optional[int] chunk_type: Optional[str] metadata: Dict[str, Any] = field(default_factory=dict) class CrossProjectRAG: """ 四层 RAG 检索器 检索时自动向下继承：小说私有 > 工作空间 > 工程 > 插件 """ # 权重：小说私有 > 工作空间 > 工程 > 插件 DEFAULT_WEIGHTS = { RAGLayer.NOVEL: 1.0, RAGLayer.WORKSPACE: 0.8, RAGLayer.PROJECT: 0.6, RAGLayer.PLUGIN: 0.3, } def __init__( self, project_root: Path, workspace_root: Optional[Path] = None, project_root_dir: Optional[Path] = None, plugin_root: Optional[Path] = None, weights: Optional[Dict[RAGLayer, float]] = None, config: Optional[Any] = None, ): self.project_root = Path(project_root).resolve() # 四层根路径 self.novel_root = self.project_root self.workspace_root = workspace_root or self._resolve_workspace_root() self.project_root_dir = project_root_dir or self._resolve_project_root_dir() self.plugin_root = plugin_root or self._resolve_plugin_root() self.weights = weights or self.DEFAULT_WEIGHTS # 加载配置 self.config = config self._load_config() # 初始化四层路径 self._init_paths() # 确保目录存在 self._ensure_dirs() def _init_paths(self): """初始化四层 RAG 路径""" # 层1：小说私有 self.novel_rag_dir = self.novel_root / ".noma" / "rag" self.novel_vectors_db = self.novel_rag_dir / "vectors.db" self.novel_learned_db = self.novel_rag_dir / "learned.db" # 层2：工作空间共享 if self.workspace_root: self.ws_rag_dir = self.workspace_root / ".noma" / "rag" self.ws_shared_dir = self.ws_rag_dir / "shared" self.ws_catharsis_dir = self.ws_shared_dir / "catharsis" self.ws_genres_dir = self.ws_shared_dir / "genres" self.ws_learned_dir = self.ws_rag_dir / "learned" else: self.ws_rag_dir = None self.ws_shared_dir = None self.ws_catharsis_dir = None self.ws_genres_dir = None self.ws_learned_dir = None # 层3：工程共享 if self.project_root_dir: self.proj_rag_dir = self.project_root_dir / ".noma" / "rag" self.proj_shared_dir = self.proj_rag_dir / "shared" self.proj_catharsis_dir = self.proj_shared_dir / "catharsis" self.proj_genres_dir = self.proj_shared_dir / "genres" self.proj_learned_dir = self.proj_rag_dir / "learned" else: self.proj_rag_dir = None self.proj_shared_dir = None self.proj_catharsis_dir = None self.proj_genres_dir = None self.proj_learned_dir = None # 层4：插件内置 self.plugin_matrices_dir = self.plugin_root / "matrices" self.plugin_catharsis_dir = self.plugin_matrices_dir / "catharsis_models" self.plugin_genres_dir = self.plugin_matrices_dir / "genres" def _resolve_workspace_root(self) -> Optional[Path]: """解析工作空间根目录""" env_path = os.environ.get("NOMA_WORKSPACE_ROOT") if env_path: p = Path(env_path).resolve() if p.exists() and (p / ".noma" / "rag").exists(): return p current = self.project_root while True: workspaces_dir = current / "workspaces" if workspaces_dir.is_dir(): for ws_dir in workspaces_dir.iterdir(): if ws_dir.is_dir() and (ws_dir / ".noma" / "rag").exists(): if self.project_root == ws_dir or str(self.project_root).startswith(str(ws_dir) + os.sep): return ws_dir parent = current.parent if parent == current: break current = parent return None def _resolve_project_root_dir(self) -> Optional[Path]: """解析工程目录根""" env_path = os.environ.get("NOMA_PROJECT_ROOT_DIR") if env_path: p = Path(env_path).resolve() if p.exists() and (p / "workspaces").is_dir() and (p / ".noma" / "rag").is_dir(): return p current = self.project_root while True: if (current / "workspaces").is_dir() and (current / ".noma" / "rag").is_dir(): return current parent = current.parent if parent == current: break current = parent return None def _resolve_plugin_root(self) -> Path: """解析插件根目录""" env_path = os.environ.get("NOMA_PLUGIN_ROOT") if env_path: p = Path(env_path) if p.exists(): return p current_file = Path(__file__).resolve() candidate = current_file.parent.parent.parent.parent if (candidate / "matrices").exists(): return candidate return current_file.parent.parent.parent def _ensure_dirs(self): """确保必要的目录存在""" self.novel_rag_dir.mkdir(parents=True, exist_ok=True) if self.ws_rag_dir: self.ws_rag_dir.mkdir(parents=True, exist_ok=True) if self.ws_shared_dir: self.ws_shared_dir.mkdir(parents=True, exist_ok=True) if self.ws_catharsis_dir: self.ws_catharsis_dir.mkdir(parents=True, exist_ok=True) if self.ws_genres_dir: self.ws_genres_dir.mkdir(parents=True, exist_ok=True) if self.ws_learned_dir: self.ws_learned_dir.mkdir(parents=True, exist_ok=True) if self.proj_rag_dir: self.proj_rag_dir.mkdir(parents=True, exist_ok=True) if self.proj_shared_dir: self.proj_shared_dir.mkdir(parents=True, exist_ok=True) if self.proj_catharsis_dir: self.proj_catharsis_dir.mkdir(parents=True, exist_ok=True) if self.proj_genres_dir: self.proj_genres_dir.mkdir(parents=True, exist_ok=True) if self.proj_learned_dir: self.proj_learned_dir.mkdir(parents=True, exist_ok=True) def _load_config(self): """加载配置""" if self.config is not None: self._extract_embed_config(self.config) return try: from data_modules.config import DataModulesConfig self.config = DataModulesConfig.from_project_root(self.project_root) self._extract_embed_config(self.config) return except Exception: pass self._embed_base_url = os.getenv("EMBED_BASE_URL", "") self._embed_model = os.getenv("EMBED_MODEL", "") self._embed_api_key = os.getenv("EMBED_API_KEY", "") def _extract_embed_config(self, config): """从配置对象提取 embedding 配置""" self._embed_base_url = getattr(config, 'embed_base_url', "") or os.getenv("EMBED_BASE_URL", "") self._embed_model = getattr(config, 'embed_model', "") or os.getenv("EMBED_MODEL", "") self._embed_api_key = getattr(config, 'embed_api_key', "") or os.getenv("EMBED_API_KEY", "") async def search( self, query: str, top_k: int = 5, layers: Optional[List[RAGLayer]] = None, chunk_type: Optional[str] = None, ) -> List[CrossProjectSearchResult]: """四层检索（向下继承）""" if layers is None: layers = [RAGLayer.NOVEL, RAGLayer.WORKSPACE, RAGLayer.PROJECT, RAGLayer.PLUGIN] all_results = [] tasks_with_layers = [] if RAGLayer.NOVEL in layers: tasks_with_layers.append((RAGLayer.NOVEL, self._search_novel(query, top_k, chunk_type))) if RAGLayer.WORKSPACE in layers and self.ws_rag_dir: tasks_with_layers.append((RAGLayer.WORKSPACE, self._search_workspace(query, top_k, chunk_type))) if RAGLayer.PROJECT in layers and self.proj_rag_dir: tasks_with_layers.append((RAGLayer.PROJECT, self._search_project(query, top_k, chunk_type))) if RAGLayer.PLUGIN in layers: tasks_with_layers.append((RAGLayer.PLUGIN, self._search_plugin(query, top_k, chunk_type))) if tasks_with_layers: tasks = [t[1] for t in tasks_with_layers] layer_results = await asyncio.gather(*tasks) for (layer, _), results in zip(tasks_with_layers, layer_results): if results: for r in results: r.score *= self.weights.get(layer, 1.0) all_results.append(r) all_results.sort(key=lambda x: x.score, reverse=True) return all_results[:top_k] async def _search_novel(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """检索小说私有 RAG""" if not self.novel_vectors_db.exists(): return [] if self._embed_api_key: return await self._novel_vector_search(query, top_k, chunk_type) return await asyncio.to_thread(self._novel_keyword_search, query, top_k, chunk_type) def _novel_keyword_search(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """小说关键词检索""" try: conn = sqlite3.connect(str(self.novel_vectors_db)) cursor = conn.cursor() if chunk_type: cursor.execute(""" SELECT chunk_id, chapter, content, chunk_type, source_file FROM vectors WHERE chunk_type = ? AND content LIKE ? ORDER BY chapter DESC LIMIT ? """, (chunk_type, f"%{query}%", top_k)) else: cursor.execute(""" SELECT chunk_id, chapter, content, chunk_type, source_file FROM vectors WHERE content LIKE ? ORDER BY chapter DESC LIMIT ? """, (f"%{query}%", top_k)) rows = cursor.fetchall() conn.close() keywords = self._extract_keywords(query) results = [] for row in rows: content = row[2] or "" matches = sum(1 for kw in keywords if kw in content) score = matches / max(len(keywords), 1) * 100 results.append(CrossProjectSearchResult( chunk_id=row[0], content=content[:500], score=score, source_layer=RAGLayer.NOVEL, source_project=self.novel_root.name, chapter=row[1], chunk_type=row[3], metadata={"source_file": row[4]} )) results.sort(key=lambda x: x.score, reverse=True) return results[:top_k] except Exception: return [] async def _novel_vector_search(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """小说向量检索""" try: embeddings = await self._embed_texts([query]) if not embeddings: return await asyncio.to_thread(self._novel_keyword_search, query, top_k, chunk_type) query_embedding = embeddings[0] conn = sqlite3.connect(str(self.novel_vectors_db)) cursor = conn.cursor() if chunk_type: cursor.execute(""" SELECT chunk_id, chapter, content, embedding, chunk_type, source_file FROM vectors WHERE chunk_type = ? """, (chunk_type,)) else: cursor.execute("SELECT chunk_id, chapter, content, embedding, chunk_type, source_file FROM vectors") rows = cursor.fetchall() conn.close() results = [] for row in rows: if not row[3]: continue embedding = self._deserialize_embedding(row[3]) score = self._cosine_similarity(query_embedding, embedding) results.append(CrossProjectSearchResult( chunk_id=row[0], content=row[2][:500] if row[2] else "", score=score, source_layer=RAGLayer.NOVEL, source_project=self.novel_root.name, chapter=row[1], chunk_type=row[4], metadata={"source_file": row[5]} )) results.sort(key=lambda x: x.score, reverse=True) return results[:top_k] except Exception: return await asyncio.to_thread(self._novel_keyword_search, query, top_k, chunk_type) async def _search_workspace(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """检索工作空间共享 RAG""" results = [] # 检索 catharsis 模板 if self.ws_catharsis_dir and self.ws_catharsis_dir.exists(): for model_file in self.ws_catharsis_dir.glob("*.md"): try: content = model_file.read_text(encoding="utf-8") keywords = self._extract_keywords(query) matches = sum(1 for kw in keywords if kw in content[:1000]) if matches > 0: score = matches / len(keywords) * 70 results.append(CrossProjectSearchResult( chunk_id=f"workspace:{model_file.stem}", content=content[:500], score=score, source_layer=RAGLayer.WORKSPACE, source_project="workspace", chapter=None, chunk_type="catharsis_model" )) except Exception: continue # 检索题材库 if self.ws_genres_dir and self.ws_genres_dir.exists(): keywords = self._extract_keywords(query) for genre_file in self.ws_genres_dir.glob("**/*.md"): try: content = genre_file.read_text(encoding="utf-8") matches = sum(1 for kw in keywords if kw in content[:1000]) if matches > 0: score = matches / len(keywords) * 50 results.append(CrossProjectSearchResult( chunk_id=f"workspace_genre:{genre_file.stem}", content=content[:500], score=score, source_layer=RAGLayer.WORKSPACE, source_project="workspace", chapter=None, chunk_type="genre_template" )) except Exception: continue # 检索学习成果 if self.ws_learned_dir and self.ws_learned_dir.exists(): results.extend(await self._search_learned_dir(self.ws_learned_dir, query, top_k, RAGLayer.WORKSPACE)) return results[:top_k] async def _search_project(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """检索工程共享 RAG""" results = [] if self.proj_catharsis_dir and self.proj_catharsis_dir.exists(): for model_file in self.proj_catharsis_dir.glob("*.md"): try: content = model_file.read_text(encoding="utf-8") keywords = self._extract_keywords(query) matches = sum(1 for kw in keywords if kw in content[:1000]) if matches > 0: score = matches / len(keywords) * 50 results.append(CrossProjectSearchResult( chunk_id=f"project:{model_file.stem}", content=content[:500], score=score, source_layer=RAGLayer.PROJECT, source_project="project", chapter=None, chunk_type="catharsis_model" )) except Exception: continue if self.proj_genres_dir and self.proj_genres_dir.exists(): keywords = self._extract_keywords(query) for genre_file in self.proj_genres_dir.glob("**/*.md"): try: content = genre_file.read_text(encoding="utf-8") matches = sum(1 for kw in keywords if kw in content[:1000]) if matches > 0: score = matches / len(keywords) * 30 results.append(CrossProjectSearchResult( chunk_id=f"project_genre:{genre_file.stem}", content=content[:500], score=score, source_layer=RAGLayer.PROJECT, source_project="project", chapter=None, chunk_type="genre_template" )) except Exception: continue if self.proj_learned_dir and self.proj_learned_dir.exists(): results.extend(await self._search_learned_dir(self.proj_learned_dir, query, top_k, RAGLayer.PROJECT)) return results[:top_k] async def _search_plugin(self, query: str, top_k: int, chunk_type: Optional[str]) -> List[CrossProjectSearchResult]: """检索插件内置 RAG""" results = [] if self.plugin_catharsis_dir and self.plugin_catharsis_dir.exists(): for model_file in self.plugin_catharsis_dir.glob("*.md"): try: content = model_file.read_text(encoding="utf-8") keywords = self._extract_keywords(query) matches = sum(1 for kw in keywords if kw in content[:1000]) if matches > 0: score = matches / len(keywords) * 20 results.append(CrossProjectSearchResult( chunk_id=f"plugin:{model_file.stem}", content=content[:500], score=score, source_layer=RAGLayer.PLUGIN, source_project="plugin", chapter=None, chunk_type="catharsis_model" )) except Exception: continue return results[:top_k] async def _search_learned_dir(self, learned_dir: Path, query: str, top_k: int, layer: RAGLayer) -> List[CrossProjectSearchResult]: """检索学习成果目录""" results = [] keywords = self._extract_keywords(query) for pattern_file in learned_dir.glob("**/*.json"): try: data = json.loads(pattern_file.read_text(encoding="utf-8")) content = json.dumps(data, ensure_ascii=False) matches = sum(1 for kw in keywords if kw in content) if matches > 0: score = matches / max(len(keywords), 1) * 50 results.append(CrossProjectSearchResult( chunk_id=f"learned:{pattern_file.stem}", content=content[:500], score=score, source_layer=layer, source_project=data.get("source_project", "unknown"), chapter=data.get("source_chapter"), chunk_type="learned_pattern", metadata=data )) except Exception: continue return results[:top_k] def _extract_keywords(self, query: str) -> List[str]: """提取关键词""" import re chinese = re.findall(r'[\u4e00-\u9fff]{2,8}', query) english = re.findall(r'[a-zA-Z]{2,}', query.lower()) return chinese + english def _cosine_similarity(self, a: List[float], b: List[float]) -> float: dot_product = sum(x * y for x, y in zip(a, b)) norm_a = sum(x * x for x in a) ** 0.5 norm_b = sum(x * x for x in b) ** 0.5 if norm_a == 0 or norm_b == 0: return 0.0 return dot_product / (norm_a * norm_b) def _deserialize_embedding(self, data: bytes) -> List[float]: import struct count = len(data) // 4 return list(struct.unpack(f"{count}f", data)) async def _embed_texts(self, texts: List[str]) -> Optional[List[List[float]]]: """调用 Embedding API""" if not self._embed_api_key or not texts: return None try: import aiohttp headers = {"Authorization": f"Bearer {self._embed_api_key}", "Content-Type": "application/json"} payload = {"model": self._embed_model, "input": texts} async with aiohttp.ClientSession() as session: async with session.post( f"{self._embed_base_url}/embeddings", headers=headers, json=payload, timeout=aiohttp.ClientTimeout(total=60) ) as resp: if resp.status == 200: result = await resp.json() return [item["embedding"] for item in result["data"]] except Exception: pass return None # ==================== 存储接口 ==================== def store_learned_pattern(self, pattern: LearnedPattern, layer: RAGLayer = RAGLayer.NOVEL) -> bool: """存储学习到的模式""" if layer == RAGLayer.NOVEL: return self._store_novel_learned(pattern) elif layer == RAGLayer.WORKSPACE: return self._store_workspace_learned(pattern) elif layer == RAGLayer.PROJECT: return self._store_project_learned(pattern) return False def _store_novel_learned(self, pattern: LearnedPattern) -> bool: """存储到小说私有库""" self._init_learned_db(self.novel_learned_db) try: conn = sqlite3.connect(str(self.novel_learned_db)) cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO learned_patterns (pattern_id, pattern_type, title, description, tension_curve, catharsis_model, structure, hot_spots, style_tags, source_project, source_chapter, learned_at, usage_count, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( pattern.pattern_id, pattern.pattern_type, pattern.title, pattern.description, json.dumps(pattern.tension_curve), pattern.catharsis_model, json.dumps(pattern.structure), json.dumps(pattern.hot_spots), json.dumps(pattern.style_tags), pattern.source_project, pattern.source_chapter, pattern.learned_at, pattern.usage_count, json.dumps(pattern.metadata) )) conn.commit() conn.close() return True except Exception: return False def _store_workspace_learned(self, pattern: LearnedPattern) -> bool: """存储到工作空间共享""" if not self.ws_learned_dir: return False pattern_file = self.ws_learned_dir / f"{pattern.pattern_id}.json" return self._write_pattern_file(pattern_file, pattern) def _store_project_learned(self, pattern: LearnedPattern) -> bool: """存储到工程共享""" if not self.proj_learned_dir: return False pattern_file = self.proj_learned_dir / f"{pattern.pattern_id}.json" return self._write_pattern_file(pattern_file, pattern) def _write_pattern_file(self, path: Path, pattern: LearnedPattern) -> bool: try: path.parent.mkdir(parents=True, exist_ok=True) data = { "pattern_id": pattern.pattern_id, "pattern_type": pattern.pattern_type, "title": pattern.title, "description": pattern.description, "tension_curve": pattern.tension_curve, "catharsis_model": pattern.catharsis_model, "structure": pattern.structure, "hot_spots": pattern.hot_spots, "style_tags": pattern.style_tags, "source_project": pattern.source_project, "source_chapter": pattern.source_chapter, "learned_at": pattern.learned_at, "usage_count": pattern.usage_count, "metadata": pattern.metadata } path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") return True except Exception: return False def _init_learned_db(self, db_path: Path): """初始化学习库""" if db_path.exists(): return conn = sqlite3.connect(str(db_path)) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS learned_patterns ( pattern_id TEXT PRIMARY KEY, pattern_type TEXT NOT NULL, title TEXT NOT NULL, description TEXT, tension_curve TEXT, catharsis_model TEXT, structure TEXT, hot_spots TEXT, style_tags TEXT, source_project TEXT, source_chapter INTEGER, learned_at TEXT, usage_count INTEGER DEFAULT 0, metadata TEXT ) """) conn.commit() conn.close() # ==================== 项目索引 ==================== @staticmethod def get_projects_index(system_root: Path) -> Dict[str, Any]: projects_file = system_root / "projects.json" if projects_file.exists(): return json.loads(projects_file.read_text(encoding="utf-8")) return {"projects": []} @staticmethod def register_project(system_root: Path, project_path: Path, project_info: Dict[str, Any]) -> bool: system_root = Path(system_root) system_root.mkdir(parents=True, exist_ok=True) projects_file = system_root / "projects.json" data = CrossProjectRAG.get_projects_index(system_root) project_path_str = str(project_path.resolve()) projects = data.get("projects", []) for i, p in enumerate(projects): if p.get("path") == project_path_str: projects[i] = project_info break else: projects.append(project_info) data["projects"] = projects data["last_updated"] = datetime.now().isoformat() try: projects_file.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") return True except Exception: return False if __name__ == "__main__": import argparse import sys if sys.platform == "win32": enable_windows_utf8_stdio() parser = argparse.ArgumentParser(description="Cross-Project RAG CLI") parser.add_argument("--project-root", type=str, required=True) parser.add_argument("--workspace-root", type=str) parser.add_argument("--project-root-dir", type=str) subparsers = parser.add_subparsers(dest="command") search_parser = subparsers.add_parser("search") search_parser.add_argument("--query", required=True) search_parser.add_argument("--top-k", type=int, default=5) search_parser.add_argument("--layers", type=str, default="novel,workspace,project,plugin") args = parser.parse_args() if not args.project_root: print("Error: --project-root is required") sys.exit(1) rag = CrossProjectRAG( project_root=Path(args.project_root).resolve(), workspace_root=Path(args.workspace_root).resolve() if args.workspace_root else None, project_root_dir=Path(args.project_root_dir).resolve() if args.project_root_dir else None, ) if args.command == "search": layer_map = {"novel": RAGLayer.NOVEL, "workspace": RAGLayer.WORKSPACE, "project": RAGLayer.PROJECT, "plugin": RAGLayer.PLUGIN} layers = [layer_map[l.strip()] for l in args.layers.split(",") if l.strip() in layer_map] results = asyncio.run(rag.search(args.query, args.top_k, layers)) print(f"\n=== Search Results ({len(results)}) ===") for r in results: print(f"\n[{r.source_layer.value}] {r.chunk_id} (score: {r.score:.2f})") print(f"Source: {r.source_project or 'unknown'}") if r.chapter: print(f"Chapter: {r.chapter}") print(f"Content: {r.content[:200]}...") else: parser.print_help()