Aman Nalakath
7503a441c2
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from kg_ocr.embeddings import create_and_index
|
|
from kg_ocr.rag import retrieve
|
|
|
|
|
|
class _FakeEmbeddings:
|
|
instances: list[_FakeEmbeddings] = []
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
self.indexed: list[str] | None = None
|
|
_FakeEmbeddings.instances.append(self)
|
|
|
|
def index(self, data):
|
|
self.indexed = data
|
|
|
|
|
|
def _install_fake_txtai(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_txtai = types.ModuleType("txtai")
|
|
fake_embeddings_mod = types.ModuleType("txtai.embeddings")
|
|
fake_embeddings_mod.Embeddings = _FakeEmbeddings # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "txtai", fake_txtai)
|
|
monkeypatch.setitem(sys.modules, "txtai.embeddings", fake_embeddings_mod)
|
|
|
|
|
|
def test_create_and_index_config_and_data(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
_install_fake_txtai(monkeypatch)
|
|
emb = create_and_index(["alpha", "beta"], model="some-model")
|
|
assert emb.config["path"] == "some-model"
|
|
assert emb.config["content"] is True
|
|
assert emb.config["hybrid"] is True
|
|
assert emb.indexed == ["alpha", "beta"]
|
|
|
|
|
|
def test_create_and_index_raises_without_txtai(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setitem(sys.modules, "txtai", None)
|
|
monkeypatch.setitem(sys.modules, "txtai.embeddings", None)
|
|
with pytest.raises(ImportError, match="txtai"):
|
|
create_and_index(["x"])
|
|
|
|
|
|
def test_retrieve_passes_through() -> None:
|
|
class FakeEmb:
|
|
def search(self, query, limit):
|
|
return [{"text": f"{query}-{i}", "score": 1.0} for i in range(limit)]
|
|
|
|
out = retrieve(FakeEmb(), "q", limit=2)
|
|
assert out == [{"text": "q-0", "score": 1.0}, {"text": "q-1", "score": 1.0}]
|