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

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

@ -166,12 +166,12 @@ class HaikuRAG:
# Create a temporary file with the appropriate extension # Create a temporary file with the appropriate extension
with tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(
mode="wb", suffix=file_extension, delete=False mode="wb", suffix=file_extension
) as temp_file: ) as temp_file:
temp_file.write(response.content) temp_file.write(response.content)
temp_file.flush() # Ensure content is written to disk
temp_path = Path(temp_file.name) temp_path = Path(temp_file.name)
try:
# Parse the content using FileReader # Parse the content using FileReader
content = FileReader.parse_file(temp_path) content = FileReader.parse_file(temp_path)
@ -186,9 +186,6 @@ class HaikuRAG:
return await self.create_document( return await self.create_document(
content=content, uri=url, metadata=metadata content=content, uri=url, metadata=metadata
) )
finally:
# Clean up temporary file
temp_path.unlink(missing_ok=True)
def _get_extension_from_content_type_or_url( def _get_extension_from_content_type_or_url(
self, url: str, content_type: str self, url: str, content_type: str

122
tests/test_client.py
View file

@ -12,9 +12,7 @@ from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_document_crud(qa_corpus: Dataset): async def test_client_document_crud(qa_corpus: Dataset):
"""Test HaikuRAG CRUD operations for documents.""" """Test HaikuRAG CRUD operations for documents."""
# Create client with in-memory database async with HaikuRAG(":memory:") as client:
client = HaikuRAG(":memory:")
# Get test data # Get test data
first_doc = qa_corpus[0] first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"] document_text = first_doc["document_extracted"]
@ -76,30 +74,22 @@ async def test_client_document_crud(qa_corpus: Dataset):
deleted_again = await client.delete_document(created_doc.id) deleted_again = await client.delete_document(created_doc.id)
assert deleted_again is False assert deleted_again is False
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source(): async def test_client_create_document_from_source():
"""Test creating a document from a file source.""" """Test creating a document from a file source."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
with tempfile.TemporaryDirectory() as temp_dir:
# 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." test_content = "This is test content from a file."
f.write(test_content) temp_path = Path(temp_dir) / "test.txt"
temp_path = Path(f.name) temp_path.write_text(test_content)
try:
# Test create_document_from_source with Path # Test create_document_from_source with Path
doc = await client.create_document_from_source( doc = await client.create_document_from_source(source=temp_path)
source=temp_path, metadata={"source_type": "file"}
)
assert doc.id is not None assert doc.id is not None
assert doc.content == test_content assert doc.content == test_content
assert doc.uri == temp_path.as_uri() assert doc.uri == temp_path.as_uri()
assert doc.metadata["source_type"] == "file"
assert "contentType" in doc.metadata assert "contentType" in doc.metadata
assert "md5" in doc.metadata assert "md5" in doc.metadata
assert doc.metadata["contentType"] == "text/plain" assert doc.metadata["contentType"] == "text/plain"
@ -113,17 +103,11 @@ async def test_client_create_document_from_source():
assert "contentType" in doc2.metadata assert "contentType" in doc2.metadata
assert "md5" in doc2.metadata assert "md5" in doc2.metadata
finally:
# Clean up
temp_path.unlink()
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(): async def test_client_create_document_from_source_unsupported():
"""Test creating a document from an unsupported file type.""" """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 # Create a temporary file with unsupported extension
with tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(
mode="w", suffix=".unsupported", delete=False mode="w", suffix=".unsupported", delete=False
@ -131,35 +115,26 @@ async def test_client_create_document_from_source_unsupported():
f.write("content") f.write("content")
temp_path = Path(f.name) temp_path = Path(f.name)
try:
# Should raise ValueError for unsupported extension # Should raise ValueError for unsupported extension
with pytest.raises(ValueError, match="Unsupported file extension"): with pytest.raises(ValueError, match="Unsupported file extension"):
await client.create_document_from_source(temp_path) await client.create_document_from_source(temp_path)
finally:
temp_path.unlink()
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_source_nonexistent(): async def test_client_create_document_from_source_nonexistent():
"""Test creating a document from a non-existent file.""" """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 # Should raise ValueError when file doesn't exist
with pytest.raises(ValueError, match="File does not exist"): with pytest.raises(ValueError, match="File does not exist"):
await client.create_document_from_source(non_existent_path) await client.create_document_from_source(non_existent_path)
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url(): async def test_client_create_document_from_url():
"""Test creating a document from a URL.""" """Test creating a document from a URL."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Mock the HTTP response # Mock the HTTP response
mock_response = AsyncMock() 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.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
@ -180,14 +155,11 @@ async def test_client_create_document_from_url():
assert "md5" in doc.metadata assert "md5" in doc.metadata
assert doc.metadata["contentType"] == "text/html" assert doc.metadata["contentType"] == "text/html"
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_with_different_content_types(): async def test_client_create_document_from_url_with_different_content_types():
"""Test creating documents from URLs with different content types.""" """Test creating documents from URLs with different content types."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Test JSON content # Test JSON content
mock_json_response = AsyncMock() mock_json_response = AsyncMock()
mock_json_response.content = ( mock_json_response.content = (
@ -215,7 +187,9 @@ async def test_client_create_document_from_url_with_different_content_types():
mock_text_response.raise_for_status = AsyncMock() mock_text_response.raise_for_status = AsyncMock()
with patch("httpx.AsyncClient.get", return_value=mock_text_response): with patch("httpx.AsyncClient.get", return_value=mock_text_response):
doc = await client.create_document_from_source("https://example.com/readme.txt") doc = await client.create_document_from_source(
"https://example.com/readme.txt"
)
assert doc.id is not None assert doc.id is not None
assert doc.content == "This is plain text content from a URL." assert doc.content == "This is plain text content from a URL."
@ -224,14 +198,11 @@ async def test_client_create_document_from_url_with_different_content_types():
assert "md5" in doc.metadata assert "md5" in doc.metadata
assert doc.metadata["contentType"] == "text/plain" assert doc.metadata["contentType"] == "text/plain"
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_unsupported_content(): async def test_client_create_document_from_url_unsupported_content():
"""Test creating a document from URL with unsupported content type.""" """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 with unsupported content type
mock_response = AsyncMock() mock_response = AsyncMock()
mock_response.content = b"binary content" mock_response.content = b"binary content"
@ -240,16 +211,15 @@ async def test_client_create_document_from_url_unsupported_content():
with patch("httpx.AsyncClient.get", return_value=mock_response): with patch("httpx.AsyncClient.get", return_value=mock_response):
with pytest.raises(ValueError, match="Unsupported content type"): with pytest.raises(ValueError, match="Unsupported content type"):
await client.create_document_from_source("https://example.com/binary.bin") await client.create_document_from_source(
"https://example.com/binary.bin"
client.close() )
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_document_from_url_http_error(): async def test_client_create_document_from_url_http_error():
"""Test handling HTTP errors when creating document from URL.""" """Test handling HTTP errors when creating document from URL."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
with patch("httpx.AsyncClient.get") as mock_get: with patch("httpx.AsyncClient.get") as mock_get:
mock_get.side_effect = httpx.HTTPStatusError( mock_get.side_effect = httpx.HTTPStatusError(
"404 Not Found", "404 Not Found",
@ -262,20 +232,22 @@ async def test_client_create_document_from_url_http_error():
"https://example.com/notfound.html" "https://example.com/notfound.html"
) )
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_get_extension_from_content_type_or_url(): async def test_get_extension_from_content_type_or_url():
"""Test the helper method for determining file extensions.""" """Test the helper method for determining file extensions."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Test content type mappings # Test content type mappings
assert client._get_extension_from_content_type_or_url("", "text/html") == ".html"
assert ( assert (
client._get_extension_from_content_type_or_url("", "application/pdf") == ".pdf" 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"
) )
assert client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
# Test URL extension detection # Test URL extension detection
assert ( assert (
@ -305,25 +277,21 @@ async def test_get_extension_from_content_type_or_url():
== ".pdf" == ".pdf"
) )
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_metadata_content_type_and_md5(): async def test_client_metadata_content_type_and_md5():
"""Test that contentType and md5 metadata are correctly set.""" """Test that contentType and md5 metadata are correctly set."""
import hashlib import hashlib
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Create a temporary file with known content # Create a temporary file with known content
test_content = "Test content for MD5 calculation." test_content = "Test content for MD5 calculation."
expected_md5 = hashlib.md5(test_content.encode()).hexdigest() expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: with tempfile.TemporaryDirectory() as temp_dir:
f.write(test_content) temp_path = Path(temp_dir) / "test.txt"
temp_path = Path(f.name) temp_path.write_text(test_content)
try:
doc = await client.create_document_from_source(temp_path) doc = await client.create_document_from_source(temp_path)
assert doc.metadata["contentType"] == "text/plain" assert doc.metadata["contentType"] == "text/plain"
@ -342,23 +310,17 @@ async def test_client_metadata_content_type_and_md5():
assert url_doc.metadata["contentType"] == "text/plain" assert url_doc.metadata["contentType"] == "text/plain"
assert url_doc.metadata["md5"] == expected_md5 assert url_doc.metadata["md5"] == expected_md5
finally:
temp_path.unlink()
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_create_update_no_op_behavior(): async def test_client_create_update_no_op_behavior():
"""Test create/update/no-op behavior based on MD5 changes.""" """Test create/update/no-op behavior based on MD5 changes."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Create a temporary file # Create a temporary file
test_content = "Original content for testing." test_content = "Original content for testing."
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: with tempfile.TemporaryDirectory() as temp_dir:
f.write(test_content) temp_path = Path(temp_dir) / "test.txt"
temp_path = Path(f.name) temp_path.write_text(test_content)
try:
# First call - should create new document # First call - should create new document
doc1 = await client.create_document_from_source(temp_path) doc1 = await client.create_document_from_source(temp_path)
assert doc1.id is not None assert doc1.id is not None
@ -384,16 +346,11 @@ async def test_client_create_update_no_op_behavior():
assert retrieved_doc is not None assert retrieved_doc is not None
assert retrieved_doc.content == updated_content assert retrieved_doc.content == updated_content
finally:
temp_path.unlink()
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_url_create_update_no_op_behavior(): async def test_client_url_create_update_no_op_behavior():
"""Test create/update/no-op behavior for URLs based on MD5 changes.""" """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" url = "https://example.com/test.txt"
original_content = b"Original URL content" original_content = b"Original URL content"
updated_content = b"Updated URL content" updated_content = b"Updated URL content"
@ -425,14 +382,11 @@ async def test_client_url_create_update_no_op_behavior():
assert doc3.id == original_id # Same document ID assert doc3.id == original_id # Same document ID
assert doc3.content == updated_content.decode() # Updated content assert doc3.content == updated_content.decode() # Updated content
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_search(): async def test_client_search():
"""Test HaikuRAG search functionality.""" """Test HaikuRAG search functionality."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
# Add multiple documents to search from # Add multiple documents to search from
doc1_text = "Python is a high-level programming language known for its simplicity and readability." 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." doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
@ -471,8 +425,6 @@ async def test_client_search():
limited_results = await client.search("programming", limit=1) limited_results = await client.search("programming", limit=1)
assert len(limited_results) <= 1 assert len(limited_results) <= 1
client.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_async_context_manager(): async def test_client_async_context_manager():

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(): async def test_file_watcher_upsert_document():
"""Test FileWatcher._upsert_document method.""" """Test FileWatcher._upsert_document method."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: with tempfile.TemporaryDirectory() as temp_dir:
f.write("Test content for file watcher") temp_path = Path(temp_dir) / "test.txt"
temp_path = Path(f.name) temp_path.write_text("Test content for file watcher")
try:
mock_client = AsyncMock(spec=HaikuRAG) mock_client = AsyncMock(spec=HaikuRAG)
mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri()) mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri())
mock_client.create_document_from_source.return_value = mock_doc 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.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)) mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
finally:
temp_path.unlink(missing_ok=True)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_file_watcher_upsert_existing_document(): async def test_file_watcher_upsert_existing_document():
"""Test FileWatcher._upsert_document with existing document.""" """Test FileWatcher._upsert_document with existing document."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: with tempfile.TemporaryDirectory() as temp_dir:
f.write("Test content for file watcher") temp_path = Path(temp_dir) / "test.txt"
temp_path = Path(f.name) temp_path.write_text("Test content for file watcher")
try:
mock_client = AsyncMock(spec=HaikuRAG) mock_client = AsyncMock(spec=HaikuRAG)
existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri()) 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()) 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.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)) mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
finally:
temp_path.unlink(missing_ok=True)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_file_watcher_delete_document(): async def test_file_watcher_delete_document():

5
tests/test_rebuild.py
View file

@ -8,8 +8,7 @@ from haiku.rag.store.models.document import Document
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_rebuild_database(qa_corpus: Dataset): async def test_rebuild_database(qa_corpus: Dataset):
"""Test rebuild functionality with existing documents.""" """Test rebuild functionality with existing documents."""
client = HaikuRAG(":memory:") async with HaikuRAG(":memory:") as client:
created_docs: list[Document] = [] created_docs: list[Document] = []
for content in qa_corpus["document_extracted"][:3]: for content in qa_corpus["document_extracted"][:3]:
doc = await client.create_document( doc = await client.create_document(
@ -48,5 +47,3 @@ async def test_rebuild_database(qa_corpus: Dataset):
chunks_after.extend(doc_chunks) chunks_after.extend(doc_chunks)
assert len(chunks_after) > 0 assert len(chunks_after) > 0
client.close()

10
tests/test_settings.py
View file

@ -45,10 +45,9 @@ def test_settings_save_and_retrieve():
async def test_config_validation_on_db_load(): async def test_config_validation_on_db_load():
"""Test that config validation fails when loading db with mismatched settings.""" """Test that config validation fails when loading db with mismatched settings."""
# Create a temporary database file # 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) db_path = Path(tmp.name)
try:
# Create store and save settings # Create store and save settings
store1 = Store(db_path) store1 = Store(db_path)
SettingsRepository(store1) SettingsRepository(store1)
@ -58,6 +57,7 @@ async def test_config_validation_on_db_load():
original_chunk_size = Config.CHUNK_SIZE original_chunk_size = Config.CHUNK_SIZE
Config.CHUNK_SIZE = 999 Config.CHUNK_SIZE = 999
try:
# Loading the database should raise ConfigMismatchError # Loading the database should raise ConfigMismatchError
with pytest.raises(ConfigMismatchError) as exc_info: with pytest.raises(ConfigMismatchError) as exc_info:
Store(db_path) Store(db_path)
@ -77,9 +77,5 @@ async def test_config_validation_on_db_load():
assert db_settings["CHUNK_SIZE"] == 999 assert db_settings["CHUNK_SIZE"] == 999
store2.close() store2.close()
Config.CHUNK_SIZE = original_chunk_size
finally: finally:
# Cleanup Config.CHUNK_SIZE = original_chunk_size
if db_path.exists():
db_path.unlink()