danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove Claude nonsense, use context manager

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-07-12 23:00:49 +03:00
parent 62a519b094
commit a2e0f2cefd
No known key found for this signature in database
5 changed files with 366 additions and 432 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
src/haiku/rag/client.py
View file

@ -166,29 +166,26 @@ class HaikuRAG:
# Create a temporary file with the appropriate extension
with tempfile.NamedTemporaryFile(
mode="wb", suffix=file_extension, delete=False
mode="wb", suffix=file_extension
) as temp_file:
temp_file.write(response.content)
temp_file.flush() # Ensure content is written to disk
temp_path = Path(temp_file.name)
try:
# Parse the content using FileReader
content = FileReader.parse_file(temp_path)
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
if existing_doc:
existing_doc.content = content
existing_doc.metadata = metadata
return await self.update_document(existing_doc)
else:
return await self.create_document(
content=content, uri=url, metadata=metadata
)
finally:
# Clean up temporary file
temp_path.unlink(missing_ok=True)
if existing_doc:
existing_doc.content = content
existing_doc.metadata = metadata
return await self.update_document(existing_doc)
else:
return await self.create_document(
content=content, uri=url, metadata=metadata
)
def _get_extension_from_content_type_or_url(
self, url: str, content_type: str

638
tests/test_client.py
View file

@ -12,300 +12,270 @@ from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio
async def test_client_document_crud(qa_corpus: Dataset):
"""Test HaikuRAG CRUD operations for documents."""
# Create client with in-memory database
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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"}
# 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
)
# 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
assert created_doc.content == document_text
assert created_doc.uri == test_uri
assert created_doc.metadata == test_metadata
assert created_doc.id is not None
assert created_doc.content == document_text
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 retrieved_doc.content == document_text
assert retrieved_doc.uri == test_uri
# 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 retrieved_doc.content == document_text
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 retrieved_by_uri.content == document_text
# 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 retrieved_by_uri.content == document_text
# 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 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
retrieved_doc.content = "Updated content"
retrieved_doc.uri = "file:///updated/path.txt"
updated_doc = await client.update_document(retrieved_doc)
assert updated_doc.content == "Updated content"
assert updated_doc.uri == "file:///updated/path.txt"
# Test update_document
retrieved_doc.content = "Updated content"
retrieved_doc.uri = "file:///updated/path.txt"
updated_doc = await client.update_document(retrieved_doc)
assert updated_doc.content == "Updated content"
assert updated_doc.uri == "file:///updated/path.txt"
# 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
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 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
# 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
# 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
client.close()
# 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_create_document_from_source():
"""Test creating a document from a file source."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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)
# Create a temporary text file
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
test_content = "This is test content from a file."
f.write(test_content)
temp_path = Path(f.name)
# Test create_document_from_source with Path
doc = await client.create_document_from_source(source=temp_path)
try:
# Test create_document_from_source with Path
doc = await client.create_document_from_source(
source=temp_path, metadata={"source_type": "file"}
)
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"
assert doc.id is not None
assert doc.content == test_content
assert doc.uri == temp_path.as_uri()
assert doc.metadata["source_type"] == "file"
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))
# Test create_document_from_source with string path
doc2 = await client.create_document_from_source(source=str(temp_path))
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
finally:
# Clean up
temp_path.unlink()
client.close()
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_unsupported():
"""Test creating a document from an unsupported file type."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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)
# 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)
try:
# Should raise ValueError for unsupported extension
with pytest.raises(ValueError, match="Unsupported file extension"):
await client.create_document_from_source(temp_path)
finally:
temp_path.unlink()
client.close()
# 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():
"""Test creating a document from a non-existent file."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") as client:
non_existent_path = Path("/non/existent/file.txt")
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)
client.close()
# 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_url():
"""Test creating a document from a URL."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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()
# 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"}
)
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 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"
client.close()
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():
"""Test creating documents from URLs with different content types."""
client = HaikuRAG(":memory:")
# 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"
async with HaikuRAG(":memory:") 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()
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"
with patch("httpx.AsyncClient.get", return_value=mock_json_response):
doc = await client.create_document_from_source(
"https://api.example.com/data.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()
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"
with patch("httpx.AsyncClient.get", return_value=mock_text_response):
doc = await client.create_document_from_source("https://example.com/readme.txt")
# 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()
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"
with patch("httpx.AsyncClient.get", return_value=mock_text_response):
doc = await client.create_document_from_source(
"https://example.com/readme.txt"
)
client.close()
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():
"""Test creating a document from URL with unsupported content type."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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()
# 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")
client.close()
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():
"""Test handling HTTP errors when creating document from URL."""
client = HaikuRAG(":memory:")
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"
async with HaikuRAG(":memory:") 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),
)
client.close()
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():
"""Test the helper method for determining file extensions."""
client = HaikuRAG(":memory:")
# 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", ""
async with HaikuRAG(":memory:") as client:
# Test content type mappings
assert (
client._get_extension_from_content_type_or_url("", "text/html") == ".html"
)
== ".pdf"
)
assert (
client._get_extension_from_content_type_or_url(
"https://example.com/data.json", ""
assert (
client._get_extension_from_content_type_or_url("", "application/pdf")
== ".pdf"
)
== ".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"
assert (
client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
)
== ".pdf"
)
client.close()
# 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
@ -313,165 +283,147 @@ async def test_client_metadata_content_type_and_md5():
"""Test that contentType and md5 metadata are correctly set."""
import hashlib
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") as client:
# Create a temporary file with known content
test_content = "Test content for MD5 calculation."
expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
# 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)
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(test_content)
temp_path = Path(f.name)
doc = await client.create_document_from_source(temp_path)
try:
doc = await client.create_document_from_source(temp_path)
assert doc.metadata["contentType"] == "text/plain"
assert doc.metadata["md5"] == expected_md5
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()
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"
)
with patch("httpx.AsyncClient.get", return_value=mock_response):
url_doc = await client.create_document_from_source(
"https://example.com/test.txt"
)
assert url_doc.metadata["contentType"] == "text/plain"
assert url_doc.metadata["md5"] == expected_md5
finally:
temp_path.unlink()
client.close()
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():
"""Test create/update/no-op behavior based on MD5 changes."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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)
# Create a temporary file
test_content = "Original content for testing."
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(test_content)
temp_path = Path(f.name)
# First call - should create new document
doc1 = await client.create_document_from_source(temp_path)
assert doc1.id is not None
assert doc1.content == test_content
original_id = doc1.id
try:
# First call - should create new document
doc1 = await client.create_document_from_source(temp_path)
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 doc2.id == original_id # Same document
assert doc2.content == test_content
# Second call with same content - should return existing document (no-op)
doc2 = await client.create_document_from_source(temp_path)
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)
# 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 doc3.id == original_id # Same document ID
assert doc3.content == updated_content # Updated content
# Third call with changed content - should update existing document
doc3 = await client.create_document_from_source(temp_path)
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
finally:
temp_path.unlink()
client.close()
# 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_url_create_update_no_op_behavior():
"""Test create/update/no-op behavior for URLs based on MD5 changes."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") as client:
url = "https://example.com/test.txt"
original_content = b"Original URL content"
updated_content = b"Updated URL content"
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()
# 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 doc1.id is not None
original_id = doc1.id
with patch("httpx.AsyncClient.get", return_value=mock_response1):
# First call - should create new document
doc1 = await client.create_document_from_source(url)
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 doc2.id == original_id # Same document
# Second call with same content - should return existing document (no-op)
doc2 = await client.create_document_from_source(url)
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()
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 doc3.id == original_id # Same document ID
assert doc3.content == updated_content.decode() # Updated content
client.close()
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 doc3.id == original_id # Same document ID
assert doc3.content == updated_content.decode() # Updated content
@pytest.mark.asyncio
async def test_client_search():
"""Test HaikuRAG search functionality."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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."
# 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"}
)
# 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)
# Test search with keyword that should match doc1
results = await client.search("Python programming", limit=3)
assert len(results) > 0
assert all(len(result) == 2 for result in results)
assert len(results) > 0
assert all(len(result) == 2 for result in results)
# Verify first result is from the Python document (doc1)
first_chunk, _ = results[0]
assert first_chunk.document_id == doc1.id
# Verify first result is from the Python document (doc1)
first_chunk, _ = results[0]
assert first_chunk.document_id == doc1.id
# Test search with different query
ml_results = await client.search("machine learning data", limit=2)
assert len(ml_results) > 0
# 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_chunk, _ = ml_results[0]
assert first_ml_chunk.document_id == doc2.id
# Verify first result is from the machine learning document (doc2)
first_ml_chunk, _ = ml_results[0]
assert first_ml_chunk.document_id == doc2.id
# Test search with limit parameter
limited_results = await client.search("programming", limit=1)
assert len(limited_results) <= 1
client.close()
# Test search with limit parameter
limited_results = await client.search("programming", limit=1)
assert len(limited_results) <= 1
@pytest.mark.asyncio

20
tests/test_monitor.py
View file

@ -13,11 +13,10 @@ from haiku.rag.store.models.document import Document
async def test_file_watcher_upsert_document():
"""Test FileWatcher._upsert_document method."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write("Test content for file watcher")
temp_path = Path(f.name)
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text("Test content for file watcher")
try:
mock_client = AsyncMock(spec=HaikuRAG)
mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri())
mock_client.create_document_from_source.return_value = mock_doc
@ -32,19 +31,15 @@ async def test_file_watcher_upsert_document():
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
finally:
temp_path.unlink(missing_ok=True)
@pytest.mark.asyncio
async def test_file_watcher_upsert_existing_document():
"""Test FileWatcher._upsert_document with existing document."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write("Test content for file watcher")
temp_path = Path(f.name)
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text("Test content for file watcher")
try:
mock_client = AsyncMock(spec=HaikuRAG)
existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri())
updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri())
@ -61,9 +56,6 @@ async def test_file_watcher_upsert_existing_document():
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
finally:
temp_path.unlink(missing_ok=True)
@pytest.mark.asyncio
async def test_file_watcher_delete_document():

71
tests/test_rebuild.py
View file

@ -8,45 +8,42 @@ from haiku.rag.store.models.document import Document
@pytest.mark.asyncio
async def test_rebuild_database(qa_corpus: Dataset):
"""Test rebuild functionality with existing documents."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") as client:
created_docs: list[Document] = []
for content in qa_corpus["document_extracted"][:3]:
doc = await client.create_document(
content=content,
)
created_docs.append(doc)
created_docs: list[Document] = []
for content in qa_corpus["document_extracted"][:3]:
doc = await client.create_document(
content=content,
)
created_docs.append(doc)
documents_before = await client.list_documents()
assert len(documents_before) == 3
documents_before = await client.list_documents()
assert len(documents_before) == 3
chunks_before = []
for doc in created_docs:
assert doc.id is not None
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_before.extend(doc_chunks)
assert len(chunks_before) > 0
# Perform rebuild
processed_doc_ids = []
async for doc_id in client.rebuild_database():
processed_doc_ids.append(doc_id)
# Verify all documents were processed
expected_doc_ids = [doc.id for doc in created_docs]
assert set(processed_doc_ids) == set(expected_doc_ids)
documents_after = await client.list_documents()
assert len(documents_after) == 3
# Verify chunks were recreated
chunks_after = []
for doc in documents_after:
if doc.id is not None:
chunks_before = []
for doc in created_docs:
assert doc.id is not None
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_after.extend(doc_chunks)
chunks_before.extend(doc_chunks)
assert len(chunks_after) > 0
assert len(chunks_before) > 0
client.close()
# Perform rebuild
processed_doc_ids = []
async for doc_id in client.rebuild_database():
processed_doc_ids.append(doc_id)
# Verify all documents were processed
expected_doc_ids = [doc.id for doc in created_docs]
assert set(processed_doc_ids) == set(expected_doc_ids)
documents_after = await client.list_documents()
assert len(documents_after) == 3
# Verify chunks were recreated
chunks_after = []
for doc in documents_after:
if doc.id is not None:
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_after.extend(doc_chunks)
assert len(chunks_after) > 0

42
tests/test_settings.py
View file

@ -45,10 +45,9 @@ def test_settings_save_and_retrieve():
async def test_config_validation_on_db_load():
"""Test that config validation fails when loading db with mismatched settings."""
# Create a temporary database file
with tempfile.NamedTemporaryFile(suffix=".sqlite", delete=False) as tmp:
with tempfile.NamedTemporaryFile(suffix=".sqlite") as tmp:
db_path = Path(tmp.name)
try:
# Create store and save settings
store1 = Store(db_path)
SettingsRepository(store1)
@ -58,28 +57,25 @@ async def test_config_validation_on_db_load():
original_chunk_size = Config.CHUNK_SIZE
Config.CHUNK_SIZE = 999
# Loading the database should raise ConfigMismatchError
with pytest.raises(ConfigMismatchError) as exc_info:
Store(db_path)
try:
# Loading the database should raise ConfigMismatchError
with pytest.raises(ConfigMismatchError) as exc_info:
Store(db_path)
assert "CHUNK_SIZE" in str(exc_info.value)
assert "Consider rebuilding" in str(exc_info.value)
assert "CHUNK_SIZE" in str(exc_info.value)
assert "Consider rebuilding" in str(exc_info.value)
# Rebuild
async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
async for _ in client.rebuild_database():
pass # Process all documents
# Rebuild
async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
async for _ in client.rebuild_database():
pass # Process all documents
# Verify we can now load the database without exception (settings were updated)
store2 = Store(db_path)
settings_repo2 = SettingsRepository(store2)
db_settings = settings_repo2.get()
assert db_settings["CHUNK_SIZE"] == 999
store2.close()
# Verify we can now load the database without exception (settings were updated)
store2 = Store(db_path)
settings_repo2 = SettingsRepository(store2)
db_settings = settings_repo2.get()
assert db_settings["CHUNK_SIZE"] == 999
store2.close()
Config.CHUNK_SIZE = original_chunk_size
finally:
# Cleanup
if db_path.exists():
db_path.unlink()
finally:
Config.CHUNK_SIZE = original_chunk_size