diff --git a/README.md b/README.md index e69de29b..fefe0106 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,124 @@ +# Haiku SQLite RAG + +A SQLite-based Retrieval-Augmented Generation (RAG) system built for efficient document storage, chunking, and hybrid search capabilities. + +## Features + +- **Document Management**: Store and manage documents with automatic content parsing +- **Smart Updates**: Intelligent file/URL monitoring with MD5-based change detection +- **Hybrid Search**: Full-text search (FTS5) combined with vector embeddings +- **Multi-format Support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, and more +- **Web Content**: Direct URL ingestion with automatic content type detection +- **Vector Embeddings**: Uses sqlite-vec for efficient similarity search +- **Automatic Chunking**: Intelligent document segmentation for better retrieval + +## Installation + +```bash +uv pip install haiku.rag +``` + +or for development, checkout the repository and then, + +```bash +# Install dependencies +uv sync + +# Activate virtual environment +source .venv/bin/activate +``` + +## Quick Start + +```python +from pathlib import Path +from haiku.rag.client import HaikuRAG + +# Initialize client with database path +client = HaikuRAG("path/to/database.db") +# Or use in-memory database for testing +client = HaikuRAG(":memory:") + +# Create document from text +doc = await client.create_document( + content="Your document content here", + uri="doc://example", + metadata={"source": "manual", "topic": "example"} +) + +# Create document from file (auto-parses content) +doc = await client.create_document_from_source("path/to/document.pdf") + +# Create document from URL +doc = await client.create_document_from_source("https://example.com/article.html") + +# Retrieve documents +doc = await client.get_document_by_id(1) +doc = await client.get_document_by_uri("file:///path/to/document.pdf") + +# List all documents with pagination +docs = await client.list_documents(limit=10, offset=0) + +# Update document content +doc.content = "Updated content" +await client.update_document(doc) + +# Delete document +await client.delete_document(doc.id) + +# Clean up +client.close() +``` + +## Smart Document Updates + +The system automatically tracks file changes using MD5 hashes: + +```python +# First call - creates new document +doc1 = await client.create_document_from_source("document.txt") + +# Second call - no changes, returns existing document (no processing) +doc2 = await client.create_document_from_source("document.txt") +assert doc1.id == doc2.id + +# After file modification - automatically updates existing document +# File content changed... +doc3 = await client.create_document_from_source("document.txt") +assert doc1.id == doc3.id # Same document +assert doc3.content != doc1.content # Updated content +``` + +## Supported File Formats + +The system supports 40+ file formats through MarkItDown: + +- **Documents**: PDF, DOCX, PPTX, XLSX +- **Web**: HTML, XML +- **Text**: TXT, MD, CSV, JSON, YAML +- **Code**: PY, JS, TS, C, CPP, JAVA, GO, RS, and more +- **Media**: MP3, WAV (transcription) + +## Document Metadata + +Documents automatically include metadata: + +```python +doc = await client.create_document_from_source("example.pdf") +print(doc.metadata) +# { +# "contentType": "application/pdf", +# "md5": "abc123...", +# "custom_field": "value" # Your custom metadata +# } +``` + +## Contributing + +1. Fork the repository +2. Create a feature branch +3. Add tests for new functionality +4. Ensure all tests pass: `pytest` +5. Run type checking: `pyright` +6. Run linting: `ruff check` +7. Submit a pull request diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 0e46b39b..1a7e8724 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -1,3 +1,5 @@ +import hashlib +import mimetypes import tempfile from pathlib import Path from typing import Literal @@ -11,7 +13,7 @@ from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.document import DocumentRepository -class RAGClient: +class HaikuRAG: """High-level haiku-rag client.""" def __init__(self, db_path: Path | Literal[":memory:"]): @@ -31,16 +33,21 @@ class RAGClient: return await self.document_repository.create(document) async def create_document_from_source( - self, source: str | Path, metadata: dict | None = None + self, source: str | Path, metadata: dict = {} ) -> Document: - """Create a document from a file path or URL. + """Create or update a document from a file path or URL. + + Checks if a document with the same URI already exists: + - If MD5 is unchanged, returns existing document + - If MD5 changed, updates the document + - If no document exists, creates a new one Args: source: File path (as string or Path) or URL to parse metadata: Optional metadata dictionary Returns: - Created Document instance + Document instance (created, updated, or existing) Raises: ValueError: If the file/URL cannot be parsed or doesn't exist @@ -51,7 +58,7 @@ class RAGClient: source_str = str(source) parsed_url = urlparse(source_str) if parsed_url.scheme in ("http", "https"): - return await self._create_document_from_url(source_str, metadata) + return await self._create_or_update_document_from_url(source_str, metadata) # Handle as file path source_path = Path(source) if isinstance(source, str) else source @@ -61,24 +68,52 @@ class RAGClient: if not source_path.exists(): raise ValueError(f"File does not exist: {source_path}") + uri = str(source_path.resolve()) + md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest() + + # Check if document already exists + existing_doc = await self.get_document_by_uri(uri) + if existing_doc and existing_doc.metadata.get("md5") == md5_hash: + # MD5 unchanged, return existing document + return existing_doc + content = FileReader.parse_file(source_path) - # Create the document - return await self.create_document( - content=content, uri=str(source_path.resolve()), metadata=metadata - ) + # Get content type from file extension + content_type, _ = mimetypes.guess_type(str(source_path)) + if not content_type: + content_type = "application/octet-stream" - async def _create_document_from_url( - self, url: str, metadata: dict | None = None + # Merge metadata with contentType and md5 + metadata.update({"contentType": content_type, "md5": md5_hash}) + + if existing_doc: + # Update existing document + existing_doc.content = content + existing_doc.metadata = metadata + return await self.update_document(existing_doc) + else: + # Create new document + return await self.create_document( + content=content, uri=uri, metadata=metadata + ) + + async def _create_or_update_document_from_url( + self, url: str, metadata: dict = {} ) -> Document: - """Create a document from a URL by downloading and parsing the content. + """Create or update a document from a URL by downloading and parsing the content. + + Checks if a document with the same URI already exists: + - If MD5 is unchanged, returns existing document + - If MD5 changed, updates the document + - If no document exists, creates a new one Args: url: URL to download and parse metadata: Optional metadata dictionary Returns: - Created Document instance + Document instance (created, updated, or existing) Raises: ValueError: If the content cannot be parsed @@ -88,10 +123,16 @@ class RAGClient: response = await client.get(url) response.raise_for_status() + md5_hash = hashlib.md5(response.content).hexdigest() + + # Check if document already exists + existing_doc = await self.get_document_by_uri(url) + if existing_doc and existing_doc.metadata.get("md5") == md5_hash: + # MD5 unchanged, return existing document + return existing_doc + # Get content type to determine file extension content_type = response.headers.get("content-type", "").lower() - - # Try to determine file extension from content type or URL file_extension = self._get_extension_from_content_type_or_url( url, content_type ) @@ -112,10 +153,17 @@ class RAGClient: # Parse the content using FileReader content = FileReader.parse_file(temp_path) - # Create the document with the original URL as URI - return await self.create_document( - content=content, uri=url, metadata=metadata - ) + # 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) diff --git a/tests/test_client.py b/tests/test_client.py index d5af9df3..31b34ee0 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -6,14 +6,14 @@ import httpx import pytest from datasets import Dataset -from haiku.rag.client import RAGClient +from haiku.rag.client import HaikuRAG @pytest.mark.asyncio async def test_client_document_crud(qa_corpus: Dataset): - """Test RAGClient CRUD operations for documents.""" + """Test HaikuRAG CRUD operations for documents.""" # Create client with in-memory database - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Get test data first_doc = qa_corpus[0] @@ -82,7 +82,7 @@ async def test_client_document_crud(qa_corpus: Dataset): @pytest.mark.asyncio async def test_client_create_document_from_source(): """Test creating a document from a file source.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Create a temporary text file with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: @@ -100,6 +100,9 @@ async def test_client_create_document_from_source(): assert doc.content == test_content assert doc.uri == str(temp_path.resolve()) 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)) @@ -107,6 +110,8 @@ async def test_client_create_document_from_source(): assert doc2.id is not None assert doc2.content == test_content assert doc2.uri == str(temp_path.resolve()) + assert "contentType" in doc2.metadata + assert "md5" in doc2.metadata finally: # Clean up @@ -117,7 +122,7 @@ async def test_client_create_document_from_source(): @pytest.mark.asyncio async def test_client_create_document_from_source_unsupported(): """Test creating a document from an unsupported file type.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Create a temporary file with unsupported extension with tempfile.NamedTemporaryFile( @@ -139,7 +144,7 @@ async def test_client_create_document_from_source_unsupported(): @pytest.mark.asyncio async def test_client_create_document_from_source_nonexistent(): """Test creating a document from a non-existent file.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") non_existent_path = Path("/non/existent/file.txt") @@ -153,7 +158,7 @@ async def test_client_create_document_from_source_nonexistent(): @pytest.mark.asyncio async def test_client_create_document_from_url(): """Test creating a document from a URL.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Mock the HTTP response mock_response = AsyncMock() @@ -171,6 +176,9 @@ async def test_client_create_document_from_url(): 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() @@ -178,7 +186,7 @@ async def test_client_create_document_from_url(): @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 = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Test JSON content mock_json_response = AsyncMock() @@ -196,6 +204,9 @@ async def test_client_create_document_from_url_with_different_content_types(): 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() @@ -209,6 +220,9 @@ async def test_client_create_document_from_url_with_different_content_types(): 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" client.close() @@ -216,7 +230,7 @@ async def test_client_create_document_from_url_with_different_content_types(): @pytest.mark.asyncio async def test_client_create_document_from_url_unsupported_content(): """Test creating a document from URL with unsupported content type.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Mock response with unsupported content type mock_response = AsyncMock() @@ -234,7 +248,7 @@ async def test_client_create_document_from_url_unsupported_content(): @pytest.mark.asyncio async def test_client_create_document_from_url_http_error(): """Test handling HTTP errors when creating document from URL.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") with patch("httpx.AsyncClient.get") as mock_get: mock_get.side_effect = httpx.HTTPStatusError( @@ -254,7 +268,7 @@ async def test_client_create_document_from_url_http_error(): @pytest.mark.asyncio async def test_get_extension_from_content_type_or_url(): """Test the helper method for determining file extensions.""" - client = RAGClient(":memory:") + client = HaikuRAG(":memory:") # Test content type mappings assert client._get_extension_from_content_type_or_url("", "text/html") == ".html" @@ -292,3 +306,123 @@ async def test_get_extension_from_content_type_or_url(): ) 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:") + + # 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) + + try: + doc = await client.create_document_from_source(temp_path) + + 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 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:") + + # 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) + + 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 + + # 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 + + # 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() + + +@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:") + + 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 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 + + 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()