From 11fd4d75addf14657abfb9e988c936e943e13f06 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 17 Jun 2025 11:26:02 +0200 Subject: [PATCH] Handle files/url with create_document_from_source in client --- pyproject.toml | 1 + src/haiku/rag/client.py | 126 +++++++++++++++++++++++ tests/test_client.py | 220 ++++++++++++++++++++++++++++++++++++++++ uv.lock | 2 + 4 files changed, 349 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 018be9e5..042ec815 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ readme = "README.md" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] requires-python = ">=3.13" dependencies = [ + "httpx>=0.28.1", "markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2", "ollama>=0.5.1", "pydantic>=2.11.7", diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 01e348e7..0e46b39b 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -1,6 +1,11 @@ +import tempfile from pathlib import Path from typing import Literal +from urllib.parse import urlparse +import httpx + +from haiku.rag.reader import FileReader from haiku.rag.store.engine import Store from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.document import DocumentRepository @@ -25,6 +30,127 @@ class RAGClient: ) return await self.document_repository.create(document) + async def create_document_from_source( + self, source: str | Path, metadata: dict | None = None + ) -> Document: + """Create a document from a file path or URL. + + Args: + source: File path (as string or Path) or URL to parse + metadata: Optional metadata dictionary + + Returns: + Created Document instance + + Raises: + ValueError: If the file/URL cannot be parsed or doesn't exist + httpx.RequestError: If URL request fails + """ + + # Check if it's a URL + 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) + + # Handle as file path + source_path = Path(source) if isinstance(source, str) else source + if source_path.suffix.lower() not in FileReader.extensions: + raise ValueError(f"Unsupported file extension: {source_path.suffix}") + + if not source_path.exists(): + raise ValueError(f"File does not exist: {source_path}") + + content = FileReader.parse_file(source_path) + + # Create the document + return await self.create_document( + content=content, uri=str(source_path.resolve()), metadata=metadata + ) + + async def _create_document_from_url( + self, url: str, metadata: dict | None = None + ) -> Document: + """Create a document from a URL by downloading and parsing the content. + + Args: + url: URL to download and parse + metadata: Optional metadata dictionary + + Returns: + Created Document instance + + Raises: + ValueError: If the content cannot be parsed + httpx.RequestError: If URL request fails + """ + async with httpx.AsyncClient() as client: + response = await client.get(url) + response.raise_for_status() + + # 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 + ) + + if file_extension not in FileReader.extensions: + raise ValueError( + f"Unsupported content type/extension: {content_type}/{file_extension}" + ) + + # Create a temporary file with the appropriate extension + with tempfile.NamedTemporaryFile( + mode="wb", suffix=file_extension, delete=False + ) as temp_file: + temp_file.write(response.content) + temp_path = Path(temp_file.name) + + try: + # 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 + ) + 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 + ) -> str: + """Determine file extension from content type or URL.""" + # Common content type mappings + content_type_map = { + "text/html": ".html", + "text/plain": ".txt", + "text/markdown": ".md", + "application/pdf": ".pdf", + "application/json": ".json", + "text/csv": ".csv", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", + } + + # Try content type first + for ct, ext in content_type_map.items(): + if ct in content_type: + return ext + + # Try URL extension + parsed_url = urlparse(url) + path = Path(parsed_url.path) + if path.suffix: + return path.suffix.lower() + + # Default to .html for web content + return ".html" + async def get_document_by_id(self, document_id: int) -> Document | None: """Get a document by its ID.""" return await self.document_repository.get_by_id(document_id) diff --git a/tests/test_client.py b/tests/test_client.py index 7d46d78b..d5af9df3 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,3 +1,8 @@ +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import httpx import pytest from datasets import Dataset @@ -72,3 +77,218 @@ async def test_client_document_crud(qa_corpus: Dataset): 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 = RAGClient(":memory:") + + # 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) + + 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 == str(temp_path.resolve()) + assert doc.metadata["source_type"] == "file" + + # 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 == str(temp_path.resolve()) + + 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 = RAGClient(":memory:") + + # 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() + + +@pytest.mark.asyncio +async def test_client_create_document_from_source_nonexistent(): + """Test creating a document from a non-existent file.""" + client = RAGClient(":memory:") + + 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 = RAGClient(":memory:") + + # Mock the HTTP response + mock_response = AsyncMock() + mock_response.content = b"

Test Page

This is test content from a webpage.

" + 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"} + ) + + 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" + + 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 = RAGClient(":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" + ) + + assert doc.id is not None + assert "Test JSON" in doc.content + assert doc.uri == "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() + + with patch("httpx.AsyncClient.get", return_value=mock_text_response): + 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." + assert doc.uri == "https://example.com/readme.txt" + + 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 = RAGClient(":memory:") + + # 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() + + +@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:") + + 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" + ) + + client.close() + + +@pytest.mark.asyncio +async def test_get_extension_from_content_type_or_url(): + """Test the helper method for determining file extensions.""" + client = RAGClient(":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", "" + ) + == ".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" + ) + + client.close() diff --git a/uv.lock b/uv.lock index e89a1eeb..a2f87e12 100644 --- a/uv.lock +++ b/uv.lock @@ -460,6 +460,7 @@ name = "haiku-rag" version = "0.1.0" source = { editable = "." } dependencies = [ + { name = "httpx" }, { name = "markitdown", extra = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"] }, { name = "ollama" }, { name = "pydantic" }, @@ -482,6 +483,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "httpx", specifier = ">=0.28.1" }, { name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" }, { name = "ollama", specifier = ">=0.5.1" }, { name = "pydantic", specifier = ">=2.11.7" },