402 lines
14 KiB
Python
402 lines
14 KiB
Python
import inspect
|
|
import json
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
from uuid import uuid4
|
|
|
|
from lancedb.rerankers import RRFReranker
|
|
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.embeddings import get_embedder
|
|
from haiku.rag.store.engine import DocumentRecord, Store
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
from haiku.rag.utils import load_callable
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ChunkRepository:
|
|
"""Repository for Chunk operations."""
|
|
|
|
def __init__(self, store: Store) -> None:
|
|
self.store = store
|
|
self.embedder = get_embedder()
|
|
|
|
def _ensure_fts_index(self) -> None:
|
|
"""Ensure FTS index exists on the content column."""
|
|
try:
|
|
self.store.chunks_table.create_fts_index(
|
|
"content", replace=True, with_position=True, remove_stop_words=False
|
|
)
|
|
except Exception as e:
|
|
# Log the error but don't fail - FTS might already exist
|
|
logger.debug(f"FTS index creation skipped: {e}")
|
|
|
|
async def create(self, entity: Chunk) -> Chunk:
|
|
"""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)
|
|
order_val = int(entity.order)
|
|
|
|
chunk_record = self.store.ChunkRecord(
|
|
id=chunk_id,
|
|
document_id=entity.document_id,
|
|
content=entity.content,
|
|
metadata=json.dumps(
|
|
{k: v for k, v in entity.metadata.items() if k != "order"}
|
|
),
|
|
order=order_val,
|
|
vector=embedding,
|
|
)
|
|
|
|
self.store.chunks_table.add([chunk_record])
|
|
|
|
entity.id = chunk_id
|
|
return entity
|
|
|
|
async def get_by_id(self, entity_id: str) -> Chunk | None:
|
|
"""Get a chunk by its ID."""
|
|
results = list(
|
|
self.store.chunks_table.search()
|
|
.where(f"id = '{entity_id}'")
|
|
.limit(1)
|
|
.to_pydantic(self.store.ChunkRecord)
|
|
)
|
|
|
|
if not results:
|
|
return None
|
|
|
|
chunk_record = results[0]
|
|
md = json.loads(chunk_record.metadata)
|
|
return Chunk(
|
|
id=chunk_record.id,
|
|
document_id=chunk_record.document_id,
|
|
content=chunk_record.content,
|
|
metadata=md,
|
|
order=chunk_record.order,
|
|
)
|
|
|
|
async def update(self, entity: Chunk) -> Chunk:
|
|
"""Update an existing chunk."""
|
|
assert entity.id, "Chunk ID is required for update"
|
|
|
|
embedding = await self.embedder.embed(entity.content)
|
|
order_val = int(entity.order)
|
|
|
|
self.store.chunks_table.update(
|
|
where=f"id = '{entity.id}'",
|
|
values={
|
|
"document_id": entity.document_id,
|
|
"content": entity.content,
|
|
"metadata": json.dumps(
|
|
{k: v for k, v in entity.metadata.items() if k != "order"}
|
|
),
|
|
"order": order_val,
|
|
"vector": embedding,
|
|
},
|
|
)
|
|
return entity
|
|
|
|
async def delete(self, entity_id: str) -> bool:
|
|
"""Delete a chunk by its ID."""
|
|
chunk = await self.get_by_id(entity_id)
|
|
if chunk is None:
|
|
return False
|
|
|
|
self.store.chunks_table.delete(f"id = '{entity_id}'")
|
|
return True
|
|
|
|
async def list_all(
|
|
self, limit: int | None = None, offset: int | None = None
|
|
) -> list[Chunk]:
|
|
"""List all chunks with optional pagination."""
|
|
query = self.store.chunks_table.search()
|
|
|
|
if offset is not None:
|
|
query = query.offset(offset)
|
|
if limit is not None:
|
|
query = query.limit(limit)
|
|
|
|
results = list(query.to_pydantic(self.store.ChunkRecord))
|
|
|
|
chunks: list[Chunk] = []
|
|
for rec in results:
|
|
md = json.loads(rec.metadata)
|
|
chunks.append(
|
|
Chunk(
|
|
id=rec.id,
|
|
document_id=rec.document_id,
|
|
content=rec.content,
|
|
metadata=md,
|
|
order=rec.order,
|
|
)
|
|
)
|
|
return chunks
|
|
|
|
async def create_chunks_for_document(
|
|
self, document_id: str, document: "DoclingDocument"
|
|
) -> list[Chunk]:
|
|
"""Create chunks and embeddings for a document from DoclingDocument."""
|
|
# Lazy imports to avoid loading docling during module import
|
|
from haiku.rag.chunker import chunker
|
|
from haiku.rag.utils import text_to_docling_document
|
|
|
|
# Optionally preprocess markdown before chunking
|
|
processed_document = document
|
|
preprocessor_path = Config.processing.markdown_preprocessor
|
|
if preprocessor_path:
|
|
try:
|
|
pre_fn = load_callable(preprocessor_path)
|
|
markdown = document.export_to_markdown()
|
|
result = pre_fn(markdown)
|
|
if inspect.isawaitable(result):
|
|
result = await result # type: ignore[assignment]
|
|
processed_markdown = result
|
|
if not isinstance(processed_markdown, str):
|
|
raise ValueError("Preprocessor must return a markdown string")
|
|
processed_document = text_to_docling_document(
|
|
processed_markdown, name="content.md"
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to apply MARKDOWN_PREPROCESSOR '{preprocessor_path}': {e}. Proceeding without preprocessing."
|
|
)
|
|
raise e
|
|
|
|
chunk_texts = await chunker.chunk(processed_document)
|
|
|
|
embeddings = await self.embedder.embed(chunk_texts)
|
|
|
|
# Prepare all chunk records for batch insertion
|
|
chunk_records = []
|
|
created_chunks = []
|
|
|
|
for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)):
|
|
chunk_id = str(uuid4())
|
|
|
|
chunk_record = self.store.ChunkRecord(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=chunk_text,
|
|
metadata=json.dumps({}),
|
|
order=order,
|
|
vector=embedding,
|
|
)
|
|
chunk_records.append(chunk_record)
|
|
|
|
chunk = Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=chunk_text,
|
|
metadata={},
|
|
order=order,
|
|
)
|
|
created_chunks.append(chunk)
|
|
|
|
# Batch insert all chunks at once
|
|
if chunk_records:
|
|
self.store.chunks_table.add(chunk_records)
|
|
|
|
return created_chunks
|
|
|
|
async def delete_all(self) -> None:
|
|
"""Delete all chunks from the database."""
|
|
# Drop and recreate table to clear all data
|
|
self.store.db.drop_table("chunks")
|
|
self.store.chunks_table = self.store.db.create_table(
|
|
"chunks", schema=self.store.ChunkRecord
|
|
)
|
|
# Create FTS index on the new table with phrase query support
|
|
self.store.chunks_table.create_fts_index(
|
|
"content", replace=True, with_position=True, remove_stop_words=False
|
|
)
|
|
|
|
async def delete_by_document_id(self, document_id: str) -> bool:
|
|
"""Delete all chunks for a document."""
|
|
chunks = await self.get_by_document_id(document_id)
|
|
|
|
if not chunks:
|
|
return False
|
|
|
|
self.store.chunks_table.delete(f"document_id = '{document_id}'")
|
|
return True
|
|
|
|
async def search(
|
|
self, query: str, limit: int = 5, search_type: str = "hybrid"
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Search for relevant chunks using the specified search method.
|
|
|
|
Args:
|
|
query: The search query string.
|
|
limit: Maximum number of results to return.
|
|
search_type: Type of search - "vector", "fts", or "hybrid" (default).
|
|
|
|
Returns:
|
|
List of (chunk, score) tuples ordered by relevance.
|
|
"""
|
|
if not query.strip():
|
|
return []
|
|
|
|
if search_type == "vector":
|
|
query_embedding = await self.embedder.embed(query)
|
|
|
|
results = self.store.chunks_table.search(
|
|
query_embedding, query_type="vector", vector_column_name="vector"
|
|
).limit(limit)
|
|
|
|
return await self._process_search_results(results)
|
|
|
|
elif search_type == "fts":
|
|
results = self.store.chunks_table.search(query, query_type="fts").limit(
|
|
limit
|
|
)
|
|
return await self._process_search_results(results)
|
|
|
|
else: # hybrid (default)
|
|
query_embedding = await self.embedder.embed(query)
|
|
|
|
# Create RRF reranker
|
|
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)
|
|
)
|
|
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."""
|
|
results = list(
|
|
self.store.chunks_table.search()
|
|
.where(f"document_id = '{document_id}'")
|
|
.to_pydantic(self.store.ChunkRecord)
|
|
)
|
|
|
|
# Get document info
|
|
doc_results = list(
|
|
self.store.documents_table.search()
|
|
.where(f"id = '{document_id}'")
|
|
.limit(1)
|
|
.to_pydantic(DocumentRecord)
|
|
)
|
|
|
|
doc_uri = doc_results[0].uri if doc_results else None
|
|
doc_title = doc_results[0].title if doc_results else None
|
|
doc_meta = doc_results[0].metadata if doc_results else "{}"
|
|
|
|
chunks: list[Chunk] = []
|
|
for rec in results:
|
|
md = json.loads(rec.metadata)
|
|
chunks.append(
|
|
Chunk(
|
|
id=rec.id,
|
|
document_id=rec.document_id,
|
|
content=rec.content,
|
|
metadata=md,
|
|
order=rec.order,
|
|
document_uri=doc_uri,
|
|
document_title=doc_title,
|
|
document_meta=json.loads(doc_meta),
|
|
)
|
|
)
|
|
|
|
chunks.sort(key=lambda c: c.order)
|
|
return chunks
|
|
|
|
async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
|
|
"""Get adjacent chunks before and after the given chunk within the same document."""
|
|
assert chunk.document_id, "Document id is required for adjacent chunk finding"
|
|
|
|
chunk_order = chunk.order
|
|
|
|
# Fetch chunks for the same document and filter by order proximity
|
|
all_chunks = await self.get_by_document_id(chunk.document_id)
|
|
|
|
adjacent_chunks: list[Chunk] = []
|
|
for c in all_chunks:
|
|
c_order = c.order
|
|
if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
|
|
adjacent_chunks.append(c)
|
|
|
|
return adjacent_chunks
|
|
|
|
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 = []
|
|
|
|
# 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")
|
|
|
|
# Collect all unique document IDs for batch lookup
|
|
document_ids = list(set(chunk.document_id for chunk in pydantic_results))
|
|
|
|
# Batch fetch all documents at once
|
|
documents_map = {}
|
|
if document_ids:
|
|
# Create a WHERE clause for all document IDs
|
|
where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
|
|
doc_results = list(
|
|
self.store.documents_table.search()
|
|
.where(where_clause)
|
|
.to_pydantic(DocumentRecord)
|
|
)
|
|
documents_map = {doc.id: doc for doc in doc_results}
|
|
|
|
for i, chunk_record in enumerate(pydantic_results):
|
|
# Get document info from pre-fetched map
|
|
doc = documents_map.get(chunk_record.document_id)
|
|
doc_uri = doc.uri if doc else None
|
|
doc_title = doc.title if doc else None
|
|
doc_meta = doc.metadata if doc else "{}"
|
|
|
|
md = json.loads(chunk_record.metadata)
|
|
|
|
chunk = Chunk(
|
|
id=chunk_record.id,
|
|
document_id=chunk_record.document_id,
|
|
content=chunk_record.content,
|
|
metadata=md,
|
|
order=chunk_record.order,
|
|
document_uri=doc_uri,
|
|
document_title=doc_title,
|
|
document_meta=json.loads(doc_meta),
|
|
)
|
|
|
|
# Get score from arrow result
|
|
score = scores[i] if i < len(scores) else 1.0
|
|
|
|
chunks_with_scores.append((chunk, score))
|
|
|
|
return chunks_with_scores
|