199 lines
6.5 KiB
Python
199 lines
6.5 KiB
Python
import pytest
|
|
from datasets import Dataset
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.converters import get_converter
|
|
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_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
|
|
"""Test ChunkRepository operations."""
|
|
# Create client
|
|
client = HaikuRAG(db_path=temp_db_path, config=Config)
|
|
|
|
# Get the first document from the corpus
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
# Create a document first with chunks
|
|
created_document = await client.create_document(
|
|
content=document_text, metadata={"source": "test"}
|
|
)
|
|
assert created_document.id is not None
|
|
|
|
# Test getting chunks by document ID
|
|
chunks = await client.chunk_repository.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 client.chunk_repository.search(
|
|
"election", limit=2, search_type="vector"
|
|
)
|
|
assert len(results) <= 2
|
|
assert all(hasattr(chunk, "content") for chunk, _ in results)
|
|
|
|
# Test deleting chunks by document ID
|
|
deleted = await client.chunk_repository.delete_by_document_id(created_document.id)
|
|
assert deleted is True
|
|
|
|
# Verify chunks are gone
|
|
chunks_after_delete = await client.chunk_repository.get_by_document_id(
|
|
created_document.id
|
|
)
|
|
assert len(chunks_after_delete) == 0
|
|
|
|
client.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_chunks_for_document(qa_corpus: Dataset, temp_db_path):
|
|
"""Test creating chunks for a document."""
|
|
# Create a store and repositories
|
|
store = Store(temp_db_path)
|
|
chunk_repo = ChunkRepository(store)
|
|
doc_repo = DocumentRepository(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"})
|
|
created_document = await doc_repo.create(document)
|
|
document_id = created_document.id
|
|
|
|
assert document_id is not None, "Document ID should not be None"
|
|
|
|
# Convert text to DoclingDocument
|
|
converter = get_converter(Config)
|
|
docling_document = converter.convert_text(document_text, name="test.md")
|
|
|
|
# Test creating chunks for the document
|
|
chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document)
|
|
|
|
# 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
|
|
for i, chunk in enumerate(chunks):
|
|
assert chunk.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(temp_db_path):
|
|
"""Test basic CRUD operations in ChunkRepository."""
|
|
# Create a store
|
|
store = Store(temp_db_path)
|
|
chunk_repo = ChunkRepository(store)
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
# First create a document to reference
|
|
document = Document(content="Test document content", metadata={})
|
|
created_document = await doc_repo.create(document)
|
|
document_id = created_document.id
|
|
|
|
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 isinstance(created_chunk, 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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_adjacent_chunks(temp_db_path):
|
|
"""Test the get_adjacent_chunks repository method."""
|
|
store = Store(temp_db_path)
|
|
doc_repo = DocumentRepository(store)
|
|
chunk_repo = ChunkRepository(store)
|
|
|
|
# Create a simple document first
|
|
document_content = "Test document for chunking"
|
|
document = Document(content=document_content)
|
|
created_document = await doc_repo.create(document)
|
|
|
|
# Manually create multiple chunks with order metadata
|
|
chunks_data = [
|
|
("First chunk content", 0),
|
|
("Second chunk content", 1),
|
|
("Third chunk content", 2),
|
|
("Fourth chunk content", 3),
|
|
("Fifth chunk content", 4),
|
|
]
|
|
|
|
created_chunks = []
|
|
for content, order in chunks_data:
|
|
chunk = Chunk(document_id=created_document.id, content=content, order=order)
|
|
created_chunk = await chunk_repo.create(chunk)
|
|
created_chunks.append(created_chunk)
|
|
|
|
# Test with the middle chunk (index 2, order 2)
|
|
middle_chunk = created_chunks[2]
|
|
|
|
# Get adjacent chunks (1 before and after)
|
|
adjacent_chunks = await chunk_repo.get_adjacent_chunks(middle_chunk, 1)
|
|
|
|
# Should have 2 chunks (one before, one after)
|
|
assert len(adjacent_chunks) == 2
|
|
|
|
# Should not include the original chunk
|
|
assert middle_chunk.id not in [chunk.id for chunk in adjacent_chunks]
|
|
|
|
# Should include chunks with order 1 and 3
|
|
orders = [chunk.order for chunk in adjacent_chunks]
|
|
assert 1 in orders
|
|
assert 3 in orders
|
|
|
|
# All adjacent chunks should be from the same document
|
|
for chunk in adjacent_chunks:
|
|
assert chunk.document_id == created_document.id
|
|
|
|
store.close()
|