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
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)
@ -186,9 +186,6 @@ class HaikuRAG:
return await self.create_document(
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(
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
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"]
@ -76,30 +74,22 @@ async def test_client_document_crud(qa_corpus: Dataset):
deleted_again = await client.delete_document(created_doc.id)
assert deleted_again is False
client.close()
@pytest.mark.asyncio
async def test_client_create_document_from_source():
"""Test creating a document from a file source."""
client = HaikuRAG(":memory:")
# Create a temporary text file
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
async with HaikuRAG(":memory:") as client:
with tempfile.TemporaryDirectory() as temp_dir:
test_content = "This is test content from a file."
f.write(test_content)
temp_path = Path(f.name)
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text(test_content)
try:
# Test create_document_from_source with Path
doc = await client.create_document_from_source(
source=temp_path, metadata={"source_type": "file"}
)
doc = await client.create_document_from_source(source=temp_path)
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"
@ -113,17 +103,11 @@ async def test_client_create_document_from_source():
assert "contentType" in doc2.metadata
assert "md5" in doc2.metadata
finally:
# Clean up
temp_path.unlink()
client.close()
@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
@ -131,35 +115,26 @@ async def test_client_create_document_from_source_unsupported():
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()
@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")
# 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()
@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>"
@ -180,14 +155,11 @@ async def test_client_create_document_from_url():
assert "md5" in doc.metadata
assert doc.metadata["contentType"] == "text/html"
client.close()
@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:")
async with HaikuRAG(":memory:") as client:
# Test JSON content
mock_json_response = AsyncMock()
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()
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.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 doc.metadata["contentType"] == "text/plain"
client.close()
@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"
@ -240,16 +211,15 @@ async def test_client_create_document_from_url_unsupported_content():
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()
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:")
async with HaikuRAG(":memory:") as client:
with patch("httpx.AsyncClient.get") as mock_get:
mock_get.side_effect = httpx.HTTPStatusError(
"404 Not Found",
@ -262,20 +232,22 @@ async def test_client_create_document_from_url_http_error():
"https://example.com/notfound.html"
)
client.close()
@pytest.mark.asyncio
async def test_get_extension_from_content_type_or_url():
"""Test the helper method for determining file extensions."""
client = HaikuRAG(":memory:")
async with HaikuRAG(":memory:") 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"
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
assert (
@ -305,25 +277,21 @@ async def test_get_extension_from_content_type_or_url():
== ".pdf"
)
client.close()
@pytest.mark.asyncio
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()
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(test_content)
temp_path = Path(f.name)
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text(test_content)
try:
doc = await client.create_document_from_source(temp_path)
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["md5"] == expected_md5
finally:
temp_path.unlink()
client.close()
@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.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(test_content)
temp_path = Path(f.name)
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text(test_content)
try:
# First call - should create new document
doc1 = await client.create_document_from_source(temp_path)
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.content == updated_content
finally:
temp_path.unlink()
client.close()
@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"
@ -425,14 +382,11 @@ async def test_client_url_create_update_no_op_behavior():
assert doc3.id == original_id # Same document ID
assert doc3.content == updated_content.decode() # Updated content
client.close()
@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."
@ -471,8 +425,6 @@ async def test_client_search():
limited_results = await client.search("programming", limit=1)
assert len(limited_results) <= 1
client.close()
@pytest.mark.asyncio
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():
"""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():

5
tests/test_rebuild.py
View file

@ -8,8 +8,7 @@ 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(
@ -48,5 +47,3 @@ async def test_rebuild_database(qa_corpus: Dataset):
chunks_after.extend(doc_chunks)
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():
"""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,6 +57,7 @@ async def test_config_validation_on_db_load():
original_chunk_size = Config.CHUNK_SIZE
Config.CHUNK_SIZE = 999
try:
# Loading the database should raise ConfigMismatchError
with pytest.raises(ConfigMismatchError) as exc_info:
Store(db_path)
@ -77,9 +77,5 @@ async def test_config_validation_on_db_load():
assert db_settings["CHUNK_SIZE"] == 999
store2.close()
Config.CHUNK_SIZE = original_chunk_size
finally:
# Cleanup
if db_path.exists():
db_path.unlink()
Config.CHUNK_SIZE = original_chunk_size