opencode-config/src/memory/index.py

import argparse
import fcntl
import hashlib
import json
import os
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"

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 _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 _run_index_locked(
    *,
    output_dir: Path,
    md_files: list[Path],
    memory_dir: Path,
    chunk_config: tuple[int, int],
) -> None:
    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]

    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:
        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:
        n_batches = (len(new_texts) + BATCH_SIZE - 1) // BATCH_SIZE
        print(f"Embedding {len(new_texts)} chunks in {n_batches} batches...", flush=True)
        new_embeddings = embed_texts(new_texts)
    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),
    )

    print(
        f"Indexed {len(md_files)} files to {index_path} "
        f"({len(changed_files)} changed, {len(unchanged_files)} unchanged, "
        f"{len(deleted_files)} deleted)"
    )