diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 83c61120..abcd1b09 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -139,124 +139,19 @@ class HaikuRAG: async def convert( self, source: Path | str, *, format: str = "md" ) -> "DoclingDocument": - """Convert a file, URL, or text to DoclingDocument. + from haiku.rag.client.processing import convert - Args: - source: One of: - - Path: Local file path to convert - - str (URL): HTTP/HTTPS URL to download and convert - - str (text): Raw text content to convert - format: The format of text content ("md", "html", or "plain"). - Defaults to "md". Use "plain" for plain text without parsing. - Only used when source is raw text (not a file path or URL). - Files and URLs determine format from extension/content-type. - - Returns: - DoclingDocument from the converted source. - - Raises: - ValueError: If the file doesn't exist or has unsupported extension. - httpx.RequestError: If URL download fails. - """ - converter = get_converter(self._config) - - # Path object - convert file directly - if isinstance(source, Path): - if not source.exists(): - raise ValueError(f"File does not exist: {source}") - if source.suffix.lower() not in converter.supported_extensions: - raise ValueError(f"Unsupported file extension: {source.suffix}") - return await converter.convert_file(source) - - # String - check if URL or text - parsed = urlparse(source) - - if parsed.scheme in ("http", "https"): - # URL - download and convert - async with httpx.AsyncClient() as http: - response = await http.get(source) - response.raise_for_status() - - content_type = response.headers.get("content-type", "").lower() - file_extension = self._get_extension_from_content_type_or_url( - source, content_type - ) - - if file_extension not in converter.supported_extensions: - raise ValueError( - f"Unsupported content type/extension: {content_type}/{file_extension}" - ) - - with tempfile.NamedTemporaryFile( - mode="wb", suffix=file_extension, delete=False - ) as temp_file: - temp_file.write(response.content) - temp_file.flush() - temp_path = Path(temp_file.name) - - try: - return await converter.convert_file(temp_path) - finally: - temp_path.unlink(missing_ok=True) - - elif parsed.scheme == "file": - # file:// URI - file_path = Path(parsed.path) - if not file_path.exists(): - raise ValueError(f"File does not exist: {file_path}") - if file_path.suffix.lower() not in converter.supported_extensions: - raise ValueError(f"Unsupported file extension: {file_path.suffix}") - return await converter.convert_file(file_path) - - else: - # Treat as text content - return await converter.convert_text(source, format=format) + return await convert(self._config, source, format=format) async def chunk(self, docling_document: "DoclingDocument") -> list[Chunk]: - """Chunk a DoclingDocument into Chunks. + from haiku.rag.client.processing import chunk - Args: - docling_document: The DoclingDocument to chunk. - - Returns: - List of Chunk objects (without embeddings, without document_id). - Each chunk has its `order` field set to its position in the list. - """ - from haiku.rag.chunkers import get_chunker - - chunker = get_chunker(self._config) - return await chunker.chunk(docling_document) + return await chunk(self._config, docling_document) async def _ensure_chunks_embedded(self, chunks: list[Chunk]) -> list[Chunk]: - """Ensure all chunks have embeddings, embedding any that don't. + from haiku.rag.client.processing import ensure_chunks_embedded - Args: - chunks: List of chunks, some may have embeddings already. - - Returns: - List of chunks with all embeddings populated. - """ - from haiku.rag.embeddings import embed_chunks - - # Find chunks that need embedding - chunks_to_embed = [c for c in chunks if c.embedding is None] - - if not chunks_to_embed: - return chunks - - # Embed chunks that don't have embeddings (returns new Chunk objects) - embedded = await embed_chunks(chunks_to_embed, self._config) - - # Build result maintaining original order - embedded_map = {(c.content, c.order): c for c in embedded} - result = [] - for chunk in chunks: - if chunk.embedding is not None: - result.append(chunk) - else: - result.append(embedded_map[(chunk.content, chunk.order)]) - - return result + return await ensure_chunks_embedded(self._config, chunks) # ========================================================================= # Title Generation @@ -768,33 +663,9 @@ class HaikuRAG: 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", - } + from haiku.rag.client.processing import get_extension_from_content_type_or_url - # 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" + return get_extension_from_content_type_or_url(url, content_type) async def get_document_by_id(self, document_id: str) -> Document | None: """Get a document by its ID. diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py new file mode 100644 index 00000000..b3aa59c5 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -0,0 +1,160 @@ +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +import httpx + +from haiku.rag.config import AppConfig +from haiku.rag.converters import get_converter +from haiku.rag.store.models.chunk import Chunk + +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + + +async def convert( + config: AppConfig, source: Path | str, *, format: str = "md" +) -> "DoclingDocument": + """Convert a file, URL, or text to DoclingDocument. + + Args: + config: Application configuration. + source: One of: + - Path: Local file path to convert + - str (URL): HTTP/HTTPS URL to download and convert + - str (text): Raw text content to convert + format: The format of text content ("md", "html", or "plain"). + Defaults to "md". Use "plain" for plain text without parsing. + Only used when source is raw text (not a file path or URL). + Files and URLs determine format from extension/content-type. + + Returns: + DoclingDocument from the converted source. + + Raises: + ValueError: If the file doesn't exist or has unsupported extension. + httpx.RequestError: If URL download fails. + """ + converter = get_converter(config) + + # Path object - convert file directly + if isinstance(source, Path): + if not source.exists(): + raise ValueError(f"File does not exist: {source}") + if source.suffix.lower() not in converter.supported_extensions: + raise ValueError(f"Unsupported file extension: {source.suffix}") + return await converter.convert_file(source) + + # String - check if URL or text + parsed = urlparse(source) + + if parsed.scheme in ("http", "https"): + # URL - download and convert + async with httpx.AsyncClient() as http: + response = await http.get(source) + response.raise_for_status() + + content_type = response.headers.get("content-type", "").lower() + file_extension = get_extension_from_content_type_or_url( + source, content_type + ) + + if file_extension not in converter.supported_extensions: + raise ValueError( + f"Unsupported content type/extension: {content_type}/{file_extension}" + ) + + with tempfile.NamedTemporaryFile( + mode="wb", suffix=file_extension, delete=False + ) as temp_file: + temp_file.write(response.content) + temp_file.flush() + temp_path = Path(temp_file.name) + + try: + return await converter.convert_file(temp_path) + finally: + temp_path.unlink(missing_ok=True) + + elif parsed.scheme == "file": + # file:// URI + file_path = Path(parsed.path) + if not file_path.exists(): + raise ValueError(f"File does not exist: {file_path}") + if file_path.suffix.lower() not in converter.supported_extensions: + raise ValueError(f"Unsupported file extension: {file_path.suffix}") + return await converter.convert_file(file_path) + + else: + # Treat as text content + return await converter.convert_text(source, format=format) + + +async def chunk(config: AppConfig, docling_document: "DoclingDocument") -> list[Chunk]: + """Chunk a DoclingDocument into Chunks. + + Returns chunks without embeddings or document_id. Each chunk's `order` + field is set to its position in the list. + """ + from haiku.rag.chunkers import get_chunker + + chunker = get_chunker(config) + return await chunker.chunk(docling_document) + + +async def ensure_chunks_embedded(config: AppConfig, chunks: list[Chunk]) -> list[Chunk]: + """Ensure all chunks have embeddings, embedding any that don't. + + Chunks that already have embeddings are passed through unchanged; missing + embeddings are filled in in-place in the returned list (preserving order). + """ + from haiku.rag.embeddings import embed_chunks + + chunks_to_embed = [c for c in chunks if c.embedding is None] + + if not chunks_to_embed: + return chunks + + embedded = await embed_chunks(chunks_to_embed, config) + + # Build result maintaining original order + embedded_map = {(c.content, c.order): c for c in embedded} + result = [] + for ch in chunks: + if ch.embedding is not None: + result.append(ch) + else: + result.append(embedded_map[(ch.content, ch.order)]) + + return result + + +def get_extension_from_content_type_or_url(url: str, content_type: str) -> str: + """Determine file extension from HTTP Content-Type header or URL suffix. + + Returns the mapped extension for known content types, falling back to the + URL path suffix, and finally `.html` for generic web content. + """ + 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", + } + + for ct, ext in content_type_map.items(): + if ct in content_type: + return ext + + parsed_url = urlparse(url) + path = Path(parsed_url.path) + if path.suffix: + return path.suffix.lower() + + return ".html"