Use db repositories
This commit is contained in:
Yiorgis Gozadinos
parent
5d1a854934
commit
5f9fb12644
9 changed files with 716 additions and 154 deletions
|
|
@ -1,11 +1,7 @@
|
|||
import functools
|
||||
|
||||
|
||||
class EmbedderBase:
|
||||
_model: str = ""
|
||||
_vector_dim: int = 0
|
||||
|
||||
@functools.lru_cache(maxsize=128)
|
||||
async def embed(self, text: str) -> list[float]:
|
||||
raise NotImplementedError(
|
||||
"Embedder is an abstract class. Please implement the embed method in a subclass."
|
||||
|
|
|
|||
|
|
@ -1,5 +1,3 @@
|
|||
import functools
|
||||
|
||||
from ollama import AsyncClient
|
||||
|
||||
from haiku.rag.config import Config
|
||||
|
|
@ -10,7 +8,6 @@ class Embedder(EmbedderBase):
|
|||
_model: str = Config.EMBEDDING_MODEL
|
||||
_vector_dim: int = 1024
|
||||
|
||||
@functools.lru_cache(maxsize=128)
|
||||
async def embed(self, text: str) -> list[float]:
|
||||
client = AsyncClient(host=Config.OLLAMA_BASE_URL)
|
||||
res = await client.embeddings(model=self._model, prompt=text)
|
||||
|
|
|
|||
|
|
@ -1,16 +1,7 @@
|
|||
import json
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.chunker import chunker
|
||||
from haiku.rag.embeddings.ollama import Embedder
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
|
||||
class Document(BaseModel):
|
||||
"""
|
||||
|
|
@ -22,124 +13,3 @@ class Document(BaseModel):
|
|||
metadata: dict = {}
|
||||
created_at: datetime = Field(default_factory=datetime.now)
|
||||
updated_at: datetime = Field(default_factory=datetime.now)
|
||||
|
||||
async def create_with_chunks(self, store: "Store") -> "Document":
|
||||
"""
|
||||
Create a document in the database along with its chunks and embeddings.
|
||||
|
||||
Args:
|
||||
store: The Store instance to use for database operations
|
||||
|
||||
Returns:
|
||||
Document: The created document with updated id
|
||||
"""
|
||||
if store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
|
||||
cursor = store._connection.cursor()
|
||||
embedder = Embedder()
|
||||
|
||||
# Insert the document
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO documents (content, metadata, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
self.content,
|
||||
json.dumps(self.metadata),
|
||||
self.created_at,
|
||||
self.updated_at,
|
||||
),
|
||||
)
|
||||
|
||||
document_id = cursor.lastrowid
|
||||
assert document_id is not None, "Failed to create document in database"
|
||||
self.id = document_id
|
||||
|
||||
# Chunk the document content
|
||||
chunk_texts = await chunker.chunk(self.content)
|
||||
|
||||
# Create chunks with embeddings
|
||||
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}
|
||||
)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO chunks (document_id, content, metadata)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(chunk.document_id, chunk.content, json.dumps(chunk.metadata)),
|
||||
)
|
||||
chunk_id = cursor.lastrowid
|
||||
|
||||
# Generate and store embedding
|
||||
embedding = await embedder.embed(chunk_text)
|
||||
serialized_embedding = store.serialize_embedding(embedding)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO chunk_embeddings (chunk_id, embedding)
|
||||
VALUES (?, ?)
|
||||
""",
|
||||
(chunk_id, serialized_embedding),
|
||||
)
|
||||
|
||||
store._connection.commit()
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
async def search_chunks(
|
||||
cls, store: "Store", query: str, limit: int = 5
|
||||
) -> list[Chunk]:
|
||||
"""
|
||||
Search for relevant chunks using vector similarity with sqlite-vec.
|
||||
|
||||
Args:
|
||||
store: The Store instance to use for database operations
|
||||
query: The text query to search for
|
||||
limit: Maximum number of chunks to return
|
||||
|
||||
Returns:
|
||||
List of relevant Chunk objects ordered by similarity
|
||||
"""
|
||||
if store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
|
||||
embedder = Embedder()
|
||||
cursor = store._connection.cursor()
|
||||
|
||||
# Generate embedding for the query
|
||||
query_embedding = await embedder.embed(query)
|
||||
serialized_query_embedding = 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
|
||||
FROM chunk_embeddings
|
||||
JOIN chunks c ON c.id = chunk_embeddings.chunk_id
|
||||
WHERE embedding MATCH ? AND k = ?
|
||||
ORDER BY distance
|
||||
""",
|
||||
(serialized_query_embedding, limit),
|
||||
)
|
||||
|
||||
results = cursor.fetchall()
|
||||
chunks = []
|
||||
|
||||
for row in results:
|
||||
chunk_id, document_id, content, metadata_json, distance = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=chunk_id,
|
||||
document_id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return chunks
|
||||
|
|
|
|||
5
src/haiku/rag/store/repositories/__init__.py
Normal file
5
src/haiku/rag/store/repositories/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from haiku.rag.store.repositories.base import BaseRepository
|
||||
from haiku.rag.store.repositories.chunk import ChunkRepository
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
|
||||
__all__ = ["BaseRepository", "DocumentRepository", "ChunkRepository"]
|
||||
40
src/haiku/rag/store/repositories/base.py
Normal file
40
src/haiku/rag/store/repositories/base.py
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class BaseRepository(ABC, Generic[T]):
|
||||
"""Base repository interface for database operations."""
|
||||
|
||||
def __init__(self, store: Store):
|
||||
self.store = store
|
||||
|
||||
@abstractmethod
|
||||
async def create(self, entity: T) -> T:
|
||||
"""Create a new entity in the database."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_by_id(self, entity_id: int) -> T | None:
|
||||
"""Get an entity by its ID."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def update(self, entity: T) -> T:
|
||||
"""Update an existing entity."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, entity_id: int) -> bool:
|
||||
"""Delete an entity by its ID."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def list_all(
|
||||
self, limit: int | None = None, offset: int | None = None
|
||||
) -> list[T]:
|
||||
"""List all entities with optional pagination."""
|
||||
pass
|
||||
270
src/haiku/rag/store/repositories/chunk.py
Normal file
270
src/haiku/rag/store/repositories/chunk.py
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
import json
|
||||
|
||||
from haiku.rag.chunker import chunker
|
||||
from haiku.rag.embeddings.ollama import 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, embedder: Embedder | None = None):
|
||||
super().__init__(store)
|
||||
self.embedder = embedder or 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 (?, ?, ?)
|
||||
""",
|
||||
(entity.document_id, entity.content, 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 (?, ?)
|
||||
""",
|
||||
(entity.id, serialized_embedding),
|
||||
)
|
||||
|
||||
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 = ?
|
||||
""",
|
||||
(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 = ?, content = ?, metadata = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
entity.document_id,
|
||||
entity.content,
|
||||
json.dumps(entity.metadata),
|
||||
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 = ?
|
||||
WHERE chunk_id = ?
|
||||
""",
|
||||
(serialized_embedding, 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 the embedding first
|
||||
cursor.execute("DELETE FROM chunk_embeddings WHERE chunk_id = ?", (entity_id,))
|
||||
|
||||
# Delete the chunk
|
||||
cursor.execute("DELETE FROM chunks WHERE 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 ?"
|
||||
params.append(limit)
|
||||
|
||||
if offset is not None:
|
||||
query += " OFFSET ?"
|
||||
params.append(offset)
|
||||
|
||||
cursor.execute(query, params)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
chunks = []
|
||||
for row in rows:
|
||||
chunk_id, document_id, content, metadata_json = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=chunk_id,
|
||||
document_id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return chunks
|
||||
|
||||
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_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[Chunk]:
|
||||
"""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
|
||||
FROM chunk_embeddings
|
||||
JOIN chunks c ON c.id = chunk_embeddings.chunk_id
|
||||
WHERE embedding MATCH ? AND k = ?
|
||||
ORDER BY distance
|
||||
""",
|
||||
(serialized_query_embedding, limit),
|
||||
)
|
||||
|
||||
results = cursor.fetchall()
|
||||
chunks = []
|
||||
|
||||
for row in results:
|
||||
chunk_id, document_id, content, metadata_json, _ = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=chunk_id,
|
||||
document_id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return chunks
|
||||
|
||||
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 id, document_id, content, metadata
|
||||
FROM chunks WHERE document_id = ?
|
||||
ORDER BY JSON_EXTRACT(metadata, '$.order')
|
||||
""",
|
||||
(document_id,),
|
||||
)
|
||||
|
||||
rows = cursor.fetchall()
|
||||
chunks = []
|
||||
|
||||
for row in rows:
|
||||
chunk_id, document_id, content, metadata_json = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=chunk_id,
|
||||
document_id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return chunks
|
||||
181
src/haiku/rag/store/repositories/document.py
Normal file
181
src/haiku/rag/store/repositories/document.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
import json
|
||||
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.repositories.base import BaseRepository
|
||||
|
||||
|
||||
class DocumentRepository(BaseRepository[Document]):
|
||||
"""Repository for Document database operations."""
|
||||
|
||||
def __init__(self, store, chunk_repository=None):
|
||||
super().__init__(store)
|
||||
# Avoid circular import by using late import if not provided
|
||||
if chunk_repository is None:
|
||||
from haiku.rag.store.repositories.chunk import ChunkRepository
|
||||
|
||||
chunk_repository = ChunkRepository(store)
|
||||
self.chunk_repository = chunk_repository
|
||||
|
||||
async def create(self, entity: Document) -> Document:
|
||||
"""Create a document with its chunks and embeddings."""
|
||||
if self.store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
|
||||
cursor = self.store._connection.cursor()
|
||||
|
||||
# Start transaction
|
||||
cursor.execute("BEGIN TRANSACTION")
|
||||
|
||||
try:
|
||||
# Insert the document
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO documents (content, metadata, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
entity.content,
|
||||
json.dumps(entity.metadata),
|
||||
entity.created_at,
|
||||
entity.updated_at,
|
||||
),
|
||||
)
|
||||
|
||||
document_id = cursor.lastrowid
|
||||
assert document_id is not None, "Failed to create document in database"
|
||||
entity.id = document_id
|
||||
|
||||
# Create chunks and embeddings using ChunkRepository
|
||||
await self.chunk_repository.create_chunks_for_document(
|
||||
document_id, entity.content, commit=False
|
||||
)
|
||||
|
||||
cursor.execute("COMMIT")
|
||||
return entity
|
||||
|
||||
except Exception:
|
||||
cursor.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
async def get_by_id(self, entity_id: int) -> Document | None:
|
||||
"""Get a document 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, content, metadata, created_at, updated_at
|
||||
FROM documents WHERE id = ?
|
||||
""",
|
||||
(entity_id,),
|
||||
)
|
||||
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
document_id, content, metadata_json, created_at, updated_at = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
|
||||
return Document(
|
||||
id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
|
||||
async def update(self, entity: Document) -> Document:
|
||||
"""Update an existing document and regenerate its chunks and embeddings."""
|
||||
if self.store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
if entity.id is None:
|
||||
raise ValueError("Document ID is required for update")
|
||||
|
||||
cursor = self.store._connection.cursor()
|
||||
|
||||
# Start transaction
|
||||
cursor.execute("BEGIN TRANSACTION")
|
||||
|
||||
try:
|
||||
# Update the document
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE documents
|
||||
SET content = ?, metadata = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
entity.content,
|
||||
json.dumps(entity.metadata),
|
||||
entity.updated_at,
|
||||
entity.id,
|
||||
),
|
||||
)
|
||||
|
||||
# Delete existing chunks and regenerate using ChunkRepository
|
||||
await self.chunk_repository.delete_by_document_id(entity.id, commit=False)
|
||||
await self.chunk_repository.create_chunks_for_document(
|
||||
entity.id, entity.content, commit=False
|
||||
)
|
||||
|
||||
cursor.execute("COMMIT")
|
||||
return entity
|
||||
|
||||
except Exception:
|
||||
cursor.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
async def delete(self, entity_id: int) -> bool:
|
||||
"""Delete a document and all its associated chunks and embeddings."""
|
||||
# Delete chunks and embeddings first
|
||||
await self.chunk_repository.delete_by_document_id(entity_id)
|
||||
|
||||
if self.store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
|
||||
cursor = self.store._connection.cursor()
|
||||
cursor.execute("DELETE FROM documents WHERE id = ?", (entity_id,))
|
||||
|
||||
deleted = cursor.rowcount > 0
|
||||
self.store._connection.commit()
|
||||
return deleted
|
||||
|
||||
async def list_all(
|
||||
self, limit: int | None = None, offset: int | None = None
|
||||
) -> list[Document]:
|
||||
"""List all documents with optional pagination."""
|
||||
if self.store._connection is None:
|
||||
raise ValueError("Store connection is not available")
|
||||
|
||||
cursor = self.store._connection.cursor()
|
||||
query = "SELECT id, content, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
|
||||
params = []
|
||||
|
||||
if limit is not None:
|
||||
query += " LIMIT ?"
|
||||
params.append(limit)
|
||||
|
||||
if offset is not None:
|
||||
query += " OFFSET ?"
|
||||
params.append(offset)
|
||||
|
||||
cursor.execute(query, params)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
documents = []
|
||||
for row in rows:
|
||||
document_id, content, metadata_json, created_at, updated_at = row
|
||||
metadata = json.loads(metadata_json) if metadata_json else {}
|
||||
documents.append(
|
||||
Document(
|
||||
id=document_id,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
)
|
||||
|
||||
return documents
|
||||
186
tests/test_chunk.py
Normal file
186
tests/test_chunk.py
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
import pytest
|
||||
from datasets import Dataset
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.repositories.chunk import ChunkRepository
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_chunks(qa_corpus: Dataset):
|
||||
"""Test vector search functionality using ChunkRepository."""
|
||||
# Create an in-memory store and repositories
|
||||
store = Store(":memory:")
|
||||
doc_repo = DocumentRepository(store)
|
||||
chunk_repo = ChunkRepository(store)
|
||||
|
||||
# Get the first document from the corpus
|
||||
first_doc = qa_corpus[0]
|
||||
document_text = first_doc["document_extracted"]
|
||||
|
||||
# Create and store a document
|
||||
document = Document(content=document_text, metadata={"source": "qa_corpus"})
|
||||
created_document = await doc_repo.create(document)
|
||||
|
||||
# Perform a search using ChunkRepository
|
||||
search_query = "news" # Simple query
|
||||
results = await chunk_repo.search_chunks(search_query, limit=3)
|
||||
|
||||
# Verify search results
|
||||
assert len(results) <= 3
|
||||
assert all(hasattr(chunk, "content") for chunk in results)
|
||||
assert all(hasattr(chunk, "document_id") for chunk in results)
|
||||
assert all(chunk.document_id == created_document.id for chunk in results)
|
||||
|
||||
store.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_repository_operations(qa_corpus: Dataset):
|
||||
"""Test ChunkRepository operations."""
|
||||
# Create an in-memory store and repositories
|
||||
store = Store(":memory:")
|
||||
doc_repo = DocumentRepository(store)
|
||||
chunk_repo = ChunkRepository(store)
|
||||
|
||||
# Get the first document from the corpus
|
||||
first_doc = qa_corpus[0]
|
||||
document_text = first_doc["document_extracted"]
|
||||
|
||||
# Create a document first
|
||||
document = Document(content=document_text, metadata={"source": "test"})
|
||||
created_document = await doc_repo.create(document)
|
||||
assert created_document.id is not None
|
||||
|
||||
# Test getting chunks by document ID
|
||||
chunks = await chunk_repo.get_by_document_id(created_document.id)
|
||||
assert len(chunks) > 0
|
||||
assert all(chunk.document_id == created_document.id for chunk in chunks)
|
||||
|
||||
# Test chunk search
|
||||
results = await chunk_repo.search_chunks("election", limit=2)
|
||||
assert len(results) <= 2
|
||||
assert all(hasattr(chunk, "content") for chunk in results)
|
||||
|
||||
# Test deleting chunks by document ID
|
||||
deleted = await chunk_repo.delete_by_document_id(created_document.id)
|
||||
assert deleted is True
|
||||
|
||||
# Verify chunks are gone
|
||||
chunks_after_delete = await chunk_repo.get_by_document_id(created_document.id)
|
||||
assert len(chunks_after_delete) == 0
|
||||
|
||||
store.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_chunks_for_document(qa_corpus: Dataset):
|
||||
"""Test creating chunks for a document."""
|
||||
# Create an in-memory store and repositories
|
||||
store = Store(":memory:")
|
||||
chunk_repo = ChunkRepository(store)
|
||||
|
||||
# Get the first document from the corpus
|
||||
first_doc = qa_corpus[0]
|
||||
document_text = first_doc["document_extracted"]
|
||||
|
||||
# Create a document first (without chunks)
|
||||
document = Document(content=document_text, metadata={"source": "test"})
|
||||
|
||||
# Insert document manually to test chunk creation independently
|
||||
document_id = None
|
||||
if store._connection is not None:
|
||||
cursor = store._connection.cursor()
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO documents (content, metadata, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(document.content, "{}", document.created_at, document.updated_at),
|
||||
)
|
||||
document_id = cursor.lastrowid
|
||||
document.id = document_id
|
||||
store._connection.commit()
|
||||
|
||||
assert document_id is not None, "Document ID should not be None"
|
||||
|
||||
# Test creating chunks for the document
|
||||
chunks = await chunk_repo.create_chunks_for_document(document_id, document_text)
|
||||
|
||||
# Verify chunks were created
|
||||
assert len(chunks) > 0
|
||||
assert all(chunk.document_id == document_id for chunk in chunks)
|
||||
assert all(chunk.id is not None for chunk in chunks)
|
||||
|
||||
# Verify chunk order metadata
|
||||
for i, chunk in enumerate(chunks):
|
||||
assert chunk.metadata.get("order") == i
|
||||
|
||||
# Verify chunks exist in database
|
||||
db_chunks = await chunk_repo.get_by_document_id(document_id)
|
||||
assert len(db_chunks) == len(chunks)
|
||||
|
||||
store.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_repository_crud():
|
||||
"""Test basic CRUD operations in ChunkRepository."""
|
||||
# Create an in-memory store
|
||||
store = Store(":memory:")
|
||||
chunk_repo = ChunkRepository(store)
|
||||
|
||||
# First create a document to reference
|
||||
document_id = None
|
||||
if store._connection is not None:
|
||||
cursor = store._connection.cursor()
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO documents (content, metadata, created_at, updated_at)
|
||||
VALUES (?, ?, datetime('now'), datetime('now'))
|
||||
""",
|
||||
("Test document content", "{}"),
|
||||
)
|
||||
document_id = cursor.lastrowid
|
||||
store._connection.commit()
|
||||
|
||||
assert document_id is not None, "Document ID should not be None"
|
||||
|
||||
# Test create chunk manually
|
||||
chunk = Chunk(
|
||||
document_id=document_id,
|
||||
content="Test chunk content",
|
||||
metadata={"test": "value"},
|
||||
)
|
||||
|
||||
created_chunk = await chunk_repo.create(chunk)
|
||||
assert created_chunk.id is not None
|
||||
assert created_chunk.content == "Test chunk content"
|
||||
|
||||
# Test get by ID
|
||||
retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
|
||||
assert retrieved_chunk is not None
|
||||
assert retrieved_chunk.content == "Test chunk content"
|
||||
assert retrieved_chunk.metadata["test"] == "value"
|
||||
|
||||
# Test update
|
||||
retrieved_chunk.content = "Updated chunk content"
|
||||
updated_chunk = await chunk_repo.update(retrieved_chunk)
|
||||
assert updated_chunk.content == "Updated chunk content"
|
||||
|
||||
# Test list all
|
||||
all_chunks = await chunk_repo.list_all()
|
||||
assert len(all_chunks) >= 1
|
||||
assert any(chunk.id == created_chunk.id for chunk in all_chunks)
|
||||
|
||||
# Test delete
|
||||
deleted = await chunk_repo.delete(created_chunk.id)
|
||||
assert deleted is True
|
||||
|
||||
# Verify chunk is gone
|
||||
retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
|
||||
assert retrieved_chunk is None
|
||||
|
||||
store.close()
|
||||
|
|
@ -3,13 +3,15 @@ from datasets import Dataset
|
|||
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_document_with_chunks(qa_corpus: Dataset):
|
||||
"""Test creating a document with chunks from the qa_corpus."""
|
||||
# Create an in-memory store
|
||||
"""Test creating a document with chunks from the qa_corpus using repository."""
|
||||
# Create an in-memory store and repository
|
||||
store = Store(":memory:")
|
||||
doc_repo = DocumentRepository(store)
|
||||
|
||||
# Get the first document from the corpus
|
||||
first_doc = qa_corpus[0]
|
||||
|
|
@ -22,7 +24,7 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
|
|||
)
|
||||
|
||||
# Create the document with chunks in the database
|
||||
created_document = await document.create_with_chunks(store)
|
||||
created_document = await doc_repo.create(document)
|
||||
|
||||
# Verify the document was created
|
||||
assert created_document.id is not None
|
||||
|
|
@ -60,30 +62,45 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_chunks(qa_corpus: Dataset):
|
||||
"""Test vector search functionality."""
|
||||
# Create an in-memory store
|
||||
async def test_document_repository_crud(qa_corpus: Dataset):
|
||||
"""Test CRUD operations in DocumentRepository."""
|
||||
# Create an in-memory store and repository
|
||||
store = Store(":memory:")
|
||||
doc_repo = DocumentRepository(store)
|
||||
|
||||
# Get the first document from the corpus
|
||||
first_doc = qa_corpus[0]
|
||||
document_text = first_doc["document_extracted"]
|
||||
|
||||
# Create and store a document
|
||||
# Create a document
|
||||
document = Document(
|
||||
content=document_text,
|
||||
metadata={"source": "qa_corpus"}
|
||||
metadata={"source": "test"}
|
||||
)
|
||||
await document.create_with_chunks(store)
|
||||
created_document = await doc_repo.create(document)
|
||||
|
||||
# Perform a search
|
||||
search_query = "news" # Simple query
|
||||
results = await Document.search_chunks(store, search_query, limit=3)
|
||||
# Test get_by_id
|
||||
assert created_document.id is not None
|
||||
retrieved_document = await doc_repo.get_by_id(created_document.id)
|
||||
assert retrieved_document is not None
|
||||
assert retrieved_document.content == document_text
|
||||
|
||||
# Verify search results
|
||||
assert len(results) <= 3
|
||||
assert all(hasattr(chunk, "content") for chunk in results)
|
||||
assert all(hasattr(chunk, "document_id") for chunk in results)
|
||||
assert all(chunk.document_id == document.id for chunk in results)
|
||||
# Test update (should regenerate chunks)
|
||||
retrieved_document.content = "Updated content for testing"
|
||||
updated_document = await doc_repo.update(retrieved_document)
|
||||
assert updated_document.content == "Updated content for testing"
|
||||
|
||||
# Test list_all
|
||||
all_documents = await doc_repo.list_all()
|
||||
assert len(all_documents) == 1
|
||||
assert all_documents[0].id == created_document.id
|
||||
|
||||
# Test delete
|
||||
deleted = await doc_repo.delete(created_document.id)
|
||||
assert deleted is True
|
||||
|
||||
# Verify document is gone
|
||||
retrieved_document = await doc_repo.get_by_id(created_document.id)
|
||||
assert retrieved_document is None
|
||||
|
||||
store.close()
|
||||
Loading…
Reference in a new issue