danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor config from flat to nested Pydantic models

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-23 13:30:03 +03:00
parent 0fbe8cc516
commit 618b3ffd5f
No known key found for this signature in database
39 changed files with 319 additions and 407 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
examples/ag-ui-research/backend/agent.py
View file

@ -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(

12
examples/ag-ui-research/backend/main.py
View file

@ -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",

2
src/evaluations/benchmark.py
View file

@ -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]](

2
src/evaluations/llm_judge.py
View file

@ -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

6
src/haiku/rag/a2a/__init__.py
View file

@ -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

12
src/haiku/rag/app.py
View file

@ -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)

2
src/haiku/rag/chunker.py
View file

@ -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(

28
src/haiku/rag/cli.py
View file

@ -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",
),

4
src/haiku/rag/client.py
View file

@ -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.

104
src/haiku/rag/config.py
View file

@ -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

67
src/haiku/rag/config/__init__.py Normal file
View file

@ -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

128
src/haiku/rag/config_loader.py → src/haiku/rag/config/loader.py
View file

@ -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"

103
src/haiku/rag/config/models.py Normal file
View file

@ -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

18
src/haiku/rag/embeddings/__init__.py
View file

@ -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}")

4
src/haiku/rag/embeddings/base.py
View file

@ -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

4
src/haiku/rag/embeddings/ollama.py
View file

@ -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(

2
src/haiku/rag/embeddings/vllm.py
View file

@ -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 []

4
src/haiku/rag/graph/common.py
View file

@ -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",
),
)

22
src/haiku/rag/mcp.py
View file

@ -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,

4
src/haiku/rag/qa/__init__.py
View file

@ -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,

6
src/haiku/rag/qa/agent.py
View file

@ -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:

4
src/haiku/rag/reranking/__init__.py
View file

@ -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

2
src/haiku/rag/reranking/base.py
View file

@ -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

2
src/haiku/rag/reranking/cohere.py
View file

@ -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

2
src/haiku/rag/reranking/mxbai.py
View file

@ -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(

2
src/haiku/rag/reranking/vllm.py
View file

@ -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

16
src/haiku/rag/store/engine.py
View file

@ -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:

2
src/haiku/rag/store/repositories/chunk.py
View file

@ -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)

23
src/haiku/rag/store/repositories/settings.py
View file

@ -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 = []

18
src/haiku/rag/utils.py
View file

@ -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):

6
tests/test_client.py
View file

@ -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 = [

51
tests/test_config_loader.py
View file

@ -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)

6
tests/test_embedder.py
View file

@ -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

10
tests/test_lancedb_connection.py
View file

@ -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

6
tests/test_preprocessor.py
View file

@ -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

6
tests/test_qa.py
View file

@ -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

8
tests/test_reranker.py
View file

@ -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")

20
tests/test_settings.py
View file

@ -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

2
tests/test_versioning.py
View file

@ -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