haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/converters/docling_serve.py

@@ -0,0 +1,229 @@
+"""docling-serve remote converter implementation."""
+
+import asyncio
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar
+
+from haiku.rag.config import AppConfig
+from haiku.rag.converters.base import DocumentConverter
+from haiku.rag.converters.text_utils import TextFileHandler
+from haiku.rag.providers.docling_serve import DoclingServeClient
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+    from haiku.rag.config.models import ModelConfig
+
+
+class DoclingServeConverter(DocumentConverter):
+    """Converter that uses docling-serve for document conversion.
+
+    This converter offloads document processing to a docling-serve instance,
+    which handles heavy operations like PDF parsing, OCR, and table extraction.
+
+    For plain text files, it reads them locally and converts to markdown format
+    before sending to docling-serve for DoclingDocument conversion.
+    """
+
+    # Extensions that docling-serve can handle
+    docling_serve_extensions: ClassVar[list[str]] = [
+        ".adoc",
+        ".asc",
+        ".asciidoc",
+        ".bmp",
+        ".csv",
+        ".docx",
+        ".html",
+        ".xhtml",
+        ".jpeg",
+        ".jpg",
+        ".md",
+        ".pdf",
+        ".png",
+        ".pptx",
+        ".tiff",
+        ".xlsx",
+        ".xml",
+        ".webp",
+    ]
+
+    def __init__(self, config: AppConfig):
+        """Initialize the converter with configuration.
+
+        Args:
+            config: Application configuration containing docling-serve settings.
+        """
+        self.config = config
+        self.client = DoclingServeClient(
+            base_url=config.providers.docling_serve.base_url,
+            api_key=config.providers.docling_serve.api_key,
+        )
+
+    @property
+    def supported_extensions(self) -> list[str]:
+        """Return list of file extensions supported by this converter."""
+        return self.docling_serve_extensions + TextFileHandler.text_extensions
+
+    def _get_vlm_api_url(self, model: "ModelConfig") -> str:
+        """Construct VLM API URL from model config."""
+        if model.base_url:
+            base = model.base_url.rstrip("/")
+            return f"{base}/v1/chat/completions"
+
+        if model.provider == "ollama":
+            base = self.config.providers.ollama.base_url.rstrip("/")
+            return f"{base}/v1/chat/completions"
+
+        if model.provider == "openai":
+            return "https://api.openai.com/v1/chat/completions"
+
+        raise ValueError(f"Unsupported VLM provider: {model.provider}")
+
+    def _build_conversion_data(self) -> dict[str, str | list[str]]:
+        """Build form data for conversion request."""
+        opts = self.config.processing.conversion_options
+        pic_desc = opts.picture_description
+
+        data: dict[str, str | list[str]] = {
+            "to_formats": "json",
+            "do_ocr": str(opts.do_ocr).lower(),
+            "force_ocr": str(opts.force_ocr).lower(),
+            "do_table_structure": str(opts.do_table_structure).lower(),
+            "table_mode": opts.table_mode,
+            "table_cell_matching": str(opts.table_cell_matching).lower(),
+            "images_scale": str(opts.images_scale),
+            "generate_picture_images": str(
+                opts.generate_picture_images or pic_desc.enabled
+            ).lower(),
+            "do_picture_description": str(pic_desc.enabled).lower(),
+        }
+
+        if opts.ocr_lang:
+            data["ocr_lang"] = opts.ocr_lang
+
+        if pic_desc.enabled:
+            prompt = self.config.prompts.picture_description
+            picture_description_api = {
+                "url": self._get_vlm_api_url(pic_desc.model),
+                "params": {
+                    "model": pic_desc.model.name,
+                    "max_completion_tokens": pic_desc.max_tokens,
+                },
+                "prompt": prompt,
+                "timeout": pic_desc.timeout,
+            }
+            data["picture_description_api"] = json.dumps(picture_description_api)
+
+        return data
+
+    async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
+        """Make an async request to docling-serve and poll for results.
+
+        Args:
+            files: Dictionary with files parameter for httpx
+            name: Name of the document being converted (for error messages)
+
+        Returns:
+            DoclingDocument representation
+
+        Raises:
+            ValueError: If conversion fails or service is unavailable
+        """
+        from docling_core.types.doc.document import DoclingDocument
+
+        data = self._build_conversion_data()
+        result = await self.client.submit_and_poll(
+            endpoint="/v1/convert/file/async",
+            files=files,
+            data=data,
+            name=name,
+        )
+
+        if result.get("status") not in ("success", "partial_success", None):
+            errors = result.get("errors", [])
+            raise ValueError(f"Conversion failed: {errors}")
+
+        json_content = result.get("document", {}).get("json_content")
+
+        if json_content is None:
+            raise ValueError(
+                f"docling-serve did not return JSON content for {name}. "
+                "This may indicate an unsupported file format."
+            )
+
+        return DoclingDocument.model_validate(json_content)
+
+    async def convert_file(self, path: Path) -> "DoclingDocument":
+        """Convert a file to DoclingDocument using docling-serve.
+
+        Args:
+            path: Path to the file to convert.
+
+        Returns:
+            DoclingDocument representation of the file.
+
+        Raises:
+            ValueError: If the file cannot be converted or service is unavailable.
+        """
+        file_extension = path.suffix.lower()
+
+        if file_extension in TextFileHandler.text_extensions:
+            try:
+                content = await asyncio.to_thread(path.read_text, encoding="utf-8")
+                prepared_content = TextFileHandler.prepare_text_content(
+                    content, file_extension
+                )
+                return await self.convert_text(prepared_content, name=f"{path.stem}.md")
+            except Exception as e:
+                raise ValueError(f"Failed to read text file {path}: {e}")
+
+        def read_file():
+            with open(path, "rb") as f:
+                return f.read()
+
+        file_content = await asyncio.to_thread(read_file)
+        files = {"files": (path.name, file_content, "application/octet-stream")}
+        return await self._make_request(files, path.name)
+
+    SUPPORTED_FORMATS = ("md", "html", "plain")
+
+    async def convert_text(
+        self, text: str, name: str = "content.md", format: str = "md"
+    ) -> "DoclingDocument":
+        """Convert text content to DoclingDocument via docling-serve.
+
+        Sends the text to docling-serve for conversion using the specified format.
+
+        Args:
+            text: The text content to convert.
+            name: The name to use for the document (defaults to "content.md").
+            format: The format of the text content ("md", "html", or "plain").
+                Defaults to "md". Use "plain" for plain text without parsing.
+
+        Returns:
+            DoclingDocument representation of the text.
+
+        Raises:
+            ValueError: If the text cannot be converted or format is unsupported.
+        """
+        from haiku.rag.converters.text_utils import TextFileHandler
+
+        if format not in self.SUPPORTED_FORMATS:
+            raise ValueError(
+                f"Unsupported format: {format}. "
+                f"Supported formats: {', '.join(self.SUPPORTED_FORMATS)}"
+            )
+
+        # Derive document name from format to tell docling which parser to use
+        doc_name = f"content.{format}" if name == "content.md" else name
+
+        # Plain text doesn't need remote parsing - create document directly
+        if format == "plain":
+            return TextFileHandler._create_simple_docling_document(text, doc_name)
+
+        mime_type = "text/html" if format == "html" else "text/markdown"
+
+        text_bytes = text.encode("utf-8")
+        files = {"files": (doc_name, text_bytes, mime_type)}
+        return await self._make_request(files, doc_name)

haiku/rag/converters/text_utils.py

@@ -0,0 +1,237 @@
+"""Shared utilities for text file handling in converters."""
+
+import asyncio
+from io import BytesIO
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+
+class TextFileHandler:
+    """Handles conversion of text files to DoclingDocument format.
+
+    This class provides shared functionality for converting plain text and code files
+    to DoclingDocument format, with proper code block wrapping for syntax highlighting.
+    """
+
+    # Plain text extensions that we'll read directly
+    text_extensions: ClassVar[list[str]] = [
+        ".astro",
+        ".bash",
+        ".c",
+        ".clj",
+        ".cljs",
+        ".cpp",
+        ".cs",
+        ".css",
+        ".dart",
+        ".elm",
+        ".ex",
+        ".exs",
+        ".fs",
+        ".fsx",
+        ".go",
+        ".gql",
+        ".graphql",
+        ".groovy",
+        ".h",
+        ".hcl",
+        ".hpp",
+        ".hs",
+        ".java",
+        ".jl",
+        ".js",
+        ".json",
+        ".kt",
+        ".less",
+        ".lua",
+        ".mdx",
+        ".mjs",
+        ".ml",
+        ".mli",
+        ".nim",
+        ".nix",
+        ".php",
+        ".pl",
+        ".pm",
+        ".proto",
+        ".ps1",
+        ".py",
+        ".r",
+        ".rb",
+        ".rs",
+        ".sass",
+        ".scala",
+        ".scss",
+        ".sh",
+        ".sql",
+        ".svelte",
+        ".swift",
+        ".tf",
+        ".toml",
+        ".ts",
+        ".tsx",
+        ".txt",
+        ".vue",
+        ".xml",
+        ".yaml",
+        ".yml",
+        ".zig",
+    ]
+
+    # Code file extensions with their markdown language identifiers
+    code_markdown_identifier: ClassVar[dict[str, str]] = {
+        ".astro": "astro",
+        ".bash": "bash",
+        ".c": "c",
+        ".clj": "clojure",
+        ".cljs": "clojure",
+        ".cpp": "cpp",
+        ".cs": "csharp",
+        ".css": "css",
+        ".dart": "dart",
+        ".elm": "elm",
+        ".ex": "elixir",
+        ".exs": "elixir",
+        ".fs": "fsharp",
+        ".fsx": "fsharp",
+        ".go": "go",
+        ".gql": "graphql",
+        ".graphql": "graphql",
+        ".groovy": "groovy",
+        ".h": "c",
+        ".hcl": "hcl",
+        ".hpp": "cpp",
+        ".hs": "haskell",
+        ".java": "java",
+        ".jl": "julia",
+        ".js": "javascript",
+        ".json": "json",
+        ".kt": "kotlin",
+        ".less": "less",
+        ".lua": "lua",
+        ".mjs": "javascript",
+        ".ml": "ocaml",
+        ".mli": "ocaml",
+        ".nim": "nim",
+        ".nix": "nix",
+        ".php": "php",
+        ".pl": "perl",
+        ".pm": "perl",
+        ".proto": "protobuf",
+        ".ps1": "powershell",
+        ".py": "python",
+        ".r": "r",
+        ".rb": "ruby",
+        ".rs": "rust",
+        ".sass": "sass",
+        ".scala": "scala",
+        ".scss": "scss",
+        ".sh": "bash",
+        ".sql": "sql",
+        ".svelte": "svelte",
+        ".swift": "swift",
+        ".tf": "hcl",
+        ".toml": "toml",
+        ".ts": "typescript",
+        ".tsx": "tsx",
+        ".vue": "vue",
+        ".xml": "xml",
+        ".yaml": "yaml",
+        ".yml": "yaml",
+        ".zig": "zig",
+    }
+
+    @staticmethod
+    def prepare_text_content(content: str, file_extension: str) -> str:
+        """Prepare text content for conversion to DoclingDocument.
+
+        Wraps code files in markdown code blocks with appropriate language identifiers.
+
+        Args:
+            content: The text content.
+            file_extension: File extension (including dot, e.g., ".py").
+
+        Returns:
+            Prepared text content, possibly wrapped in code blocks.
+        """
+        if file_extension in TextFileHandler.code_markdown_identifier:
+            language = TextFileHandler.code_markdown_identifier[file_extension]
+            return f"```{language}\n{content}\n```"
+        return content
+
+    SUPPORTED_FORMATS = ("md", "html", "plain")
+
+    @staticmethod
+    def _create_simple_docling_document(text: str, name: str) -> "DoclingDocument":
+        """Create a simple DoclingDocument directly from text.
+
+        Used as fallback when docling's format detection fails for plain text
+        that doesn't contain markdown syntax.
+        """
+        from docling_core.types.doc.document import DoclingDocument
+        from docling_core.types.doc.labels import DocItemLabel
+
+        doc_name = name.rsplit(".", 1)[0] if "." in name else name
+        doc = DoclingDocument(name=doc_name)
+        doc.add_text(label=DocItemLabel.TEXT, text=text)
+        return doc
+
+    @staticmethod
+    def _sync_text_to_docling_document(
+        text: str, name: str = "content.md", format: str = "md"
+    ) -> "DoclingDocument":
+        """Synchronous implementation of text to DoclingDocument conversion."""
+        from docling.document_converter import DocumentConverter as DoclingDocConverter
+        from docling.exceptions import ConversionError
+        from docling_core.types.io import DocumentStream
+
+        if format not in TextFileHandler.SUPPORTED_FORMATS:
+            raise ValueError(
+                f"Unsupported format: {format}. "
+                f"Supported formats: {', '.join(TextFileHandler.SUPPORTED_FORMATS)}"
+            )
+
+        # Derive document name from format to tell docling which parser to use
+        doc_name = f"content.{format}" if name == "content.md" else name
+
+        # Plain text doesn't need parsing - create document directly
+        if format == "plain":
+            return TextFileHandler._create_simple_docling_document(text, doc_name)
+
+        bytes_io = BytesIO(text.encode("utf-8"))
+        doc_stream = DocumentStream(name=doc_name, stream=bytes_io)
+        converter = DoclingDocConverter()
+        try:
+            result = converter.convert(doc_stream)
+            return result.document
+        except ConversionError:
+            # Docling's format detection fails for plain text without markdown syntax.
+            # Fall back to creating a simple document directly.
+            return TextFileHandler._create_simple_docling_document(text, doc_name)
+
+    @staticmethod
+    async def text_to_docling_document(
+        text: str, name: str = "content.md", format: str = "md"
+    ) -> "DoclingDocument":
+        """Convert text to DoclingDocument using docling's parser.
+
+        Args:
+            text: The text content to convert.
+            name: The name to use for the document.
+            format: The format of the text content ("md", "html", or "plain").
+                Defaults to "md". Use "plain" for plain text without parsing.
+
+        Returns:
+            DoclingDocument representation of the text.
+
+        Raises:
+            ValueError: If the conversion fails or format is unsupported.
+        """
+        try:
+            return await asyncio.to_thread(
+                TextFileHandler._sync_text_to_docling_document, text, name, format
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to convert text to DoclingDocument: {e}")

haiku/rag/embeddings/__init__.py

@@ -1,11 +1,100 @@
+from typing import TYPE_CHECKING
+
+from pydantic_ai.embeddings import Embedder
+from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.embeddings.base import EmbedderBase
-from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
 
+if TYPE_CHECKING:
+    from haiku.rag.store.models.chunk import Chunk
+
+
+class EmbedderWrapper:
+    """Wrapper around pydantic-ai Embedder with explicit query/document methods."""
+
+    def __init__(self, embedder: Embedder, vector_dim: int):
+        self._embedder = embedder
+        self._vector_dim = vector_dim
+
+    async def embed_query(self, text: str) -> list[float]:
+        """Embed a search query."""
+        result = await self._embedder.embed_query(text)
+        return list(result.embeddings[0])
+
+    async def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents/chunks for indexing."""
+        if not texts:
+            return []
+        result = await self._embedder.embed_documents(texts)
+        return [list(e) for e in result.embeddings]
+
+
+def contextualize(chunks: list["Chunk"]) -> list[str]:
+    """Prepare chunk content for embedding/FTS by adding context.
 
-def get_embedder(config: AppConfig = Config) -> EmbedderBase:
+    Prepends section headings to chunk content for better semantic search.
+
+    Args:
+        chunks: List of chunks to contextualize.
+
+    Returns:
+        List of contextualized text strings.
     """
-    Factory function to get the appropriate embedder based on the configuration.
+    texts = []
+    for chunk in chunks:
+        meta = chunk.get_chunk_metadata()
+        if meta.headings:
+            text = "\n".join(meta.headings) + "\n" + chunk.content
+        else:
+            text = chunk.content
+        texts.append(text)
+    return texts
+
+
+async def embed_chunks(
+    chunks: list["Chunk"], config: AppConfig = Config
+) -> list["Chunk"]:
+    """Generate embeddings for chunks.
+
+    Contextualizes chunks (prepends headings) before embedding for better
+    semantic search. Returns new Chunk objects with embeddings set.
+
+    Args:
+        chunks: List of chunks to embed.
+        config: Configuration for embedder selection.
+
+    Returns:
+        New list of Chunk objects with embedding field populated.
+    """
+    if not chunks:
+        return []
+
+    from haiku.rag.store.models.chunk import Chunk
+
+    embedder = get_embedder(config)
+    texts = contextualize(chunks)
+    embeddings = await embedder.embed_documents(texts)
+
+    return [
+        Chunk(
+            id=chunk.id,
+            document_id=chunk.document_id,
+            content=chunk.content,
+            metadata=chunk.metadata,
+            order=chunk.order,
+            document_uri=chunk.document_uri,
+            document_title=chunk.document_title,
+            document_meta=chunk.document_meta,
+            embedding=embedding,
+        )
+        for chunk, embedding in zip(chunks, embeddings)
+    ]
+
+
+def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
+    """Factory function to get the appropriate embedder based on the configuration.
 
     Args:
         config: Configuration to use. Defaults to global Config.
@@ -13,37 +102,47 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase:
     Returns:
         An embedder instance configured according to the config.
     """
+    embedding_model = config.embeddings.model
+    provider = embedding_model.provider
+    model_name = embedding_model.name
+    vector_dim = embedding_model.vector_dim
 
-    if config.embeddings.provider == "ollama":
-        return OllamaEmbedder(
-            config.embeddings.model, config.embeddings.vector_dim, config
+    if provider == "ollama":
+        # Use model-level base_url if set, otherwise fall back to providers config
+        base_url = embedding_model.base_url or f"{config.providers.ollama.base_url}/v1"
+        model = OpenAIEmbeddingModel(
+            model_name,
+            provider=OllamaProvider(base_url=base_url),
         )
+        return EmbedderWrapper(Embedder(model), vector_dim)
+
+    if provider == "openai":
+        if embedding_model.base_url:
+            model = OpenAIEmbeddingModel(
+                model_name,
+                provider=OpenAIProvider(base_url=embedding_model.base_url),
+            )
+            return EmbedderWrapper(Embedder(model), vector_dim)
+        return EmbedderWrapper(Embedder(f"openai:{model_name}"), vector_dim)
 
-    if config.embeddings.provider == "voyageai":
+    if provider == "voyageai":
         try:
-            from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
+            from haiku.rag.embeddings.voyageai import VoyageAIEmbeddingModel
         except ImportError:
             raise ImportError(
                 "VoyageAI embedder requires the 'voyageai' package. "
                 "Please install haiku.rag with the 'voyageai' extra: "
                 "uv pip install haiku.rag[voyageai]"
             )
-        return VoyageAIEmbedder(
-            config.embeddings.model, config.embeddings.vector_dim, config
-        )
-
-    if config.embeddings.provider == "openai":
-        from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
-
-        return OpenAIEmbedder(
-            config.embeddings.model, config.embeddings.vector_dim, config
-        )
+        model = VoyageAIEmbeddingModel(model_name)
+        return EmbedderWrapper(Embedder(model), vector_dim)
 
-    if config.embeddings.provider == "vllm":
-        from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
+    if provider == "cohere":
+        return EmbedderWrapper(Embedder(f"cohere:{model_name}"), vector_dim)
 
-        return VllmEmbedder(
-            config.embeddings.model, config.embeddings.vector_dim, config
+    if provider == "sentence-transformers":
+        return EmbedderWrapper(
+            Embedder(f"sentence-transformers:{model_name}"), vector_dim
         )
 
-    raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")
+    raise ValueError(f"Unsupported embedding provider: {provider}")