diff --git a/src/genesis/rag/impact_rag.py b/src/genesis/rag/impact_rag.py new file mode 100644 index 0000000..fa10be8 --- /dev/null +++ b/src/genesis/rag/impact_rag.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from pathlib import Path +from typing import List, Tuple + +from genesis.rag.embeddings import Embedder +from genesis.rag.store import RagStore + +_CHUNK = 800 +_TEXT_EXTS = { + ".java", ".py", ".js", ".ts", ".go", ".kt", ".scala", ".xml", ".yml", + ".yaml", ".md", ".txt", ".json", ".csv", ".sql", ".html", ".css", ".sh", +} + + +def _split(text: str, size: int = _CHUNK) -> List[str]: + paras = [p.strip() for p in text.split("\n") if p.strip()] + out, buf = [], "" + for p in paras: + if len(buf) + len(p) > size and buf: + out.append(buf) + buf = p + else: + buf = (buf + "\n" + p).strip() + if buf: + out.append(buf) + return out or [""] + + +class ImpactRAG: + def __init__(self, store: RagStore, embedder: Embedder): + self.store = store + self.embedder = embedder + + def index(self, scope: str, sources: List[Tuple[str, str]]) -> None: + chunks = [] + for name, text in sources: + for piece in _split(text): + chunks.append(f"[{name}]\n{piece}") + embs = self.embedder.embed(chunks) + self.store.add(scope, chunks, embs) + + def index_dir(self, scope: str, root: str) -> int: + """D1:索引目录内文本源文件(既有系统源码)。返回索引的片段数。""" + root_path = Path(root) + sources: List[Tuple[str, str]] = [] + if root_path.is_dir(): + for p in root_path.rglob("*"): + if p.is_file() and p.suffix.lower() in _TEXT_EXTS: + try: + text = p.read_text(encoding="utf-8", errors="ignore") + except Exception: + continue + if text.strip(): + sources.append((str(p.relative_to(root_path)), text)) + if sources: + self.index(scope, sources) + return len(sources) + + def retrieve(self, scope: str, query: str, k: int = 5) -> List[str]: + if not query.strip(): + return [] + qv = self.embedder.embed([query])[0] + return self.store.search(scope, qv, k) diff --git a/tests/test_impact_rag.py b/tests/test_impact_rag.py new file mode 100644 index 0000000..e6deabd --- /dev/null +++ b/tests/test_impact_rag.py @@ -0,0 +1,35 @@ +from genesis.rag.embeddings import FakeEmbedder +from genesis.rag.store import RagStore +from genesis.rag.impact_rag import ImpactRAG + + +def test_retrieve_returns_relevant_chunk(): + store = RagStore(":memory:") + try: + rag = ImpactRAG(store, FakeEmbedder()) + sources = [ + ("OrderController.java", "public class OrderController { 创建订单 }"), + ("UserAuth.java", "public class UserAuth { 用户登录认证 }"), + ] + rag.index("p1", sources) + res = rag.retrieve("p1", "OrderController 的影响范围", k=1) + assert res and "OrderController" in res[0] + finally: + store.close() + + +def test_index_dir_reads_text_files(tmp_path): + (tmp_path / "src").mkdir() + (tmp_path / "src" / "OrderController.java").write_text( + "public class OrderController { 创建订单 }", encoding="utf-8" + ) + (tmp_path / "src" / "binary.bin").write_bytes(b"\x00\x01") + store = RagStore(":memory:") + try: + rag = ImpactRAG(store, FakeEmbedder()) + n = rag.index_dir("p1", str(tmp_path / "src")) + assert n == 1 + res = rag.retrieve("p1", "OrderController 的影响范围", k=1) + assert res and "OrderController" in res[0] + finally: + store.close()