forked from innovacion/searchbox
63 lines
1.6 KiB
Python
63 lines
1.6 KiB
Python
from uuid import uuid4
|
|
import pytest
|
|
from searchbox.engine.qdrant_engine import QdrantEngine
|
|
from searchbox.models import Chunk, Match, MatchAny, MatchExclude
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def engine():
|
|
return QdrantEngine(":memory:")
|
|
|
|
|
|
async def test_create_index(engine: QdrantEngine):
|
|
result = await engine.create_index("test_index", 3)
|
|
|
|
assert result is True
|
|
|
|
|
|
async def test_upload_chunk(engine: QdrantEngine):
|
|
chunk = Chunk.model_validate(
|
|
{
|
|
"id": uuid4(),
|
|
"vector": [0.0, 0.1, 0.3],
|
|
"payload": {
|
|
"page_content": "This is a test page content.",
|
|
"filename": "test.txt",
|
|
"page": 1,
|
|
},
|
|
}
|
|
)
|
|
|
|
result = await engine.upload_chunk("test_index", chunk)
|
|
|
|
assert result is True
|
|
|
|
|
|
async def test_search_chunk(engine: QdrantEngine):
|
|
result = await engine.semantic_search([0.0, 0.1, 0.3], "test_index")
|
|
|
|
assert len(result) == 1
|
|
|
|
first_result = result[0]
|
|
assert first_result.chunk_id is not None
|
|
assert first_result.score > 0.9
|
|
assert first_result.payload == {
|
|
"page_content": "This is a test page content.",
|
|
"filename": "test.txt",
|
|
"page": 1,
|
|
}
|
|
|
|
|
|
async def test_search_chunk_with_conditions(engine: QdrantEngine):
|
|
conditions = [
|
|
Match(key="filename", value="test.md"),
|
|
MatchAny(key="filename", any=["test.md", "test.docx"]),
|
|
MatchExclude(key="filename", exclude=["test.txt"]),
|
|
]
|
|
|
|
result = await engine.semantic_search(
|
|
[0.0, 0.1, 0.3], "test_index", conditions=conditions
|
|
)
|
|
|
|
assert len(result) == 0
|