a6491df5a6
* refactor(memory): lazy-init embedder env reads * feat(memory): reindex log to .rag/reindex.log * refactor(memory): search handles None embedder * test(memory): cover None embedder in search and index * docs(handoff): add handoff and ADR for lazy-init embedder * docs(handoff): set PR number * docs(handoff): rename to pr-100 prefix for pipeline detection * refactor(memory): split run_search to satisfy xenon rank A * fix(memory): cast search result fields to satisfy mypy * style(memory): ruff format search.py --------- Co-authored-by: opencode-agent <agent@opencode.local>
315 lines
10 KiB
Python
315 lines
10 KiB
Python
import argparse
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from src.memory.embedder import BATCH_SIZE, EMBEDDING_MODEL, embed_texts
|
|
|
|
INDEX_FILENAME = "index.json"
|
|
META_FILENAME = "meta.json"
|
|
LOCK_FILENAME = ".lock"
|
|
LOG_FILENAME = "reindex.log"
|
|
|
|
FileEntry = dict[str, str | int]
|
|
IndexedEntry = dict[str, str | int | list[float]]
|
|
FileMap = dict[str, tuple[str, list[FileEntry]]]
|
|
FilesMeta = dict[str, dict[str, str | int]]
|
|
Meta = dict[str, str | FilesMeta]
|
|
|
|
|
|
def _extract_text(content: str) -> str:
|
|
if content.startswith("---"):
|
|
parts = content.split("---", 2)
|
|
if len(parts) >= 3:
|
|
return parts[2].strip()
|
|
return content.strip()
|
|
|
|
|
|
def _chunk_text(text: str, size: int, overlap: int) -> list[tuple[str, int]]:
|
|
if not text or size <= overlap or len(text) <= size:
|
|
return [(text, 0)] if text else []
|
|
step = size - overlap
|
|
return [(text[i : i + size], i) for i in range(0, len(text), step)]
|
|
|
|
|
|
def _content_hash(text: str) -> str:
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _file_entry(rel: Path, chunk_idx: int, offset: int, chunk: str) -> FileEntry:
|
|
return {
|
|
"source": str(rel),
|
|
"chunk_idx": chunk_idx,
|
|
"offset": offset,
|
|
"text": chunk[:500],
|
|
}
|
|
|
|
|
|
def _atomic_write(path: Path, content: str) -> None:
|
|
tmp = path.with_suffix(path.suffix + ".tmp")
|
|
tmp.write_text(content, encoding="utf-8")
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def _log_line(log_path: Path, message: str) -> None:
|
|
ts = datetime.now(tz=UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
with open(log_path, "a", encoding="utf-8") as f:
|
|
f.write(f"{ts} {message}\n")
|
|
|
|
|
|
def _load_meta(meta_path: Path) -> Meta:
|
|
if not meta_path.exists():
|
|
return {}
|
|
data: object = json.loads(meta_path.read_text(encoding="utf-8"))
|
|
if isinstance(data, dict):
|
|
return data
|
|
return {}
|
|
|
|
|
|
def _build_file_map(
|
|
memory_dir: Path, md_files: list[Path], chunk_size: int, chunk_overlap: int
|
|
) -> FileMap:
|
|
file_map: FileMap = {}
|
|
for fpath in md_files:
|
|
rel = fpath.relative_to(memory_dir)
|
|
content = fpath.read_text(encoding="utf-8")
|
|
text = _extract_text(content)
|
|
chunks = _chunk_text(text, chunk_size, chunk_overlap) or [("", 0)]
|
|
entries = [
|
|
_file_entry(rel, chunk_idx, offset, chunk)
|
|
for chunk_idx, (chunk, offset) in enumerate(chunks)
|
|
]
|
|
file_map[str(rel)] = (_content_hash(text), entries)
|
|
return file_map
|
|
|
|
|
|
def _current_version(chunk_size: int, chunk_overlap: int) -> str:
|
|
return f"{EMBEDDING_MODEL}:{chunk_size}:{chunk_overlap}"
|
|
|
|
|
|
def _split_changed_unchanged(
|
|
file_map: FileMap,
|
|
old_files_meta: FilesMeta,
|
|
needs_full_reindex: bool,
|
|
) -> tuple[list[str], list[str]]:
|
|
changed: list[str] = []
|
|
unchanged: list[str] = []
|
|
for rel, (sha, _entries) in file_map.items():
|
|
old = old_files_meta.get(rel)
|
|
if old is None or old.get("sha256") != sha or needs_full_reindex:
|
|
changed.append(rel)
|
|
else:
|
|
unchanged.append(rel)
|
|
return changed, unchanged
|
|
|
|
|
|
def _load_old_entries_by_source(index_path: Path) -> dict[str, list[FileEntry]]:
|
|
if not index_path.exists():
|
|
return {}
|
|
old_index: object = json.loads(index_path.read_text(encoding="utf-8"))
|
|
if not isinstance(old_index, dict):
|
|
return {}
|
|
files: object = old_index.get("files", [])
|
|
if not isinstance(files, list):
|
|
return {}
|
|
by_source: dict[str, list[FileEntry]] = {}
|
|
for entry in files:
|
|
if isinstance(entry, dict):
|
|
src = str(entry.get("source", ""))
|
|
by_source.setdefault(src, []).append(entry)
|
|
return by_source
|
|
|
|
|
|
def _collect_kept_entries(
|
|
unchanged: list[str], old_entries_by_source: dict[str, list[FileEntry]]
|
|
) -> list[FileEntry]:
|
|
kept: list[FileEntry] = []
|
|
for rel in unchanged:
|
|
kept.extend(old_entries_by_source.get(rel, []))
|
|
return kept
|
|
|
|
|
|
def _prepare_changed(
|
|
changed: list[str],
|
|
file_map: FileMap,
|
|
memory_dir: Path,
|
|
chunk_size: int,
|
|
chunk_overlap: int,
|
|
) -> tuple[list[str], list[FileEntry], FilesMeta]:
|
|
texts: list[str] = []
|
|
entries: list[FileEntry] = []
|
|
files_meta: FilesMeta = {}
|
|
for rel in changed:
|
|
sha, file_entries = file_map[rel]
|
|
content = Path(memory_dir, rel).read_text(encoding="utf-8")
|
|
text = _extract_text(content)
|
|
chunks = _chunk_text(text, chunk_size, chunk_overlap) or [("", 0)]
|
|
for entry in file_entries:
|
|
idx = int(entry["chunk_idx"])
|
|
texts.append(chunks[idx][0] if 0 <= idx < len(chunks) else "")
|
|
entries.append(entry)
|
|
files_meta[rel] = {"sha256": sha, "chunks": len(file_entries)}
|
|
return texts, entries, files_meta
|
|
|
|
|
|
def _build_full_meta(
|
|
file_map: FileMap,
|
|
changed_meta: FilesMeta,
|
|
unchanged: list[str],
|
|
) -> FilesMeta:
|
|
full: FilesMeta = dict(changed_meta)
|
|
for rel in unchanged:
|
|
sha, entries = file_map[rel]
|
|
full[rel] = {"sha256": sha, "chunks": len(entries)}
|
|
return full
|
|
|
|
|
|
def _merge_and_sort(
|
|
kept: list[FileEntry],
|
|
new_entries: list[FileEntry],
|
|
new_embeddings: list[list[float]],
|
|
) -> list[IndexedEntry]:
|
|
merged: list[IndexedEntry] = [dict(e) for e in kept]
|
|
for entry, emb in zip(new_entries, new_embeddings, strict=False):
|
|
merged.append({**entry, "embedding": emb})
|
|
|
|
def _sort_key(e: IndexedEntry) -> tuple[str, int]:
|
|
chunk_idx = e.get("chunk_idx", 0)
|
|
return str(e["source"]), int(chunk_idx) if isinstance(chunk_idx, int) else 0
|
|
|
|
merged.sort(key=_sort_key)
|
|
return merged
|
|
|
|
|
|
def run_index(args: argparse.Namespace) -> None:
|
|
memory_dir = Path(args.memory_dir)
|
|
output_dir = Path(args.output)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
md_files = [f for f in sorted(memory_dir.rglob("*.md")) if ".rag" not in f.parts]
|
|
if not md_files:
|
|
print("No .md files found")
|
|
return
|
|
|
|
chunk_size = int(os.environ.get("MEMORY_CHUNK_SIZE", "512"))
|
|
chunk_overlap = int(os.environ.get("MEMORY_CHUNK_OVERLAP", "64"))
|
|
|
|
lock_path = output_dir / LOCK_FILENAME
|
|
lock_path.touch()
|
|
with open(lock_path, "w") as lock_file:
|
|
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
|
|
_run_index_locked(
|
|
output_dir=output_dir,
|
|
md_files=md_files,
|
|
memory_dir=memory_dir,
|
|
chunk_config=(chunk_size, chunk_overlap),
|
|
)
|
|
|
|
|
|
def _embed_in_batches(
|
|
texts: list[str],
|
|
log_path: Path,
|
|
) -> list[list[float]] | None:
|
|
n_batches = (len(texts) + BATCH_SIZE - 1) // BATCH_SIZE
|
|
results: list[list[float]] = []
|
|
for batch_idx, i in enumerate(range(0, len(texts), BATCH_SIZE), start=1):
|
|
batch = texts[i : i + BATCH_SIZE]
|
|
_log_line(log_path, f"batch {batch_idx}/{n_batches}, model={EMBEDDING_MODEL}, start")
|
|
chunk = embed_texts(batch)
|
|
if chunk is None:
|
|
_log_line(
|
|
log_path,
|
|
f"batch {batch_idx}/{n_batches}, model={EMBEDDING_MODEL}, status=error",
|
|
)
|
|
return None
|
|
_log_line(log_path, f"batch {batch_idx}/{n_batches}, model={EMBEDDING_MODEL}, status=ok")
|
|
results.extend(chunk)
|
|
return results
|
|
|
|
|
|
def _run_index_locked(
|
|
*,
|
|
output_dir: Path,
|
|
md_files: list[Path],
|
|
memory_dir: Path,
|
|
chunk_config: tuple[int, int],
|
|
) -> None:
|
|
log_path = output_dir / LOG_FILENAME
|
|
start_ts = time.monotonic()
|
|
chunk_size, chunk_overlap = chunk_config
|
|
current_version = _current_version(chunk_size, chunk_overlap)
|
|
index_path = output_dir / INDEX_FILENAME
|
|
meta_path = output_dir / META_FILENAME
|
|
meta = _load_meta(meta_path)
|
|
meta_version = meta.get("version")
|
|
needs_full_reindex = meta_path.exists() and meta_version != current_version
|
|
if needs_full_reindex:
|
|
print("Index version mismatch, full reindex")
|
|
|
|
file_map = _build_file_map(memory_dir, md_files, chunk_size, chunk_overlap)
|
|
raw_files = meta.get("files", {})
|
|
old_files_meta: FilesMeta = raw_files if isinstance(raw_files, dict) else {}
|
|
|
|
changed_files, unchanged_files = _split_changed_unchanged(
|
|
file_map, old_files_meta, needs_full_reindex
|
|
)
|
|
deleted_files = [rel for rel in old_files_meta if rel not in file_map]
|
|
|
|
_log_line(
|
|
log_path,
|
|
f"start: changed={len(changed_files)} total={len(md_files)}",
|
|
)
|
|
|
|
old_entries_by_source = _load_old_entries_by_source(index_path)
|
|
kept_entries = _collect_kept_entries(unchanged_files, old_entries_by_source)
|
|
|
|
if not changed_files and not deleted_files:
|
|
_log_line(log_path, f"done: changed=0 total={len(md_files)} took=0s noop")
|
|
print(f"No changes detected ({len(md_files)} files, {len(unchanged_files)} unchanged)")
|
|
return
|
|
|
|
new_texts, new_file_entries, changed_meta = _prepare_changed(
|
|
changed_files, file_map, memory_dir, chunk_size, chunk_overlap
|
|
)
|
|
|
|
if changed_files:
|
|
print(f"Embedding {len(new_texts)} chunks...", flush=True)
|
|
new_embeddings = _embed_in_batches(new_texts, log_path)
|
|
if new_embeddings is None:
|
|
_log_line(
|
|
log_path,
|
|
"failed: embeddings unavailable (OPENAI_BASE_URL not set or API error)",
|
|
)
|
|
print(
|
|
"embeddings unavailable (OPENAI_BASE_URL not set or API error), "
|
|
"skipping semantic index"
|
|
)
|
|
return
|
|
else:
|
|
new_embeddings = []
|
|
|
|
full_meta = _build_full_meta(file_map, changed_meta, unchanged_files)
|
|
merged = _merge_and_sort(kept_entries, new_file_entries, new_embeddings)
|
|
|
|
_atomic_write(index_path, json.dumps({"files": merged}, ensure_ascii=False))
|
|
_atomic_write(
|
|
meta_path,
|
|
json.dumps({"version": current_version, "files": full_meta}, ensure_ascii=False),
|
|
)
|
|
|
|
elapsed = int(time.monotonic() - start_ts)
|
|
_log_line(
|
|
log_path,
|
|
f"done: total={len(md_files)} took={elapsed}s "
|
|
f"changed={len(changed_files)} unchanged={len(unchanged_files)} "
|
|
f"deleted={len(deleted_files)}",
|
|
)
|
|
print(
|
|
f"Indexed {len(md_files)} files to {index_path} "
|
|
f"({len(changed_files)} changed, {len(unchanged_files)} unchanged, "
|
|
f"{len(deleted_files)} deleted)"
|
|
)
|