diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index 418307e5..a4607a0b 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -38,13 +38,13 @@ def _as_state_snapshot(ctx: RunContext[ResearchDeps]) -> StateSnapshotEvent: def create_agent( - qa_provider: str = Config.QA_PROVIDER, qa_model: str = Config.QA_MODEL + qa_provider: str = Config.qa.provider, qa_model: str = Config.qa.model ) -> Agent[ResearchDeps, str]: """Create and configure the research agent. Args: - qa_provider: QA provider for the agent (default: from Config.QA_PROVIDER) - qa_model: Model name to use (default: from Config.QA_MODEL) + qa_provider: QA provider for the agent (default: from Config.qa.provider) + qa_model: Model name to use (default: from Config.qa.model) """ print(f"[AGENT SETUP] Creating agent with provider={qa_provider}, model={qa_model}") agent = Agent( diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 1c696445..41b39945 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -33,7 +33,7 @@ async def lifespan(app): logger.info(f"Initializing HaikuRAG client with database: {db_path}") client = HaikuRAG(db_path) logger.info("Research assistant backend ready") - logger.info(f"QA Provider: {Config.QA_PROVIDER}, Model: {Config.QA_MODEL}") + logger.info(f"QA Provider: {Config.qa.provider}, Model: {Config.qa.model}") yield @@ -51,9 +51,9 @@ async def health(request): { "status": "healthy", "agent_model": str(agent.model), - "qa_provider": Config.QA_PROVIDER, - "qa_model": Config.QA_MODEL, - "ollama_base_url": Config.OLLAMA_BASE_URL, + "qa_provider": Config.qa.provider, + "qa_model": Config.qa.model, + "ollama_base_url": Config.providers.ollama.base_url, "db_path": db_path_str, "db_exists": Path(db_path_str).exists(), } @@ -100,8 +100,8 @@ if __name__ == "__main__": print("Starting haiku.rag research assistant backend...") print(f"Agent model: {agent.model}") - print(f"QA provider: {Config.QA_PROVIDER}") - print(f"QA model: {Config.QA_MODEL}") + print(f"QA provider: {Config.qa.provider}") + print(f"QA model: {Config.qa.model}") uvicorn.run( "main:app", diff --git a/src/evaluations/benchmark.py b/src/evaluations/benchmark.py index 793f9abf..954bdabe 100644 --- a/src/evaluations/benchmark.py +++ b/src/evaluations/benchmark.py @@ -174,7 +174,7 @@ async def run_qa_benchmark( judge_model = OpenAIChatModel( model_name=QA_JUDGE_MODEL, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), + provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"), ) evaluation_dataset = EvalDataset[str, str, dict[str, str]]( diff --git a/src/evaluations/llm_judge.py b/src/evaluations/llm_judge.py index 7dcfccac..9b2cc8aa 100644 --- a/src/evaluations/llm_judge.py +++ b/src/evaluations/llm_judge.py @@ -41,7 +41,7 @@ class LLMJudge: # Create Ollama model ollama_model = 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"), ) # Create Pydantic AI agent diff --git a/src/haiku/rag/a2a/__init__.py b/src/haiku/rag/a2a/__init__.py index f4481fbc..892056eb 100644 --- a/src/haiku/rag/a2a/__init__.py +++ b/src/haiku/rag/a2a/__init__.py @@ -57,12 +57,12 @@ def create_a2a_app( """ base_storage = InMemoryStorage() storage = LRUMemoryStorage( - storage=base_storage, max_contexts=Config.A2A_MAX_CONTEXTS + storage=base_storage, max_contexts=Config.a2a.max_contexts ) broker = InMemoryBroker() # Create the agent with native search tool - model = get_model(Config.QA_PROVIDER, Config.QA_MODEL) + model = get_model(Config.qa.provider, Config.qa.model) agent = Agent( model=model, deps_type=AgentDependencies, @@ -120,7 +120,7 @@ def create_a2a_app( # Create FastA2A app with custom worker lifecycle @asynccontextmanager async def lifespan(app): - logger.info(f"Started A2A server (max contexts: {Config.A2A_MAX_CONTEXTS})") + logger.info(f"Started A2A server (max contexts: {Config.a2a.max_contexts})") async with app.task_manager: async with worker.run(): yield diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index f701eeec..ac789d13 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -231,8 +231,8 @@ class HaikuRAGApp: ) start_node = DeepQAPlanNode( - provider=Config.QA_PROVIDER, - model=Config.QA_MODEL, + provider=Config.qa.provider, + model=Config.qa.model, ) result = await graph.run( @@ -278,8 +278,8 @@ class HaikuRAGApp: ) start = 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, ) report = None async for event in stream_research_graph(graph, start, state, deps): @@ -474,7 +474,9 @@ class HaikuRAGApp: # Start file monitor if enabled if enable_monitor: - monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES, client=client) + monitor = FileWatcher( + paths=Config.storage.monitor_directories, client=client + ) monitor_task = asyncio.create_task(monitor.observe()) tasks.append(monitor_task) diff --git a/src/haiku/rag/chunker.py b/src/haiku/rag/chunker.py index 31cb2717..f478d64c 100644 --- a/src/haiku/rag/chunker.py +++ b/src/haiku/rag/chunker.py @@ -22,7 +22,7 @@ class Chunker: def __init__( self, - chunk_size: int = Config.CHUNK_SIZE, + chunk_size: int = Config.processing.chunk_size, ): self.chunk_size = chunk_size tokenizer = OpenAITokenizer( diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 078a2178..1d778487 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -56,7 +56,7 @@ def main( os.environ["HAIKU_RAG_CONFIG_PATH"] = str(config.absolute()) # Configure logging minimally for CLI context - if Config.ENV == "development": + if Config.environment == "development": # Lazy import logfire only in development try: import logfire # type: ignore @@ -80,7 +80,7 @@ def main( @cli.command("list", help="List all stored documents") def list_documents( db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -127,7 +127,7 @@ def add_document_text( metavar="KEY=VALUE", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -156,7 +156,7 @@ def add_document_src( metavar="KEY=VALUE", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -178,7 +178,7 @@ def get_document( help="The ID of the document to get", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -195,7 +195,7 @@ def delete_document( help="The ID of the document to delete", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -222,7 +222,7 @@ def search( help="Maximum number of results to return", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -239,7 +239,7 @@ def ask( help="The question to ask", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -287,7 +287,7 @@ def research( help="Max concurrent searches per iteration (planned)", ), db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -334,7 +334,7 @@ def init_config( """Generate a YAML configuration file with defaults or from .env.""" import yaml - from haiku.rag.config_loader import generate_default_config, load_config_from_env + from haiku.rag.config.loader import generate_default_config, load_config_from_env if output.exists(): typer.echo( @@ -373,7 +373,7 @@ def init_config( ) def rebuild( db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -387,7 +387,7 @@ def rebuild( @cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage") def vacuum( db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -401,7 +401,7 @@ def vacuum( @cli.command("info", help="Show read-only database info (no upgrades or writes)") def info( db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), @@ -430,7 +430,7 @@ def download_models_cmd(): ) def serve( db: Path = typer.Option( - Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + Config.storage.data_dir / "haiku.rag.lancedb", "--db", help="Path to the LanceDB database file", ), diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 8c6c5ba8..40b81f0f 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -25,7 +25,7 @@ class HaikuRAG: def __init__( self, - db_path: Path = Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", + db_path: Path = Config.storage.data_dir / "haiku.rag.lancedb", skip_validation: bool = False, ): """Initialize the RAG client with a database path. @@ -452,7 +452,7 @@ class HaikuRAG: async def expand_context( self, search_results: list[tuple[Chunk, float]], - radius: int = Config.CONTEXT_CHUNK_RADIUS, + radius: int = Config.processing.context_chunk_radius, ) -> list[tuple[Chunk, float]]: """Expand search results with adjacent chunks, merging overlapping chunks. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py deleted file mode 100644 index 2d25523b..00000000 --- a/src/haiku/rag/config.py +++ /dev/null @@ -1,104 +0,0 @@ -import os -from pathlib import Path - -from pydantic import BaseModel, field_validator - -from haiku.rag.config_loader import ( - check_for_deprecated_env, - find_config_file, - flatten_yaml_to_env_dict, - load_yaml_config, -) -from haiku.rag.utils import get_default_data_dir - - -class AppConfig(BaseModel): - ENV: str = "production" - - LANCEDB_API_KEY: str = "" - LANCEDB_URI: str = "" - LANCEDB_REGION: str = "" - - DEFAULT_DATA_DIR: Path = get_default_data_dir() - MONITOR_DIRECTORIES: list[Path] = [] - - EMBEDDINGS_PROVIDER: str = "ollama" - EMBEDDINGS_MODEL: str = "qwen3-embedding" - EMBEDDINGS_VECTOR_DIM: int = 4096 - - RERANK_PROVIDER: str = "" - RERANK_MODEL: str = "" - - QA_PROVIDER: str = "ollama" - QA_MODEL: str = "gpt-oss" - - # Research defaults (fallback to QA if not provided via env) - RESEARCH_PROVIDER: str = "ollama" - RESEARCH_MODEL: str = "gpt-oss" - - CHUNK_SIZE: int = 256 - CONTEXT_CHUNK_RADIUS: int = 0 - - # Optional dotted path or file path to a callable that preprocesses - # markdown content before chunking. Examples: - MARKDOWN_PREPROCESSOR: str = "" - - OLLAMA_BASE_URL: str = "http://localhost:11434" - - VLLM_EMBEDDINGS_BASE_URL: str = "" - VLLM_RERANK_BASE_URL: str = "" - VLLM_QA_BASE_URL: str = "" - VLLM_RESEARCH_BASE_URL: str = "" - - # Provider keys - VOYAGE_API_KEY: str = "" - OPENAI_API_KEY: str = "" - ANTHROPIC_API_KEY: str = "" - COHERE_API_KEY: str = "" - - # If true, refuse to auto-create a new LanceDB database or tables - # and error out when the database does not already exist. - DISABLE_DB_AUTOCREATE: bool = False - - # Vacuum retention threshold in seconds. Only versions older than this - # threshold will be removed during vacuum operations. Default is 60 seconds - # to allow concurrent connections to safely use recent versions. - VACUUM_RETENTION_SECONDS: int = 60 - - # Maximum number of A2A contexts to keep in memory. When exceeded, least - # recently used contexts will be evicted. Default is 1000. - A2A_MAX_CONTEXTS: int = 1000 - - @field_validator("MONITOR_DIRECTORIES", mode="before") - @classmethod - def parse_monitor_directories(cls, v): - if isinstance(v, str): - if not v.strip(): - return [] - return [ - Path(path.strip()).absolute() for path in v.split(",") if path.strip() - ] - return v - - -# 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_dict = flatten_yaml_to_env_dict(yaml_data) -else: - config_dict = {} - -# Check for deprecated .env file -check_for_deprecated_env() - -# Expose Config object for app to import -Config = AppConfig.model_validate(config_dict) -if Config.OPENAI_API_KEY: - os.environ["OPENAI_API_KEY"] = Config.OPENAI_API_KEY -if Config.VOYAGE_API_KEY: - os.environ["VOYAGE_API_KEY"] = Config.VOYAGE_API_KEY -if Config.ANTHROPIC_API_KEY: - os.environ["ANTHROPIC_API_KEY"] = Config.ANTHROPIC_API_KEY -if Config.COHERE_API_KEY: - os.environ["CO_API_KEY"] = Config.COHERE_API_KEY diff --git a/src/haiku/rag/config/__init__.py b/src/haiku/rag/config/__init__.py new file mode 100644 index 00000000..7af9c6f2 --- /dev/null +++ b/src/haiku/rag/config/__init__.py @@ -0,0 +1,67 @@ +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, + APIKeysConfig, + AppConfig, + EmbeddingsConfig, + LanceDBConfig, + OllamaConfig, + ProcessingConfig, + ProvidersConfig, + QAConfig, + RerankingConfig, + ResearchConfig, + StorageConfig, + VLLMConfig, +) + +__all__ = [ + "Config", + "AppConfig", + "StorageConfig", + "LanceDBConfig", + "EmbeddingsConfig", + "RerankingConfig", + "QAConfig", + "ResearchConfig", + "ProcessingConfig", + "OllamaConfig", + "VLLMConfig", + "APIKeysConfig", + "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() + +# Export API keys to os.environ for provider libraries +if Config.providers.api_keys.openai: + os.environ["OPENAI_API_KEY"] = Config.providers.api_keys.openai +if Config.providers.api_keys.voyage: + os.environ["VOYAGE_API_KEY"] = Config.providers.api_keys.voyage +if Config.providers.api_keys.anthropic: + os.environ["ANTHROPIC_API_KEY"] = Config.providers.api_keys.anthropic +if Config.providers.api_keys.cohere: + # Cohere SDK expects CO_API_KEY (not COHERE_API_KEY) + os.environ["CO_API_KEY"] = Config.providers.api_keys.cohere diff --git a/src/haiku/rag/config_loader.py b/src/haiku/rag/config/loader.py similarity index 53% rename from src/haiku/rag/config_loader.py rename to src/haiku/rag/config/loader.py index b5a8e272..147bb542 100644 --- a/src/haiku/rag/config_loader.py +++ b/src/haiku/rag/config/loader.py @@ -45,134 +45,6 @@ def load_yaml_config(path: Path) -> dict: return data or {} -def flatten_yaml_to_env_dict(yaml_dict: dict) -> dict: - """Convert nested YAML structure to flat environment variable dict. - - Maps YAML structure like: - embeddings: - provider: ollama - model: qwen3 - - To flat dict like: - EMBEDDINGS_PROVIDER: ollama - EMBEDDINGS_MODEL: qwen3 - """ - result = {} - - # Top-level simple fields - if "environment" in yaml_dict: - result["ENV"] = yaml_dict["environment"] - - # Storage section - if "storage" in yaml_dict: - storage = yaml_dict["storage"] - if "data_dir" in storage: - result["DEFAULT_DATA_DIR"] = storage["data_dir"] - if "monitor_directories" in storage: - dirs = storage["monitor_directories"] - if isinstance(dirs, list): - result["MONITOR_DIRECTORIES"] = ",".join(str(d) for d in dirs) - else: - result["MONITOR_DIRECTORIES"] = str(dirs) - if "disable_autocreate" in storage: - result["DISABLE_DB_AUTOCREATE"] = storage["disable_autocreate"] - if "vacuum_retention_seconds" in storage: - result["VACUUM_RETENTION_SECONDS"] = storage["vacuum_retention_seconds"] - - # LanceDB section - if "lancedb" in yaml_dict: - lancedb = yaml_dict["lancedb"] - if "uri" in lancedb: - result["LANCEDB_URI"] = lancedb["uri"] - if "api_key" in lancedb: - result["LANCEDB_API_KEY"] = lancedb["api_key"] - if "region" in lancedb: - result["LANCEDB_REGION"] = lancedb["region"] - - # Embeddings section - if "embeddings" in yaml_dict: - embeddings = yaml_dict["embeddings"] - if "provider" in embeddings: - result["EMBEDDINGS_PROVIDER"] = embeddings["provider"] - if "model" in embeddings: - result["EMBEDDINGS_MODEL"] = embeddings["model"] - if "vector_dim" in embeddings: - result["EMBEDDINGS_VECTOR_DIM"] = embeddings["vector_dim"] - - # Reranking section - if "reranking" in yaml_dict: - reranking = yaml_dict["reranking"] - if "provider" in reranking: - result["RERANK_PROVIDER"] = reranking["provider"] - if "model" in reranking: - result["RERANK_MODEL"] = reranking["model"] - - # QA section - if "qa" in yaml_dict: - qa = yaml_dict["qa"] - if "provider" in qa: - result["QA_PROVIDER"] = qa["provider"] - if "model" in qa: - result["QA_MODEL"] = qa["model"] - - # Research section - if "research" in yaml_dict: - research = yaml_dict["research"] - if "provider" in research: - result["RESEARCH_PROVIDER"] = research["provider"] - if "model" in research: - result["RESEARCH_MODEL"] = research["model"] - - # Processing section - if "processing" in yaml_dict: - processing = yaml_dict["processing"] - if "chunk_size" in processing: - result["CHUNK_SIZE"] = processing["chunk_size"] - if "context_chunk_radius" in processing: - result["CONTEXT_CHUNK_RADIUS"] = processing["context_chunk_radius"] - if "markdown_preprocessor" in processing: - result["MARKDOWN_PREPROCESSOR"] = processing["markdown_preprocessor"] - - # Providers section - if "providers" in yaml_dict: - providers = yaml_dict["providers"] - - if "ollama" in providers: - ollama = providers["ollama"] - if "base_url" in ollama: - result["OLLAMA_BASE_URL"] = ollama["base_url"] - - if "vllm" in providers: - vllm = providers["vllm"] - if "embeddings_base_url" in vllm: - result["VLLM_EMBEDDINGS_BASE_URL"] = vllm["embeddings_base_url"] - if "rerank_base_url" in vllm: - result["VLLM_RERANK_BASE_URL"] = vllm["rerank_base_url"] - if "qa_base_url" in vllm: - result["VLLM_QA_BASE_URL"] = vllm["qa_base_url"] - if "research_base_url" in vllm: - result["VLLM_RESEARCH_BASE_URL"] = vllm["research_base_url"] - - if "api_keys" in providers: - api_keys = providers["api_keys"] - if "voyage" in api_keys: - result["VOYAGE_API_KEY"] = api_keys["voyage"] - if "openai" in api_keys: - result["OPENAI_API_KEY"] = api_keys["openai"] - if "anthropic" in api_keys: - result["ANTHROPIC_API_KEY"] = api_keys["anthropic"] - if "cohere" in api_keys: - result["COHERE_API_KEY"] = api_keys["cohere"] - - # A2A section - if "a2a" in yaml_dict: - a2a = yaml_dict["a2a"] - if "max_contexts" in a2a: - result["A2A_MAX_CONTEXTS"] = a2a["max_contexts"] - - return result - - def check_for_deprecated_env() -> None: """Check for .env file and warn if found.""" env_file = Path.cwd() / ".env" diff --git a/src/haiku/rag/config/models.py b/src/haiku/rag/config/models.py new file mode 100644 index 00000000..51bb55a5 --- /dev/null +++ b/src/haiku/rag/config/models.py @@ -0,0 +1,103 @@ +from pathlib import Path + +from pydantic import BaseModel, Field, field_validator + +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 APIKeysConfig(BaseModel): + voyage: str = "" + openai: str = "" + anthropic: str = "" + cohere: str = "" + + +class ProvidersConfig(BaseModel): + ollama: OllamaConfig = Field(default_factory=OllamaConfig) + vllm: VLLMConfig = Field(default_factory=VLLMConfig) + api_keys: APIKeysConfig = Field(default_factory=APIKeysConfig) + + +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) + + @field_validator("storage", mode="before") + @classmethod + def parse_storage(cls, v): + """Parse storage config, handling comma-separated monitor directories.""" + if isinstance(v, dict) and "monitor_directories" in v: + dirs = v["monitor_directories"] + if isinstance(dirs, str): + if not dirs.strip(): + v["monitor_directories"] = [] + else: + v["monitor_directories"] = [ + Path(path.strip()).absolute() + for path in dirs.split(",") + if path.strip() + ] + return v diff --git a/src/haiku/rag/embeddings/__init__.py b/src/haiku/rag/embeddings/__init__.py index 0964ea9e..97c20341 100644 --- a/src/haiku/rag/embeddings/__init__.py +++ b/src/haiku/rag/embeddings/__init__.py @@ -8,10 +8,10 @@ def get_embedder() -> EmbedderBase: Factory function to get the appropriate embedder based on the configuration. """ - 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 +20,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/embeddings/base.py b/src/haiku/rag/embeddings/base.py index 8a8ea4e2..fe04e811 100644 --- a/src/haiku/rag/embeddings/base.py +++ b/src/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 diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index ea565541..e9fbe8a5 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/src/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( diff --git a/src/haiku/rag/embeddings/vllm.py b/src/haiku/rag/embeddings/vllm.py index d7e2c608..4a55bf7d 100644 --- a/src/haiku/rag/embeddings/vllm.py +++ b/src/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 [] diff --git a/src/haiku/rag/graph/common.py b/src/haiku/rag/graph/common.py index 2b2b721a..31877206 100644 --- a/src/haiku/rag/graph/common.py +++ b/src/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", ), ) diff --git a/src/haiku/rag/mcp.py b/src/haiku/rag/mcp.py index 379086c8..7248d744 100644 --- a/src/haiku/rag/mcp.py +++ b/src/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, diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index 0b6e3dd3..91adf8f5 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -8,8 +8,8 @@ def get_qa_agent( use_citations: bool = False, system_prompt: str | None = None, ) -> QuestionAnswerAgent: - provider = Config.QA_PROVIDER - model_name = Config.QA_MODEL + provider = Config.qa.provider + model_name = Config.qa.model return QuestionAnswerAgent( client=client, diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index f1abe7ab..863296b8 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/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: diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index f4753d50..09635430 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -15,7 +15,7 @@ def get_reranker() -> RerankerBase | None: if _reranker is not None: return _reranker - if Config.RERANK_PROVIDER == "mxbai": + if Config.reranking.provider == "mxbai": try: from haiku.rag.reranking.mxbai import MxBAIReranker @@ -25,7 +25,7 @@ def get_reranker() -> RerankerBase | None: except ImportError: return None - if Config.RERANK_PROVIDER == "cohere": + if Config.reranking.provider == "cohere": try: from haiku.rag.reranking.cohere import CohereReranker diff --git a/src/haiku/rag/reranking/base.py b/src/haiku/rag/reranking/base.py index 0e95e26b..4f93d1a2 100644 --- a/src/haiku/rag/reranking/base.py +++ b/src/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 diff --git a/src/haiku/rag/reranking/cohere.py b/src/haiku/rag/reranking/cohere.py index 6d30952d..69f4f7d6 100644 --- a/src/haiku/rag/reranking/cohere.py +++ b/src/haiku/rag/reranking/cohere.py @@ -12,7 +12,7 @@ except ImportError as e: class CohereReranker(RerankerBase): def __init__(self): - self._client = cohere.ClientV2(api_key=Config.COHERE_API_KEY) + self._client = cohere.ClientV2(api_key=Config.providers.api_keys.cohere) async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 diff --git a/src/haiku/rag/reranking/mxbai.py b/src/haiku/rag/reranking/mxbai.py index df39ac4a..f4b7834c 100644 --- a/src/haiku/rag/reranking/mxbai.py +++ b/src/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( diff --git a/src/haiku/rag/reranking/vllm.py b/src/haiku/rag/reranking/vllm.py index c99ddbb4..52035589 100644 --- a/src/haiku/rag/reranking/vllm.py +++ b/src/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 diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 558f8c0d..e3d56e17 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -59,7 +59,7 @@ class Store: # Local filesystem handling for DB directory if not self._has_cloud_config(): - if Config.DISABLE_DB_AUTOCREATE: + if Config.storage.disable_autocreate: # LanceDB uses a directory path for local databases; enforce presence if not db_path.exists(): raise FileNotFoundError( @@ -85,13 +85,13 @@ 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(Config.lancedb.uri).startswith("db://"): return # Skip if already running (non-blocking) @@ -102,7 +102,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 = Config.storage.vacuum_retention_seconds # Perform maintenance per table using optimize() with configurable retention retention = timedelta(seconds=retention_seconds) for table in [ @@ -120,9 +120,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=Config.lancedb.uri, + api_key=Config.lancedb.api_key, + region=Config.lancedb.region, ) else: # Local file system connection @@ -131,7 +131,7 @@ 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 + Config.lancedb.uri and Config.lancedb.api_key and Config.lancedb.region ) def _validate_configuration(self) -> None: diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index bfd2e624..cfe9cbe3 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -153,7 +153,7 @@ class ChunkRepository: # Optionally preprocess markdown before chunking processed_document = document - preprocessor_path = Config.MARKDOWN_PREPROCESSOR + preprocessor_path = 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 e085ef61..e23a7122 100644 --- a/src/haiku/rag/store/repositories/settings.py +++ b/src/haiku/rag/store/repositories/settings.py @@ -119,14 +119,25 @@ class SettingsRepository: current_config = Config.model_dump(mode="json") # Check if embedding provider or model has changed - stored_provider = stored_settings.get("EMBEDDINGS_PROVIDER") - current_provider = current_config.get("EMBEDDINGS_PROVIDER") + # Support both old flat structure and new nested structure for backward compatibility + stored_embeddings = stored_settings.get("embeddings", {}) + current_embeddings = current_config.get("embeddings", {}) - stored_model = stored_settings.get("EMBEDDINGS_MODEL") - current_model = current_config.get("EMBEDDINGS_MODEL") + # Try nested structure first, fall back to flat for old databases + stored_provider = stored_embeddings.get("provider") or stored_settings.get( + "EMBEDDINGS_PROVIDER" + ) + current_provider = current_embeddings.get("provider") - stored_vector_dim = stored_settings.get("EMBEDDINGS_VECTOR_DIM") - current_vector_dim = current_config.get("EMBEDDINGS_VECTOR_DIM") + stored_model = stored_embeddings.get("model") or stored_settings.get( + "EMBEDDINGS_MODEL" + ) + current_model = current_embeddings.get("model") + + stored_vector_dim = stored_embeddings.get("vector_dim") or stored_settings.get( + "EMBEDDINGS_VECTOR_DIM" + ) + current_vector_dim = current_embeddings.get("vector_dim") # Check for incompatible changes incompatible_changes = [] diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py index 433c6f1e..228aa7c2 100644 --- a/src/haiku/rag/utils.py +++ b/src/haiku/rag/utils.py @@ -176,19 +176,19 @@ def prefetch_models(): # 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.RERANK_PROVIDER == "ollama": - required_models.add(Config.RERANK_MODEL) + 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.OLLAMA_BASE_URL + base_url = Config.providers.ollama.base_url with httpx.Client(timeout=None) as client: for model in sorted(required_models): diff --git a/tests/test_client.py b/tests/test_client.py index 25bd1249..a7889912 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -557,7 +557,7 @@ async def test_client_create_document_with_custom_chunks(temp_db_path): Chunk( content="This is the second chunk", metadata={"custom": "metadata2"}, - embedding=[0.1] * Config.EMBEDDINGS_VECTOR_DIM, + embedding=[0.1] * Config.embeddings.vector_dim, order=1, ), # With embedding Chunk( @@ -641,7 +641,7 @@ async def test_client_ask_with_cite(monkeypatch, temp_db_path): async def test_client_expand_context(temp_db_path): """Test expanding search results with adjacent chunks.""" # Mock Config to have CONTEXT_CHUNK_RADIUS = 2 - with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2): + with patch("haiku.rag.client.Config.processing.context_chunk_radius", 2): async with HaikuRAG(temp_db_path) as client: # Create chunks manually with precomputed embeddings to avoid network dim = client.chunk_repository.embedder._vector_dim @@ -710,7 +710,7 @@ async def test_client_expand_context_radius_zero(temp_db_path): @pytest.mark.asyncio async def test_client_expand_context_multiple_chunks(temp_db_path): """Test expand_context with multiple search results.""" - with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): + with patch("haiku.rag.client.Config.processing.context_chunk_radius", 1): async with HaikuRAG(temp_db_path) as client: # Create first document with manual chunks doc1_chunks = [ diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 9edee0f8..a2b0869f 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -3,9 +3,8 @@ from pathlib import Path import pytest -from haiku.rag.config_loader import ( +from haiku.rag.config.loader import ( find_config_file, - flatten_yaml_to_env_dict, generate_default_config, load_config_from_env, load_yaml_config, @@ -30,54 +29,6 @@ embeddings: assert config["embeddings"]["vector_dim"] == 1024 -def test_flatten_yaml_to_env_dict(): - """Test converting nested YAML to flat env dict.""" - yaml_dict = { - "environment": "development", - "storage": { - "data_dir": "/tmp/data", - "monitor_directories": ["/path/one", "/path/two"], - "disable_autocreate": True, - "vacuum_retention_seconds": 30, - }, - "embeddings": { - "provider": "openai", - "model": "text-embedding-3-small", - "vector_dim": 1536, - }, - "qa": {"provider": "anthropic", "model": "claude-3-haiku"}, - "processing": {"chunk_size": 512, "context_chunk_radius": 1}, - "providers": { - "ollama": {"base_url": "http://localhost:11434"}, - "api_keys": {"openai": "test-key", "anthropic": "test-key-2"}, - }, - } - - result = flatten_yaml_to_env_dict(yaml_dict) - - assert result["ENV"] == "development" - assert result["DEFAULT_DATA_DIR"] == "/tmp/data" - assert result["MONITOR_DIRECTORIES"] == "/path/one,/path/two" - assert result["DISABLE_DB_AUTOCREATE"] is True - assert result["VACUUM_RETENTION_SECONDS"] == 30 - assert result["EMBEDDINGS_PROVIDER"] == "openai" - assert result["EMBEDDINGS_MODEL"] == "text-embedding-3-small" - assert result["EMBEDDINGS_VECTOR_DIM"] == 1536 - assert result["QA_PROVIDER"] == "anthropic" - assert result["QA_MODEL"] == "claude-3-haiku" - assert result["CHUNK_SIZE"] == 512 - assert result["CONTEXT_CHUNK_RADIUS"] == 1 - assert result["OLLAMA_BASE_URL"] == "http://localhost:11434" - assert result["OPENAI_API_KEY"] == "test-key" - assert result["ANTHROPIC_API_KEY"] == "test-key-2" - - -def test_flatten_yaml_empty(): - """Test flattening empty YAML dict returns empty dict.""" - result = flatten_yaml_to_env_dict({}) - assert result == {} - - def test_find_config_file_cwd(tmp_path, monkeypatch): """Test finding config in current directory.""" monkeypatch.chdir(tmp_path) diff --git a/tests/test_embedder.py b/tests/test_embedder.py index fbde114a..1391fe35 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -6,9 +6,9 @@ from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder from haiku.rag.embeddings.vllm import Embedder as VLLMEmbedder -OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY) -VOYAGEAI_AVAILABLE = bool(Config.VOYAGE_API_KEY) -VLLM_EMBEDDINGS_AVAILABLE = bool(Config.VLLM_EMBEDDINGS_BASE_URL) +OPENAI_AVAILABLE = bool(Config.providers.api_keys.openai) +VOYAGEAI_AVAILABLE = bool(Config.providers.api_keys.voyage) +VLLM_EMBEDDINGS_AVAILABLE = bool(Config.providers.vllm.embeddings_base_url) # Calculate cosine similarity diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index c61eae0d..d1de0877 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -14,9 +14,9 @@ async def test_lancedb_cloud_skips_optimization(temp_db_path): # Mock all cloud config to simulate LanceDB Cloud usage with ( - patch.object(Config, "LANCEDB_URI", "db://test-database"), - patch.object(Config, "LANCEDB_API_KEY", "test-api-key"), - patch.object(Config, "LANCEDB_REGION", "us-east-1"), + patch.object(Config.lancedb, "uri", "db://test-database"), + patch.object(Config.lancedb, "api_key", "test-api-key"), + patch.object(Config.lancedb, "region", "us-east-1"), ): # Mock the optimize method to track if it's called with patch.object(store.chunks_table, "optimize") as mock_optimize: @@ -35,8 +35,8 @@ async def test_local_storage_calls_optimization(temp_db_path): # Create a store store = Store(temp_db_path) - # Ensure LANCEDB_URI is empty (local storage) - with patch.object(Config, "LANCEDB_URI", ""): + # Ensure uri is empty (local storage) + with patch.object(Config.lancedb, "uri", ""): # Mock the optimize method to track if it's called with patch.object(store.chunks_table, "optimize") as mock_optimize: # Call vacuum - this should optimize all tables for local storage diff --git a/tests/test_preprocessor.py b/tests/test_preprocessor.py index 5fecdac7..c79d1ed7 100644 --- a/tests/test_preprocessor.py +++ b/tests/test_preprocessor.py @@ -41,9 +41,9 @@ def add_marker(text: str) -> str: """ ) - original_pre = Config.MARKDOWN_PREPROCESSOR + original_pre = Config.processing.markdown_preprocessor try: - Config.MARKDOWN_PREPROCESSOR = f"{pre_file}:add_marker" + Config.processing.markdown_preprocessor = f"{pre_file}:add_marker" store = Store(temp_db_path) chunk_repo = ChunkRepository(store) @@ -68,4 +68,4 @@ def add_marker(text: str) -> str: assert any(marker in c.content for c in chunks) finally: - Config.MARKDOWN_PREPROCESSOR = original_pre + Config.processing.markdown_preprocessor = original_pre diff --git a/tests/test_qa.py b/tests/test_qa.py index 1f2f4022..2bcf20a5 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -6,9 +6,9 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.qa.agent import QuestionAnswerAgent -OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY) -ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY) -VLLM_QA_AVAILABLE = bool(Config.VLLM_QA_BASE_URL) +OPENAI_AVAILABLE = bool(Config.providers.api_keys.openai) +ANTHROPIC_AVAILABLE = bool(Config.providers.api_keys.anthropic) +VLLM_QA_AVAILABLE = bool(Config.providers.vllm.qa_base_url) @pytest.mark.asyncio diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 1d7dc980..bc98b5b6 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -5,8 +5,8 @@ from haiku.rag.reranking.base import RerankerBase from haiku.rag.reranking.vllm import VLLMReranker from haiku.rag.store.models.chunk import Chunk -COHERE_AVAILABLE = bool(Config.COHERE_API_KEY) -VLLM_RERANK_AVAILABLE = bool(Config.VLLM_RERANK_BASE_URL) +COHERE_AVAILABLE = bool(Config.providers.api_keys.cohere) +VLLM_RERANK_AVAILABLE = bool(Config.providers.vllm.rerank_base_url) chunks = [ Chunk(content=content, document_id=str(i)) @@ -37,7 +37,7 @@ async def test_mxbai_reranker(): try: from haiku.rag.reranking.mxbai import MxBAIReranker - Config.RERANK_MODEL = "mixedbread-ai/mxbai-rerank-base-v2" + Config.reranking.model = "mixedbread-ai/mxbai-rerank-base-v2" reranker = MxBAIReranker() # reranker._model = "mixedbread-ai/mxbai-rerank-base-v2" reranked = await reranker.rerank( @@ -45,7 +45,7 @@ async def test_mxbai_reranker(): ) assert [chunk.document_id for chunk, score in reranked] == ["0", "2"] assert all(isinstance(score, float) for chunk, score in reranked) - Config.RERANK_MODEL = "" + Config.reranking.model = "" except ImportError: pytest.skip("MxBAI package not installed") diff --git a/tests/test_settings.py b/tests/test_settings.py index afb86f57..4e6d9548 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -35,14 +35,14 @@ def test_settings_save_and_retrieve(temp_db_path): store = Store(temp_db_path) settings_repo = SettingsRepository(store) - original_chunk_size = Config.CHUNK_SIZE - Config.CHUNK_SIZE = 2 * original_chunk_size + original_chunk_size = Config.processing.chunk_size + Config.processing.chunk_size = 2 * original_chunk_size settings_repo.save_current_settings() retrieved_settings = settings_repo.get_current_settings() - assert retrieved_settings["CHUNK_SIZE"] == 2 * original_chunk_size + assert retrieved_settings["processing"]["chunk_size"] == 2 * original_chunk_size - Config.CHUNK_SIZE = original_chunk_size + Config.processing.chunk_size = original_chunk_size store.close() @@ -57,16 +57,16 @@ async def test_config_validation_on_db_load(temp_db_path): store1.close() # Change config - original_chunk_size = Config.CHUNK_SIZE - Config.CHUNK_SIZE = 999 + original_chunk_size = Config.processing.chunk_size + Config.processing.chunk_size = 999 try: # Loading the database should raise ConfigMismatchError with pytest.raises(ConfigMismatchError) as exc_info: Store(temp_db_path) - assert "CHUNK_SIZE" in str(exc_info.value) - assert "Consider rebuilding" in str(exc_info.value) + assert "chunk_size" in str(exc_info.value) + assert "rebuild" in str(exc_info.value).lower() # Rebuild async with HaikuRAG(db_path=temp_db_path, skip_validation=True) as client: @@ -77,8 +77,8 @@ async def test_config_validation_on_db_load(temp_db_path): store2 = Store(temp_db_path) settings_repo2 = SettingsRepository(store2) db_settings = settings_repo2.get_current_settings() - assert db_settings["CHUNK_SIZE"] == 999 + assert db_settings["processing"]["chunk_size"] == 999 store2.close() finally: - Config.CHUNK_SIZE = original_chunk_size + Config.processing.chunk_size = original_chunk_size diff --git a/tests/test_versioning.py b/tests/test_versioning.py index 874696cd..be0bdd1a 100644 --- a/tests/test_versioning.py +++ b/tests/test_versioning.py @@ -204,7 +204,7 @@ async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch): from haiku.rag.utils import text_to_docling_document # Set aggressive vacuum retention for this test - monkeypatch.setattr(Config, "VACUUM_RETENTION_SECONDS", 0) + monkeypatch.setattr(Config.storage, "vacuum_retention_seconds", 0) async with HaikuRAG(db_path=temp_db_path) as client: # Create multiple documents - each creation triggers automatic vacuum with retention=0