460 lines
16 KiB
Python
460 lines
16 KiB
Python
import json
|
|
import re
|
|
|
|
from haiku.rag.chunker import chunker
|
|
from haiku.rag.embeddings import get_embedder
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
from haiku.rag.store.repositories.base import BaseRepository
|
|
|
|
|
|
class ChunkRepository(BaseRepository[Chunk]):
|
|
"""Repository for Chunk database operations."""
|
|
|
|
def __init__(self, store):
|
|
super().__init__(store)
|
|
self.embedder = get_embedder()
|
|
|
|
async def create(self, entity: Chunk, commit: bool = True) -> Chunk:
|
|
"""Create a chunk in the database."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO chunks (document_id, content, metadata)
|
|
VALUES (:document_id, :content, :metadata)
|
|
""",
|
|
{
|
|
"document_id": entity.document_id,
|
|
"content": entity.content,
|
|
"metadata": json.dumps(entity.metadata),
|
|
},
|
|
)
|
|
|
|
entity.id = cursor.lastrowid
|
|
|
|
# Generate and store embedding
|
|
embedding = await self.embedder.embed(entity.content)
|
|
serialized_embedding = self.store.serialize_embedding(embedding)
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO chunk_embeddings (chunk_id, embedding)
|
|
VALUES (:chunk_id, :embedding)
|
|
""",
|
|
{"chunk_id": entity.id, "embedding": serialized_embedding},
|
|
)
|
|
|
|
# Insert into FTS5 table for full-text search
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO chunks_fts(rowid, content)
|
|
VALUES (:rowid, :content)
|
|
""",
|
|
{"rowid": entity.id, "content": entity.content},
|
|
)
|
|
|
|
if commit:
|
|
self.store._connection.commit()
|
|
return entity
|
|
|
|
async def get_by_id(self, entity_id: int) -> Chunk | None:
|
|
"""Get a chunk by its ID."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, document_id, content, metadata
|
|
FROM chunks WHERE id = :id
|
|
""",
|
|
{"id": entity_id},
|
|
)
|
|
|
|
row = cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
|
|
chunk_id, document_id, content, metadata_json = row
|
|
metadata = json.loads(metadata_json) if metadata_json else {}
|
|
|
|
return Chunk(
|
|
id=chunk_id, document_id=document_id, content=content, metadata=metadata
|
|
)
|
|
|
|
async def update(self, entity: Chunk) -> Chunk:
|
|
"""Update an existing chunk."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
if entity.id is None:
|
|
raise ValueError("Chunk ID is required for update")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
cursor.execute(
|
|
"""
|
|
UPDATE chunks
|
|
SET document_id = :document_id, content = :content, metadata = :metadata
|
|
WHERE id = :id
|
|
""",
|
|
{
|
|
"document_id": entity.document_id,
|
|
"content": entity.content,
|
|
"metadata": json.dumps(entity.metadata),
|
|
"id": entity.id,
|
|
},
|
|
)
|
|
|
|
# Regenerate and update embedding
|
|
embedding = await self.embedder.embed(entity.content)
|
|
serialized_embedding = self.store.serialize_embedding(embedding)
|
|
cursor.execute(
|
|
"""
|
|
UPDATE chunk_embeddings
|
|
SET embedding = :embedding
|
|
WHERE chunk_id = :chunk_id
|
|
""",
|
|
{"embedding": serialized_embedding, "chunk_id": entity.id},
|
|
)
|
|
|
|
# Update FTS5 table
|
|
cursor.execute(
|
|
"""
|
|
UPDATE chunks_fts
|
|
SET content = :content
|
|
WHERE rowid = :rowid
|
|
""",
|
|
{"content": entity.content, "rowid": entity.id},
|
|
)
|
|
|
|
self.store._connection.commit()
|
|
return entity
|
|
|
|
async def delete(self, entity_id: int, commit: bool = True) -> bool:
|
|
"""Delete a chunk by its ID."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
|
|
# Delete from FTS5 table first
|
|
cursor.execute(
|
|
"DELETE FROM chunks_fts WHERE rowid = :rowid", {"rowid": entity_id}
|
|
)
|
|
|
|
# Delete the embedding
|
|
cursor.execute(
|
|
"DELETE FROM chunk_embeddings WHERE chunk_id = :chunk_id",
|
|
{"chunk_id": entity_id},
|
|
)
|
|
|
|
# Delete the chunk
|
|
cursor.execute("DELETE FROM chunks WHERE id = :id", {"id": entity_id})
|
|
|
|
deleted = cursor.rowcount > 0
|
|
if commit:
|
|
self.store._connection.commit()
|
|
return deleted
|
|
|
|
async def list_all(
|
|
self, limit: int | None = None, offset: int | None = None
|
|
) -> list[Chunk]:
|
|
"""List all chunks with optional pagination."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
query = "SELECT id, document_id, content, metadata FROM chunks ORDER BY document_id, id"
|
|
params = {}
|
|
|
|
if limit is not None:
|
|
query += " LIMIT :limit"
|
|
params["limit"] = limit
|
|
|
|
if offset is not None:
|
|
query += " OFFSET :offset"
|
|
params["offset"] = offset
|
|
|
|
cursor.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
|
|
return [
|
|
Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=content,
|
|
metadata=json.loads(metadata_json) if metadata_json else {},
|
|
)
|
|
for chunk_id, document_id, content, metadata_json in rows
|
|
]
|
|
|
|
async def create_chunks_for_document(
|
|
self, document_id: int, content: str, commit: bool = True
|
|
) -> list[Chunk]:
|
|
"""Create chunks and embeddings for a document."""
|
|
# Chunk the document content
|
|
chunk_texts = await chunker.chunk(content)
|
|
created_chunks = []
|
|
|
|
# Create chunks with embeddings using the create method
|
|
for order, chunk_text in enumerate(chunk_texts):
|
|
# Create chunk with order in metadata
|
|
chunk = Chunk(
|
|
document_id=document_id, content=chunk_text, metadata={"order": order}
|
|
)
|
|
|
|
created_chunk = await self.create(chunk, commit=commit)
|
|
created_chunks.append(created_chunk)
|
|
|
|
return created_chunks
|
|
|
|
async def delete_all(self, commit: bool = True) -> bool:
|
|
"""Delete all chunks from the database."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
|
|
cursor.execute("DELETE FROM chunks_fts")
|
|
cursor.execute("DELETE FROM chunk_embeddings")
|
|
cursor.execute("DELETE FROM chunks")
|
|
|
|
deleted = cursor.rowcount > 0
|
|
if commit:
|
|
self.store._connection.commit()
|
|
return deleted
|
|
|
|
async def delete_by_document_id(
|
|
self, document_id: int, commit: bool = True
|
|
) -> bool:
|
|
"""Delete all chunks for a document."""
|
|
chunks = await self.get_by_document_id(document_id)
|
|
|
|
deleted_any = False
|
|
for chunk in chunks:
|
|
if chunk.id is not None:
|
|
deleted = await self.delete(chunk.id, commit=False)
|
|
deleted_any = deleted_any or deleted
|
|
|
|
if commit and deleted_any and self.store._connection:
|
|
self.store._connection.commit()
|
|
return deleted_any
|
|
|
|
async def search_chunks(
|
|
self, query: str, limit: int = 5
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Search for relevant chunks using vector similarity."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
|
|
# Generate embedding for the query
|
|
query_embedding = await self.embedder.embed(query)
|
|
serialized_query_embedding = self.store.serialize_embedding(query_embedding)
|
|
|
|
# Search for similar chunks using sqlite-vec
|
|
cursor.execute(
|
|
"""
|
|
SELECT c.id, c.document_id, c.content, c.metadata, distance, d.uri, d.metadata as document_metadata
|
|
FROM chunk_embeddings
|
|
JOIN chunks c ON c.id = chunk_embeddings.chunk_id
|
|
JOIN documents d ON c.document_id = d.id
|
|
WHERE embedding MATCH :embedding AND k = :k
|
|
ORDER BY distance
|
|
""",
|
|
{"embedding": serialized_query_embedding, "k": limit},
|
|
)
|
|
|
|
results = cursor.fetchall()
|
|
return [
|
|
(
|
|
Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=content,
|
|
metadata=json.loads(metadata_json) if metadata_json else {},
|
|
document_uri=document_uri,
|
|
document_meta=json.loads(document_metadata_json)
|
|
if document_metadata_json
|
|
else {},
|
|
),
|
|
1.0 / (1.0 + distance),
|
|
)
|
|
for chunk_id, document_id, content, metadata_json, distance, document_uri, document_metadata_json in results
|
|
]
|
|
|
|
async def search_chunks_fts(
|
|
self, query: str, limit: int = 5
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Search for chunks using FTS5 full-text search."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
|
|
# Clean the query for FTS5 - extract keywords for better matching
|
|
# Remove special characters and split into words
|
|
words = re.findall(r"\b\w+\b", query.lower())
|
|
# Join with OR to find chunks containing any of the keywords
|
|
fts_query = " OR ".join(words) if words else query
|
|
|
|
# Search using FTS5
|
|
cursor.execute(
|
|
"""
|
|
SELECT c.id, c.document_id, c.content, c.metadata, rank, d.uri, d.metadata as document_metadata
|
|
FROM chunks_fts
|
|
JOIN chunks c ON c.id = chunks_fts.rowid
|
|
JOIN documents d ON c.document_id = d.id
|
|
WHERE chunks_fts MATCH :query
|
|
ORDER BY rank
|
|
LIMIT :limit
|
|
""",
|
|
{"query": fts_query, "limit": limit},
|
|
)
|
|
|
|
results = cursor.fetchall()
|
|
|
|
return [
|
|
(
|
|
Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=content,
|
|
metadata=json.loads(metadata_json) if metadata_json else {},
|
|
document_uri=document_uri,
|
|
document_meta=json.loads(document_metadata_json)
|
|
if document_metadata_json
|
|
else {},
|
|
),
|
|
-rank,
|
|
)
|
|
for chunk_id, document_id, content, metadata_json, rank, document_uri, document_metadata_json in results
|
|
# FTS5 rank is negative BM25 score
|
|
]
|
|
|
|
async def search_chunks_hybrid(
|
|
self, query: str, limit: int = 5, k: int = 60
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Hybrid search using Reciprocal Rank Fusion (RRF) combining vector similarity and FTS5 full-text search."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
|
|
# Generate embedding for the query
|
|
query_embedding = await self.embedder.embed(query)
|
|
serialized_query_embedding = self.store.serialize_embedding(query_embedding)
|
|
|
|
# Clean the query for FTS5 - extract keywords for better matching
|
|
# Remove special characters and split into words
|
|
words = re.findall(r"\b\w+\b", query.lower())
|
|
# Join with OR to find chunks containing any of the keywords
|
|
fts_query = " OR ".join(words) if words else query
|
|
# Perform hybrid search using RRF (Reciprocal Rank Fusion)
|
|
cursor.execute(
|
|
"""
|
|
WITH vector_search AS (
|
|
SELECT
|
|
c.id,
|
|
c.document_id,
|
|
c.content,
|
|
c.metadata,
|
|
ROW_NUMBER() OVER (ORDER BY ce.distance) as vector_rank
|
|
FROM chunk_embeddings ce
|
|
JOIN chunks c ON c.id = ce.chunk_id
|
|
WHERE ce.embedding MATCH :embedding AND k = :k_vector
|
|
ORDER BY ce.distance
|
|
),
|
|
fts_search AS (
|
|
SELECT
|
|
c.id,
|
|
c.document_id,
|
|
c.content,
|
|
c.metadata,
|
|
ROW_NUMBER() OVER (ORDER BY chunks_fts.rank) as fts_rank
|
|
FROM chunks_fts
|
|
JOIN chunks c ON c.id = chunks_fts.rowid
|
|
WHERE chunks_fts MATCH :fts_query
|
|
ORDER BY chunks_fts.rank
|
|
),
|
|
all_chunks AS (
|
|
SELECT id, document_id, content, metadata FROM vector_search
|
|
UNION
|
|
SELECT id, document_id, content, metadata FROM fts_search
|
|
),
|
|
rrf_scores AS (
|
|
SELECT
|
|
a.id,
|
|
a.document_id,
|
|
a.content,
|
|
a.metadata,
|
|
COALESCE(1.0 / (:k + v.vector_rank), 0) + COALESCE(1.0 / (:k + f.fts_rank), 0) as rrf_score
|
|
FROM all_chunks a
|
|
LEFT JOIN vector_search v ON a.id = v.id
|
|
LEFT JOIN fts_search f ON a.id = f.id
|
|
)
|
|
SELECT r.id, r.document_id, r.content, r.metadata, r.rrf_score, d.uri, d.metadata as document_metadata
|
|
FROM rrf_scores r
|
|
JOIN documents d ON r.document_id = d.id
|
|
ORDER BY r.rrf_score DESC
|
|
LIMIT :limit
|
|
""",
|
|
{
|
|
"embedding": serialized_query_embedding,
|
|
"k_vector": limit * 3,
|
|
"fts_query": fts_query,
|
|
"k": k,
|
|
"limit": limit,
|
|
},
|
|
)
|
|
|
|
results = cursor.fetchall()
|
|
return [
|
|
(
|
|
Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=content,
|
|
metadata=json.loads(metadata_json) if metadata_json else {},
|
|
document_uri=document_uri,
|
|
document_meta=json.loads(document_metadata_json)
|
|
if document_metadata_json
|
|
else {},
|
|
),
|
|
rrf_score,
|
|
)
|
|
for chunk_id, document_id, content, metadata_json, rrf_score, document_uri, document_metadata_json in results
|
|
]
|
|
|
|
async def get_by_document_id(self, document_id: int) -> list[Chunk]:
|
|
"""Get all chunks for a specific document."""
|
|
if self.store._connection is None:
|
|
raise ValueError("Store connection is not available")
|
|
|
|
cursor = self.store._connection.cursor()
|
|
cursor.execute(
|
|
"""
|
|
SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata
|
|
FROM chunks c
|
|
JOIN documents d ON c.document_id = d.id
|
|
WHERE c.document_id = :document_id
|
|
ORDER BY JSON_EXTRACT(c.metadata, '$.order')
|
|
""",
|
|
{"document_id": document_id},
|
|
)
|
|
|
|
rows = cursor.fetchall()
|
|
return [
|
|
Chunk(
|
|
id=chunk_id,
|
|
document_id=document_id,
|
|
content=content,
|
|
metadata=json.loads(metadata_json) if metadata_json else {},
|
|
document_uri=document_uri,
|
|
document_meta=json.loads(document_metadata_json)
|
|
if document_metadata_json
|
|
else {},
|
|
)
|
|
for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows
|
|
]
|