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>
170 lines
4.3 KiB
Python
170 lines
4.3 KiB
Python
import importlib
|
|
import os
|
|
|
|
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
|
|
|
from typing import NoReturn
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
import pytest
|
|
import src.memory.embedder as embedder_mod
|
|
from src.memory.embedder import embed_texts
|
|
|
|
|
|
def test_embed_texts_success() -> None:
|
|
fake_embedding = [0.1, 0.2, 0.3]
|
|
|
|
class FakeResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {
|
|
"data": [{"embedding": fake_embedding, "index": 0}],
|
|
"model": "gemini-embedding-2-preview",
|
|
}
|
|
|
|
def raise_for_status(self) -> None:
|
|
pass
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
assert "embeddings" in url
|
|
return FakeResponse()
|
|
|
|
with patch.object(httpx.Client, "post", mock_post):
|
|
result = embed_texts(["test text"])
|
|
|
|
assert len(result) == 1
|
|
assert result[0] == fake_embedding
|
|
|
|
|
|
def test_embed_texts_empty() -> None:
|
|
assert embed_texts([]) == []
|
|
|
|
|
|
def test_embed_texts_api_error() -> None:
|
|
class FakeErrorResponse:
|
|
status_code = 401
|
|
|
|
def raise_for_status(self) -> NoReturn:
|
|
msg = "Unauthorized"
|
|
raise httpx.HTTPStatusError(
|
|
msg,
|
|
request=None,
|
|
response=self,
|
|
)
|
|
|
|
def json(self):
|
|
return {}
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
return FakeErrorResponse()
|
|
|
|
with patch.object(httpx.Client, "post", mock_post), pytest.raises(httpx.HTTPStatusError):
|
|
embed_texts(["test"])
|
|
|
|
|
|
def test_embed_texts_default_model() -> None:
|
|
captured: dict[str, str | list[str]] = {}
|
|
|
|
class FakeResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {"data": [{"embedding": [0.1], "index": 0}]}
|
|
|
|
def raise_for_status(self) -> None:
|
|
pass
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
captured.update(kwargs["json"])
|
|
return FakeResponse()
|
|
|
|
with patch.object(httpx.Client, "post", mock_post):
|
|
embed_texts(["text"])
|
|
|
|
assert captured["model"] == "gemini-embedding-2-preview"
|
|
|
|
|
|
def test_embed_texts_custom_model(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
|
|
importlib.reload(embedder_mod)
|
|
|
|
captured: dict[str, str | list[str]] = {}
|
|
|
|
class FakeResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {"data": [{"embedding": [0.1], "index": 0}]}
|
|
|
|
def raise_for_status(self) -> None:
|
|
pass
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
captured.update(kwargs["json"])
|
|
return FakeResponse()
|
|
|
|
with patch.object(httpx.Client, "post", mock_post):
|
|
embedder_mod.embed_texts(["text"])
|
|
|
|
assert captured["model"] == "text-embedding-3-small"
|
|
|
|
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
|
|
importlib.reload(embedder_mod)
|
|
|
|
|
|
def test_embed_texts_trailing_slash(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENAI_BASE_URL", "https://api.test.com/v1/")
|
|
importlib.reload(embedder_mod)
|
|
|
|
captured_url: dict[str, str] = {}
|
|
|
|
class FakeResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {"data": [{"embedding": [0.1], "index": 0}]}
|
|
|
|
def raise_for_status(self) -> None:
|
|
pass
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
captured_url["url"] = url
|
|
return FakeResponse()
|
|
|
|
with patch.object(httpx.Client, "post", mock_post):
|
|
embedder_mod.embed_texts(["text"])
|
|
|
|
assert captured_url["url"] == "https://api.test.com/v1/embeddings"
|
|
|
|
monkeypatch.setenv("OPENAI_BASE_URL", "http://test/v1")
|
|
importlib.reload(embedder_mod)
|
|
|
|
|
|
def test_embed_texts_batches() -> None:
|
|
calls: list[int] = []
|
|
|
|
class FakeResponse:
|
|
def __init__(self, count: int) -> None:
|
|
self.count = count
|
|
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {"data": [{"embedding": [0.1], "index": i} for i in range(self.count)]}
|
|
|
|
def raise_for_status(self) -> None:
|
|
pass
|
|
|
|
def mock_post(self, url, **kwargs):
|
|
count = len(kwargs["json"]["input"])
|
|
calls.append(count)
|
|
return FakeResponse(count)
|
|
|
|
with patch.object(httpx.Client, "post", mock_post):
|
|
result = embed_texts(["text"] * 3000)
|
|
|
|
assert len(calls) == 2
|
|
assert calls == [2048, 952]
|
|
assert len(result) == 3000
|