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

haiku/rag/utils.py

@@ -1,51 +1,379 @@
-import asyncio
-import importlib
-import importlib.util
 import sys
-from collections.abc import Callable
-from functools import wraps
+from datetime import UTC, datetime
 from importlib import metadata
-from io import BytesIO
 from pathlib import Path
-from types import ModuleType
+from typing import TYPE_CHECKING, Any
 
+from dateutil import parser as dateutil_parser
 from packaging.version import Version, parse
 
+if TYPE_CHECKING:
+    from rich.console import RenderableType
 
-def debounce(wait: float) -> Callable:
+    from haiku.rag.config.models import AppConfig, ModelConfig
+    from haiku.rag.graph.research.models import Citation
+
+
+def parse_datetime(s: str) -> datetime:
+    """Parse a datetime string into a datetime object.
+
+    Supports:
+    - ISO 8601 format: "2025-01-15T14:30:00", "2025-01-15T14:30:00Z", "2025-01-15T14:30:00+00:00"
+    - Date only: "2025-01-15" (interpreted as 00:00:00)
+    - Various other formats via dateutil
+
+    Args:
+        s: String to parse
+
+    Returns:
+        Parsed datetime object
+
+    Raises:
+        ValueError: If the string cannot be parsed
+    """
+    try:
+        return dateutil_parser.parse(s)
+    except (ValueError, TypeError) as e:
+        raise ValueError(
+            f"Could not parse datetime: {s}. "
+            "Use ISO 8601 format (e.g., 2025-01-15T14:30:00) or date (e.g., 2025-01-15)"
+        ) from e
+
+
+def to_utc(dt: datetime) -> datetime:
+    """Convert a datetime to UTC.
+
+    - Naive datetimes are assumed to be local time and converted to UTC
+    - Datetimes with timezone info are converted to UTC
+    - UTC datetimes are returned as-is
+
+    Args:
+        dt: Datetime to convert
+
+    Returns:
+        Datetime in UTC timezone
+    """
+    if dt.tzinfo is None:
+        # Naive datetime - assume local time
+        local_dt = dt.astimezone()  # Adds local timezone
+        return local_dt.astimezone(UTC)
+    elif dt.tzinfo == UTC:
+        return dt
+    else:
+        return dt.astimezone(UTC)
+
+
+def apply_common_settings(
+    settings: Any | None,
+    settings_class: type[Any],
+    model_config: Any,
+) -> Any | None:
+    """Apply common settings (temperature, max_tokens) to model settings.
+
+    Args:
+        settings: Existing settings instance or None
+        settings_class: Settings class to instantiate if needed
+        model_config: ModelConfig with temperature and max_tokens
+
+    Returns:
+        Updated settings instance or None if no settings to apply
+    """
+    if model_config.temperature is None and model_config.max_tokens is None:
+        return settings
+
+    if settings is None:
+        settings_dict = settings_class()
+    else:
+        settings_dict = settings
+
+    if model_config.temperature is not None:
+        settings_dict["temperature"] = model_config.temperature
+
+    if model_config.max_tokens is not None:
+        settings_dict["max_tokens"] = model_config.max_tokens
+
+    return settings_dict
+
+
+def get_model(
+    model_config: "ModelConfig",
+    app_config: "AppConfig | None" = None,
+) -> Any:
     """
-    A decorator to debounce a function, ensuring it is called only after a specified delay
-    and always executes after the last call.
+    Get a model instance for the specified configuration.
 
     Args:
-        wait (float): The debounce delay in seconds.
+        model_config: ModelConfig with provider, model, and settings
+        app_config: AppConfig for provider base URLs (defaults to global Config)
 
     Returns:
-        Callable: The decorated function.
+        A configured model instance
     """
+    from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
+    from pydantic_ai.providers.ollama import OllamaProvider
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+    if app_config is None:
+        from haiku.rag.config import Config
+
+        app_config = Config
+
+    provider = model_config.provider
+    model = model_config.name
+
+    if provider == "ollama":
+        model_settings = None
+
+        # Apply thinking control for gpt-oss
+        if model == "gpt-oss" and model_config.enable_thinking is not None:
+            if model_config.enable_thinking is False:
+                model_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
+            else:
+                model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
+
+        model_settings = apply_common_settings(
+            model_settings, OpenAIChatModelSettings, model_config
+        )
+
+        # Use model-level base_url if set, otherwise fall back to providers config
+        base_url = model_config.base_url or f"{app_config.providers.ollama.base_url}/v1"
+
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=base_url),
+            settings=model_settings,
+        )
+
+    elif provider == "openai":
+        openai_settings: Any = None
+
+        # Apply thinking control
+        if model_config.enable_thinking is not None:
+            if model_config.enable_thinking is False:
+                openai_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
+            else:
+                openai_settings = OpenAIChatModelSettings(
+                    openai_reasoning_effort="high"
+                )
+
+        openai_settings = apply_common_settings(
+            openai_settings, OpenAIChatModelSettings, model_config
+        )
+
+        # Use model-level base_url if set (for vLLM, LM Studio, etc.)
+        if model_config.base_url:
+            return OpenAIChatModel(
+                model_name=model,
+                provider=OpenAIProvider(base_url=model_config.base_url),
+                settings=openai_settings,
+            )
+
+        return OpenAIChatModel(model_name=model, settings=openai_settings)
+
+    elif provider == "anthropic":
+        from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+
+        anthropic_settings: Any = None
+
+        # Apply thinking control
+        if model_config.enable_thinking is not None:
+            if model_config.enable_thinking:
+                anthropic_settings = AnthropicModelSettings(
+                    anthropic_thinking={"type": "enabled", "budget_tokens": 4096}
+                )
+            else:
+                anthropic_settings = AnthropicModelSettings(
+                    anthropic_thinking={"type": "disabled"}
+                )
+
+        anthropic_settings = apply_common_settings(
+            anthropic_settings, AnthropicModelSettings, model_config
+        )
+
+        return AnthropicModel(model_name=model, settings=anthropic_settings)
+
+    elif provider == "gemini":
+        from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
+
+        gemini_settings: Any = None
 
-    def decorator(func: Callable) -> Callable:
-        last_call = None
-        task = None
+        # Apply thinking control
+        if model_config.enable_thinking is not None:
+            gemini_settings = GoogleModelSettings(
+                google_thinking_config={
+                    "include_thoughts": model_config.enable_thinking
+                }
+            )
 
-        @wraps(func)
-        async def debounced(*args, **kwargs):
-            nonlocal last_call, task
-            last_call = asyncio.get_event_loop().time()
+        gemini_settings = apply_common_settings(
+            gemini_settings, GoogleModelSettings, model_config
+        )
 
-            if task:
-                task.cancel()
+        return GoogleModel(model_name=model, settings=gemini_settings)
 
-            async def call_func():
-                await asyncio.sleep(wait)
-                if asyncio.get_event_loop().time() - last_call >= wait:  # type: ignore
-                    await func(*args, **kwargs)
+    elif provider == "groq":
+        from pydantic_ai.models.groq import GroqModel, GroqModelSettings
 
-            task = asyncio.create_task(call_func())
+        groq_settings: Any = None
 
-        return debounced
+        # Apply thinking control
+        if model_config.enable_thinking is not None:
+            if model_config.enable_thinking:
+                groq_settings = GroqModelSettings(groq_reasoning_format="parsed")
+            else:
+                groq_settings = GroqModelSettings(groq_reasoning_format="hidden")
 
-    return decorator
+        groq_settings = apply_common_settings(
+            groq_settings, GroqModelSettings, model_config
+        )
+
+        return GroqModel(model_name=model, settings=groq_settings)
+
+    elif provider == "bedrock":
+        from pydantic_ai.models.bedrock import (
+            BedrockConverseModel,
+            BedrockModelSettings,
+        )
+
+        bedrock_settings: Any = None
+
+        # Apply thinking control for Claude models
+        if model_config.enable_thinking is not None:
+            additional_fields: dict[str, Any] = {}
+            if model.startswith("anthropic.claude"):
+                if model_config.enable_thinking:
+                    additional_fields = {
+                        "thinking": {"type": "enabled", "budget_tokens": 4096}
+                    }
+                else:
+                    additional_fields = {"thinking": {"type": "disabled"}}
+            elif "gpt" in model or "o1" in model or "o3" in model:
+                # OpenAI models on Bedrock
+                additional_fields = {
+                    "reasoning_effort": "high"
+                    if model_config.enable_thinking
+                    else "low"
+                }
+            elif "qwen" in model:
+                # Qwen models on Bedrock
+                additional_fields = {
+                    "reasoning_config": "high"
+                    if model_config.enable_thinking
+                    else "low"
+                }
+
+            if additional_fields:
+                bedrock_settings = BedrockModelSettings(
+                    bedrock_additional_model_requests_fields=additional_fields
+                )
+
+        bedrock_settings = apply_common_settings(
+            bedrock_settings, BedrockModelSettings, model_config
+        )
+
+        return BedrockConverseModel(model_name=model, settings=bedrock_settings)
+
+    else:
+        # For any other provider, use string format and let Pydantic AI handle it
+        return f"{provider}:{model}"
+
+
+def format_bytes(num_bytes: int) -> str:
+    """Format bytes as human-readable string."""
+    size = float(num_bytes)
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if size < 1024.0:
+            return f"{size:.1f} {unit}"
+        size /= 1024.0
+    return f"{size:.1f} PB"
+
+
+def format_citations(citations: "list[Citation]") -> str:
+    """Format citations as plain text with preserved formatting.
+
+    Used by things like the MCP server where Rich renderables are not available.
+    """
+    if not citations:
+        return ""
+
+    lines = ["## Citations\n"]
+
+    for c in citations:
+        # Header line
+        header = f"[{c.document_id}:{c.chunk_id}]"
+
+        # Location info
+        location_parts = []
+        if c.page_numbers:
+            if len(c.page_numbers) == 1:
+                location_parts.append(f"p. {c.page_numbers[0]}")
+            else:
+                location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}")
+        if c.headings:
+            location_parts.append(f"Section: {c.headings[-1]}")
+
+        source = c.document_uri
+        if c.document_title:
+            source = f"{c.document_title} ({c.document_uri})"
+        if location_parts:
+            source += f" - {', '.join(location_parts)}"
+
+        lines.append(f"{header} {source}")
+        lines.append(c.content)
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]":
+    """Format citations as Rich renderables.
+
+    Returns a list of Rich Panel objects for direct console printing,
+    with content rendered as markdown for syntax highlighting.
+    """
+    from rich.markdown import Markdown
+    from rich.panel import Panel
+    from rich.text import Text
+
+    if not citations:
+        return []
+
+    renderables: list[RenderableType] = []
+    renderables.append(Text("Citations", style="bold"))
+
+    for c in citations:
+        # Build header with IDs
+        header = Text()
+        header.append("doc: ", style="dim")
+        header.append(c.document_id, style="cyan")
+        header.append("  chunk: ", style="dim")
+        header.append(c.chunk_id, style="cyan")
+
+        # Location info for subtitle
+        location_parts = []
+        if c.page_numbers:
+            if len(c.page_numbers) == 1:
+                location_parts.append(f"p. {c.page_numbers[0]}")
+            else:
+                location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}")
+        if c.headings:
+            location_parts.append(f"Section: {c.headings[-1]}")
+
+        subtitle = c.document_uri
+        if c.document_title:
+            subtitle = f"{c.document_title} ({c.document_uri})"
+        if location_parts:
+            subtitle += f" - {', '.join(location_parts)}"
+        panel = Panel(
+            Markdown(c.content),
+            title=header,
+            subtitle=subtitle,
+            subtitle_align="left",
+            border_style="dim",
+        )
+        renderables.append(panel)
+
+    return renderables
 
 
 def get_default_data_dir() -> Path:
@@ -70,6 +398,21 @@ def get_default_data_dir() -> Path:
     return data_path
 
 
+def build_prompt(base_prompt: str, config: "AppConfig") -> str:
+    """Build a prompt with domain_preamble prepended if configured.
+
+    Args:
+        base_prompt: The base prompt to use
+        config: AppConfig with prompts.domain_preamble
+
+    Returns:
+        Prompt with domain_preamble prepended if configured
+    """
+    if config.prompts.domain_preamble:
+        return f"{config.prompts.domain_preamble}\n\n{base_prompt}"
+    return base_prompt
+
+
 async def is_up_to_date() -> tuple[bool, Version, Version]:
     """Check whether haiku.rag is current.
 
@@ -91,121 +434,3 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
             # If no network connection, do not raise alarms.
             pypi_version = running_version
     return running_version >= pypi_version, running_version, pypi_version
-
-
-def text_to_docling_document(text: str, name: str = "content.md"):
-    """Convert text content to a DoclingDocument.
-
-    Args:
-        text: The text content to convert.
-        name: The name to use for the document stream (defaults to "content.md").
-
-    Returns:
-        A DoclingDocument created from the text content.
-    """
-    try:
-        import docling  # noqa: F401
-    except ImportError as e:
-        raise ImportError(
-            "Docling is required for document conversion. "
-            "Install with: pip install haiku.rag-slim[docling]"
-        ) from e
-
-    from docling.document_converter import DocumentConverter
-    from docling_core.types.io import DocumentStream
-
-    bytes_io = BytesIO(text.encode("utf-8"))
-    doc_stream = DocumentStream(name=name, stream=bytes_io)
-    converter = DocumentConverter()
-    result = converter.convert(doc_stream)
-    return result.document
-
-
-def load_callable(path: str):
-    """Load a callable from a dotted path or file path.
-
-    Supported formats:
-    - "package.module:func" or "package.module.func"
-    - "path/to/file.py:func"
-
-    Returns the loaded callable. Raises ValueError on failure.
-    """
-    if not path:
-        raise ValueError("Empty callable path provided")
-
-    module_part = None
-    func_name = None
-
-    if ":" in path:
-        module_part, func_name = path.split(":", 1)
-    else:
-        # split by last dot for module.attr
-        if "." in path:
-            module_part, func_name = path.rsplit(".", 1)
-        else:
-            raise ValueError(
-                "Invalid callable path format. Use 'module:func' or 'module.func' or 'file.py:func'."
-            )
-
-    # Try file path first
-    mod: ModuleType | None = None
-    module_path = Path(module_part)
-    if module_path.suffix == ".py" and module_path.exists():
-        spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
-        if spec and spec.loader:
-            mod = importlib.util.module_from_spec(spec)
-            spec.loader.exec_module(mod)
-    else:
-        # Import as a module path
-        try:
-            mod = importlib.import_module(module_part)
-        except Exception as e:
-            raise ValueError(f"Failed to import module '{module_part}': {e}")
-
-    if not hasattr(mod, func_name):
-        raise ValueError(f"Callable '{func_name}' not found in module '{module_part}'")
-    func = getattr(mod, func_name)
-    if not callable(func):
-        raise ValueError(
-            f"Attribute '{func_name}' in module '{module_part}' is not callable"
-        )
-    return func
-
-
-def prefetch_models():
-    """Prefetch runtime models (Docling + Ollama as configured)."""
-    import httpx
-
-    from haiku.rag.config import Config
-
-    try:
-        from docling.utils.model_downloader import download_models
-
-        download_models()
-    except ImportError:
-        # Docling not installed, skip downloading docling models
-        pass
-
-    # Collect Ollama models from config
-    required_models: set[str] = set()
-    if Config.embeddings.provider == "ollama":
-        required_models.add(Config.embeddings.model)
-    if Config.qa.provider == "ollama":
-        required_models.add(Config.qa.model)
-    if Config.research.provider == "ollama":
-        required_models.add(Config.research.model)
-    if Config.reranking.provider == "ollama":
-        required_models.add(Config.reranking.model)
-
-    if not required_models:
-        return
-
-    base_url = Config.providers.ollama.base_url
-
-    with httpx.Client(timeout=None) as client:
-        for model in sorted(required_models):
-            with client.stream(
-                "POST", f"{base_url}/api/pull", json={"model": model}
-            ) as r:
-                for _ in r.iter_lines():
-                    pass

{haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: haiku.rag-slim
-Version: 0.16.0
-Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB - Minimal dependencies
+Version: 0.24.0
+Summary: Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
 License-File: LICENSE
@@ -17,16 +17,15 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Typing :: Typed
 Requires-Python: >=3.12
-Requires-Dist: docling-core==2.50.1
+Requires-Dist: docling-core==2.57.0
 Requires-Dist: httpx>=0.28.1
-Requires-Dist: lancedb==0.25.2
+Requires-Dist: lancedb==0.26.0
 Requires-Dist: pathspec>=0.12.1
-Requires-Dist: pydantic-ai-slim[ag-ui,fastmcp,logfire,openai]>=1.11.1
-Requires-Dist: pydantic>=2.12.3
+Requires-Dist: pydantic-ai-slim[ag-ui,fastmcp,logfire,openai]>=1.39.0
+Requires-Dist: pydantic>=2.12.5
 Requires-Dist: python-dotenv>=1.2.1
 Requires-Dist: pyyaml>=6.0.3
 Requires-Dist: rich>=14.2.0
-Requires-Dist: tiktoken>=0.12.0
 Requires-Dist: typer<0.20.0,>=0.19.2
 Requires-Dist: watchfiles>=1.1.1
 Provides-Extra: anthropic
@@ -34,13 +33,17 @@ Requires-Dist: pydantic-ai-slim[anthropic]; extra == 'anthropic'
 Provides-Extra: bedrock
 Requires-Dist: pydantic-ai-slim[bedrock]; extra == 'bedrock'
 Provides-Extra: cohere
-Requires-Dist: cohere>=5.0.0; extra == 'cohere'
+Requires-Dist: cohere>=5.20.1; extra == 'cohere'
 Provides-Extra: docling
-Requires-Dist: docling==2.61.1; extra == 'docling'
+Requires-Dist: docling==2.65.0; extra == 'docling'
+Requires-Dist: opencv-python-headless>=4.12.0.88; extra == 'docling'
 Provides-Extra: google
 Requires-Dist: pydantic-ai-slim[google]; extra == 'google'
 Provides-Extra: groq
 Requires-Dist: pydantic-ai-slim[groq]; extra == 'groq'
+Provides-Extra: inspector
+Requires-Dist: textual-image>=0.8.5; extra == 'inspector'
+Requires-Dist: textual>=7.0.0; extra == 'inspector'
 Provides-Extra: mistral
 Requires-Dist: pydantic-ai-slim[mistral]; extra == 'mistral'
 Provides-Extra: mxbai
@@ -48,14 +51,14 @@ Requires-Dist: mxbai-rerank>=0.1.6; extra == 'mxbai'
 Provides-Extra: vertexai
 Requires-Dist: pydantic-ai-slim[vertexai]; extra == 'vertexai'
 Provides-Extra: voyageai
-Requires-Dist: voyageai>=0.3.5; extra == 'voyageai'
+Requires-Dist: voyageai>=0.3.7; extra == 'voyageai'
 Provides-Extra: zeroentropy
-Requires-Dist: zeroentropy>=0.1.0a6; extra == 'zeroentropy'
+Requires-Dist: zeroentropy>=0.1.0a7; extra == 'zeroentropy'
 Description-Content-Type: text/markdown
 
 # haiku.rag-slim
 
-Retrieval-Augmented Generation (RAG) library built on LanceDB - Core package with minimal dependencies.
+Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Core package with minimal dependencies.
 
 `haiku.rag-slim` is the core package for users who want to install only the dependencies they need. Document processing (docling), and reranker support are all optional extras.