272 lines
9.5 KiB
Python
272 lines
9.5 KiB
Python
import pytest
|
|
from datasets import Dataset
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.store.models import SearchResult
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
|
|
"""Test that documents can be found by searching with their associated questions."""
|
|
# Create client
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
# Load unique documents (limited to 10)
|
|
seen_documents = set()
|
|
documents = []
|
|
|
|
for doc_data in qa_corpus:
|
|
if len(seen_documents) >= 10:
|
|
break
|
|
document_text = doc_data["document_extracted"] # type: ignore
|
|
document_id = doc_data.get("document_id", "") # type: ignore
|
|
|
|
if document_id in seen_documents:
|
|
continue
|
|
seen_documents.add(document_id)
|
|
|
|
# Create the document with chunks and embeddings
|
|
created_document = await client.create_document(content=document_text)
|
|
documents.append((created_document, doc_data))
|
|
|
|
# Test with first few unique documents
|
|
|
|
for target_document, doc_data in documents:
|
|
question = doc_data["question"]
|
|
|
|
# Test vector search
|
|
vector_results = await client.chunk_repository.search(
|
|
question, limit=5, search_type="vector"
|
|
)
|
|
target_document_ids = {chunk.document_id for chunk, _ in vector_results}
|
|
assert target_document.id in target_document_ids
|
|
|
|
# Test FTS search
|
|
fts_results = await client.chunk_repository.search(
|
|
question, limit=5, search_type="fts"
|
|
)
|
|
target_document_ids = {chunk.document_id for chunk, _ in fts_results}
|
|
assert target_document.id in target_document_ids
|
|
|
|
# Test hybrid search
|
|
hybrid_results = await client.chunk_repository.search(
|
|
question, limit=5, search_type="hybrid"
|
|
)
|
|
target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
|
|
assert target_document.id in target_document_ids
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chunks_include_document_info(temp_db_path):
|
|
"""Test that search results include document URI and metadata."""
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
# Create a document with URI and metadata
|
|
created_document = await client.create_document(
|
|
content="This is a test document with some content for searching.",
|
|
uri="https://example.com/test.html",
|
|
metadata={"title": "Test Document", "author": "Test Author"},
|
|
)
|
|
|
|
# Search for chunks
|
|
results = await client.chunk_repository.search(
|
|
"test document", limit=1, search_type="hybrid"
|
|
)
|
|
|
|
assert len(results) > 0
|
|
chunk, score = results[0]
|
|
|
|
# Test that score is valid
|
|
assert isinstance(score, int | float), f"Score should be numeric, got {type(score)}"
|
|
assert score >= 0, f"Score should be non-negative, got {score}"
|
|
|
|
# Verify the chunk includes document information
|
|
assert chunk.document_uri == "https://example.com/test.html"
|
|
assert chunk.document_meta == {"title": "Test Document", "author": "Test Author"}
|
|
assert chunk.document_id == created_document.id
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chunks_include_document_title(temp_db_path):
|
|
"""Test that search results include the parent document title when present."""
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
# Create a document with URI and title
|
|
await client.create_document(
|
|
content="This is a test document with a custom title to verify enrichment.",
|
|
uri="file:///tmp/title-test.md",
|
|
title="My Custom Title",
|
|
)
|
|
|
|
# Perform a search that should find this document
|
|
results = await client.chunk_repository.search(
|
|
"custom title", limit=3, search_type="hybrid"
|
|
)
|
|
|
|
assert results, "Expected at least one search result"
|
|
for chunk, _ in results:
|
|
# All returned chunks for this doc should carry the document title
|
|
if chunk.document_uri == "file:///tmp/title-test.md":
|
|
assert chunk.document_title == "My Custom Title"
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_score_types(temp_db_path):
|
|
"""Test that different search types return appropriate score ranges."""
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
# Create multiple documents with different content
|
|
documents_content = [
|
|
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
|
|
"Deep learning neural networks can process complex datasets and identify hidden patterns.",
|
|
"Natural language processing enables computers to understand and generate human text.",
|
|
"Computer vision systems can interpret and analyze visual information from images.",
|
|
]
|
|
|
|
for content in documents_content:
|
|
await client.create_document(content=content)
|
|
|
|
query = "machine learning"
|
|
|
|
# Test vector search scores (should be converted from distances)
|
|
vector_results = await client.chunk_repository.search(
|
|
query, limit=3, search_type="vector"
|
|
)
|
|
assert len(vector_results) > 0
|
|
vector_scores = [score for _, score in vector_results]
|
|
|
|
# Test FTS search scores (should be native LanceDB FTS scores)
|
|
fts_results = await client.chunk_repository.search(
|
|
query, limit=3, search_type="fts"
|
|
)
|
|
assert len(fts_results) > 0
|
|
fts_scores = [score for _, score in fts_results]
|
|
|
|
# Test hybrid search scores (should be native LanceDB relevance scores)
|
|
hybrid_results = await client.chunk_repository.search(
|
|
query, limit=3, search_type="hybrid"
|
|
)
|
|
assert len(hybrid_results) > 0
|
|
hybrid_scores = [score for _, score in hybrid_results]
|
|
|
|
# All scores should be numeric and non-negative
|
|
for scores, search_type in [
|
|
(vector_scores, "vector"),
|
|
(fts_scores, "fts"),
|
|
(hybrid_scores, "hybrid"),
|
|
]:
|
|
for score in scores:
|
|
assert isinstance(score, int | float), (
|
|
f"{search_type} score should be numeric"
|
|
)
|
|
assert score >= 0, f"{search_type} score should be non-negative"
|
|
|
|
# Vector scores should typically be small (0-1 range due to distance conversion)
|
|
assert all(0 <= score <= 1 for score in vector_scores), (
|
|
"Vector scores should be in 0-1 range"
|
|
)
|
|
|
|
# Scores should be sorted in descending order (most relevant first)
|
|
for scores, search_type in [
|
|
(vector_scores, "vector"),
|
|
(fts_scores, "fts"),
|
|
(hybrid_scores, "hybrid"),
|
|
]:
|
|
for i in range(len(scores) - 1):
|
|
assert scores[i] >= scores[i + 1], (
|
|
f"{search_type} results should be sorted by score descending"
|
|
)
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_returns_search_result(temp_db_path):
|
|
"""Test that client.search() returns SearchResult with provenance info."""
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
await client.create_document(
|
|
content="Machine learning models can classify images with high accuracy.",
|
|
uri="https://example.com/ml.html",
|
|
title="ML Guide",
|
|
)
|
|
|
|
results = await client.search("machine learning", limit=3)
|
|
|
|
assert len(results) > 0
|
|
result = results[0]
|
|
assert isinstance(result, SearchResult)
|
|
assert result.content
|
|
assert result.score > 0
|
|
assert result.document_uri == "https://example.com/ml.html"
|
|
assert result.document_title == "ML Guide"
|
|
# page_numbers and headings come from chunk metadata
|
|
assert isinstance(result.page_numbers, list)
|
|
assert isinstance(result.labels, list)
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_graceful_degradation(temp_db_path):
|
|
"""Test search works when docling data is unavailable."""
|
|
from haiku.rag.store.models import Chunk
|
|
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
|
|
|
|
# Import document with custom chunks (no docling document)
|
|
custom_chunks = [
|
|
Chunk(content="Custom chunk without docling metadata", metadata={}),
|
|
]
|
|
await client.import_document(
|
|
content="Document with custom chunks",
|
|
chunks=custom_chunks,
|
|
uri="https://example.com/custom.html",
|
|
)
|
|
|
|
results = await client.search("custom chunk", limit=3)
|
|
|
|
assert len(results) > 0
|
|
result = results[0]
|
|
assert isinstance(result, SearchResult)
|
|
assert result.content
|
|
# Metadata defaults should still work
|
|
assert result.page_numbers == []
|
|
assert result.labels == []
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_result_format_includes_metadata(temp_db_path):
|
|
"""Test that formatted search results include document metadata."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="Important information about machine learning algorithms.",
|
|
title="ML Guide",
|
|
uri="https://example.com/ml-guide",
|
|
)
|
|
|
|
results = await client.search("machine learning", limit=1)
|
|
assert len(results) > 0
|
|
|
|
formatted = results[0].format_for_agent()
|
|
|
|
# Should include chunk ID and score
|
|
assert "[" in formatted and "]" in formatted
|
|
assert "score:" in formatted
|
|
|
|
# Should include document title in Source
|
|
assert "ML Guide" in formatted
|
|
assert "Source:" in formatted
|
|
|
|
# Should include content
|
|
assert "Content:" in formatted
|
|
assert "machine learning" in formatted.lower()
|