73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
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):
|
|
# 跳过空片段，避免产生空向量噪声 chunk
|
|
if piece.strip():
|
|
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 not text.strip():
|
|
continue
|
|
# 含 NUL 字节的疑似二进制文件（误带白名单扩展名）跳过
|
|
if "\x00" in text:
|
|
continue
|
|
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)
|