feat: initial commit
This commit is contained in:
@@ -0,0 +1,726 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user