150 lines
4.9 KiB
Python
150 lines
4.9 KiB
Python
import pytest
|
|
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, temp_db_path):
|
|
"""Test creating a document with chunks from the qa_corpus using repository."""
|
|
# Create a store and repository
|
|
store = Store(temp_db_path)
|
|
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 instance
|
|
document = Document(
|
|
content=document_text,
|
|
metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")},
|
|
)
|
|
|
|
# Convert text to DoclingDocument for chunk creation
|
|
from haiku.rag.utils import text_to_docling_document
|
|
|
|
docling_document = text_to_docling_document(document_text, name="test.md")
|
|
|
|
# Create the document with chunks in the database
|
|
created_document = await doc_repo._create_and_chunk(document, docling_document)
|
|
|
|
# Verify the document was created
|
|
assert created_document.id is not None
|
|
assert created_document.content == document_text
|
|
|
|
# Check that chunks were created using repository
|
|
from haiku.rag.store.repositories.chunk import ChunkRepository
|
|
|
|
chunk_repo = ChunkRepository(store)
|
|
chunks = await chunk_repo.get_by_document_id(created_document.id)
|
|
|
|
assert len(chunks) > 0
|
|
|
|
# Verify chunk order is set correctly
|
|
for i, chunk in enumerate(chunks):
|
|
assert chunk.order == i
|
|
|
|
store.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_repository_crud(qa_corpus: Dataset, temp_db_path):
|
|
"""Test CRUD operations in DocumentRepository."""
|
|
# Create a store and repository
|
|
store = Store(temp_db_path)
|
|
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 with URI
|
|
test_uri = "file:///path/to/test.txt"
|
|
document = Document(
|
|
content=document_text, uri=test_uri, metadata={"source": "test"}
|
|
)
|
|
created_document = await doc_repo.create(document)
|
|
|
|
# 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
|
|
assert retrieved_document.uri == test_uri
|
|
|
|
# Test get_by_uri
|
|
retrieved_by_uri = await doc_repo.get_by_uri(test_uri)
|
|
assert retrieved_by_uri is not None
|
|
assert retrieved_by_uri.id == created_document.id
|
|
assert retrieved_by_uri.content == document_text
|
|
assert retrieved_by_uri.uri == test_uri
|
|
|
|
# Test get_by_uri with non-existent URI
|
|
non_existent = await doc_repo.get_by_uri("file:///non/existent.txt")
|
|
assert non_existent is None
|
|
|
|
# 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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
|
|
"""Test listing documents with filter clause."""
|
|
store = Store(temp_db_path)
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
doc1 = Document(
|
|
content=document_text,
|
|
uri="https://example.com/doc1.txt",
|
|
metadata={"source": "test", "category": "A"},
|
|
)
|
|
doc2 = Document(
|
|
content=document_text,
|
|
uri="https://arxiv.org/paper.pdf",
|
|
metadata={"source": "test", "category": "B"},
|
|
)
|
|
doc3 = Document(
|
|
content=document_text,
|
|
uri="https://example.com/doc3.txt",
|
|
metadata={"source": "test", "category": "A"},
|
|
)
|
|
|
|
created_doc1 = await doc_repo.create(doc1)
|
|
created_doc2 = await doc_repo.create(doc2)
|
|
created_doc3 = await doc_repo.create(doc3)
|
|
|
|
all_documents = await doc_repo.list_all()
|
|
assert len(all_documents) == 3
|
|
|
|
arxiv_documents = await doc_repo.list_all(filter="uri LIKE '%arxiv%'")
|
|
assert len(arxiv_documents) == 1
|
|
assert arxiv_documents[0].id == created_doc2.id
|
|
|
|
example_documents = await doc_repo.list_all(filter="uri LIKE '%example.com%'")
|
|
assert len(example_documents) == 2
|
|
assert {doc.id for doc in example_documents} == {created_doc1.id, created_doc3.id}
|
|
|
|
store.close()
|