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

haiku/rag/config/__init__.py

@@ -0,0 +1,74 @@
+import os
+
+from haiku.rag.config.loader import (
+    find_config_file,
+    generate_default_config,
+    load_yaml_config,
+)
+from haiku.rag.config.models import (
+    AGUIConfig,
+    AppConfig,
+    EmbeddingsConfig,
+    LanceDBConfig,
+    MonitorConfig,
+    OllamaConfig,
+    ProcessingConfig,
+    ProvidersConfig,
+    QAConfig,
+    RerankingConfig,
+    ResearchConfig,
+    StorageConfig,
+    VLLMConfig,
+)
+
+__all__ = [
+    "Config",
+    "AGUIConfig",
+    "AppConfig",
+    "StorageConfig",
+    "MonitorConfig",
+    "LanceDBConfig",
+    "EmbeddingsConfig",
+    "RerankingConfig",
+    "QAConfig",
+    "ResearchConfig",
+    "ProcessingConfig",
+    "OllamaConfig",
+    "VLLMConfig",
+    "ProvidersConfig",
+    "find_config_file",
+    "load_yaml_config",
+    "generate_default_config",
+    "get_config",
+    "set_config",
+]
+
+# Global config instance - initially loads from default locations
+_config: AppConfig | None = None
+
+
+def _load_default_config() -> AppConfig:
+    """Load config from default locations (used at import time)."""
+    config_path = find_config_file(None)
+    if config_path:
+        yaml_data = load_yaml_config(config_path)
+        return AppConfig.model_validate(yaml_data)
+    return AppConfig()
+
+
+def set_config(config: AppConfig) -> None:
+    """Set the global config instance (used by CLI to override)."""
+    global _config
+    _config = config
+
+
+def get_config() -> AppConfig:
+    """Get the current config instance."""
+    global _config
+    if _config is None:
+        _config = _load_default_config()
+    return _config
+
+
+# Legacy compatibility - Config is the default instance
+Config = _load_default_config()

haiku/rag/config/loader.py

@@ -0,0 +1,94 @@
+import os
+from pathlib import Path
+
+import yaml
+
+
+def find_config_file(cli_path: Path | None = None) -> Path | None:
+    """Find the YAML config file using the search path.
+
+    Search order:
+    1. CLI-provided path (via HAIKU_RAG_CONFIG_PATH env var or parameter)
+    2. ./haiku.rag.yaml (current directory)
+    3. Platform-specific user config directory
+
+    Returns None if no config file is found.
+    """
+    # Check environment variable first (set by CLI --config flag)
+    if not cli_path:
+        env_path = os.getenv("HAIKU_RAG_CONFIG_PATH")
+        if env_path:
+            cli_path = Path(env_path)
+
+    if cli_path:
+        if cli_path.exists():
+            return cli_path
+        raise FileNotFoundError(f"Config file not found: {cli_path}")
+
+    cwd_config = Path.cwd() / "haiku.rag.yaml"
+    if cwd_config.exists():
+        return cwd_config
+
+    # Use same directory as data storage for config
+    from haiku.rag.utils import get_default_data_dir
+
+    data_dir = get_default_data_dir()
+    user_config = data_dir / "haiku.rag.yaml"
+    if user_config.exists():
+        return user_config
+
+    return None
+
+
+def load_yaml_config(path: Path) -> dict:
+    """Load and parse a YAML config file."""
+    with open(path) as f:
+        data = yaml.safe_load(f)
+    return data or {}
+
+
+def generate_default_config() -> dict:
+    """Generate a default YAML config structure with documentation."""
+    return {
+        "environment": "production",
+        "storage": {
+            "data_dir": "",
+            "vacuum_retention_seconds": 86400,
+        },
+        "monitor": {
+            "directories": [],
+            "ignore_patterns": [],
+            "include_patterns": [],
+        },
+        "lancedb": {"uri": "", "api_key": "", "region": ""},
+        "embeddings": {
+            "provider": "ollama",
+            "model": "qwen3-embedding",
+            "vector_dim": 4096,
+        },
+        "reranking": {"provider": "", "model": ""},
+        "qa": {"provider": "ollama", "model": "gpt-oss"},
+        "research": {"provider": "", "model": ""},
+        "processing": {
+            "chunk_size": 256,
+            "context_chunk_radius": 0,
+            "markdown_preprocessor": "",
+        },
+        "providers": {
+            "ollama": {"base_url": "http://localhost:11434"},
+            "vllm": {
+                "embeddings_base_url": "",
+                "rerank_base_url": "",
+                "qa_base_url": "",
+                "research_base_url": "",
+            },
+        },
+        "agui": {
+            "host": "0.0.0.0",
+            "port": 8000,
+            "cors_origins": ["*"],
+            "cors_credentials": True,
+            "cors_methods": ["GET", "POST", "OPTIONS"],
+            "cors_headers": ["*"],
+        },
+    }

haiku/rag/config/models.py

@@ -0,0 +1,99 @@
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from haiku.rag.utils import get_default_data_dir
+
+
+class StorageConfig(BaseModel):
+    data_dir: Path = Field(default_factory=get_default_data_dir)
+    vacuum_retention_seconds: int = 86400
+
+
+class MonitorConfig(BaseModel):
+    directories: list[Path] = []
+    ignore_patterns: list[str] = []
+    include_patterns: list[str] = []
+    delete_orphans: bool = False
+
+
+class LanceDBConfig(BaseModel):
+    uri: str = ""
+    api_key: str = ""
+    region: str = ""
+
+
+class EmbeddingsConfig(BaseModel):
+    provider: str = "ollama"
+    model: str = "qwen3-embedding"
+    vector_dim: int = 4096
+
+
+class RerankingConfig(BaseModel):
+    provider: str = ""
+    model: str = ""
+
+
+class QAConfig(BaseModel):
+    provider: str = "ollama"
+    model: str = "gpt-oss"
+    max_sub_questions: int = 3
+    max_iterations: int = 2
+    max_concurrency: int = 1
+
+
+class ResearchConfig(BaseModel):
+    provider: str = "ollama"
+    model: str = "gpt-oss"
+    max_iterations: int = 3
+    confidence_threshold: float = 0.8
+    max_concurrency: int = 1
+
+
+class ProcessingConfig(BaseModel):
+    chunk_size: int = 256
+    context_chunk_radius: int = 0
+    markdown_preprocessor: str = ""
+
+
+class OllamaConfig(BaseModel):
+    base_url: str = Field(
+        default_factory=lambda: __import__("os").environ.get(
+            "OLLAMA_BASE_URL", "http://localhost:11434"
+        )
+    )
+
+
+class VLLMConfig(BaseModel):
+    embeddings_base_url: str = ""
+    rerank_base_url: str = ""
+    qa_base_url: str = ""
+    research_base_url: str = ""
+
+
+class ProvidersConfig(BaseModel):
+    ollama: OllamaConfig = Field(default_factory=OllamaConfig)
+    vllm: VLLMConfig = Field(default_factory=VLLMConfig)
+
+
+class AGUIConfig(BaseModel):
+    host: str = "0.0.0.0"
+    port: int = 8000
+    cors_origins: list[str] = ["*"]
+    cors_credentials: bool = True
+    cors_methods: list[str] = ["GET", "POST", "OPTIONS"]
+    cors_headers: list[str] = ["*"]
+
+
+class AppConfig(BaseModel):
+    environment: str = "production"
+    storage: StorageConfig = Field(default_factory=StorageConfig)
+    monitor: MonitorConfig = Field(default_factory=MonitorConfig)
+    lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
+    embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
+    reranking: RerankingConfig = Field(default_factory=RerankingConfig)
+    qa: QAConfig = Field(default_factory=QAConfig)
+    research: ResearchConfig = Field(default_factory=ResearchConfig)
+    processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
+    providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
+    agui: AGUIConfig = Field(default_factory=AGUIConfig)

haiku/rag/embeddings/__init__.py

@@ -0,0 +1,49 @@
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.embeddings.base import EmbedderBase
+from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
+
+
+def get_embedder(config: AppConfig = Config) -> EmbedderBase:
+    """
+    Factory function to get the appropriate embedder based on the configuration.
+
+    Args:
+        config: Configuration to use. Defaults to global Config.
+
+    Returns:
+        An embedder instance configured according to the config.
+    """
+
+    if config.embeddings.provider == "ollama":
+        return OllamaEmbedder(
+            config.embeddings.model, config.embeddings.vector_dim, config
+        )
+
+    if config.embeddings.provider == "voyageai":
+        try:
+            from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
+        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
+        )
+
+    if config.embeddings.provider == "vllm":
+        from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
+
+        return VllmEmbedder(
+            config.embeddings.model, config.embeddings.vector_dim, config
+        )
+
+    raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")

haiku/rag/embeddings/base.py

@@ -0,0 +1,25 @@
+from typing import overload
+
+from haiku.rag.config import AppConfig, Config
+
+
+class EmbedderBase:
+    _model: str = Config.embeddings.model
+    _vector_dim: int = Config.embeddings.vector_dim
+    _config: AppConfig = Config
+
+    def __init__(self, model: str, vector_dim: int, config: AppConfig = Config):
+        self._model = model
+        self._vector_dim = vector_dim
+        self._config = config
+
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        raise NotImplementedError(
+            "Embedder is an abstract class. Please implement the embed method in a subclass."
+        )

haiku/rag/embeddings/ollama.py

@@ -0,0 +1,28 @@
+from typing import overload
+
+from openai import AsyncOpenAI
+
+from haiku.rag.embeddings.base import EmbedderBase
+
+
+class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        client = AsyncOpenAI(
+            base_url=f"{self._config.providers.ollama.base_url}/v1", api_key="dummy"
+        )
+        if not text:
+            return []
+        response = await client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/openai.py

@@ -0,0 +1,26 @@
+from typing import overload
+
+from openai import AsyncOpenAI
+
+from haiku.rag.embeddings.base import EmbedderBase
+
+
+class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        client = AsyncOpenAI()
+        if not text:
+            return []
+        response = await client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/vllm.py

@@ -0,0 +1,29 @@
+from typing import overload
+
+from openai import AsyncOpenAI
+
+from haiku.rag.embeddings.base import EmbedderBase
+
+
+class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        client = AsyncOpenAI(
+            base_url=f"{self._config.providers.vllm.embeddings_base_url}/v1",
+            api_key="dummy",
+        )
+        if not text:
+            return []
+        response = await client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/voyageai.py

@@ -0,0 +1,27 @@
+try:
+    from typing import overload
+
+    from voyageai.client import Client  # type: ignore
+
+    from haiku.rag.embeddings.base import EmbedderBase
+
+    class Embedder(EmbedderBase):
+        @overload
+        async def embed(self, text: str) -> list[float]: ...
+
+        @overload
+        async def embed(self, text: list[str]) -> list[list[float]]: ...
+
+        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+            client = Client()
+            if not text:
+                return []
+            if isinstance(text, str):
+                res = client.embed([text], model=self._model, output_dtype="float")
+                return res.embeddings[0]  # type: ignore[return-value]
+            else:
+                res = client.embed(text, model=self._model, output_dtype="float")
+                return res.embeddings  # type: ignore[return-value]
+
+except ImportError:
+    pass

haiku/rag/graph/__init__.py

@@ -0,0 +1,26 @@
+"""Graph module for haiku.rag.
+
+This module contains all graph-related functionality including:
+- AG-UI protocol for graph streaming
+- Common graph utilities and models
+- Research graph implementation
+- Deep QA graph implementation
+"""
+
+from haiku.rag.graph.agui import (
+    AGUIConsoleRenderer,
+    AGUIEmitter,
+    create_agui_server,
+    stream_graph,
+)
+from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
+from haiku.rag.graph.research.graph import build_research_graph
+
+__all__ = [
+    "AGUIConsoleRenderer",
+    "AGUIEmitter",
+    "build_deep_qa_graph",
+    "build_research_graph",
+    "create_agui_server",
+    "stream_graph",
+]

haiku/rag/graph/agui/__init__.py

@@ -0,0 +1,53 @@
+"""Generic AG-UI protocol support for haiku.rag graphs."""
+
+from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
+from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.events import (
+    AGUIEvent,
+    emit_activity,
+    emit_activity_delta,
+    emit_run_error,
+    emit_run_finished,
+    emit_run_started,
+    emit_state_delta,
+    emit_state_snapshot,
+    emit_step_finished,
+    emit_step_started,
+    emit_text_message,
+    emit_text_message_content,
+    emit_text_message_end,
+    emit_text_message_start,
+)
+from haiku.rag.graph.agui.server import (
+    RunAgentInput,
+    create_agui_app,
+    create_agui_server,
+    format_sse_event,
+)
+from haiku.rag.graph.agui.state import compute_state_delta
+from haiku.rag.graph.agui.stream import stream_graph
+
+__all__ = [
+    "AGUIConsoleRenderer",
+    "AGUIEmitter",
+    "AGUIEvent",
+    "RunAgentInput",
+    "compute_state_delta",
+    "create_agui_app",
+    "create_agui_server",
+    "emit_activity",
+    "emit_activity_delta",
+    "emit_run_error",
+    "emit_run_finished",
+    "emit_run_started",
+    "emit_state_delta",
+    "emit_state_snapshot",
+    "emit_step_finished",
+    "emit_step_started",
+    "emit_text_message",
+    "emit_text_message_content",
+    "emit_text_message_end",
+    "emit_text_message_start",
+    "format_sse_event",
+    "stream_graph",
+]

haiku/rag/graph/agui/cli_renderer.py

@@ -0,0 +1,135 @@
+"""Generic CLI renderer for AG-UI events with Rich console output."""
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from rich.console import Console
+
+from haiku.rag.graph.agui.events import AGUIEvent
+
+
+class AGUIConsoleRenderer:
+    """Renders AG-UI events to Rich console with formatted output.
+
+    Generic renderer that processes AG-UI protocol events and renders them
+    with Rich formatting. Works with any graph that emits AG-UI events.
+    """
+
+    def __init__(self, console: Console | None = None):
+        """Initialize the renderer.
+
+        Args:
+            console: Optional Rich console instance (creates new one if not provided)
+        """
+        self.console = console or Console()
+
+    async def render(self, events: AsyncIterator[AGUIEvent]) -> Any | None:
+        """Process events and render to console, return final result.
+
+        Args:
+            events: Async iterator of AG-UI events
+
+        Returns:
+            The final result from RunFinished event, or None
+        """
+        result = None
+
+        async for event in events:
+            event_type = event.get("type")
+
+            if event_type == "RUN_STARTED":
+                self._render_run_started(event)
+            elif event_type == "RUN_FINISHED":
+                result = event.get("result")
+                self._render_run_finished()
+            elif event_type == "RUN_ERROR":
+                self._render_error(event)
+            elif event_type == "STEP_STARTED":
+                self._render_step_started(event)
+            elif event_type == "STEP_FINISHED":
+                self._render_step_finished(event)
+            elif event_type == "TEXT_MESSAGE_CHUNK":
+                self._render_text_message(event)
+            elif event_type == "TEXT_MESSAGE_START":
+                pass  # Start of streaming message, no output needed
+            elif event_type == "TEXT_MESSAGE_CONTENT":
+                self._render_text_content(event)
+            elif event_type == "TEXT_MESSAGE_END":
+                pass  # End of streaming message, no output needed
+            elif event_type == "STATE_SNAPSHOT":
+                self._render_state_snapshot(event)
+            elif event_type == "STATE_DELTA":
+                self._render_state_delta(event)
+            elif event_type == "ACTIVITY_SNAPSHOT":
+                self._render_activity(event)
+            elif event_type == "ACTIVITY_DELTA":
+                pass  # Activity deltas don't need separate rendering
+
+        return result
+
+    def _render_run_started(self, event: AGUIEvent) -> None:
+        """Render run start event."""
+        run_id = event.get("runId", "")
+        if run_id:
+            # Show shortened run ID (first 8 chars like our UUIDs)
+            short_id = run_id[:8] if len(run_id) > 8 else run_id
+            self.console.print(f"[bold green][RUN_STARTED][/bold green] Run {short_id}")
+
+    def _render_run_finished(self) -> None:
+        """Render run completion."""
+        self.console.print("[bold green][RUN_FINISHED][/bold green] Completed")
+
+    def _render_error(self, event: AGUIEvent) -> None:
+        """Render error event."""
+        message = event.get("message", "Unknown error")
+        self.console.print(f"[bold red][RUN_ERROR][/bold red] {message}")
+
+    def _render_step_started(self, event: AGUIEvent) -> None:
+        """Render step start event."""
+        step_name = event.get("stepName", "")
+        if step_name:
+            display_name = step_name.replace("_", " ").title()
+            self.console.print(
+                f"\n[bold cyan][STEP_STARTED][/bold cyan] {display_name}"
+            )
+
+    def _render_step_finished(self, event: AGUIEvent) -> None:
+        """Render step finish event."""
+        step_name = event.get("stepName", "")
+        if step_name:
+            display_name = step_name.replace("_", " ").title()
+            self.console.print(f"[cyan][STEP_FINISHED][/cyan] {display_name}")
+
+    def _render_text_message(self, event: AGUIEvent) -> None:
+        """Render complete text message."""
+        delta = event.get("delta", "")
+        self.console.print(f"[magenta][TEXT_MESSAGE][/magenta] {delta}")
+
+    def _render_text_content(self, event: AGUIEvent) -> None:
+        """Render streaming text content delta."""
+        delta = event.get("delta", "")
+        self.console.print(delta, end="")
+
+    def _render_activity(self, event: AGUIEvent) -> None:
+        """Render activity update."""
+        content = event.get("content", "")
+        if content:
+            self.console.print(f"[yellow][ACTIVITY][/yellow] {content}")
+
+    def _render_state_snapshot(self, event: AGUIEvent) -> None:
+        """Render full state snapshot."""
+        snapshot = event.get("snapshot")
+        if not snapshot:
+            return
+
+        self.console.print("[blue][STATE_SNAPSHOT][/blue]")
+        self.console.print(snapshot, style="dim")
+
+    def _render_state_delta(self, event: AGUIEvent) -> None:
+        """Render state delta operations."""
+        delta = event.get("delta", [])
+        if not delta:
+            return
+
+        self.console.print("[blue][STATE_DELTA][/blue]")
+        self.console.print(delta, style="dim")