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

Cleanup

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-01 11:21:42 +03:00
parent 0592a7d122
commit 628bdc8151
No known key found for this signature in database
8 changed files with 58 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/installation.md
View file

@ -31,5 +31,4 @@ uv pip install haiku.rag[mxbai]
## Requirements
- Python 3.10+
- SQLite 3.38+
- Ollama (for default embeddings)

2
mkdocs.yml
View file

@ -1,5 +1,5 @@
site_name: haiku.rag
site_description: Retrieval-Augmented Generation (RAG) library on SQLite.
site_description: Retrieval-Augmented Generation (RAG) library on LanceDB.
site_url: https://ggozad.github.io/haiku.rag/
theme:
name: material

24
src/haiku/rag/client.py
View file

@ -334,27 +334,13 @@ class HaikuRAG:
if reranker is None:
# No reranking - return direct search results
if search_type == "vector":
return await self.chunk_repository.search_chunks(query, limit)
elif search_type == "fts":
return await self.chunk_repository.search_chunks_fts(query, limit)
else: # hybrid (default)
return await self.chunk_repository.search_chunks_hybrid(query, limit)
return await self.chunk_repository.search(query, limit, search_type)
# Get more initial results (3X) for reranking
search_limit = limit * 3
if search_type == "vector":
search_results = await self.chunk_repository.search_chunks(
query, search_limit
)
elif search_type == "fts":
search_results = await self.chunk_repository.search_chunks_fts(
query, search_limit
)
else: # hybrid (default)
search_results = await self.chunk_repository.search_chunks_hybrid(
query, search_limit
)
search_results = await self.chunk_repository.search(
query, search_limit, search_type
)
# Apply reranking
chunks = [chunk for chunk, _ in search_results]
@ -564,8 +550,6 @@ class HaikuRAG:
)
yield doc.id
# LanceDB doesn't need explicit commits
def close(self):
"""Close the underlying store connection."""
self.store.close()

4
src/haiku/rag/store/engine.py
View file

@ -47,9 +47,7 @@ class Store:
# Create the ChunkRecord model with the correct vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
# For file paths, create a LanceDB directory structure
lance_path = str(db_path).replace(".sqlite", ".lancedb")
self.db = lancedb.connect(lance_path)
self.db = lancedb.connect(db_path)
self.create_or_update_db()

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

@ -28,16 +28,14 @@ class ChunkRepository:
"""Create a chunk in the database."""
assert entity.document_id, "Chunk must have a document_id to be created"
chunk_id = str(uuid4())
# Generate embedding if not provided
if entity.embedding is not None:
embedding = entity.embedding
else:
embedding = await self.embedder.embed(entity.content)
# Generate new UUID
chunk_id = str(uuid4())
# Create chunk record
chunk_record = self.store.ChunkRecord(
id=chunk_id,
document_id=entity.document_id,
@ -46,7 +44,6 @@ class ChunkRepository:
vector=embedding,
)
# Add to table
self.store.chunks_table.add([chunk_record])
entity.id = chunk_id
@ -76,10 +73,8 @@ class ChunkRepository:
"""Update an existing chunk."""
assert entity.id, "Chunk ID is required for update"
# Generate new embedding
embedding = await self.embedder.embed(entity.content)
# Update the record
self.store.chunks_table.update(
where=f"id = '{entity.id}'",
values={
@ -94,12 +89,10 @@ class ChunkRepository:
async def delete(self, entity_id: str) -> bool:
"""Delete a chunk by its ID."""
# Check if chunk exists
chunk = await self.get_by_id(entity_id)
if chunk is None:
return False
# Delete the chunk
self.store.chunks_table.delete(f"id = '{entity_id}'")
return True
@ -145,7 +138,6 @@ class ChunkRepository:
async def delete_all(self) -> bool:
"""Delete all chunks from the database."""
try:
# Get count before deletion
count = len(
list(
self.store.chunks_table.search()
@ -171,74 +163,66 @@ class ChunkRepository:
if not chunks:
return False
# Delete chunks by document_id
self.store.chunks_table.delete(f"document_id = '{document_id}'")
return True
async def search_chunks(
self, query: str, limit: int = 5
async def search(
self, query: str, limit: int = 5, search_type: str = "hybrid"
) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using vector similarity."""
# Generate embedding for the query
query_embedding = await self.embedder.embed(query)
"""Search for relevant chunks using the specified search method.
# Perform vector search with proper query type
results = (
self.store.chunks_table.search(query_embedding, query_type="vector")
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
Args:
query: The search query string.
limit: Maximum number of results to return.
search_type: Type of search - "vector", "fts", or "hybrid" (default).
return await self._process_search_results(results)
async def search_chunks_fts(
self, query: str, limit: int = 5
) -> list[tuple[Chunk, float]]:
"""Search for chunks using full-text search."""
Returns:
List of (chunk, score) tuples ordered by relevance.
"""
if not query.strip():
return []
# Ensure FTS index exists
self._ensure_fts_index()
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)
)
return await self._process_search_results(results)
elif search_type == "fts":
# Ensure FTS index exists
self._ensure_fts_index()
# Use LanceDB's native full-text search
try:
results = (
self.store.chunks_table.search(query, query_type="fts")
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
return await self._process_search_results(results)
except Exception:
# Fallback to vector search if FTS is not available or fails
return await self.search_chunks(query, limit)
async def search_chunks_hybrid(
self, query: str, limit: int = 5
) -> list[tuple[Chunk, float]]:
"""Hybrid search combining vector and full-text search with native LanceDB RRFReranker."""
if not query.strip():
return []
else: # hybrid (default)
# Ensure FTS index exists for hybrid search
self._ensure_fts_index()
# Ensure FTS index exists for hybrid search
self._ensure_fts_index()
query_embedding = await self.embedder.embed(query)
# Generate embedding for the query since LanceDB doesn't have embedding function configured
query_embedding = await self.embedder.embed(query)
# Create RRF reranker
reranker = RRFReranker()
# Create RRF reranker (k parameter is handled internally)
reranker = RRFReranker()
# Perform native hybrid search with RRF reranking
results = (
self.store.chunks_table.search(query_type="hybrid")
.vector(query_embedding)
.text(query)
.rerank(reranker)
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
return await self._process_search_results(results)
# Perform native hybrid search with RRF reranking
results = (
self.store.chunks_table.search(query_type="hybrid")
.vector(query_embedding)
.text(query)
.rerank(reranker)
.limit(limit)
.to_pydantic(self.store.ChunkRecord)
)
return await self._process_search_results(results)
async def get_by_document_id(self, document_id: str) -> list[Chunk]:
"""Get all chunks for a specific document."""
@ -272,7 +256,6 @@ class ChunkRepository:
for chunk in results
]
# Sort by order if available
chunks.sort(key=lambda c: c.metadata.get("order", 0))
return chunks

2
tests/generate_benchmark_db.py
View file

@ -11,7 +11,7 @@ from haiku.rag.qa import get_qa_agent
console = Console()
db_path = Path(__file__).parent / "data" / "benchmark.sqlite"
db_path = Path(__file__).parent / "data" / "benchmark.lancedb"
async def populate_db():

2
tests/test_chunk.py
View file

@ -35,7 +35,7 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
assert all(chunk.document_id == created_document.id for chunk in chunks)
# Test chunk search
results = await chunk_repo.search_chunks("election", limit=2)
results = await chunk_repo.search("election", limit=2, search_type="vector")
assert len(results) <= 2
assert all(hasattr(chunk, "content") for chunk, _ in results)

12
tests/test_search.py
View file

@ -47,17 +47,21 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
question = doc_data["question"]
# Test vector search
vector_results = await chunk_repo.search_chunks(question, limit=5)
vector_results = await chunk_repo.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 chunk_repo.search_chunks_fts(question, limit=5)
fts_results = await chunk_repo.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 chunk_repo.search_chunks_hybrid(question, limit=5)
hybrid_results = await chunk_repo.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
@ -85,7 +89,7 @@ async def test_chunks_include_document_info(temp_db_path):
created_document = await doc_repo._create_with_docling(document, docling_document)
# Search for chunks
results = await chunk_repo.search_chunks_hybrid("test document", limit=1)
results = await chunk_repo.search("test document", limit=1, search_type="hybrid")
assert len(results) > 0
chunk, _ = results[0]