haiku.rag 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

haiku/rag/client.py

@@ -8,8 +8,7 @@ from urllib.parse import urlparse
 
 import httpx
 
-from haiku.rag.config import Config
-from haiku.rag.reader import FileReader
+from haiku.rag.config import AppConfig, Config
 from haiku.rag.reranking import get_reranker
 from haiku.rag.store.engine import Store
 from haiku.rag.store.models.chunk import Chunk
@@ -17,7 +16,6 @@ from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
 from haiku.rag.store.repositories.settings import SettingsRepository
-from haiku.rag.utils import text_to_docling_document
 
 logger = logging.getLogger(__name__)
 
@@ -27,16 +25,23 @@ class HaikuRAG:
 
     def __init__(
         self,
-        db_path: Path = Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        db_path: Path | None = None,
+        config: AppConfig = Config,
         skip_validation: bool = False,
     ):
         """Initialize the RAG client with a database path.
 
         Args:
-            db_path: Path to the database file.
+            db_path: Path to the database file. If None, uses config.storage.data_dir.
+            config: Configuration to use. Defaults to global Config.
             skip_validation: Whether to skip configuration validation on database load.
         """
-        self.store = Store(db_path, skip_validation=skip_validation)
+        self._config = config
+        if db_path is None:
+            db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
+        self.store = Store(
+            db_path, config=self._config, skip_validation=skip_validation
+        )
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
 
@@ -91,6 +96,9 @@ class HaikuRAG:
         Returns:
             The created Document instance.
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         # Convert content to DoclingDocument for processing
         docling_document = text_to_docling_document(content)
 
@@ -106,8 +114,8 @@ class HaikuRAG:
 
     async def create_document_from_source(
         self, source: str | Path, title: str | None = None, metadata: dict | None = None
-    ) -> Document:
-        """Create or update a document from a file path or URL.
+    ) -> Document | list[Document]:
+        """Create or update document(s) from a file path, directory, or URL.
 
         Checks if a document with the same URI already exists:
         - If MD5 is unchanged, returns existing document
@@ -115,16 +123,20 @@ class HaikuRAG:
         - If no document exists, creates a new one
 
         Args:
-            source: File path (as string or Path) or URL to parse
+            source: File path, directory (as string or Path), or URL to parse
+            title: Optional title (only used for single files, not directories)
             metadata: Optional metadata dictionary
 
         Returns:
-            Document instance (created, updated, or existing)
+            Document instance (created, updated, or existing) for single files/URLs
+            List of Document instances for directories
 
         Raises:
             ValueError: If the file/URL cannot be parsed or doesn't exist
             httpx.RequestError: If URL request fails
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
 
         # Normalize metadata
         metadata = metadata or {}
@@ -142,6 +154,48 @@ class HaikuRAG:
         else:
             # Handle as regular file path
             source_path = Path(source) if isinstance(source, str) else source
+
+        # Handle directories
+        if source_path.is_dir():
+            documents = []
+            supported_extensions = set(FileReader.extensions)
+            for file_path in source_path.rglob("*"):
+                if (
+                    file_path.is_file()
+                    and file_path.suffix.lower() in supported_extensions
+                ):
+                    doc = await self._create_document_from_file(
+                        file_path, title=None, metadata=metadata
+                    )
+                    documents.append(doc)
+            return documents
+
+        # Handle single file
+        return await self._create_document_from_file(
+            source_path, title=title, metadata=metadata
+        )
+
+    async def _create_document_from_file(
+        self, source_path: Path, title: str | None = None, metadata: dict | None = None
+    ) -> Document:
+        """Create or update a document from a single file path.
+
+        Args:
+            source_path: Path to the file
+            title: Optional title
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Document instance (created, updated, or existing)
+
+        Raises:
+            ValueError: If the file cannot be parsed or doesn't exist
+        """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
+        metadata = metadata or {}
+
         if source_path.suffix.lower() not in FileReader.extensions:
             raise ValueError(f"Unsupported file extension: {source_path.suffix}")
 
@@ -215,6 +269,9 @@ class HaikuRAG:
             ValueError: If the content cannot be parsed
             httpx.RequestError: If URL request fails
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         metadata = metadata or {}
 
         async with httpx.AsyncClient() as client:
@@ -338,6 +395,9 @@ class HaikuRAG:
 
     async def update_document(self, document: Document) -> Document:
         """Update an existing document."""
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         # Convert content to DoclingDocument
         docling_document = text_to_docling_document(document.content)
 
@@ -377,7 +437,7 @@ class HaikuRAG:
             List of (chunk, score) tuples ordered by relevance.
         """
         # Get reranker if available
-        reranker = get_reranker()
+        reranker = get_reranker(config=self._config)
 
         if reranker is None:
             # No reranking - return direct search results
@@ -399,18 +459,20 @@ class HaikuRAG:
     async def expand_context(
         self,
         search_results: list[tuple[Chunk, float]],
-        radius: int = Config.CONTEXT_CHUNK_RADIUS,
+        radius: int | None = None,
     ) -> list[tuple[Chunk, float]]:
         """Expand search results with adjacent chunks, merging overlapping chunks.
 
         Args:
             search_results: List of (chunk, score) tuples from search.
             radius: Number of adjacent chunks to include before/after each chunk.
-                   Defaults to CONTEXT_CHUNK_RADIUS config setting.
+                   If None, uses config.processing.context_chunk_radius.
 
         Returns:
             List of (chunk, score) tuples with expanded and merged context chunks.
         """
+        if radius is None:
+            radius = self._config.processing.context_chunk_radius
         if radius == 0:
             return search_results
 
@@ -540,7 +602,9 @@ class HaikuRAG:
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
+        qa_agent = get_qa_agent(
+            self, config=self._config, use_citations=cite, system_prompt=system_prompt
+        )
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[str, None]:
@@ -556,6 +620,9 @@ class HaikuRAG:
         Yields:
             int: The ID of the document currently being processed
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         await self.chunk_repository.delete_all()
         self.store.recreate_embeddings_table()
 
@@ -592,6 +659,8 @@ class HaikuRAG:
                         new_doc = await self.create_document_from_source(
                             source=doc.uri, metadata=doc.metadata or {}
                         )
+                        # URIs always point to single files/URLs, never directories
+                        assert isinstance(new_doc, Document)
                         assert new_doc.id is not None, (
                             "New document ID should not be None"
                         )

haiku/rag/config/__init__.py

@@ -0,0 +1,54 @@
+import os
+
+from haiku.rag.config.loader import (
+    check_for_deprecated_env,
+    find_config_file,
+    generate_default_config,
+    load_config_from_env,
+    load_yaml_config,
+)
+from haiku.rag.config.models import (
+    A2AConfig,
+    AppConfig,
+    EmbeddingsConfig,
+    LanceDBConfig,
+    OllamaConfig,
+    ProcessingConfig,
+    ProvidersConfig,
+    QAConfig,
+    RerankingConfig,
+    ResearchConfig,
+    StorageConfig,
+    VLLMConfig,
+)
+
+__all__ = [
+    "Config",
+    "AppConfig",
+    "StorageConfig",
+    "LanceDBConfig",
+    "EmbeddingsConfig",
+    "RerankingConfig",
+    "QAConfig",
+    "ResearchConfig",
+    "ProcessingConfig",
+    "OllamaConfig",
+    "VLLMConfig",
+    "ProvidersConfig",
+    "A2AConfig",
+    "find_config_file",
+    "load_yaml_config",
+    "generate_default_config",
+    "load_config_from_env",
+]
+
+# Load config from YAML file or use defaults
+config_path = find_config_file(None)
+if config_path:
+    yaml_data = load_yaml_config(config_path)
+    Config = AppConfig.model_validate(yaml_data)
+else:
+    Config = AppConfig()
+
+# Check for deprecated .env file
+check_for_deprecated_env()

haiku/rag/config/loader.py

@@ -0,0 +1,151 @@
+import os
+import warnings
+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. ~/.config/haiku.rag/config.yaml (user config)
+
+    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
+
+    user_config_dir = Path.home() / ".config" / "haiku.rag"
+    user_config = user_config_dir / "config.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 check_for_deprecated_env() -> None:
+    """Check for .env file and warn if found."""
+    env_file = Path.cwd() / ".env"
+    if env_file.exists():
+        warnings.warn(
+            ".env file detected but YAML configuration is now preferred. "
+            "Environment variable configuration is deprecated and will be removed in future versions."
+            "Run 'haiku-rag init-config' to generate a YAML config file.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+
+def generate_default_config() -> dict:
+    """Generate a default YAML config structure with documentation."""
+    return {
+        "environment": "production",
+        "storage": {
+            "data_dir": "",
+            "monitor_directories": [],
+            "disable_autocreate": False,
+            "vacuum_retention_seconds": 60,
+        },
+        "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": "",
+            },
+        },
+        "a2a": {"max_contexts": 1000},
+    }
+
+
+def load_config_from_env() -> dict:
+    """Load current config from environment variables (for migration)."""
+    result = {}
+
+    env_mappings = {
+        "ENV": "environment",
+        "DEFAULT_DATA_DIR": ("storage", "data_dir"),
+        "MONITOR_DIRECTORIES": ("storage", "monitor_directories"),
+        "DISABLE_DB_AUTOCREATE": ("storage", "disable_autocreate"),
+        "VACUUM_RETENTION_SECONDS": ("storage", "vacuum_retention_seconds"),
+        "LANCEDB_URI": ("lancedb", "uri"),
+        "LANCEDB_API_KEY": ("lancedb", "api_key"),
+        "LANCEDB_REGION": ("lancedb", "region"),
+        "EMBEDDINGS_PROVIDER": ("embeddings", "provider"),
+        "EMBEDDINGS_MODEL": ("embeddings", "model"),
+        "EMBEDDINGS_VECTOR_DIM": ("embeddings", "vector_dim"),
+        "RERANK_PROVIDER": ("reranking", "provider"),
+        "RERANK_MODEL": ("reranking", "model"),
+        "QA_PROVIDER": ("qa", "provider"),
+        "QA_MODEL": ("qa", "model"),
+        "RESEARCH_PROVIDER": ("research", "provider"),
+        "RESEARCH_MODEL": ("research", "model"),
+        "CHUNK_SIZE": ("processing", "chunk_size"),
+        "CONTEXT_CHUNK_RADIUS": ("processing", "context_chunk_radius"),
+        "MARKDOWN_PREPROCESSOR": ("processing", "markdown_preprocessor"),
+        "OLLAMA_BASE_URL": ("providers", "ollama", "base_url"),
+        "VLLM_EMBEDDINGS_BASE_URL": ("providers", "vllm", "embeddings_base_url"),
+        "VLLM_RERANK_BASE_URL": ("providers", "vllm", "rerank_base_url"),
+        "VLLM_QA_BASE_URL": ("providers", "vllm", "qa_base_url"),
+        "VLLM_RESEARCH_BASE_URL": ("providers", "vllm", "research_base_url"),
+        "A2A_MAX_CONTEXTS": ("a2a", "max_contexts"),
+    }
+
+    for env_var, path in env_mappings.items():
+        value = os.getenv(env_var)
+        if value is not None:
+            # Special handling for MONITOR_DIRECTORIES - parse comma-separated list
+            if env_var == "MONITOR_DIRECTORIES":
+                if value.strip():
+                    value = [p.strip() for p in value.split(",") if p.strip()]
+                else:
+                    value = []
+
+            if isinstance(path, tuple):
+                current = result
+                for key in path[:-1]:
+                    if key not in current:
+                        current[key] = {}
+                    current = current[key]
+                current[path[-1]] = value
+            else:
+                result[path] = value
+
+    return result

haiku/rag/config/models.py

@@ -0,0 +1,78 @@
+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)
+    monitor_directories: list[Path] = []
+    disable_autocreate: bool = False
+    vacuum_retention_seconds: int = 60
+
+
+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"
+
+
+class ResearchConfig(BaseModel):
+    provider: str = "ollama"
+    model: str = "gpt-oss"
+
+
+class ProcessingConfig(BaseModel):
+    chunk_size: int = 256
+    context_chunk_radius: int = 0
+    markdown_preprocessor: str = ""
+
+
+class OllamaConfig(BaseModel):
+    base_url: str = "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 A2AConfig(BaseModel):
+    max_contexts: int = 1000
+
+
+class AppConfig(BaseModel):
+    environment: str = "production"
+    storage: StorageConfig = Field(default_factory=StorageConfig)
+    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)
+    a2a: A2AConfig = Field(default_factory=A2AConfig)

haiku/rag/embeddings/__init__.py

@@ -1,17 +1,23 @@
-from haiku.rag.config import Config
+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() -> EmbedderBase:
+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)
+    if config.embeddings.provider == "ollama":
+        return OllamaEmbedder(config.embeddings.model, config.embeddings.vector_dim)
 
-    if Config.EMBEDDINGS_PROVIDER == "voyageai":
+    if config.embeddings.provider == "voyageai":
         try:
             from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
         except ImportError:
@@ -20,16 +26,16 @@ def get_embedder() -> EmbedderBase:
                 "Please install haiku.rag with the 'voyageai' extra: "
                 "uv pip install haiku.rag[voyageai]"
             )
-        return VoyageAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
+        return VoyageAIEmbedder(config.embeddings.model, config.embeddings.vector_dim)
 
-    if Config.EMBEDDINGS_PROVIDER == "openai":
+    if config.embeddings.provider == "openai":
         from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
 
-        return OpenAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
+        return OpenAIEmbedder(config.embeddings.model, config.embeddings.vector_dim)
 
-    if Config.EMBEDDINGS_PROVIDER == "vllm":
+    if config.embeddings.provider == "vllm":
         from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
 
-        return VllmEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
+        return VllmEmbedder(config.embeddings.model, config.embeddings.vector_dim)
 
-    raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDINGS_PROVIDER}")
+    raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")

haiku/rag/embeddings/base.py

@@ -1,14 +1,22 @@
+from typing import overload
+
 from haiku.rag.config import Config
 
 
 class EmbedderBase:
-    _model: str = Config.EMBEDDINGS_MODEL
-    _vector_dim: int = Config.EMBEDDINGS_VECTOR_DIM
+    _model: str = Config.embeddings.model
+    _vector_dim: int = Config.embeddings.vector_dim
 
     def __init__(self, model: str, vector_dim: int):
         self._model = model
         self._vector_dim = vector_dim
 
+    @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

@@ -1,3 +1,5 @@
+from typing import overload
+
 from openai import AsyncOpenAI
 
 from haiku.rag.config import Config
@@ -5,8 +7,16 @@ 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"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy")
+        client = AsyncOpenAI(
+            base_url=f"{Config.providers.ollama.base_url}/v1", api_key="dummy"
+        )
         if not text:
             return []
         response = await client.embeddings.create(

haiku/rag/embeddings/openai.py

@@ -1,9 +1,17 @@
+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:

haiku/rag/embeddings/vllm.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from openai import AsyncOpenAI
 
 from haiku.rag.config import Config
@@ -5,9 +7,15 @@ 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"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy"
+            base_url=f"{Config.providers.vllm.embeddings_base_url}/v1", api_key="dummy"
         )
         if not text:
             return []

haiku/rag/embeddings/voyageai.py

@@ -1,9 +1,17 @@
 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:

haiku/rag/graph/common.py

@@ -15,13 +15,13 @@ def get_model(provider: str, model: str) -> Any:
     if provider == "ollama":
         return OpenAIChatModel(
             model_name=model,
-            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+            provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
         )
     elif provider == "vllm":
         return OpenAIChatModel(
             model_name=model,
             provider=OpenAIProvider(
-                base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
+                base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1",
                 api_key="none",
             ),
         )