Function test_vector_stores_context_window

tests/main/test_vector_stores.py:401–454

Test whether retrieving context-window around matches is working.

(vecdb)

    indirect=True,
)
def test_vector_stores_context_window(vecdb):
    """Test whether retrieving context-window around matches is working."""

    phrases = SimpleNamespace(
        CATS="Cats are quiet and clean.",
        DOGS="Dogs are noisy and messy.",
        GIRAFFES="Giraffes are tall and quiet.",
        ELEPHANTS="Elephants are big and noisy.",
        OWLS="Owls are quiet and nocturnal.",
        BATS="Bats are nocturnal and noisy.",
    )
    text = "\n\n".join(vars(phrases).values())
    doc = Document(content=text, metadata=DocMetaData(id="0"))
    cfg = ParsingConfig(
        splitter=Splitter.SIMPLE,
        n_neighbor_ids=2,
        chunk_size=1,
        max_chunks=20,
        min_chunk_chars=3,
        discard_chunk_chars=1,
    )

    parser = Parser(cfg)
    splits = parser.split([doc])

    vecdb.create_collection(
        collection_name=f"testcw-{os.environ.get('PYTEST_XDIST_WORKER', 'main')}",
        replace=True,
    )
    vecdb.add_documents(splits)

    # Test context window retrieval
    docs_scores = vecdb.similar_texts_with_scores("What are Giraffes like?", k=1)
    docs_scores = vecdb.add_context_window(docs_scores, neighbors=2)

    assert len(docs_scores) == 1
    giraffes, score = docs_scores[0]
    assert all(
        p in giraffes.content
        for p in [
            phrases.CATS,
            phrases.DOGS,
            phrases.GIRAFFES,
            phrases.ELEPHANTS,
            phrases.OWLS,
        ]
    )
    # check they are in the right sequence
    indices = [
        giraffes.content.index(p)
        for p in ["Cats", "Dogs", "Giraffes", "Elephants", "Owls"]
    ]

    assert indices == sorted(indices)


@pytest.mark.parametrize(