diff --git a/docs/configuration.md b/docs/configuration.md index 9e522a44..dadf6015 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -104,10 +104,10 @@ a2a: ## Programmatic Configuration -When using haiku.rag as a Python library, you can configure it programmatically using `set_config()`: +When using haiku.rag as a Python library, you can pass configuration directly to the `HaikuRAG` client: ```python -from haiku.rag.config import set_config, AppConfig +from haiku.rag.config import AppConfig from haiku.rag.client import HaikuRAG # Create custom configuration @@ -117,18 +117,17 @@ custom_config = AppConfig( processing={"chunk_size": 512} ) -# Set the configuration globally -set_config(custom_config) - -# All subsequent operations use this configuration -client = HaikuRAG(db_path) +# Pass configuration to the client +client = HaikuRAG(config=custom_config) ``` +If you don't pass a config, the client uses the global configuration loaded from your YAML file or defaults. + This is useful for: - Jupyter notebooks - Python scripts - Testing with different configurations -- Applications that need runtime configuration +- Applications that need multiple clients with different configurations ## API Keys diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 40b81f0f..b9d4a1b5 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -8,7 +8,7 @@ from urllib.parse import urlparse import httpx -from haiku.rag.config import Config +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 @@ -25,16 +25,23 @@ class HaikuRAG: def __init__( self, - db_path: Path = Config.storage.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) @@ -430,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 @@ -452,18 +459,20 @@ class HaikuRAG: async def expand_context( self, search_results: list[tuple[Chunk, float]], - radius: int = Config.processing.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 @@ -593,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]: diff --git a/src/haiku/rag/config/__init__.py b/src/haiku/rag/config/__init__.py index b5717060..e0299a6f 100644 --- a/src/haiku/rag/config/__init__.py +++ b/src/haiku/rag/config/__init__.py @@ -40,53 +40,15 @@ __all__ = [ "load_yaml_config", "generate_default_config", "load_config_from_env", - "set_config", ] - -class ConfigProxy: - """Proxy for the global configuration that allows runtime updates.""" - - def __init__(self): - # Load config from YAML file or use defaults - config_path = find_config_file(None) - if config_path: - yaml_data = load_yaml_config(config_path) - self._config = AppConfig.model_validate(yaml_data) - else: - self._config = AppConfig() - - def __getattr__(self, name): - """Proxy attribute access to the underlying config.""" - return getattr(self._config, name) - - def set(self, config: AppConfig) -> None: - """Replace the current configuration.""" - self._config = config - - -# Create the global Config instance -Config = ConfigProxy() +# 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() - - -def set_config(config: AppConfig) -> None: - """Set the global configuration programmatically. - - This allows library users to configure haiku.rag without needing - a YAML file or environment variables. - - Args: - config: The AppConfig instance to use globally. - - Example: - >>> from haiku.rag.config import set_config, AppConfig - >>> custom_config = AppConfig( - ... qa={"provider": "openai", "model": "gpt-4o"}, - ... embeddings={"provider": "voyage", "model": "voyage-3"} - ... ) - >>> set_config(custom_config) - """ - Config.set(config) diff --git a/src/haiku/rag/embeddings/__init__.py b/src/haiku/rag/embeddings/__init__.py index 97c20341..17e85a3e 100644 --- a/src/haiku/rag/embeddings/__init__.py +++ b/src/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}") diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index 91adf8f5..4e8e7f92 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/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, diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index 09635430..797aa0f7 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/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. - """ - global _reranker - if _reranker is not None: - return _reranker + Returns None if reranking is disabled. - if Config.reranking.provider == "mxbai": + Args: + config: Configuration to use. Defaults to global Config. + + Returns: + A reranker instance if configured, None otherwise. + """ + # 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.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.reranking.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 diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index e3d56e17..6cbe328e 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/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.storage.disable_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.storage.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.storage.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))] ) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index cfe9cbe3..b92c96de 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -6,8 +6,6 @@ from uuid import uuid4 from lancedb.rerankers import RRFReranker -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 @@ -23,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.""" @@ -153,7 +151,7 @@ class ChunkRepository: # Optionally preprocess markdown before chunking processed_document = document - preprocessor_path = Config.processing.markdown_preprocessor + preprocessor_path = self.store._config.processing.markdown_preprocessor if preprocessor_path: try: pre_fn = load_callable(preprocessor_path) diff --git a/src/haiku/rag/store/repositories/settings.py b/src/haiku/rag/store/repositories/settings.py index e23a7122..4d5bba43 100644 --- a/src/haiku/rag/store/repositories/settings.py +++ b/src/haiku/rag/store/repositories/settings.py @@ -1,6 +1,5 @@ import json -from haiku.rag.config import Config from haiku.rag.store.engine import SettingsRecord, Store @@ -73,7 +72,7 @@ class SettingsRepository: def save_current_settings(self) -> None: """Save the current configuration to the database.""" - current_config = Config.model_dump(mode="json") + current_config = self.store._config.model_dump(mode="json") # Check if settings exist existing = list( @@ -116,7 +115,7 @@ class SettingsRepository: self.save_current_settings() return - current_config = Config.model_dump(mode="json") + current_config = self.store._config.model_dump(mode="json") # Check if embedding provider or model has changed # Support both old flat structure and new nested structure for backward compatibility