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>
103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
import json
|
|
import os
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
|
|
|
import src.memory.search as search_mod
|
|
from src.memory.search import _cosine_sim
|
|
|
|
|
|
class TestCosineSim:
|
|
def test_identical(self) -> None:
|
|
v = np.array([1.0, 2.0, 3.0])
|
|
assert _cosine_sim(v, v) == pytest.approx(1.0)
|
|
|
|
def test_orthogonal(self) -> None:
|
|
a = np.array([1.0, 0.0])
|
|
b = np.array([0.0, 1.0])
|
|
assert _cosine_sim(a, b) == pytest.approx(0.0)
|
|
|
|
def test_zero_vector(self) -> None:
|
|
a = np.array([0.0, 0.0])
|
|
b = np.array([1.0, 1.0])
|
|
assert _cosine_sim(a, b) == pytest.approx(0.0)
|
|
|
|
def test_parallel(self) -> None:
|
|
a = np.array([1.0, 2.0])
|
|
b = np.array([2.0, 4.0])
|
|
assert _cosine_sim(a, b) == pytest.approx(1.0)
|
|
|
|
def test_opposite(self) -> None:
|
|
a = np.array([1.0, 1.0])
|
|
b = np.array([-1.0, -1.0])
|
|
assert _cosine_sim(a, b) == pytest.approx(-1.0)
|
|
|
|
|
|
class TestSearchOutput:
|
|
def test_search_json_format(self, tmp_path: Path) -> None:
|
|
index_dir = tmp_path / ".rag"
|
|
index_dir.mkdir()
|
|
index = {
|
|
"files": [
|
|
{
|
|
"source": "test.md",
|
|
"text": "hello world",
|
|
"embedding": [1.0, 0.0, 0.0],
|
|
},
|
|
],
|
|
}
|
|
(index_dir / "index.json").write_text(json.dumps(index))
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[1.0, 0.0, 0.0]]
|
|
|
|
with patch.object(search_mod, "embed_texts", fake_embed_texts):
|
|
args = Namespace(index_dir=str(index_dir), query="hello", k=5, json=True)
|
|
search_mod.run_search(args)
|
|
|
|
def test_dedup_by_source(self, tmp_path: Path, capsys) -> None:
|
|
index_dir = tmp_path / ".rag"
|
|
index_dir.mkdir()
|
|
index = {
|
|
"files": [
|
|
{
|
|
"source": "doc.md",
|
|
"text": "chunk 0",
|
|
"embedding": [1.0, 0.0, 0.0],
|
|
},
|
|
{
|
|
"source": "doc.md",
|
|
"text": "chunk 1",
|
|
"embedding": [0.9, 0.1, 0.0],
|
|
},
|
|
{
|
|
"source": "other.md",
|
|
"text": "chunk 0",
|
|
"embedding": [0.5, 0.5, 0.0],
|
|
},
|
|
],
|
|
}
|
|
(index_dir / "index.json").write_text(json.dumps(index))
|
|
|
|
def fake_embed_texts(texts):
|
|
return [[1.0, 0.0, 0.0]]
|
|
|
|
with patch.object(search_mod, "embed_texts", fake_embed_texts):
|
|
args = Namespace(index_dir=str(index_dir), query="q", k=5, json=True)
|
|
search_mod.run_search(args)
|
|
|
|
captured = capsys.readouterr()
|
|
result = json.loads(captured.out)
|
|
assert len(result) == 2
|
|
sources = [r["source"] for r in result]
|
|
assert "doc.md" in sources
|
|
assert "other.md" in sources
|
|
assert sources.count("doc.md") == 1
|
|
doc_entry = next(r for r in result if r["source"] == "doc.md")
|
|
assert doc_entry["text"] == "chunk 0"
|