danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix scores depending on search type

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-01 12:40:46 +03:00
parent 628bdc8151
commit a5e087aa53
No known key found for this signature in database
2 changed files with 120 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
src/haiku/rag/store/repositories/chunk.py
View file

@ -185,11 +185,9 @@ class ChunkRepository:
if search_type == "vector":
query_embedding = await self.embedder.embed(query)
results = (
self.store.chunks_table.search(query_embedding, query_type="vector")
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
results = self.store.chunks_table.search(
query_embedding, query_type="vector"
).limit(limit)
return await self._process_search_results(results)
@ -197,10 +195,8 @@ class ChunkRepository:
# Ensure FTS index exists
self._ensure_fts_index()
results = (
self.store.chunks_table.search(query, query_type="fts")
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
results = self.store.chunks_table.search(query, query_type="fts").limit(
limit
)
return await self._process_search_results(results)
@ -220,7 +216,6 @@ class ChunkRepository:
.text(query)
.rerank(reranker)
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
return await self._process_search_results(results)
@ -279,11 +274,31 @@ class ChunkRepository:
return adjacent_chunks
async def _process_search_results(self, results) -> list[tuple[Chunk, float]]:
async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
"""Process search results into chunks with document info and scores."""
chunks_with_scores = []
for chunk_record in results:
# Get both arrow and pydantic results to access scores
arrow_result = query_result.to_arrow()
pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
# Extract scores from arrow result based on search type
scores = []
column_names = arrow_result.column_names
if "_distance" in column_names:
# Vector search - distance (lower is better, convert to similarity)
distances = arrow_result.column("_distance").to_pylist()
scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
elif "_relevance_score" in column_names:
# Hybrid search - relevance score (higher is better)
scores = arrow_result.column("_relevance_score").to_pylist()
elif "_score" in column_names:
# FTS search - score (higher is better)
scores = arrow_result.column("_score").to_pylist()
else:
raise ValueError("Unknown search result format, cannot extract scores")
for i, chunk_record in enumerate(pydantic_results):
# Get document info
doc_results = list(
self.store.documents_table.search()
@ -306,12 +321,8 @@ class ChunkRepository:
document_meta=json.loads(doc_meta) if doc_meta else {},
)
# Get distance score - LanceDB returns _distance (lower is better)
distance = getattr(chunk_record, "_distance", 1.0)
# Convert distance to similarity score (higher is better)
# Using exponential decay to convert distance to similarity
score = max(0.0, 1.0 / (1.0 + distance))
# Get score from arrow result
score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score))

108
tests/test_search.py
View file

@ -15,23 +15,22 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# Load first 20 documents with embeddings (reduced for faster testing)
num_documents = 20
# Load unique documents (limited to 10)
seen_documents = set()
documents = []
for i in range(num_documents):
doc_data = qa_corpus[i]
document_text = doc_data["document_extracted"]
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 a Document instance
document = Document(
content=document_text,
metadata={
"source": "qa_corpus",
"topic": doc_data.get("document_topic", ""),
"document_id": doc_data.get("document_id", ""),
"question": doc_data["question"],
},
)
document = Document(content=document_text)
# Create the document with chunks and embeddings
from haiku.rag.utils import text_to_docling_document
@ -42,8 +41,9 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
)
documents.append((created_document, doc_data))
for i in range(3): # Test with first few documents
target_document, doc_data = documents[i]
# Test with first few unique documents
for target_document, doc_data in documents:
question = doc_data["question"]
# Test vector search
@ -92,7 +92,11 @@ async def test_chunks_include_document_info(temp_db_path):
results = await chunk_repo.search("test document", limit=1, search_type="hybrid")
assert len(results) > 0
chunk, _ = 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"
@ -100,3 +104,73 @@ async def test_chunks_include_document_info(temp_db_path):
assert chunk.document_id == created_document.id
store.close()
@pytest.mark.asyncio
async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# 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:
document = Document(content=content)
from haiku.rag.utils import text_to_docling_document
docling_document = text_to_docling_document(content, name="test.md")
await doc_repo._create_with_docling(document, docling_document)
query = "machine learning"
# Test vector search scores (should be converted from distances)
vector_results = await chunk_repo.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 chunk_repo.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 chunk_repo.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"
)
store.close()