Add a context manager to the client

README.md

@@ -34,11 +34,8 @@ source .venv/bin/activate
 from pathlib import Path
 from haiku.rag.client import HaikuRAG
 
-# Initialize client with database path
+# Use as async context manager (recommended)
-client = HaikuRAG("path/to/database.db")
+async with HaikuRAG("path/to/database.db") as client:
-# Or use in-memory database for testing
-client = HaikuRAG(":memory:")
-
     # Create document from text
     doc = await client.create_document(
         content="Your document content here",
@@ -74,7 +71,12 @@ for chunk, score in results:
         print(f"Document ID: {chunk.document_id}")
         print("---")
 
-# Clean up
+
+# Or use without the context manager.
+client = HaikuRAG(":memory:")
+try:
+    # ... operations ...
+finally:
     client.close()
 ```
 
@@ -87,6 +89,7 @@ client.close()
 4. **Chunked Results**: Returns relevant document chunks with scores
 
 ```python
+async with HaikuRAG("database.db") as client:
     # Basic search
     results = await client.search("your query here")
 
@@ -109,6 +112,7 @@ for chunk, relevance_score in results:
 The system automatically tracks file changes using MD5 hashes:
 
 ```python
+async with HaikuRAG("database.db") as client:
     # First call - creates new document
     doc1 = await client.create_document_from_source("document.txt")
 

src/haiku/rag/client.py

@@ -24,6 +24,15 @@ class HaikuRAG:
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
 
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        self.close()
+        return False
+
     async def create_document(
         self, content: str, uri: str | None = None, metadata: dict | None = None
     ) -> Document:

tests/test_client.py

@@ -472,3 +472,28 @@ async def test_client_search():
     assert len(limited_results) <= 1
 
     client.close()
+
+
+@pytest.mark.asyncio
+async def test_client_async_context_manager():
+    """Test HaikuRAG as async context manager."""
+
+    # Test that context manager works and auto-closes
+    async with HaikuRAG(":memory:") 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