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

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

@@ -4,8 +4,8 @@ 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

haiku/rag/embeddings/ollama.py

@@ -14,7 +14,9 @@ class Embedder(EmbedderBase):
     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/vllm.py

@@ -15,7 +15,7 @@ class Embedder(EmbedderBase):
 
     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/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",
             ),
         )

haiku/rag/mcp.py

@@ -38,10 +38,13 @@ def create_mcp_server(db_path: Path) -> FastMCP:
         """Add a document to the RAG system from a file path."""
         try:
             async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document_from_source(
+                result = await rag.create_document_from_source(
                     Path(file_path), title=title, metadata=metadata or {}
                 )
-                return document.id
+                # Handle both single document and list of documents (directories)
+                if isinstance(result, list):
+                    return result[0].id if result else None
+                return result.id
         except Exception:
             return None
 
@@ -52,10 +55,13 @@ def create_mcp_server(db_path: Path) -> FastMCP:
         """Add a document to the RAG system from a URL."""
         try:
             async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document_from_source(
+                result = await rag.create_document_from_source(
                     url, title=title, metadata=metadata or {}
                 )
-                return document.id
+                # Handle both single document and list of documents
+                if isinstance(result, list):
+                    return result[0].id if result else None
+                return result.id
         except Exception:
             return None
 
@@ -188,8 +194,8 @@ def create_mcp_server(db_path: Path) -> FastMCP:
                     deps = DeepQADeps(client=rag)
 
                     start_node = DeepQAPlanNode(
-                        provider=Config.QA_PROVIDER,
-                        model=Config.QA_MODEL,
+                        provider=Config.qa.provider,
+                        model=Config.qa.model,
                     )
 
                     result = await graph.run(
@@ -241,8 +247,8 @@ def create_mcp_server(db_path: Path) -> FastMCP:
 
                 result = await graph.run(
                     PlanNode(
-                        provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
-                        model=Config.RESEARCH_MODEL or Config.QA_MODEL,
+                        provider=Config.research.provider or Config.qa.provider,
+                        model=Config.research.model or Config.qa.model,
                     ),
                     state=state,
                     deps=deps,

haiku/rag/monitor.py

@@ -1,21 +1,27 @@
 import logging
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from watchfiles import Change, DefaultFilter, awatch
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.reader import FileReader
 from haiku.rag.store.models.document import Document
 
+if TYPE_CHECKING:
+    pass
+
 logger = logging.getLogger(__name__)
 
 
 class FileFilter(DefaultFilter):
     def __init__(self, *, ignore_paths: list[Path] | None = None) -> None:
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         self.extensions = tuple(FileReader.extensions)
         super().__init__(ignore_paths=ignore_paths)
 
-    def __call__(self, change: "Change", path: str) -> bool:
+    def __call__(self, change: Change, path: str) -> bool:
         return path.endswith(self.extensions) and super().__call__(change, path)
 
 
@@ -40,6 +46,9 @@ class FileWatcher:
                 await self._delete_document(Path(path))
 
     async def refresh(self):
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         for path in self.paths:
             for f in Path(path).rglob("**/*"):
                 if f.is_file() and f.suffix in FileReader.extensions:
@@ -50,11 +59,15 @@ class FileWatcher:
             uri = file.as_uri()
             existing_doc = await self.client.get_document_by_uri(uri)
             if existing_doc:
-                doc = await self.client.create_document_from_source(str(file))
+                result = await self.client.create_document_from_source(str(file))
+                # Since we're passing a file (not directory), result should be a single Document
+                doc = result if isinstance(result, Document) else result[0]
                 logger.info(f"Updated document {existing_doc.id} from {file}")
                 return doc
             else:
-                doc = await self.client.create_document_from_source(str(file))
+                result = await self.client.create_document_from_source(str(file))
+                # Since we're passing a file (not directory), result should be a single Document
+                doc = result if isinstance(result, Document) else result[0]
                 logger.info(f"Created new document {doc.id} from {file}")
                 return doc
         except Exception as e:

haiku/rag/qa/__init__.py

@@ -1,15 +1,28 @@
 from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
+from haiku.rag.config import AppConfig, Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
 def get_qa_agent(
     client: HaikuRAG,
+    config: AppConfig = Config,
     use_citations: bool = False,
     system_prompt: str | None = None,
 ) -> QuestionAnswerAgent:
-    provider = Config.QA_PROVIDER
-    model_name = Config.QA_MODEL
+    """
+    Factory function to get a QA agent based on the configuration.
+
+    Args:
+        client: HaikuRAG client instance.
+        config: Configuration to use. Defaults to global Config.
+        use_citations: Whether to include citations in responses.
+        system_prompt: Optional custom system prompt.
+
+    Returns:
+        A configured QuestionAnswerAgent instance.
+    """
+    provider = config.qa.provider
+    model_name = config.qa.model
 
     return QuestionAnswerAgent(
         client=client,

haiku/rag/qa/agent.py

@@ -71,13 +71,15 @@ class QuestionAnswerAgent:
         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_QA_BASE_URL}/v1", api_key="none"
+                    base_url=f"{Config.providers.vllm.qa_base_url}/v1", api_key="none"
                 ),
             )
         else:

haiku/rag/reranking/__init__.py

@@ -1,37 +1,45 @@
 import os
 
-from haiku.rag.config import Config
+from haiku.rag.config import AppConfig, Config
 from haiku.rag.reranking.base import RerankerBase
 
-_reranker: RerankerBase | None = None
+_reranker_cache: dict[int, RerankerBase | None] = {}
 
 
-def get_reranker() -> RerankerBase | None:
+def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
     """
     Factory function to get the appropriate reranker based on the configuration.
-    Returns None if if reranking is disabled.
+    Returns None if reranking is disabled.
+
+    Args:
+        config: Configuration to use. Defaults to global Config.
+
+    Returns:
+        A reranker instance if configured, None otherwise.
     """
-    global _reranker
-    if _reranker is not None:
-        return _reranker
+    # Use config id as cache key to support multiple configs
+    config_id = id(config)
+    if config_id in _reranker_cache:
+        return _reranker_cache[config_id]
+
+    reranker: RerankerBase | None = None
 
-    if Config.RERANK_PROVIDER == "mxbai":
+    if config.reranking.provider == "mxbai":
         try:
             from haiku.rag.reranking.mxbai import MxBAIReranker
 
             os.environ["TOKENIZERS_PARALLELISM"] = "true"
-            _reranker = MxBAIReranker()
-            return _reranker
+            reranker = MxBAIReranker()
         except ImportError:
-            return None
+            reranker = None
 
-    if Config.RERANK_PROVIDER == "cohere":
+    elif config.reranking.provider == "cohere":
         try:
             from haiku.rag.reranking.cohere import CohereReranker
 
-            _reranker = CohereReranker()
-            return _reranker
+            reranker = CohereReranker()
         except ImportError:
-            return None
+            reranker = None
 
-    return None
+    _reranker_cache[config_id] = reranker
+    return reranker

haiku/rag/reranking/base.py

@@ -3,7 +3,7 @@ from haiku.rag.store.models.chunk import Chunk
 
 
 class RerankerBase:
-    _model: str = Config.RERANK_MODEL
+    _model: str = Config.reranking.model
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/reranking/cohere.py

@@ -1,4 +1,3 @@
-from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
@@ -12,7 +11,8 @@ except ImportError as e:
 
 class CohereReranker(RerankerBase):
     def __init__(self):
-        self._client = cohere.ClientV2(api_key=Config.COHERE_API_KEY)
+        # Cohere SDK reads CO_API_KEY from environment by default
+        self._client = cohere.ClientV2()
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/reranking/mxbai.py

@@ -8,7 +8,7 @@ from haiku.rag.store.models.chunk import Chunk
 class MxBAIReranker(RerankerBase):
     def __init__(self):
         self._client = MxbaiRerankV2(
-            Config.RERANK_MODEL, disable_transformers_warnings=True
+            Config.reranking.model, disable_transformers_warnings=True
         )
 
     async def rerank(

haiku/rag/reranking/vllm.py

@@ -8,7 +8,7 @@ from haiku.rag.store.models.chunk import Chunk
 class VLLMReranker(RerankerBase):
     def __init__(self, model: str):
         self._model = model
-        self._base_url = Config.VLLM_RERANK_BASE_URL
+        self._base_url = Config.providers.vllm.rerank_base_url
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/store/engine.py

@@ -10,7 +10,7 @@ import lancedb
 from lancedb.pydantic import LanceModel, Vector
 from pydantic import Field
 
-from haiku.rag.config import Config
+from haiku.rag.config import AppConfig, Config
 from haiku.rag.embeddings import get_embedder
 
 logger = logging.getLogger(__name__)
@@ -49,9 +49,12 @@ class SettingsRecord(LanceModel):
 
 
 class Store:
-    def __init__(self, db_path: Path, skip_validation: bool = False):
+    def __init__(
+        self, db_path: Path, config: AppConfig = Config, skip_validation: bool = False
+    ):
         self.db_path: Path = db_path
-        self.embedder = get_embedder()
+        self._config = config
+        self.embedder = get_embedder(config=self._config)
         self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
@@ -59,7 +62,7 @@ class Store:
 
         # Local filesystem handling for DB directory
         if not self._has_cloud_config():
-            if Config.DISABLE_DB_AUTOCREATE:
+            if self._config.storage.disable_autocreate:
                 # LanceDB uses a directory path for local databases; enforce presence
                 if not db_path.exists():
                     raise FileNotFoundError(
@@ -85,13 +88,15 @@ class Store:
 
         Args:
             retention_seconds: Retention threshold in seconds. Only versions older
-                              than this will be removed. If None, uses Config.VACUUM_RETENTION_SECONDS.
+                              than this will be removed. If None, uses config.storage.vacuum_retention_seconds.
 
         Note:
             If vacuum is already running, this method returns immediately without blocking.
             Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
         """
-        if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
+        if self._has_cloud_config() and str(self._config.lancedb.uri).startswith(
+            "db://"
+        ):
             return
 
         # Skip if already running (non-blocking)
@@ -102,7 +107,7 @@ class Store:
             try:
                 # Evaluate config at runtime to allow dynamic changes
                 if retention_seconds is None:
-                    retention_seconds = Config.VACUUM_RETENTION_SECONDS
+                    retention_seconds = self._config.storage.vacuum_retention_seconds
                 # Perform maintenance per table using optimize() with configurable retention
                 retention = timedelta(seconds=retention_seconds)
                 for table in [
@@ -120,9 +125,9 @@ class Store:
         # Check if we have cloud configuration
         if self._has_cloud_config():
             return lancedb.connect(
-                uri=Config.LANCEDB_URI,
-                api_key=Config.LANCEDB_API_KEY,
-                region=Config.LANCEDB_REGION,
+                uri=self._config.lancedb.uri,
+                api_key=self._config.lancedb.api_key,
+                region=self._config.lancedb.region,
             )
         else:
             # Local file system connection
@@ -131,7 +136,9 @@ class Store:
     def _has_cloud_config(self) -> bool:
         """Check if cloud configuration is complete."""
         return bool(
-            Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION
+            self._config.lancedb.uri
+            and self._config.lancedb.api_key
+            and self._config.lancedb.region
         )
 
     def _validate_configuration(self) -> None:
@@ -173,7 +180,7 @@ class Store:
                 "settings", schema=SettingsRecord
             )
             # Save current settings to the new database
-            settings_data = Config.model_dump(mode="json")
+            settings_data = self._config.model_dump(mode="json")
             self.settings_table.add(
                 [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
             )

haiku/rag/store/repositories/chunk.py

@@ -1,17 +1,17 @@
 import inspect
 import json
 import logging
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
-from docling_core.types.doc.document import DoclingDocument
 from lancedb.rerankers import RRFReranker
 
-from haiku.rag.chunker import chunker
-from haiku.rag.config import Config
-from haiku.rag.embeddings import get_embedder
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.chunk import Chunk
-from haiku.rag.utils import load_callable, text_to_docling_document
+from haiku.rag.utils import load_callable
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +21,7 @@ class ChunkRepository:
 
     def __init__(self, store: Store) -> None:
         self.store = store
-        self.embedder = get_embedder()
+        self.embedder = store.embedder
 
     def _ensure_fts_index(self) -> None:
         """Ensure FTS index exists on the content column."""
@@ -142,12 +142,16 @@ class ChunkRepository:
         return chunks
 
     async def create_chunks_for_document(
-        self, document_id: str, document: DoclingDocument
+        self, document_id: str, document: "DoclingDocument"
     ) -> list[Chunk]:
         """Create chunks and embeddings for a document from DoclingDocument."""
+        # Lazy imports to avoid loading docling during module import
+        from haiku.rag.chunker import chunker
+        from haiku.rag.utils import text_to_docling_document
+
         # Optionally preprocess markdown before chunking
         processed_document = document
-        preprocessor_path = Config.MARKDOWN_PREPROCESSOR
+        preprocessor_path = self.store._config.processing.markdown_preprocessor
         if preprocessor_path:
             try:
                 pre_fn = load_callable(preprocessor_path)

haiku/rag/store/repositories/document.py

@@ -4,12 +4,12 @@ from datetime import datetime
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
-from docling_core.types.doc.document import DoclingDocument
-
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.document import Document
 
 if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
     from haiku.rag.store.models.chunk import Chunk
 
 
@@ -171,7 +171,7 @@ class DocumentRepository:
     async def _create_with_docling(
         self,
         entity: Document,
-        docling_document: DoclingDocument,
+        docling_document: "DoclingDocument",
         chunks: list["Chunk"] | None = None,
     ) -> Document:
         """Create a document with its chunks and embeddings."""
@@ -211,7 +211,7 @@ class DocumentRepository:
             raise
 
     async def _update_with_docling(
-        self, entity: Document, docling_document: DoclingDocument
+        self, entity: Document, docling_document: "DoclingDocument"
     ) -> Document:
         """Update a document and regenerate its chunks."""
         assert entity.id is not None, "Document ID is required for update"