80a4be1d21
* refactor(memory): rename package second-brain to memory * refactor(memory): use OpenAI env naming and fix trailing slash * feat(memory): add chunking with env-configurable size and overlap * feat(memory): dedup search results by source in top-K * test(memory): add chunking, batching, dedup, live tests * docs(memory): update README and project map after rename * docs(handoff): add handoff and ADR-032 for memory refactor * docs(handoff): set PR number * fix(ci): reduce index.py complexity to rank A --------- Co-authored-by: opencode-agent <agent@opencode.local>
130 lines
4.6 KiB
Python
130 lines
4.6 KiB
Python
import json
|
|
import os
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
|
|
|
from src.memory.index import _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)
|