1432 lines
56 KiB
Python
1432 lines
56 KiB
Python
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from datasets import Dataset
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
from haiku.rag.store.models.document import Document
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
|
|
"""Test HaikuRAG CRUD operations for documents."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Get test data
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
test_uri = "file:///path/to/test.txt"
|
|
test_metadata = {"source": "test", "topic": "testing"}
|
|
|
|
# Test create_document
|
|
created_doc = await client.create_document(
|
|
content=document_text, uri=test_uri, metadata=test_metadata
|
|
)
|
|
|
|
assert created_doc.id is not None
|
|
# Content is stored as markdown export, check key text is preserved
|
|
assert "Jakarta" in created_doc.content
|
|
assert created_doc.uri == test_uri
|
|
assert created_doc.metadata == test_metadata
|
|
|
|
# Test get_document_by_id
|
|
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
|
assert retrieved_doc is not None
|
|
assert retrieved_doc.id == created_doc.id
|
|
assert "Jakarta" in retrieved_doc.content
|
|
assert retrieved_doc.uri == test_uri
|
|
|
|
# Test get_document_by_uri
|
|
retrieved_by_uri = await client.get_document_by_uri(test_uri)
|
|
assert retrieved_by_uri is not None
|
|
assert retrieved_by_uri.id == created_doc.id
|
|
assert "Jakarta" in retrieved_by_uri.content
|
|
|
|
# Test get_document_by_uri with non-existent URI
|
|
non_existent = await client.get_document_by_uri("file:///non/existent.txt")
|
|
assert non_existent is None
|
|
|
|
# Test update_document
|
|
updated_doc = await client.update_document(
|
|
document_id=retrieved_doc.id, # type: ignore[arg-type]
|
|
content="Updated content",
|
|
)
|
|
assert updated_doc.content == "Updated content"
|
|
|
|
# Test list_documents
|
|
all_docs = await client.list_documents()
|
|
assert len(all_docs) == 1
|
|
assert all_docs[0].id == created_doc.id
|
|
|
|
# Test list_documents with pagination
|
|
limited_docs = await client.list_documents(limit=10, offset=0)
|
|
assert len(limited_docs) == 1
|
|
|
|
# Test delete_document
|
|
deleted = await client.delete_document(created_doc.id)
|
|
assert deleted is True
|
|
|
|
# Verify document is gone
|
|
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
|
assert retrieved_doc is None
|
|
|
|
# Test delete non-existent document
|
|
deleted_again = await client.delete_document(created_doc.id)
|
|
assert deleted_again is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document(qa_corpus: Dataset, temp_db_path):
|
|
"""Test updating document with individual parameters."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Get test data
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
test_uri = "file:///path/to/test.txt"
|
|
test_metadata = {"source": "test", "topic": "testing"}
|
|
|
|
# Create a document
|
|
created_doc = await client.create_document(
|
|
content=document_text,
|
|
uri=test_uri,
|
|
title="Original Title",
|
|
metadata=test_metadata,
|
|
)
|
|
assert created_doc.id is not None
|
|
original_id = created_doc.id
|
|
|
|
# Test updating only content
|
|
updated_doc = await client.update_document(
|
|
document_id=original_id, content="Updated content only"
|
|
)
|
|
assert updated_doc.id == original_id
|
|
assert updated_doc.content == "Updated content only"
|
|
assert updated_doc.title == "Original Title"
|
|
assert updated_doc.uri == test_uri
|
|
|
|
# Test updating only metadata
|
|
new_metadata = {"source": "updated", "version": "2.0"}
|
|
updated_doc = await client.update_document(
|
|
document_id=original_id, metadata=new_metadata
|
|
)
|
|
assert updated_doc.metadata == new_metadata
|
|
assert (
|
|
updated_doc.content == "Updated content only"
|
|
) # Should keep previous update
|
|
|
|
# Test updating only title
|
|
updated_doc = await client.update_document(
|
|
document_id=original_id, title="New Title"
|
|
)
|
|
assert updated_doc.title == "New Title"
|
|
assert updated_doc.content == "Updated content only"
|
|
assert updated_doc.metadata == new_metadata
|
|
|
|
# Test updating multiple fields at once
|
|
custom_chunks = [
|
|
Chunk(content="Custom chunk 1", order=0),
|
|
Chunk(content="Custom chunk 2", order=1),
|
|
]
|
|
updated_doc = await client.update_document(
|
|
document_id=original_id,
|
|
content="Content with custom chunks",
|
|
title="Final Title",
|
|
metadata={"final": "true"},
|
|
chunks=custom_chunks,
|
|
)
|
|
assert updated_doc.id == original_id
|
|
assert updated_doc.content == "Content with custom chunks"
|
|
assert updated_doc.title == "Final Title"
|
|
assert updated_doc.metadata == {"final": "true"}
|
|
|
|
# Verify the custom chunks were created
|
|
doc_chunks = await client.chunk_repository.get_by_document_id(original_id)
|
|
assert len(doc_chunks) == 2
|
|
assert doc_chunks[0].content == "Custom chunk 1"
|
|
assert doc_chunks[1].content == "Custom chunk 2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_source(temp_db_path):
|
|
"""Test creating a document from a file source."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
test_content = "This is test content from a file."
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text(test_content)
|
|
|
|
# Test create_document_from_source with Path
|
|
doc = await client.create_document_from_source(source=temp_path)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.id is not None
|
|
assert doc.content == test_content
|
|
assert doc.uri == temp_path.as_uri()
|
|
assert "contentType" in doc.metadata
|
|
assert "md5" in doc.metadata
|
|
assert doc.metadata["contentType"] == "text/plain"
|
|
|
|
# Test create_document_from_source with string path
|
|
doc2 = await client.create_document_from_source(source=str(temp_path))
|
|
assert isinstance(doc2, Document)
|
|
|
|
assert doc2.id is not None
|
|
assert doc2.content == test_content
|
|
assert doc2.uri == temp_path.as_uri()
|
|
assert "contentType" in doc2.metadata
|
|
assert "md5" in doc2.metadata
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_source_with_title(temp_db_path):
|
|
"""Test creating a document from a file source with a title."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
test_content = "This is test content from a file."
|
|
temp_path = Path(temp_dir) / "test_title.txt"
|
|
temp_path.write_text(test_content)
|
|
|
|
doc = await client.create_document_from_source(
|
|
source=temp_path, title="My Doc"
|
|
)
|
|
assert isinstance(doc, Document)
|
|
assert doc.id is not None
|
|
assert doc.title == "My Doc"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_title_noop_behavior(temp_db_path):
|
|
"""When content is unchanged, updating title should update document without re-chunking."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test_update_title.txt"
|
|
temp_path.write_text("Original content")
|
|
|
|
doc1 = await client.create_document_from_source(temp_path, title="Title A")
|
|
assert isinstance(doc1, Document)
|
|
assert doc1.id is not None
|
|
|
|
# Re-add with same content but new title
|
|
doc2 = await client.create_document_from_source(temp_path, title="Title B")
|
|
assert isinstance(doc2, Document)
|
|
assert doc2.id == doc1.id
|
|
# Fetch and verify title updated
|
|
got = await client.get_document_by_id(doc1.id)
|
|
assert got is not None
|
|
assert got.title == "Title B"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_source_unsupported(temp_db_path):
|
|
"""Test creating a document from an unsupported file type."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a temporary file with unsupported extension
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".unsupported", delete=False
|
|
) as f:
|
|
f.write("content")
|
|
temp_path = Path(f.name)
|
|
|
|
# Should raise ValueError for unsupported extension
|
|
with pytest.raises(ValueError, match="Unsupported file extension"):
|
|
await client.create_document_from_source(temp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_source_nonexistent(temp_db_path):
|
|
"""Test creating a document from a non-existent file."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
non_existent_path = Path("/non/existent/file.txt")
|
|
|
|
# Should raise ValueError when file doesn't exist
|
|
with pytest.raises(ValueError, match="File does not exist"):
|
|
await client.create_document_from_source(non_existent_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_directory(temp_db_path):
|
|
"""Test creating documents from a directory recursively."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
test_dir = Path(temp_dir) / "test_docs"
|
|
test_dir.mkdir()
|
|
|
|
(test_dir / "doc1.txt").write_text("Content of doc1")
|
|
(test_dir / "doc2.md").write_text("# Content of doc2")
|
|
|
|
subdir = test_dir / "subdir"
|
|
subdir.mkdir()
|
|
(subdir / "doc3.py").write_text("print('hello')")
|
|
|
|
(test_dir / "unsupported.xyz").write_text("unsupported file")
|
|
|
|
result = await client.create_document_from_source(test_dir)
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
|
|
for doc in result:
|
|
assert doc.id is not None
|
|
assert doc.uri is not None
|
|
assert "md5" in doc.metadata
|
|
assert "contentType" in doc.metadata
|
|
|
|
uris = [doc.uri for doc in result if doc.uri]
|
|
assert any("doc1.txt" in uri for uri in uris)
|
|
assert any("doc2.md" in uri for uri in uris)
|
|
assert any("doc3.py" in uri for uri in uris)
|
|
assert not any("unsupported.xyz" in uri for uri in uris)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_directory_with_filters(
|
|
monkeypatch, temp_db_path
|
|
):
|
|
"""Test creating documents from a directory with ignore and include patterns."""
|
|
# Mock config to have ignore and include patterns
|
|
monkeypatch.setattr(
|
|
"haiku.rag.client.Config.monitor.ignore_patterns", ["**/ignore_me/**", "*.log"]
|
|
)
|
|
monkeypatch.setattr(
|
|
"haiku.rag.client.Config.monitor.include_patterns", ["**/include/**/*.txt"]
|
|
)
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
test_dir = Path(temp_dir) / "test_docs"
|
|
test_dir.mkdir()
|
|
|
|
# Create files in include directory - should be included
|
|
include_dir = test_dir / "include"
|
|
include_dir.mkdir()
|
|
(include_dir / "doc1.txt").write_text("Content of doc1")
|
|
(include_dir / "doc2.txt").write_text("Content of doc2")
|
|
|
|
# Create files outside include directory - should be excluded by include pattern
|
|
(test_dir / "doc3.txt").write_text("Content of doc3")
|
|
|
|
# Create files in ignore directory - should be excluded by ignore pattern
|
|
ignore_dir = test_dir / "ignore_me"
|
|
ignore_dir.mkdir()
|
|
(ignore_dir / "doc4.txt").write_text("Content of doc4")
|
|
|
|
# Create log file - should be excluded by ignore pattern
|
|
(test_dir / "debug.log").write_text("log content")
|
|
|
|
result = await client.create_document_from_source(test_dir)
|
|
|
|
assert isinstance(result, list)
|
|
# Should only include doc1.txt and doc2.txt from include directory
|
|
assert len(result) == 2
|
|
|
|
uris = [doc.uri for doc in result if doc.uri]
|
|
assert any("doc1.txt" in uri for uri in uris)
|
|
assert any("doc2.txt" in uri for uri in uris)
|
|
assert not any("doc3.txt" in uri for uri in uris)
|
|
assert not any("doc4.txt" in uri for uri in uris)
|
|
assert not any("debug.log" in uri for uri in uris)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_url(temp_db_path):
|
|
"""Test creating a document from a URL."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Mock the HTTP response
|
|
mock_response = AsyncMock()
|
|
mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
|
|
mock_response.headers = {"content-type": "text/html"}
|
|
mock_response.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
|
doc = await client.create_document_from_source(
|
|
source="https://example.com/test.html", metadata={"source_type": "web"}
|
|
)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.id is not None
|
|
assert "Test Page" in doc.content
|
|
assert "test content" in doc.content
|
|
assert doc.uri == "https://example.com/test.html"
|
|
assert doc.metadata["source_type"] == "web"
|
|
assert "contentType" in doc.metadata
|
|
assert "md5" in doc.metadata
|
|
assert doc.metadata["contentType"] == "text/html"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_url_with_different_content_types(
|
|
temp_db_path,
|
|
):
|
|
"""Test creating documents from URLs with different content types."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Test JSON content
|
|
mock_json_response = AsyncMock()
|
|
mock_json_response.content = (
|
|
b'{"title": "Test JSON", "content": "This is JSON content"}'
|
|
)
|
|
mock_json_response.headers = {"content-type": "application/json"}
|
|
mock_json_response.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_json_response):
|
|
doc = await client.create_document_from_source(
|
|
"https://api.example.com/data.json"
|
|
)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.id is not None
|
|
assert "Test JSON" in doc.content
|
|
assert doc.uri == "https://api.example.com/data.json"
|
|
assert "contentType" in doc.metadata
|
|
assert "md5" in doc.metadata
|
|
assert doc.metadata["contentType"] == "application/json"
|
|
|
|
# Test plain text content
|
|
mock_text_response = AsyncMock()
|
|
mock_text_response.content = b"This is plain text content from a URL."
|
|
mock_text_response.headers = {"content-type": "text/plain"}
|
|
mock_text_response.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_text_response):
|
|
doc = await client.create_document_from_source(
|
|
"https://example.com/readme.txt"
|
|
)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.id is not None
|
|
assert doc.content == "This is plain text content from a URL."
|
|
assert doc.uri == "https://example.com/readme.txt"
|
|
assert "contentType" in doc.metadata
|
|
assert "md5" in doc.metadata
|
|
assert doc.metadata["contentType"] == "text/plain"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_url_unsupported_content(temp_db_path):
|
|
"""Test creating a document from URL with unsupported content type."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Mock response with unsupported content type
|
|
mock_response = AsyncMock()
|
|
mock_response.content = b"binary content"
|
|
mock_response.headers = {"content-type": "application/octet-stream"}
|
|
mock_response.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
|
with pytest.raises(ValueError, match="Unsupported content type"):
|
|
await client.create_document_from_source(
|
|
"https://example.com/binary.bin"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_url_http_error(temp_db_path):
|
|
"""Test handling HTTP errors when creating document from URL."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with patch("httpx.AsyncClient.get") as mock_get:
|
|
mock_get.side_effect = httpx.HTTPStatusError(
|
|
"404 Not Found",
|
|
request=httpx.Request("GET", "https://example.com/notfound.html"),
|
|
response=httpx.Response(404),
|
|
)
|
|
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await client.create_document_from_source(
|
|
"https://example.com/notfound.html"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_extension_from_content_type_or_url(temp_db_path):
|
|
"""Test the helper method for determining file extensions."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Test content type mappings
|
|
assert (
|
|
client._get_extension_from_content_type_or_url("", "text/html") == ".html"
|
|
)
|
|
assert (
|
|
client._get_extension_from_content_type_or_url("", "application/pdf")
|
|
== ".pdf"
|
|
)
|
|
assert (
|
|
client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
|
|
)
|
|
|
|
# Test URL extension detection
|
|
assert (
|
|
client._get_extension_from_content_type_or_url(
|
|
"https://example.com/doc.pdf", ""
|
|
)
|
|
== ".pdf"
|
|
)
|
|
assert (
|
|
client._get_extension_from_content_type_or_url(
|
|
"https://example.com/data.json", ""
|
|
)
|
|
== ".json"
|
|
)
|
|
|
|
# Test default fallback
|
|
assert (
|
|
client._get_extension_from_content_type_or_url("https://example.com/", "")
|
|
== ".html"
|
|
)
|
|
|
|
# Test content type priority over URL extension
|
|
assert (
|
|
client._get_extension_from_content_type_or_url(
|
|
"https://example.com/file.txt", "application/pdf"
|
|
)
|
|
== ".pdf"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_metadata_content_type_and_md5(temp_db_path):
|
|
"""Test that contentType and md5 metadata are correctly set."""
|
|
import hashlib
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a temporary file with known content
|
|
test_content = "Test content for MD5 calculation."
|
|
expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text(test_content)
|
|
|
|
doc = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.metadata["contentType"] == "text/plain"
|
|
assert doc.metadata["md5"] == expected_md5
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.content = test_content.encode()
|
|
mock_response.headers = {"content-type": "text/plain"}
|
|
mock_response.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
|
url_doc = await client.create_document_from_source(
|
|
"https://example.com/test.txt"
|
|
)
|
|
assert isinstance(url_doc, Document)
|
|
|
|
assert url_doc.metadata["contentType"] == "text/plain"
|
|
assert url_doc.metadata["md5"] == expected_md5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_update_no_op_behavior(temp_db_path):
|
|
"""Test create/update/no-op behavior based on MD5 changes."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a temporary file
|
|
test_content = "Original content for testing."
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text(test_content)
|
|
|
|
# First call - should create new document
|
|
doc1 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc1, Document)
|
|
assert doc1.id is not None
|
|
assert doc1.content == test_content
|
|
original_id = doc1.id
|
|
|
|
# Second call with same content - should return existing document (no-op)
|
|
doc2 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc2, Document)
|
|
assert doc2.id == original_id # Same document
|
|
assert doc2.content == test_content
|
|
|
|
# Modify file content
|
|
updated_content = "Updated content for testing."
|
|
temp_path.write_text(updated_content)
|
|
|
|
# Third call with changed content - should update existing document
|
|
doc3 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc3, Document)
|
|
assert doc3.id == original_id # Same document ID
|
|
assert doc3.content == updated_content # Updated content
|
|
|
|
# Verify the document was actually updated in database
|
|
retrieved_doc = await client.get_document_by_id(original_id)
|
|
assert retrieved_doc is not None
|
|
assert retrieved_doc.content == updated_content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_unchanged_file_keeps_timestamp(temp_db_path):
|
|
"""Test that unchanged files don't update the updated_at timestamp."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a temporary file
|
|
test_content = "Test content for timestamp check."
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text(test_content)
|
|
|
|
# First call - create document
|
|
doc1 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc1, Document)
|
|
original_updated_at = doc1.updated_at
|
|
|
|
# Second call with same content - should not update timestamp
|
|
doc2 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc2, Document)
|
|
assert doc2.id == doc1.id
|
|
assert doc2.updated_at == original_updated_at # Timestamp should not change
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_url_create_update_no_op_behavior(temp_db_path):
|
|
"""Test create/update/no-op behavior for URLs based on MD5 changes."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
url = "https://example.com/test.txt"
|
|
original_content = b"Original URL content"
|
|
updated_content = b"Updated URL content"
|
|
|
|
# Mock first response
|
|
mock_response1 = AsyncMock()
|
|
mock_response1.content = original_content
|
|
mock_response1.headers = {"content-type": "text/plain"}
|
|
mock_response1.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_response1):
|
|
# First call - should create new document
|
|
doc1 = await client.create_document_from_source(url)
|
|
assert isinstance(doc1, Document)
|
|
assert doc1.id is not None
|
|
original_id = doc1.id
|
|
|
|
# Second call with same content - should return existing document (no-op)
|
|
doc2 = await client.create_document_from_source(url)
|
|
assert isinstance(doc2, Document)
|
|
assert doc2.id == original_id # Same document
|
|
|
|
mock_response2 = AsyncMock()
|
|
mock_response2.content = updated_content
|
|
mock_response2.headers = {"content-type": "text/plain"}
|
|
mock_response2.raise_for_status = AsyncMock()
|
|
|
|
with patch("httpx.AsyncClient.get", return_value=mock_response2):
|
|
# Third call with changed content - should update existing document
|
|
doc3 = await client.create_document_from_source(url)
|
|
assert isinstance(doc3, Document)
|
|
assert doc3.id == original_id # Same document ID
|
|
assert doc3.content == updated_content.decode() # Updated content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_search(temp_db_path):
|
|
"""Test HaikuRAG search functionality."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Add multiple documents to search from
|
|
doc1_text = "Python is a high-level programming language known for its simplicity and readability."
|
|
doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
|
|
doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."
|
|
|
|
# Create documents
|
|
doc1 = await client.create_document(
|
|
content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
|
|
)
|
|
doc2 = await client.create_document(
|
|
content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"}
|
|
)
|
|
await client.create_document(
|
|
content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"}
|
|
)
|
|
|
|
# Test search with keyword that should match doc1
|
|
results = await client.search("Python programming", limit=3)
|
|
|
|
assert len(results) > 0
|
|
# Verify results are SearchResult objects with expected fields
|
|
first_result = results[0]
|
|
assert first_result.content
|
|
assert first_result.score >= 0
|
|
assert first_result.document_id == doc1.id
|
|
|
|
# Test search with different query
|
|
ml_results = await client.search("machine learning data", limit=2)
|
|
assert len(ml_results) > 0
|
|
|
|
# Verify first result is from the machine learning document (doc2)
|
|
first_ml_result = ml_results[0]
|
|
assert first_ml_result.document_id == doc2.id
|
|
|
|
# Test search with limit parameter
|
|
limited_results = await client.search("programming", limit=1)
|
|
assert len(limited_results) <= 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_async_context_manager(temp_db_path):
|
|
"""Test HaikuRAG as async context manager."""
|
|
|
|
# Test that context manager works and auto-closes
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a document to ensure the client works
|
|
doc = await client.create_document(
|
|
content="Test content for context manager",
|
|
uri="test://context",
|
|
metadata={"test": "context_manager"},
|
|
)
|
|
|
|
assert doc.id is not None
|
|
assert doc.content == "Test content for context manager"
|
|
|
|
# Test search works within context
|
|
results = await client.search("Test content", limit=1)
|
|
assert len(results) > 0
|
|
|
|
# Context manager should have automatically closed the connection
|
|
# We can't easily test that the connection is closed without accessing internals,
|
|
# but the test passing means the context manager methods work correctly
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_import_document_with_custom_chunks(temp_db_path):
|
|
"""Test importing a document with pre-created chunks."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create some custom chunks with and without embeddings
|
|
chunks = [
|
|
Chunk(
|
|
content="This is the first chunk",
|
|
metadata={"custom": "metadata1"},
|
|
order=0,
|
|
),
|
|
Chunk(
|
|
content="This is the second chunk",
|
|
metadata={"custom": "metadata2"},
|
|
embedding=[0.1] * Config.embeddings.model.vector_dim,
|
|
order=1,
|
|
), # With embedding
|
|
Chunk(
|
|
content="This is the third chunk",
|
|
metadata={"custom": "metadata3"},
|
|
order=2,
|
|
),
|
|
]
|
|
|
|
# Import document with custom chunks
|
|
document = await client.import_document(
|
|
content="Full document content", chunks=chunks
|
|
)
|
|
|
|
assert document.id is not None
|
|
assert document.content == "Full document content"
|
|
|
|
# Verify the chunks were created correctly
|
|
doc_chunks = await client.chunk_repository.get_by_document_id(document.id)
|
|
assert len(doc_chunks) == 3
|
|
|
|
# Check chunks have correct content, document_id, and order from list position
|
|
for i, chunk in enumerate(doc_chunks):
|
|
assert chunk.document_id == document.id
|
|
assert chunk.content == chunks[i].content
|
|
assert chunk.order == i # Order should be set from list position
|
|
assert (
|
|
chunk.metadata["custom"] == f"metadata{i + 1}"
|
|
) # Original metadata preserved
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_ask(monkeypatch, temp_db_path):
|
|
"""Test asking questions returns answer and citations."""
|
|
from pydantic_ai.models.test import TestModel
|
|
|
|
# Mock get_model to return TestModel
|
|
monkeypatch.setattr(
|
|
"haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel()
|
|
)
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a test document for the agent to search
|
|
await client.create_document(
|
|
content="Python is a high-level programming language.", uri="test.txt"
|
|
)
|
|
|
|
# Use real QA agent with TestModel
|
|
answer, citations = await client.ask("What is Python?")
|
|
|
|
# TestModel will generate a valid string response
|
|
assert answer is not None
|
|
assert isinstance(answer, str)
|
|
assert isinstance(citations, list)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_expand_context(temp_db_path):
|
|
"""Test that expand_context method exists and works with basic input."""
|
|
from haiku.rag.store.models import SearchResult
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(content="Simple test content")
|
|
assert doc.id is not None
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
|
|
search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
|
|
expanded_results = await client.expand_context(search_results)
|
|
|
|
assert len(expanded_results) == 1
|
|
assert expanded_results[0].score == 0.9
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_stores_docling_json(temp_db_path):
|
|
"""Test that create_document stores DoclingDocument JSON."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="Test content for docling storage",
|
|
uri="test://docling",
|
|
metadata={"test": "docling_storage"},
|
|
)
|
|
|
|
assert doc.id is not None
|
|
assert doc.docling_document_json is not None
|
|
assert doc.docling_version is not None
|
|
|
|
# Verify JSON is valid and can be parsed
|
|
import json
|
|
|
|
parsed = json.loads(doc.docling_document_json)
|
|
assert "version" in parsed
|
|
assert parsed["version"] == doc.docling_version
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_import_document_without_docling(temp_db_path):
|
|
"""Test that import_document without docling params does not store docling JSON."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
custom_chunks = [Chunk(content="Custom chunk", order=0)]
|
|
|
|
doc = await client.import_document(content="Test content", chunks=custom_chunks)
|
|
|
|
assert doc.id is not None
|
|
# When no docling params provided, they remain None
|
|
assert doc.docling_document_json is None
|
|
assert doc.docling_version is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_import_document_validates_docling_params(temp_db_path):
|
|
"""Test that import_document validates docling parameters."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
custom_chunks = [Chunk(content="Custom chunk", order=0)]
|
|
|
|
# Should fail if only one docling param is provided
|
|
with pytest.raises(ValueError, match="must both be provided"):
|
|
await client.import_document(
|
|
content="Test content",
|
|
chunks=custom_chunks,
|
|
docling_document_json='{"some": "json"}',
|
|
# Missing docling_version
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="must both be provided"):
|
|
await client.import_document(
|
|
content="Test content",
|
|
chunks=custom_chunks,
|
|
docling_version="1.0.0",
|
|
# Missing docling_document_json
|
|
)
|
|
|
|
# Should fail with invalid JSON
|
|
with pytest.raises(ValueError, match="Invalid docling_document_json"):
|
|
await client.import_document(
|
|
content="Test content",
|
|
chunks=custom_chunks,
|
|
docling_document_json='{"invalid": "not a docling document"}',
|
|
docling_version="1.0.0",
|
|
)
|
|
|
|
# Should fail if neither content nor docling_document_json provided
|
|
with pytest.raises(ValueError, match="Either content or docling_document_json"):
|
|
await client.import_document(chunks=custom_chunks)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_import_document_extracts_content_from_docling(temp_db_path):
|
|
"""Test that import_document extracts content from DoclingDocument when not provided."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a docling document with some content
|
|
docling_doc = DoclingDocument(name="test")
|
|
docling_doc.add_text(
|
|
label=DocItemLabel.TEXT, text="Content from docling document"
|
|
)
|
|
|
|
custom_chunks = [Chunk(content="Chunk content", order=0)]
|
|
|
|
# Import without content - should extract from docling
|
|
doc = await client.import_document(
|
|
chunks=custom_chunks,
|
|
docling_document_json=docling_doc.model_dump_json(),
|
|
docling_version=docling_doc.version,
|
|
)
|
|
|
|
assert doc.id is not None
|
|
assert "Content from docling document" in doc.content
|
|
assert doc.docling_document_json == docling_doc.model_dump_json()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_from_file_stores_docling_json(temp_db_path):
|
|
"""Test that create_document_from_source stores DoclingDocument JSON for files."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text("Test file content")
|
|
|
|
doc = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc, Document)
|
|
|
|
assert doc.id is not None
|
|
assert doc.docling_document_json is not None
|
|
assert doc.docling_version is not None
|
|
|
|
# Verify the stored document also has the JSON
|
|
retrieved = await client.get_document_by_id(doc.id)
|
|
assert retrieved is not None
|
|
assert retrieved.docling_document_json == doc.docling_document_json
|
|
assert retrieved.docling_version == doc.docling_version
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document_stores_docling_json(temp_db_path):
|
|
"""Test that update_document stores DoclingDocument JSON when content changes."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create initial document
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
original_json = doc.docling_document_json
|
|
|
|
# Update content via update_document
|
|
updated_doc = await client.update_document(
|
|
document_id=doc.id, content="New content via fields update"
|
|
)
|
|
|
|
assert updated_doc.docling_document_json is not None
|
|
assert updated_doc.docling_version is not None
|
|
# JSON should be different because content changed
|
|
assert updated_doc.docling_document_json != original_json
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document_with_custom_chunks_no_docling_json(
|
|
temp_db_path,
|
|
):
|
|
"""Test that update_document with custom chunks does not update docling JSON."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create initial document
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
original_json = doc.docling_document_json
|
|
|
|
# Update with custom chunks
|
|
custom_chunks = [Chunk(content="Custom chunk", order=0)]
|
|
updated_doc = await client.update_document(
|
|
document_id=doc.id, content="New content", chunks=custom_chunks
|
|
)
|
|
|
|
# Docling JSON should remain unchanged (no conversion when custom chunks provided)
|
|
assert updated_doc.docling_document_json == original_json
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document_content_docling_mutually_exclusive(
|
|
temp_db_path,
|
|
):
|
|
"""Test that content and docling_document_json cannot both be provided."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
|
|
# Create a docling document
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
docling_doc = DoclingDocument(name="test")
|
|
docling_doc.add_text(label=DocItemLabel.TEXT, text="Some text")
|
|
|
|
with pytest.raises(ValueError, match="mutually exclusive"):
|
|
await client.update_document(
|
|
document_id=doc.id,
|
|
content="New content",
|
|
docling_document_json=docling_doc.model_dump_json(),
|
|
docling_version=docling_doc.version,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document_with_docling_rechunks(temp_db_path):
|
|
"""Test that providing docling_document_json without chunks triggers rechunk."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create initial document
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
|
|
# Create a new docling document with different content
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
docling_doc = DoclingDocument(name="updated")
|
|
docling_doc.add_text(
|
|
label=DocItemLabel.TEXT,
|
|
text="Completely different text from docling document",
|
|
)
|
|
|
|
# Update with docling document only - should rechunk from it
|
|
updated_doc = await client.update_document(
|
|
document_id=doc.id,
|
|
docling_document_json=docling_doc.model_dump_json(),
|
|
docling_version=docling_doc.version,
|
|
)
|
|
|
|
# Content should be extracted from docling document
|
|
assert "Completely different text" in updated_doc.content
|
|
assert updated_doc.docling_document_json == docling_doc.model_dump_json()
|
|
assert updated_doc.docling_version == docling_doc.version
|
|
|
|
# Chunks should be regenerated
|
|
new_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(new_chunks) > 0
|
|
# Content should differ from original
|
|
assert new_chunks[0].content != original_chunks[0].content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_update_document_docling_with_chunks(temp_db_path):
|
|
"""Test that providing both docling_document_json and chunks stores both."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create initial document
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
|
|
# Create a docling document
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
docling_doc = DoclingDocument(name="custom")
|
|
docling_doc.add_text(label=DocItemLabel.TEXT, text="Text from docling")
|
|
|
|
# Provide both docling and custom chunks
|
|
custom_chunks = [
|
|
Chunk(content="Custom chunk 1", order=0),
|
|
Chunk(content="Custom chunk 2", order=1),
|
|
]
|
|
|
|
updated_doc = await client.update_document(
|
|
document_id=doc.id,
|
|
chunks=custom_chunks,
|
|
docling_document_json=docling_doc.model_dump_json(),
|
|
docling_version=docling_doc.version,
|
|
)
|
|
|
|
# Content should be extracted from docling (since content wasn't provided)
|
|
assert "Text from docling" in updated_doc.content
|
|
assert updated_doc.docling_document_json == docling_doc.model_dump_json()
|
|
|
|
# Custom chunks should be used (not rechunked from docling)
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks) == 2
|
|
assert chunks[0].content == "Custom chunk 1"
|
|
assert chunks[1].content == "Custom chunk 2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_file_update_stores_docling_json(temp_db_path):
|
|
"""Test that updating a file re-stores DoclingDocument JSON."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text("Original content")
|
|
|
|
# Create initial document
|
|
doc1 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc1, Document)
|
|
original_json = doc1.docling_document_json
|
|
original_version = doc1.docling_version
|
|
|
|
# Modify file
|
|
temp_path.write_text("Modified content")
|
|
|
|
# Update document from source
|
|
doc2 = await client.create_document_from_source(temp_path)
|
|
assert isinstance(doc2, Document)
|
|
assert doc2.id == doc1.id # Same document
|
|
|
|
# Docling JSON should be updated
|
|
assert doc2.docling_document_json is not None
|
|
assert doc2.docling_document_json != original_json
|
|
assert doc2.docling_version == original_version # Version stays same
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_visualize_chunk_no_document(temp_db_path):
|
|
"""Test visualize_chunk returns empty list when chunk has no document_id."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
chunk = Chunk(content="Orphan chunk", order=0)
|
|
images = await client.visualize_chunk(chunk)
|
|
assert images == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_visualize_chunk_no_docling_document(temp_db_path):
|
|
"""Test visualize_chunk returns empty list when document has no DoclingDocument."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Import document with custom chunks (no DoclingDocument)
|
|
custom_chunks = [Chunk(content="Custom chunk", order=0)]
|
|
doc = await client.import_document(content="Test content", chunks=custom_chunks)
|
|
|
|
assert doc.id is not None
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks) == 1
|
|
|
|
images = await client.visualize_chunk(chunks[0])
|
|
assert images == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path):
|
|
"""Test visualize_chunk returns empty list when chunk has no bounding boxes."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create document from text (will have DoclingDocument but no page images)
|
|
doc = await client.create_document(
|
|
content="Simple text content without structure",
|
|
uri="test://simple",
|
|
)
|
|
|
|
assert doc.id is not None
|
|
assert doc.docling_document_json is not None
|
|
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks) >= 1
|
|
|
|
# Text documents converted via markdown won't have page images
|
|
# so visualize_chunk should return empty list
|
|
images = await client.visualize_chunk(chunks[0])
|
|
assert images == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_visualize_chunk_returns_list(temp_db_path):
|
|
"""Test visualize_chunk returns a list (empty or with images)."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create a structured document
|
|
markdown_content = """# Chapter 1
|
|
|
|
This is paragraph one about topic A.
|
|
|
|
This is paragraph two about topic A continued.
|
|
|
|
# Chapter 2
|
|
|
|
This is paragraph four about topic C.
|
|
"""
|
|
doc = await client.create_document(
|
|
content=markdown_content,
|
|
uri="test://structured",
|
|
)
|
|
|
|
assert doc.id is not None
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
|
|
# Find a chunk with doc_item_refs
|
|
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
|
|
|
|
if chunks_with_refs:
|
|
# visualize_chunk should return a list (possibly empty if no page images)
|
|
images = await client.visualize_chunk(chunks_with_refs[0])
|
|
assert isinstance(images, list)
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.asyncio
|
|
async def test_client_visualize_chunk_with_pdf(temp_db_path):
|
|
"""Test visualize_chunk returns images with bounding boxes for PDF documents."""
|
|
from PIL.Image import Image as PILImage
|
|
|
|
pdf_path = Path("tests/data/doclaynet.pdf")
|
|
if not pdf_path.exists():
|
|
pytest.skip("doclaynet.pdf not found")
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document_from_source(pdf_path)
|
|
assert isinstance(doc, Document)
|
|
assert doc.id is not None
|
|
assert doc.docling_document_json is not None
|
|
|
|
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(chunks) > 0
|
|
|
|
# Find a chunk with doc_item_refs (bounding box info)
|
|
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
|
|
assert len(chunks_with_refs) > 0, "PDF should have chunks with doc_item_refs"
|
|
|
|
# Visualize a chunk - should return images with bounding boxes drawn
|
|
images = await client.visualize_chunk(chunks_with_refs[0])
|
|
|
|
assert isinstance(images, list)
|
|
assert len(images) > 0, "PDF with page images should return visualizations"
|
|
|
|
# Verify returned objects are PIL Images
|
|
for img in images:
|
|
assert isinstance(img, PILImage)
|
|
|
|
|
|
# =============================================================================
|
|
# convert() method tests
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_text(temp_db_path):
|
|
"""Test convert() with plain text content."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
text = "This is some test content for conversion."
|
|
docling_doc = await client.convert(text)
|
|
|
|
assert isinstance(docling_doc, DoclingDocument)
|
|
# Check the content is preserved in markdown export
|
|
markdown = docling_doc.export_to_markdown()
|
|
assert "test content" in markdown
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_file(temp_db_path):
|
|
"""Test convert() with a file path."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text("File content for conversion test.")
|
|
|
|
docling_doc = await client.convert(temp_path)
|
|
|
|
assert isinstance(docling_doc, DoclingDocument)
|
|
markdown = docling_doc.export_to_markdown()
|
|
assert "File content" in markdown
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_file_not_found(temp_db_path):
|
|
"""Test convert() raises ValueError for non-existent file."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with pytest.raises(ValueError, match="File does not exist"):
|
|
await client.convert(Path("/nonexistent/path/file.txt"))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_unsupported_extension(temp_db_path):
|
|
"""Test convert() raises ValueError for unsupported file extension."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.xyz"
|
|
temp_path.write_text("content")
|
|
|
|
with pytest.raises(ValueError, match="Unsupported file extension"):
|
|
await client.convert(temp_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_file_uri(temp_db_path):
|
|
"""Test convert() with a file:// URI string."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_path = Path(temp_dir) / "test.txt"
|
|
temp_path.write_text("URI file content.")
|
|
file_uri = temp_path.as_uri()
|
|
|
|
docling_doc = await client.convert(file_uri)
|
|
|
|
assert isinstance(docling_doc, DoclingDocument)
|
|
markdown = docling_doc.export_to_markdown()
|
|
assert "URI file content" in markdown
|
|
|
|
|
|
# =============================================================================
|
|
# chunk() method tests
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_chunk_basic(temp_db_path):
|
|
"""Test chunk() produces Chunk objects from DoclingDocument."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# First convert some text
|
|
docling_doc = await client.convert("This is test content for chunking.")
|
|
|
|
# Then chunk it
|
|
chunks = await client.chunk(docling_doc)
|
|
|
|
assert isinstance(chunks, list)
|
|
assert len(chunks) > 0
|
|
assert all(isinstance(c, Chunk) for c in chunks)
|
|
# Chunks should have content but no embedding yet
|
|
assert all(c.content for c in chunks)
|
|
assert all(c.embedding is None for c in chunks)
|
|
# Chunks should not have document_id yet (not stored)
|
|
assert all(c.document_id is None for c in chunks)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_chunk_preserves_metadata(temp_db_path):
|
|
"""Test chunk() preserves structured metadata from DoclingDocument."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Convert structured markdown
|
|
markdown = """# Chapter 1
|
|
|
|
This is the first paragraph.
|
|
|
|
## Section 1.1
|
|
|
|
This is a subsection.
|
|
"""
|
|
docling_doc = await client.convert(markdown)
|
|
chunks = await client.chunk(docling_doc)
|
|
|
|
assert len(chunks) > 0
|
|
|
|
# Check that at least some chunks have metadata
|
|
has_metadata = False
|
|
for chunk in chunks:
|
|
meta = chunk.get_chunk_metadata()
|
|
if meta.doc_item_refs or meta.headings:
|
|
has_metadata = True
|
|
break
|
|
|
|
assert has_metadata, "Chunks should have structured metadata"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_chunk_empty_document(temp_db_path):
|
|
"""Test chunk() with empty DoclingDocument."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create an empty DoclingDocument
|
|
empty_doc = DoclingDocument(name="empty")
|
|
|
|
chunks = await client.chunk(empty_doc)
|
|
|
|
assert isinstance(chunks, list)
|
|
assert len(chunks) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_import_document_embeds_chunks_without_embeddings(temp_db_path):
|
|
"""Test that import_document embeds chunks that don't have embeddings."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create chunks without embeddings
|
|
chunks = [
|
|
Chunk(content="First chunk without embedding", order=0),
|
|
Chunk(content="Second chunk without embedding", order=1),
|
|
]
|
|
|
|
# Import document with chunks that have no embeddings
|
|
doc = await client.import_document(
|
|
content="Document with unembedded chunks",
|
|
chunks=chunks,
|
|
)
|
|
assert doc.id is not None
|
|
|
|
# Verify chunks were stored
|
|
stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(stored_chunks) == 2
|
|
|
|
# Verify vector search works (proves embeddings were generated)
|
|
results = await client.search("First chunk", search_type="vector")
|
|
assert len(results) > 0
|
|
assert results[0].content == "First chunk without embedding"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_document_embeds_chunks_without_embeddings(temp_db_path):
|
|
"""Test that update_document embeds chunks that don't have embeddings."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Create initial document
|
|
doc = await client.create_document(content="Initial content")
|
|
assert doc.id is not None
|
|
|
|
# Update with chunks that have no embeddings
|
|
new_chunks = [
|
|
Chunk(content="Updated chunk without embedding", order=0),
|
|
]
|
|
await client.update_document(
|
|
document_id=doc.id,
|
|
content="Updated content",
|
|
chunks=new_chunks,
|
|
)
|
|
|
|
# Verify chunks were stored
|
|
stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert len(stored_chunks) == 1
|
|
assert stored_chunks[0].content == "Updated chunk without embedding"
|
|
|
|
# Verify vector search works (proves embeddings were generated)
|
|
results = await client.search("Updated chunk", search_type="vector")
|
|
assert len(results) > 0
|
|
assert results[0].content == "Updated chunk without embedding"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_create_document_with_html_format(temp_db_path):
|
|
"""Test create_document with HTML format preserves document structure."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
html_content = """
|
|
<h1>Main Title</h1>
|
|
<p>Introduction paragraph.</p>
|
|
<h2>Section Header</h2>
|
|
<ul>
|
|
<li>Item 1</li>
|
|
<li>Item 2</li>
|
|
</ul>
|
|
"""
|
|
|
|
doc = await client.create_document(
|
|
content=html_content,
|
|
uri="test://html-doc",
|
|
format="html",
|
|
)
|
|
|
|
assert doc.id is not None
|
|
assert doc.docling_document_json is not None
|
|
|
|
# Verify the DoclingDocument has proper structure
|
|
docling_doc = doc.get_docling_document()
|
|
assert docling_doc is not None
|
|
|
|
items = list(docling_doc.iterate_items())
|
|
labels = [str(getattr(item, "label", "")) for item, _ in items]
|
|
|
|
# HTML format should preserve headers and list items
|
|
assert "title" in labels or "section_header" in labels
|
|
assert "list_item" in labels
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_convert_with_html_format(temp_db_path):
|
|
"""Test convert with HTML format."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
html_content = "<h1>Title</h1><p>Text</p>"
|
|
|
|
docling_doc = await client.convert(html_content, format="html")
|
|
|
|
items = list(docling_doc.iterate_items())
|
|
labels = [str(getattr(item, "label", "")) for item, _ in items]
|
|
|
|
assert "title" in labels
|