Handle files/url with create_document_from_source in client
This commit is contained in:
parent
a91249c575
commit
11fd4d75ad
4 changed files with 349 additions and 0 deletions
|
|
@ -6,6 +6,7 @@ readme = "README.md"
|
||||||
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
||||||
requires-python = ">=3.13"
|
requires-python = ">=3.13"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"httpx>=0.28.1",
|
||||||
"markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2",
|
"markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2",
|
||||||
"ollama>=0.5.1",
|
"ollama>=0.5.1",
|
||||||
"pydantic>=2.11.7",
|
"pydantic>=2.11.7",
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,11 @@
|
||||||
|
import tempfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Literal
|
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.engine import Store
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
from haiku.rag.store.repositories.document import DocumentRepository
|
from haiku.rag.store.repositories.document import DocumentRepository
|
||||||
|
|
@ -25,6 +30,127 @@ class RAGClient:
|
||||||
)
|
)
|
||||||
return await self.document_repository.create(document)
|
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:
|
async def get_document_by_id(self, document_id: int) -> Document | None:
|
||||||
"""Get a document by its ID."""
|
"""Get a document by its ID."""
|
||||||
return await self.document_repository.get_by_id(document_id)
|
return await self.document_repository.get_by_id(document_id)
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,8 @@
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import AsyncMock, patch
|
||||||
|
|
||||||
|
import httpx
|
||||||
import pytest
|
import pytest
|
||||||
from datasets import Dataset
|
from datasets import Dataset
|
||||||
|
|
||||||
|
|
@ -72,3 +77,218 @@ async def test_client_document_crud(qa_corpus: Dataset):
|
||||||
assert deleted_again is False
|
assert deleted_again is False
|
||||||
|
|
||||||
client.close()
|
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"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
|
||||||
|
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()
|
||||||
|
|
|
||||||
2
uv.lock
2
uv.lock
|
|
@ -460,6 +460,7 @@ name = "haiku-rag"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
source = { editable = "." }
|
source = { editable = "." }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
{ name = "httpx" },
|
||||||
{ name = "markitdown", extra = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"] },
|
{ name = "markitdown", extra = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"] },
|
||||||
{ name = "ollama" },
|
{ name = "ollama" },
|
||||||
{ name = "pydantic" },
|
{ name = "pydantic" },
|
||||||
|
|
@ -482,6 +483,7 @@ dev = [
|
||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
|
{ name = "httpx", specifier = ">=0.28.1" },
|
||||||
{ name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" },
|
{ name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" },
|
||||||
{ name = "ollama", specifier = ">=0.5.1" },
|
{ name = "ollama", specifier = ">=0.5.1" },
|
||||||
{ name = "pydantic", specifier = ">=2.11.7" },
|
{ name = "pydantic", specifier = ">=2.11.7" },
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue