dkma-labor-arbitration/tests/test_retrieval.py

from dkma_labor_arbitration.kb.vector_store import SimpleVectorStore
from dkma_labor_arbitration.schemas import LegalProvision


def test_simple_vector_store_search():
    store = SimpleVectorStore()
    store.add([
        LegalProvision("劳动合同法", "第八十二条", "未签订书面劳动合同，支付二倍工资。"),
        LegalProvision("劳动合同法", "第三十条", "及时足额支付劳动报酬。"),
    ])
    result = store.search("未签劳动合同", top_k=1)
    assert result
    assert "合同" in result[0].provision.content