25cf10baa6
* feat(memory): add SHA256 incremental index with atomic writes, versioning, flock * feat(memory): parse Retry-After header and add batch progress logging * test(memory): add incremental, atomic, versioning, retry, batch tests * docs(memory): update .env.example with OpenRouter defaults * docs(handoff): add handoff and ADR-036 for incremental index * docs(handoff): set PR number * docs(project-map): update index.py and embedder.py descriptions for PR#83 * fix(ci): skip index rewrite on no-op + versioning first-run * fix(ci): ruff format index.py --------- Co-authored-by: opencode-agent <agent@opencode.local>
352 lines
13 KiB
Python
352 lines
13 KiB
Python
import json
|
|
import os
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
|
|
|
from src.memory.index import _atomic_write, _content_hash, _extract_text, run_index
|
|
|
|
|
|
class TestExtractText:
|
|
def test_with_frontmatter(self) -> None:
|
|
md = "---\ntitle: test\n---\n\nbody content"
|
|
assert _extract_text(md) == "body content"
|
|
|
|
def test_no_frontmatter(self) -> None:
|
|
md = "just content"
|
|
assert _extract_text(md) == "just content"
|
|
|
|
def test_empty(self) -> None:
|
|
assert _extract_text("") == ""
|
|
|
|
|
|
class TestRunIndex:
|
|
def test_index_creates_json(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
|
|
(memory_dir / "test.md").write_text("---\ntitle: test\n---\n\nhello world")
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[0.1, 0.2, 0.3]] * len(texts)
|
|
|
|
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
|
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
|
run_index(args)
|
|
|
|
assert (index_dir / "index.json").exists()
|
|
index = json.loads((index_dir / "index.json").read_text())
|
|
assert len(index["files"]) == 1
|
|
assert index["files"][0]["source"] == "test.md"
|
|
assert index["files"][0]["embedding"] == [0.1, 0.2, 0.3]
|
|
assert index["files"][0]["chunk_idx"] == 0
|
|
assert index["files"][0]["offset"] == 0
|
|
|
|
def test_no_md_files(self, tmp_path: Path, capsys) -> None:
|
|
memory_dir = tmp_path / "empty"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
|
|
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
|
run_index(args)
|
|
|
|
captured = capsys.readouterr()
|
|
assert "No .md files found" in captured.out
|
|
|
|
def test_chunking_long_file(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
|
|
long_body = "a" * 1500
|
|
(memory_dir / "long.md").write_text(long_body)
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[0.1, 0.2, 0.3]] * len(texts)
|
|
|
|
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
|
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
|
run_index(args)
|
|
|
|
index = json.loads((index_dir / "index.json").read_text())
|
|
assert len(index["files"]) >= 3
|
|
chunk_idxs = [f["chunk_idx"] for f in index["files"]]
|
|
assert chunk_idxs == sorted(chunk_idxs)
|
|
assert chunk_idxs[0] == 0
|
|
assert index["files"][0]["offset"] == 0
|
|
offsets = [f["offset"] for f in index["files"]]
|
|
assert offsets == sorted(offsets)
|
|
|
|
def test_chunking_size_env_override(
|
|
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
|
|
monkeypatch.setenv("MEMORY_CHUNK_SIZE", "256")
|
|
monkeypatch.setenv("MEMORY_CHUNK_OVERLAP", "0")
|
|
(memory_dir / "f.md").write_text("a" * 600)
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[0.1]] * len(texts)
|
|
|
|
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
|
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
|
run_index(args)
|
|
|
|
index = json.loads((index_dir / "index.json").read_text())
|
|
assert len(index["files"]) == 3
|
|
for i, f in enumerate(index["files"]):
|
|
assert f["chunk_idx"] == i
|
|
assert f["offset"] == i * 256
|
|
|
|
def test_chunking_skip_rag_dir(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = memory_dir / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
|
|
(memory_dir / "real.md").write_text("real content")
|
|
(index_dir).mkdir()
|
|
(index_dir / "index.json").write_text('{"files": []}')
|
|
(index_dir / "ignore.md").write_text("should be ignored")
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[0.1]] * len(texts)
|
|
|
|
output_dir = tmp_path / "out"
|
|
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
|
args = Namespace(memory_dir=str(memory_dir), output=str(output_dir))
|
|
run_index(args)
|
|
|
|
index = json.loads((output_dir / "index.json").read_text())
|
|
sources = [f["source"] for f in index["files"]]
|
|
assert "real.md" in sources
|
|
assert all(".rag" not in s for s in sources)
|
|
|
|
|
|
class TestIncrementalIndex:
|
|
def _make_args(self, memory_dir: Path, output_dir: Path) -> Namespace:
|
|
return Namespace(memory_dir=str(memory_dir), output=str(output_dir))
|
|
|
|
def _fake_embed(self):
|
|
def fake_embed_texts(texts):
|
|
return [[0.1, 0.2] for _ in texts]
|
|
|
|
return fake_embed_texts
|
|
|
|
def test_incremental_add(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("file A content")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
(memory_dir / "b.md").write_text("file B content")
|
|
calls: list[list[str]] = []
|
|
|
|
def tracking_embed(texts):
|
|
calls.append(list(texts))
|
|
return [[0.1, 0.2] for _ in texts]
|
|
|
|
with patch("src.memory.index.embed_texts", tracking_embed):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
assert len(calls) == 1
|
|
embedded_text = calls[0][0]
|
|
assert "file B content" in embedded_text
|
|
assert "file A content" not in embedded_text
|
|
|
|
index = json.loads((index_dir / "index.json").read_text())
|
|
sources = {f["source"] for f in index["files"]}
|
|
assert {"a.md", "b.md"} <= sources
|
|
|
|
def test_incremental_edit(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("original content")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
(memory_dir / "a.md").write_text("edited content here")
|
|
calls: list[list[str]] = []
|
|
|
|
def tracking_embed(texts):
|
|
calls.append(list(texts))
|
|
return [[0.1, 0.2] for _ in texts]
|
|
|
|
with patch("src.memory.index.embed_texts", tracking_embed):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
assert len(calls) == 1
|
|
assert "edited content here" in calls[0][0]
|
|
|
|
def test_incremental_delete(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("file A")
|
|
(memory_dir / "b.md").write_text("file B")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
(memory_dir / "b.md").unlink()
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
index = json.loads((index_dir / "index.json").read_text())
|
|
sources = [f["source"] for f in index["files"]]
|
|
assert "a.md" in sources
|
|
assert "b.md" not in sources
|
|
|
|
meta = json.loads((index_dir / "meta.json").read_text())
|
|
assert "b.md" not in meta["files"]
|
|
assert "a.md" in meta["files"]
|
|
|
|
def test_no_changes_noop(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("stable content")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
index_mtime_before = (index_dir / "index.json").stat().st_mtime_ns
|
|
|
|
embed_mock = Mock(return_value=[[0.1, 0.2]])
|
|
with patch("src.memory.index.embed_texts", embed_mock):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
embed_mock.assert_not_called()
|
|
index_mtime_after = (index_dir / "index.json").stat().st_mtime_ns
|
|
assert index_mtime_after == index_mtime_before
|
|
|
|
def test_meta_json_persistence(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("file A")
|
|
(memory_dir / "b.md").write_text("file B")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
meta_path = index_dir / "meta.json"
|
|
assert meta_path.exists()
|
|
meta = json.loads(meta_path.read_text())
|
|
assert "a.md" in meta["files"]
|
|
assert "b.md" in meta["files"]
|
|
assert "version" in meta
|
|
assert meta["files"]["a.md"]["sha256"]
|
|
assert meta["files"]["a.md"]["chunks"] == 1
|
|
|
|
(memory_dir / "c.md").write_text("file C")
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
meta2 = json.loads(meta_path.read_text())
|
|
assert "a.md" in meta2["files"]
|
|
assert "b.md" in meta2["files"]
|
|
assert "c.md" in meta2["files"]
|
|
|
|
def test_sha256_content_hash(self, tmp_path: Path) -> None:
|
|
assert _content_hash("hello") == _content_hash("hello")
|
|
assert _content_hash("hello") != _content_hash("world")
|
|
assert len(_content_hash("hello")) == 64
|
|
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("---\ntitle: T\n---\n\nbody content")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
meta_before = json.loads((index_dir / "meta.json").read_text())
|
|
sha_before = meta_before["files"]["a.md"]["sha256"]
|
|
|
|
os.utime(memory_dir / "a.md", None)
|
|
|
|
embed_mock = Mock(return_value=[[0.1, 0.2]])
|
|
with patch("src.memory.index.embed_texts", embed_mock):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
embed_mock.assert_not_called()
|
|
meta_after = json.loads((index_dir / "meta.json").read_text())
|
|
assert meta_after["files"]["a.md"]["sha256"] == sha_before
|
|
|
|
def test_atomic_writes(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("file A")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
original_index = (index_dir / "index.json").read_text()
|
|
|
|
(memory_dir / "a.md").write_text("file A edited")
|
|
|
|
with (
|
|
patch("src.memory.index.embed_texts", self._fake_embed()),
|
|
patch("src.memory.index.os.replace", side_effect=OSError("simulated crash")),
|
|
pytest.raises(OSError),
|
|
):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
assert (index_dir / "index.json").read_text() == original_index
|
|
|
|
def test_index_versioning(self, tmp_path: Path) -> None:
|
|
memory_dir = tmp_path / "memory"
|
|
index_dir = tmp_path / ".rag"
|
|
memory_dir.mkdir(parents=True)
|
|
(memory_dir / "a.md").write_text("file A")
|
|
|
|
with patch("src.memory.index.embed_texts", self._fake_embed()):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
meta = json.loads((index_dir / "meta.json").read_text())
|
|
meta["version"] = "old-model:512:64"
|
|
(index_dir / "meta.json").write_text(json.dumps(meta))
|
|
|
|
calls: list[int] = []
|
|
|
|
def tracking_embed(texts):
|
|
calls.append(len(texts))
|
|
return [[0.1, 0.2] for _ in texts]
|
|
|
|
with patch("src.memory.index.embed_texts", tracking_embed):
|
|
run_index(self._make_args(memory_dir, index_dir))
|
|
|
|
assert len(calls) == 1
|
|
assert calls[0] >= 1
|
|
|
|
meta2 = json.loads((index_dir / "meta.json").read_text())
|
|
assert meta2["version"] != "old-model:512:64"
|
|
|
|
|
|
class TestAtomicWrite:
|
|
def test_atomic_write_replaces(self, tmp_path: Path) -> None:
|
|
target = tmp_path / "file.json"
|
|
target.write_text("old")
|
|
_atomic_write(target, "new content")
|
|
assert target.read_text() == "new content"
|
|
assert not (tmp_path / "file.json.tmp").exists()
|
|
|
|
def test_atomic_write_creates(self, tmp_path: Path) -> None:
|
|
target = tmp_path / "new.json"
|
|
_atomic_write(target, "fresh")
|
|
assert target.read_text() == "fresh"
|