From 1bb7b6d5bdb71a8e7b0a8edcf6c23e64164587c4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 25 Nov 2025 11:58:57 +0200 Subject: [PATCH 1/7] Add support for per-model configuration settings including thinking, temperature and max_tokens --- CHANGELOG.md | 4 + evaluations/evaluations/benchmark.py | 16 +- evaluations/evaluations/evaluators/judge.py | 14 +- .../a2a-server/haiku_rag_a2a/a2a/__init__.py | 2 +- examples/ag-ui-research/backend/agent.py | 2 +- examples/ag-ui-research/backend/main.py | 6 +- haiku_rag_slim/haiku/rag/config/models.py | 48 +++- .../haiku/rag/embeddings/__init__.py | 20 +- haiku_rag_slim/haiku/rag/embeddings/base.py | 2 +- .../haiku/rag/graph/common/__init__.py | 2 +- .../haiku/rag/graph/common/nodes.py | 24 +- .../haiku/rag/graph/common/utils.py | 48 ---- .../haiku/rag/graph/deep_qa/graph.py | 13 +- .../haiku/rag/graph/research/graph.py | 15 +- haiku_rag_slim/haiku/rag/qa/__init__.py | 6 +- haiku_rag_slim/haiku/rag/qa/agent.py | 30 +-- .../haiku/rag/reranking/__init__.py | 12 +- haiku_rag_slim/haiku/rag/reranking/base.py | 4 +- haiku_rag_slim/haiku/rag/reranking/cohere.py | 3 +- haiku_rag_slim/haiku/rag/reranking/mxbai.py | 7 +- .../haiku/rag/reranking/zeroentropy.py | 3 +- haiku_rag_slim/haiku/rag/utils.py | 246 +++++++++++++++++- tests/graph/test_deep_qa.py | 10 +- tests/graph/test_research_graph.py | 13 +- tests/test_chunk.py | 1 + tests/test_client.py | 8 +- tests/test_embedder_config.py | 29 ++- tests/test_qa.py | 17 +- tests/test_reranker.py | 8 +- tests/test_utils.py | 181 +++++++++++++ 30 files changed, 597 insertions(+), 197 deletions(-) delete mode 100644 haiku_rag_slim/haiku/rag/graph/common/utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 04b0534d..e68ddb51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,10 @@ ### Added +- **Model Customization**: Added support for per-model configuration settings + - New `enable_thinking` parameter to control reasoning behavior (true/false/None) + - Support for `temperature` and `max_tokens` settings on QA and research models + - All settings apply to any provider that supports them - **Database Inspector**: New `inspect` CLI command launches interactive TUI for browsing documents and chunks & searching - **Evaluations**: Added `evaluations` CLI script for running benchmarks (replaces `python -m evaluations.benchmark`) - **Evaluations**: Added `--db` option to override evaluation database path diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 44da7b9b..bba02402 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -43,15 +43,19 @@ def build_experiment_metadata( return { "dataset": dataset_key, "test_cases": test_cases, - "embedder_provider": config.embeddings.provider, - "embedder_model": config.embeddings.model, + "embedder_provider": config.embeddings.model.provider, + "embedder_model": config.embeddings.model.model, "embedder_dim": config.embeddings.vector_dim, "chunk_size": config.processing.chunk_size, "context_chunk_radius": config.processing.context_chunk_radius, - "rerank_provider": config.reranking.provider, - "rerank_model": config.reranking.model, - "qa_provider": config.qa.provider, - "qa_model": config.qa.model, + "rerank_provider": config.reranking.model.provider + if config.reranking.model + else None, + "rerank_model": config.reranking.model.model + if config.reranking.model + else None, + "qa_provider": config.qa.model.provider, + "qa_model": config.qa.model.model, "judge_provider": "ollama", "judge_model": judge_model, } diff --git a/evaluations/evaluations/evaluators/judge.py b/evaluations/evaluations/evaluators/judge.py index d1f7beb9..ccc80dd9 100644 --- a/evaluations/evaluations/evaluators/judge.py +++ b/evaluations/evaluations/evaluators/judge.py @@ -1,9 +1,9 @@ from pydantic import BaseModel from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider from haiku.rag.config import Config +from haiku.rag.config.models import ModelConfig +from haiku.rag.utils import get_model ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent. @@ -37,15 +37,15 @@ class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" def __init__(self, model: str = "gpt-oss"): - # Create Ollama model - ollama_model = OpenAIChatModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"), + # Create model using get_model with thinking disabled + model_config = ModelConfig( + provider="ollama", model=model, enable_thinking=False ) + model_obj = get_model(model_config, Config) # Create Pydantic AI agent self._agent = Agent( - model=ollama_model, + model=model_obj, output_type=LLMJudgeResponseSchema, system_prompt=ANSWER_EQUIVALENCE_RUBRIC, retries=3, diff --git a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py index d06eaae9..524365c0 100644 --- a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py +++ b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py @@ -65,7 +65,7 @@ def create_a2a_app( broker = InMemoryBroker() # Create the agent with native search tool - model = get_model(config.qa.provider, config.qa.model) + model = get_model(config.qa.model, config) agent = Agent( model=model, deps_type=AgentDependencies, diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index a557d78d..0182f1e6 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -35,7 +35,7 @@ class AgentDeps: agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None -model = get_model(Config.research.provider, Config.research.model) +model = get_model(Config.research.model, Config) agent = Agent( model, diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 4865caba..9d65c932 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -47,7 +47,7 @@ if not db_path.exists(): logger.info(f"Initializing research assistant with database: {db_path}") logger.info( - f"Research Provider: {Config.research.provider}, Model: {Config.research.model}" + f"Research Provider: {Config.research.model.provider}, Model: {Config.research.model.model}" ) # Store client reference for proper lifecycle management @@ -153,8 +153,8 @@ async def health_check(_: Request) -> JSONResponse: { "status": "healthy", "agent_model": str(agent.model), - "research_provider": Config.research.provider, - "research_model": Config.research.model, + "research_provider": Config.research.model.provider, + "research_model": Config.research.model.model, "db_path": str(db_path), "db_exists": db_path.exists(), } diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 84d8bd03..1c1007cf 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -6,6 +6,25 @@ from pydantic import BaseModel, Field from haiku.rag.utils import get_default_data_dir +class ModelConfig(BaseModel): + """Configuration for a language model. + + Attributes: + provider: Model provider (ollama, openai, anthropic, etc.) + model: Model name/identifier + enable_thinking: Control reasoning behavior (true/false/None for default) + temperature: Sampling temperature (0.0 to 1.0+) + max_tokens: Maximum tokens to generate + """ + + provider: str = "ollama" + model: str = "gpt-oss" + + enable_thinking: bool | None = None + temperature: float | None = None + max_tokens: int | None = None + + class StorageConfig(BaseModel): data_dir: Path = Field(default_factory=get_default_data_dir) vacuum_retention_seconds: int = 86400 @@ -25,27 +44,40 @@ class LanceDBConfig(BaseModel): class EmbeddingsConfig(BaseModel): - provider: str = "ollama" - model: str = "qwen3-embedding:4b" + model: ModelConfig = Field( + default_factory=lambda: ModelConfig( + provider="ollama", + model="qwen3-embedding:4b", + ) + ) vector_dim: int = 2560 class RerankingConfig(BaseModel): - provider: str = "" - model: str = "" + model: ModelConfig | None = None class QAConfig(BaseModel): - provider: str = "ollama" - model: str = "gpt-oss" + model: ModelConfig = Field( + default_factory=lambda: ModelConfig( + provider="ollama", + model="gpt-oss", + enable_thinking=False, + ) + ) max_sub_questions: int = 3 max_iterations: int = 2 max_concurrency: int = 1 class ResearchConfig(BaseModel): - provider: str = "ollama" - model: str = "gpt-oss" + model: ModelConfig = Field( + default_factory=lambda: ModelConfig( + provider="ollama", + model="gpt-oss", + enable_thinking=True, + ) + ) max_iterations: int = 3 confidence_threshold: float = 0.8 max_concurrency: int = 1 diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index c8e00fce..44464d9b 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -14,12 +14,12 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase: An embedder instance configured according to the config. """ - if config.embeddings.provider == "ollama": + if config.embeddings.model.provider == "ollama": return OllamaEmbedder( - config.embeddings.model, config.embeddings.vector_dim, config + config.embeddings.model.model, config.embeddings.vector_dim, config ) - if config.embeddings.provider == "voyageai": + if config.embeddings.model.provider == "voyageai": try: from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder except ImportError: @@ -29,21 +29,23 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase: "uv pip install haiku.rag[voyageai]" ) return VoyageAIEmbedder( - config.embeddings.model, config.embeddings.vector_dim, config + config.embeddings.model.model, config.embeddings.vector_dim, config ) - if config.embeddings.provider == "openai": + if config.embeddings.model.provider == "openai": from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder return OpenAIEmbedder( - config.embeddings.model, config.embeddings.vector_dim, config + config.embeddings.model.model, config.embeddings.vector_dim, config ) - if config.embeddings.provider == "vllm": + if config.embeddings.model.provider == "vllm": from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder return VllmEmbedder( - config.embeddings.model, config.embeddings.vector_dim, config + config.embeddings.model.model, config.embeddings.vector_dim, config ) - raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}") + raise ValueError( + f"Unsupported embedding provider: {config.embeddings.model.provider}" + ) diff --git a/haiku_rag_slim/haiku/rag/embeddings/base.py b/haiku_rag_slim/haiku/rag/embeddings/base.py index bcd80f91..61452992 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/base.py +++ b/haiku_rag_slim/haiku/rag/embeddings/base.py @@ -4,7 +4,7 @@ from haiku.rag.config import AppConfig, Config class EmbedderBase: - _model: str = Config.embeddings.model + _model: str = Config.embeddings.model.model _vector_dim: int = Config.embeddings.vector_dim _config: AppConfig = Config diff --git a/haiku_rag_slim/haiku/rag/graph/common/__init__.py b/haiku_rag_slim/haiku/rag/graph/common/__init__.py index e2a53f16..199aabe6 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/__init__.py +++ b/haiku_rag_slim/haiku/rag/graph/common/__init__.py @@ -1,5 +1,5 @@ """Common utilities for graph implementations.""" -from haiku.rag.graph.common.utils import get_model +from haiku.rag.utils import get_model __all__ = ["get_model"] diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 4597244a..7aee94a7 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -11,7 +11,7 @@ from pydantic_graph.beta import StepContext from haiku.rag.client import HaikuRAG from haiku.rag.config import Config -from haiku.rag.config.models import AppConfig +from haiku.rag.config.models import AppConfig, ModelConfig from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.common import get_model from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer @@ -52,8 +52,7 @@ class GraphAgentDeps(Protocol): def create_plan_node[AgentDepsT: GraphAgentDeps]( - provider: str, - model: str, + model_config: ModelConfig, deps_type: type[AgentDepsT], activity_message: str = "Creating plan", output_retries: int | None = None, @@ -62,8 +61,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( """Create a plan node for any graph. Args: - provider: Model provider (e.g., 'openai', 'anthropic') - model: Model name + model_config: ModelConfig with provider, model, and settings deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies) activity_message: Message to show during planning activity output_retries: Number of output retries for the agent (optional) @@ -86,7 +84,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( try: # Build agent configuration agent_config = { - "model": get_model(provider, model, config), + "model": get_model(model_config, config), "output_type": ResearchPlan, "instructions": ( PLAN_PROMPT @@ -141,8 +139,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( def create_search_node[AgentDepsT: GraphAgentDeps]( - provider: str, - model: str, + model_config: ModelConfig, deps_type: type[AgentDepsT], with_step_wrapper: bool = True, success_message_format: str = "Answered: {sub_q}", @@ -152,8 +149,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps]( """Create a search_one node for any graph. Args: - provider: Model provider - model: Model name + model_config: ModelConfig with provider, model, and settings deps_type: Type of dependencies for the agent with_step_wrapper: Whether to wrap with agui_emitter start/finish step success_message_format: Format string for success activity message @@ -186,8 +182,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps]( state, deps, sub_q, - provider, - model, + model_config, deps_type, success_message_format, handle_exceptions, @@ -204,8 +199,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( state: GraphState, deps: GraphDeps, sub_q: str, - provider: str, - model: str, + model_config: ModelConfig, deps_type: type[AgentDepsT], success_message_format: str, handle_exceptions: bool, @@ -223,7 +217,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( ) agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=ToolOutput(SearchAnswer, max_retries=3), instructions=SEARCH_AGENT_PROMPT, retries=3, diff --git a/haiku_rag_slim/haiku/rag/graph/common/utils.py b/haiku_rag_slim/haiku/rag/graph/common/utils.py deleted file mode 100644 index c2711915..00000000 --- a/haiku_rag_slim/haiku/rag/graph/common/utils.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Common utilities for all graph implementations.""" - -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider - -from haiku.rag.config import Config -from haiku.rag.config.models import AppConfig - - -def get_model( - provider: str, model: str, config: AppConfig = Config -) -> OpenAIChatModel | str: - """ - Get a model instance for the specified provider and model name. - - Args: - provider: The model provider ("ollama", "vllm", or other) - model: The model name - config: AppConfig object (defaults to global Config) - - Returns: - A configured model instance - - Raises: - ValueError: If the provider is unknown - """ - if provider == "ollama": - return OpenAIChatModel( - model_name=model, - 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.providers.vllm.research_base_url or config.providers.vllm.qa_base_url}/v1", - api_key="none", - ), - ) - elif provider in ("openai", "anthropic", "gemini", "groq", "bedrock"): - # These providers use string format - return f"{provider}:{model}" - else: - raise ValueError( - f"Unknown model provider: {provider}. " - f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock" - ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py index b28a98e7..567c0cf3 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py @@ -29,8 +29,7 @@ def build_deep_qa_graph( Returns: Configured Deep QA graph """ - provider = config.qa.provider - model = config.qa.model + model_config = config.qa.model g = GraphBuilder( state_type=DeepQAState, deps_type=DeepQADeps, @@ -40,8 +39,7 @@ def build_deep_qa_graph( # Create and register the plan node using the factory plan = g.step( create_plan_node( - provider=provider, - model=model, + model_config=model_config, deps_type=DeepQADependencies, # type: ignore[arg-type] activity_message="Planning approach", output_retries=None, # Deep QA doesn't use output_retries @@ -52,8 +50,7 @@ def build_deep_qa_graph( # Create and register the search_one node using the factory search_one = g.step( create_search_node( - provider=provider, - model=model, + model_config=model_config, deps_type=DeepQADependencies, # type: ignore[arg-type] with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step success_message_format="Answered: {sub_q}", @@ -92,7 +89,7 @@ def build_deep_qa_graph( try: agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=DeepQAEvaluation, instructions=DECISION_PROMPT, retries=3, @@ -173,7 +170,7 @@ def build_deep_qa_graph( ) agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=DeepQAAnswer, instructions=prompt_template, retries=3, diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 77c4860a..32000dce 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -36,8 +36,7 @@ def build_research_graph( Returns: Configured Research graph """ - provider = config.research.provider - model = config.research.model + model_config = config.research.model g = GraphBuilder( state_type=ResearchState, deps_type=ResearchDeps, @@ -47,8 +46,7 @@ def build_research_graph( # Create and register the plan node using the factory plan = g.step( create_plan_node( - provider=provider, - model=model, + model_config=model_config, deps_type=ResearchDependencies, # type: ignore[arg-type] activity_message="Creating research plan", output_retries=3, @@ -59,8 +57,7 @@ def build_research_graph( # Create and register the search_one node using the factory search_one = g.step( create_search_node( - provider=provider, - model=model, + model_config=model_config, deps_type=ResearchDependencies, # type: ignore[arg-type] with_step_wrapper=True, success_message_format="Found answer with {confidence:.0%} confidence", @@ -99,7 +96,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=InsightAnalysis, instructions=INSIGHT_AGENT_PROMPT, retries=3, @@ -168,7 +165,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=EvaluationResult, instructions=DECISION_AGENT_PROMPT, retries=3, @@ -247,7 +244,7 @@ def build_research_graph( try: agent = Agent( - model=get_model(provider, model, config), + model=get_model(model_config, config), output_type=ResearchReport, instructions=SYNTHESIS_AGENT_PROMPT, retries=3, diff --git a/haiku_rag_slim/haiku/rag/qa/__init__.py b/haiku_rag_slim/haiku/rag/qa/__init__.py index 4e8e7f92..faf1fca8 100644 --- a/haiku_rag_slim/haiku/rag/qa/__init__.py +++ b/haiku_rag_slim/haiku/rag/qa/__init__.py @@ -21,13 +21,9 @@ def get_qa_agent( Returns: A configured QuestionAnswerAgent instance. """ - provider = config.qa.provider - model_name = config.qa.model - return QuestionAnswerAgent( client=client, - provider=provider, - model=model_name, + model_config=config.qa.model, use_citations=use_citations, system_prompt=system_prompt, ) diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index d277ae08..2f308008 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -1,11 +1,10 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.client import HaikuRAG from haiku.rag.config import Config +from haiku.rag.config.models import ModelConfig +from haiku.rag.graph.common import get_model from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS @@ -26,8 +25,7 @@ class QuestionAnswerAgent: def __init__( self, client: HaikuRAG, - provider: str, - model: str, + model_config: ModelConfig, use_citations: bool = False, q: float = 0.0, system_prompt: str | None = None, @@ -38,7 +36,7 @@ class QuestionAnswerAgent: system_prompt = ( QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT ) - model_obj = self._get_model(provider, model) + model_obj = get_model(model_config, Config) self._agent = Agent( model=model_obj, @@ -66,26 +64,6 @@ class QuestionAnswerAgent: for chunk, score in expanded_results ] - def _get_model(self, provider: str, model: str): - """Get the appropriate model object for the provider.""" - if provider == "ollama": - return OpenAIChatModel( - model_name=model, - 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.providers.vllm.qa_base_url}/v1", api_key="none" - ), - ) - else: - # For all other providers, use the provider:model format - return f"{provider}:{model}" - async def answer(self, question: str) -> str: """Answer a question using the RAG system.""" deps = Dependencies(client=self._client) diff --git a/haiku_rag_slim/haiku/rag/reranking/__init__.py b/haiku_rag_slim/haiku/rag/reranking/__init__.py index f43e619e..b9f32694 100644 --- a/haiku_rag_slim/haiku/rag/reranking/__init__.py +++ b/haiku_rag_slim/haiku/rag/reranking/__init__.py @@ -24,7 +24,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None: reranker: RerankerBase | None = None - if config.reranking.provider == "mxbai": + if config.reranking.model and config.reranking.model.provider == "mxbai": try: from haiku.rag.reranking.mxbai import MxBAIReranker @@ -33,7 +33,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None: except ImportError: reranker = None - elif config.reranking.provider == "cohere": + elif config.reranking.model and config.reranking.model.provider == "cohere": try: from haiku.rag.reranking.cohere import CohereReranker @@ -41,20 +41,20 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None: except ImportError: reranker = None - elif config.reranking.provider == "vllm": + elif config.reranking.model and config.reranking.model.provider == "vllm": try: from haiku.rag.reranking.vllm import VLLMReranker - reranker = VLLMReranker(config.reranking.model) + reranker = VLLMReranker(config.reranking.model.model) except ImportError: reranker = None - elif config.reranking.provider == "zeroentropy": + elif config.reranking.model and config.reranking.model.provider == "zeroentropy": try: from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker # Use configured model or default to zerank-1 - model = config.reranking.model or "zerank-1" + model = config.reranking.model.model or "zerank-1" reranker = ZeroEntropyReranker(model) except ImportError: reranker = None diff --git a/haiku_rag_slim/haiku/rag/reranking/base.py b/haiku_rag_slim/haiku/rag/reranking/base.py index 4f93d1a2..808371fd 100644 --- a/haiku_rag_slim/haiku/rag/reranking/base.py +++ b/haiku_rag_slim/haiku/rag/reranking/base.py @@ -3,7 +3,9 @@ from haiku.rag.store.models.chunk import Chunk class RerankerBase: - _model: str = Config.reranking.model + _model: str | None = ( + Config.reranking.model.model if Config.reranking.model else None + ) async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 diff --git a/haiku_rag_slim/haiku/rag/reranking/cohere.py b/haiku_rag_slim/haiku/rag/reranking/cohere.py index 439cbca2..c05daab1 100644 --- a/haiku_rag_slim/haiku/rag/reranking/cohere.py +++ b/haiku_rag_slim/haiku/rag/reranking/cohere.py @@ -22,8 +22,9 @@ class CohereReranker(RerankerBase): documents = [chunk.content for chunk in chunks] + model_name = self._model or "rerank-v3.5" response = self._client.rerank( - model=self._model, query=query, documents=documents, top_n=top_n + model=model_name, query=query, documents=documents, top_n=top_n ) reranked_chunks = [] diff --git a/haiku_rag_slim/haiku/rag/reranking/mxbai.py b/haiku_rag_slim/haiku/rag/reranking/mxbai.py index f4b7834c..517c0dc7 100644 --- a/haiku_rag_slim/haiku/rag/reranking/mxbai.py +++ b/haiku_rag_slim/haiku/rag/reranking/mxbai.py @@ -7,9 +7,12 @@ from haiku.rag.store.models.chunk import Chunk class MxBAIReranker(RerankerBase): def __init__(self): - self._client = MxbaiRerankV2( - Config.reranking.model, disable_transformers_warnings=True + model_name = ( + Config.reranking.model.model + if Config.reranking.model + else "mxbai-rerank-base-v2" ) + self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True) async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 diff --git a/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py b/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py index 0c66fb08..ad37fccc 100644 --- a/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py +++ b/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py @@ -37,8 +37,9 @@ class ZeroEntropyReranker(RerankerBase): documents = [chunk.content for chunk in chunks] # Call Zero Entropy reranking API + model_name = self._model or "zerank-1" response = self._client.models.rerank( - model=self._model, + model=model_name, query=query, documents=documents, ) diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index ae9d6e04..074e01f0 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -4,10 +4,240 @@ import sys from importlib import metadata from pathlib import Path from types import ModuleType +from typing import Any from packaging.version import Version, parse +def apply_common_settings( + settings: Any | None, + settings_class: type[Any], + model_config: Any, +) -> Any | None: + """Apply common settings (temperature, max_tokens) to model settings. + + Args: + settings: Existing settings instance or None + settings_class: Settings class to instantiate if needed + model_config: ModelConfig with temperature and max_tokens + + Returns: + Updated settings instance or None if no settings to apply + """ + if model_config.temperature is None and model_config.max_tokens is None: + return settings + + if settings is None: + settings_dict = settings_class() + else: + settings_dict = settings + + if model_config.temperature is not None: + settings_dict["temperature"] = model_config.temperature + + if model_config.max_tokens is not None: + settings_dict["max_tokens"] = model_config.max_tokens + + return settings_dict + + +def get_model( + model_config: Any, + app_config: Any | None = None, +) -> Any: + """ + Get a model instance for the specified configuration. + + Args: + model_config: ModelConfig with provider, model, and settings + app_config: AppConfig for provider base URLs (defaults to global Config) + + Returns: + A configured model instance + """ + from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings + from pydantic_ai.providers.ollama import OllamaProvider + from pydantic_ai.providers.openai import OpenAIProvider + + if app_config is None: + from haiku.rag.config import Config + + app_config = Config + + provider = model_config.provider + model = model_config.model + + if provider == "ollama": + model_settings = None + + # Apply thinking control for gpt-oss + if model == "gpt-oss" and model_config.enable_thinking is not None: + if model_config.enable_thinking is False: + model_settings = OpenAIChatModelSettings(openai_reasoning_effort="low") + else: + model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high") + + model_settings = apply_common_settings( + model_settings, OpenAIChatModelSettings, model_config + ) + + return OpenAIChatModel( + model_name=model, + provider=OllamaProvider( + base_url=f"{app_config.providers.ollama.base_url}/v1" + ), + settings=model_settings, + ) + + elif provider == "openai": + openai_settings: Any = None + + # Apply thinking control + if model_config.enable_thinking is not None: + if model_config.enable_thinking is False: + openai_settings = OpenAIChatModelSettings(openai_reasoning_effort="low") + else: + openai_settings = OpenAIChatModelSettings( + openai_reasoning_effort="high" + ) + + openai_settings = apply_common_settings( + openai_settings, OpenAIChatModelSettings, model_config + ) + + return OpenAIChatModel(model_name=model, settings=openai_settings) + + elif provider == "anthropic": + from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings + + anthropic_settings: Any = None + + # Apply thinking control + if model_config.enable_thinking is not None: + if model_config.enable_thinking: + anthropic_settings = AnthropicModelSettings( + anthropic_thinking={"type": "enabled", "budget_tokens": 4096} + ) + else: + anthropic_settings = AnthropicModelSettings( + anthropic_thinking={"type": "disabled"} + ) + + anthropic_settings = apply_common_settings( + anthropic_settings, AnthropicModelSettings, model_config + ) + + return AnthropicModel(model_name=model, settings=anthropic_settings) + + elif provider == "gemini": + from pydantic_ai.models.google import GoogleModel, GoogleModelSettings + + gemini_settings: Any = None + + # Apply thinking control + if model_config.enable_thinking is not None: + gemini_settings = GoogleModelSettings( + google_thinking_config={ + "include_thoughts": model_config.enable_thinking + } + ) + + gemini_settings = apply_common_settings( + gemini_settings, GoogleModelSettings, model_config + ) + + return GoogleModel(model_name=model, settings=gemini_settings) + + elif provider == "groq": + from pydantic_ai.models.groq import GroqModel, GroqModelSettings + + groq_settings: Any = None + + # Apply thinking control + if model_config.enable_thinking is not None: + if model_config.enable_thinking: + groq_settings = GroqModelSettings(groq_reasoning_format="parsed") + else: + groq_settings = GroqModelSettings(groq_reasoning_format="hidden") + + groq_settings = apply_common_settings( + groq_settings, GroqModelSettings, model_config + ) + + return GroqModel(model_name=model, settings=groq_settings) + + elif provider == "bedrock": + from pydantic_ai.models.bedrock import ( + BedrockConverseModel, + BedrockModelSettings, + ) + + bedrock_settings: Any = None + + # Apply thinking control for Claude models + if model_config.enable_thinking is not None: + additional_fields: dict[str, Any] = {} + if model.startswith("anthropic.claude"): + if model_config.enable_thinking: + additional_fields = { + "thinking": {"type": "enabled", "budget_tokens": 4096} + } + else: + additional_fields = {"thinking": {"type": "disabled"}} + elif "gpt" in model or "o1" in model or "o3" in model: + # OpenAI models on Bedrock + additional_fields = { + "reasoning_effort": "high" + if model_config.enable_thinking + else "low" + } + elif "qwen" in model: + # Qwen models on Bedrock + additional_fields = { + "reasoning_config": "high" + if model_config.enable_thinking + else "low" + } + + if additional_fields: + bedrock_settings = BedrockModelSettings( + bedrock_additional_model_requests_fields=additional_fields + ) + + bedrock_settings = apply_common_settings( + bedrock_settings, BedrockModelSettings, model_config + ) + + return BedrockConverseModel(model_name=model, settings=bedrock_settings) + + elif provider == "vllm": + vllm_settings = None + + # Apply thinking control for gpt-oss + if model == "gpt-oss" and model_config.enable_thinking is not None: + if model_config.enable_thinking is False: + vllm_settings = OpenAIChatModelSettings(openai_reasoning_effort="low") + else: + vllm_settings = OpenAIChatModelSettings(openai_reasoning_effort="high") + + vllm_settings = apply_common_settings( + vllm_settings, OpenAIChatModelSettings, model_config + ) + + return OpenAIChatModel( + model_name=model, + provider=OpenAIProvider( + base_url=f"{app_config.providers.vllm.research_base_url or app_config.providers.vllm.qa_base_url}/v1", + api_key="none", + ), + settings=vllm_settings, + ) + + else: + # For any other provider, use string format and let Pydantic AI handle it + return f"{provider}:{model}" + + def format_bytes(num_bytes: int) -> str: """Format bytes as human-readable string.""" size = float(num_bytes) @@ -135,14 +365,14 @@ 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.reranking.provider == "ollama": - required_models.add(Config.reranking.model) + if Config.embeddings.model.provider == "ollama": + required_models.add(Config.embeddings.model.model) + if Config.qa.model.provider == "ollama": + required_models.add(Config.qa.model.model) + if Config.research.model.provider == "ollama": + required_models.add(Config.research.model.model) + if Config.reranking.model and Config.reranking.model.provider == "ollama": + required_models.add(Config.reranking.model.model) if not required_models: return diff --git a/tests/graph/test_deep_qa.py b/tests/graph/test_deep_qa.py index 3c7e4e13..c96cda33 100644 --- a/tests/graph/test_deep_qa.py +++ b/tests/graph/test_deep_qa.py @@ -16,7 +16,10 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): def test_model_factory(provider, model, config=None): return TestModel() - monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory) + # Patch all locations where get_model is imported + monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory) graph = build_deep_qa_graph() @@ -50,7 +53,10 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path): def test_model_factory(provider, model, config=None): return TestModel() - monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory) + # Patch all locations where get_model is imported + monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory) graph = build_deep_qa_graph() diff --git a/tests/graph/test_research_graph.py b/tests/graph/test_research_graph.py index 889ae92b..6311ba4e 100644 --- a/tests/graph/test_research_graph.py +++ b/tests/graph/test_research_graph.py @@ -1,5 +1,3 @@ -import asyncio - import pytest from pydantic_ai.models.test import TestModel @@ -25,12 +23,6 @@ def test_build_graph_and_state(): assert state.context.sub_questions == [] -def test_async_loop_available(): - # Ensure an event loop can be created in test env - loop = asyncio.new_event_loop() - loop.close() - - @pytest.mark.asyncio async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): """Test research graph with mocked LLM using AG-UI events.""" @@ -39,7 +31,10 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): def test_model_factory(_provider, _model, _config=None): return TestModel() - monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory) + # Patch all locations where get_model is imported + monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory) graph = build_research_graph() diff --git a/tests/test_chunk.py b/tests/test_chunk.py index bc90cb87..8da09bfa 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -115,6 +115,7 @@ async def test_chunk_repository_crud(temp_db_path): ) created_chunk = await chunk_repo.create(chunk) + assert isinstance(created_chunk, Chunk) assert created_chunk.id is not None assert created_chunk.content == "Test chunk content" diff --git a/tests/test_client.py b/tests/test_client.py index fb50acbf..0f34cd82 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -736,9 +736,9 @@ async def test_client_ask_without_cite(monkeypatch, temp_db_path): """Test asking questions without citations.""" from pydantic_ai.models.test import TestModel - # Mock OpenAIChatModel to return TestModel + # Mock get_model to return TestModel monkeypatch.setattr( - "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel() + "haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel() ) async with HaikuRAG(temp_db_path) as client: @@ -760,9 +760,9 @@ async def test_client_ask_with_cite(monkeypatch, temp_db_path): """Test asking questions with citations.""" from pydantic_ai.models.test import TestModel - # Mock OpenAIChatModel to return TestModel + # Mock get_model to return TestModel monkeypatch.setattr( - "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel() + "haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel() ) async with HaikuRAG(temp_db_path) as client: diff --git a/tests/test_embedder_config.py b/tests/test_embedder_config.py index a386d4b3..bfce4190 100644 --- a/tests/test_embedder_config.py +++ b/tests/test_embedder_config.py @@ -7,6 +7,7 @@ from haiku.rag.config import ( ProvidersConfig, VLLMConfig, ) +from haiku.rag.config.models import ModelConfig from haiku.rag.embeddings import get_embedder @@ -14,8 +15,10 @@ def test_embedder_uses_config_from_get_embedder(): """Test that embedders use the config passed to get_embedder.""" custom_config = AppConfig( embeddings=EmbeddingsConfig( - provider="ollama", - model="custom-model", + model=ModelConfig( + provider="ollama", + model="custom-model", + ), vector_dim=512, ), providers=ProvidersConfig( @@ -33,10 +36,14 @@ def test_embedder_uses_config_from_get_embedder(): def test_vllm_embedder_uses_config(): """Test that vllm embedder uses the config passed to get_embedder.""" + from haiku.rag.config.models import ModelConfig + custom_config = AppConfig( embeddings=EmbeddingsConfig( - provider="vllm", - model="custom-vllm-model", + model=ModelConfig( + provider="vllm", + model="custom-vllm-model", + ), vector_dim=768, ), providers=ProvidersConfig( @@ -55,10 +62,14 @@ def test_vllm_embedder_uses_config(): def test_openai_embedder_uses_config(): """Test that openai embedder uses the config passed to get_embedder.""" + from haiku.rag.config.models import ModelConfig + custom_config = AppConfig( embeddings=EmbeddingsConfig( - provider="openai", - model="text-embedding-3-large", + model=ModelConfig( + provider="openai", + model="text-embedding-3-large", + ), vector_dim=3072, ), ) @@ -77,8 +88,10 @@ def test_voyageai_embedder_uses_config(): """Test that voyageai embedder uses the config passed to get_embedder.""" custom_config = AppConfig( embeddings=EmbeddingsConfig( - provider="voyageai", - model="voyage-large-2", + model=ModelConfig( + provider="voyageai", + model="voyage-large-2", + ), vector_dim=1536, ), ) diff --git a/tests/test_qa.py b/tests/test_qa.py index 39053d23..f4fcfee3 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -6,6 +6,7 @@ from evaluations.evaluators import LLMJudge from haiku.rag.client import HaikuRAG from haiku.rag.config import Config +from haiku.rag.config.models import ModelConfig from haiku.rag.qa.agent import QuestionAnswerAgent OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY")) @@ -17,7 +18,9 @@ VLLM_QA_AVAILABLE = bool(Config.providers.vllm.qa_base_url) async def test_qa_ollama(qa_corpus: Dataset, temp_db_path): """Test Ollama QA with LLM judge.""" client = HaikuRAG(temp_db_path) - qa = QuestionAnswerAgent(client, "ollama", "qwen3") + qa = QuestionAnswerAgent( + client, ModelConfig(provider="ollama", model="gpt-oss", enable_thinking=False) + ) llm_judge = LLMJudge() doc = qa_corpus[1] @@ -41,7 +44,9 @@ async def test_qa_ollama(qa_corpus: Dataset, temp_db_path): async def test_qa_openai(qa_corpus: Dataset, temp_db_path): """Test OpenAI QA with LLM judge.""" client = HaikuRAG(temp_db_path) - qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini") + qa = QuestionAnswerAgent( + client, ModelConfig(provider="openai", model="gpt-4o-mini") + ) llm_judge = LLMJudge() doc = qa_corpus[1] @@ -65,7 +70,9 @@ async def test_qa_openai(qa_corpus: Dataset, temp_db_path): async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path): """Test Anthropic QA with LLM judge.""" client = HaikuRAG(temp_db_path) - qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022") + qa = QuestionAnswerAgent( + client, ModelConfig(provider="anthropic", model="claude-3-5-haiku-20241022") + ) llm_judge = LLMJudge() doc = qa_corpus[1] @@ -89,7 +96,9 @@ async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path): async def test_qa_vllm(qa_corpus: Dataset, temp_db_path): """Test vLLM QA with LLM judge.""" client = HaikuRAG(temp_db_path) - qa = QuestionAnswerAgent(client, "vllm", "Qwen/Qwen3-4B") + qa = QuestionAnswerAgent( + client, ModelConfig(provider="vllm", model="Qwen/Qwen3-4B") + ) llm_judge = LLMJudge() doc = qa_corpus[1] diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 399234b3..8dc0105e 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -40,17 +40,19 @@ async def test_reranker_base(): @pytest.mark.asyncio async def test_mxbai_reranker(): try: + from haiku.rag.config.models import ModelConfig from haiku.rag.reranking.mxbai import MxBAIReranker - Config.reranking.model = "mixedbread-ai/mxbai-rerank-base-v2" + Config.reranking.model = ModelConfig( + provider="mxbai", model="mixedbread-ai/mxbai-rerank-base-v2" + ) reranker = MxBAIReranker() - # reranker._model = "mixedbread-ai/mxbai-rerank-base-v2" reranked = await reranker.rerank( "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2 ) assert [chunk.document_id for chunk, score in reranked] == ["0", "2"] assert all(isinstance(score, float) for chunk, score in reranked) - Config.reranking.model = "" + Config.reranking.model = None except ImportError: pytest.skip("MxBAI package not installed") diff --git a/tests/test_utils.py b/tests/test_utils.py index c412242f..0a194582 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,5 +1,18 @@ +import importlib.util + +import pytest +from pydantic_ai.models.openai import OpenAIChatModel + from haiku.rag.config import Config +from haiku.rag.config.models import ModelConfig from haiku.rag.converters import get_converter +from haiku.rag.utils import get_model + +# Check for optional dependencies +HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None +HAS_GOOGLE = importlib.util.find_spec("google.generativeai") is not None +HAS_GROQ = importlib.util.find_spec("groq") is not None +HAS_BEDROCK = importlib.util.find_spec("botocore") is not None def test_text_to_docling_document(): @@ -119,3 +132,171 @@ Emoji test: 🚀 ✅ 📝""" assert "测试文档" in result_markdown assert "¡Hola mundo!" in result_markdown assert "🚀" in result_markdown + + +def test_get_model_ollama(): + """Test get_model returns OpenAIChatModel for Ollama.""" + model_config = ModelConfig(provider="ollama", model="llama3") + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_ollama_with_thinking(): + """Test get_model configures thinking for gpt-oss on Ollama.""" + model_config = ModelConfig( + provider="ollama", model="gpt-oss", enable_thinking=False + ) + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_ollama_with_settings(): + """Test get_model applies temperature and max_tokens for Ollama.""" + model_config = ModelConfig( + provider="ollama", model="llama3", temperature=0.5, max_tokens=100 + ) + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_openai(): + """Test get_model returns OpenAIChatModel for OpenAI.""" + model_config = ModelConfig(provider="openai", model="gpt-4o") + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_openai_with_thinking(): + """Test get_model configures thinking for OpenAI reasoning models.""" + model_config = ModelConfig(provider="openai", model="o1", enable_thinking=True) + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed") +def test_get_model_anthropic(): + """Test get_model returns AnthropicModel for Anthropic.""" + from pydantic_ai.models.anthropic import AnthropicModel + + model_config = ModelConfig(provider="anthropic", model="claude-3-5-sonnet-20241022") + result = get_model(model_config) + assert isinstance(result, AnthropicModel) + + +@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed") +def test_get_model_anthropic_with_thinking(): + """Test get_model configures thinking for Anthropic.""" + from pydantic_ai.models.anthropic import AnthropicModel + + model_config = ModelConfig( + provider="anthropic", + model="claude-3-5-sonnet-20241022", + enable_thinking=True, + ) + result = get_model(model_config) + assert isinstance(result, AnthropicModel) + + +@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed") +def test_get_model_gemini(): + """Test get_model returns GoogleModel for Gemini.""" + from pydantic_ai.models.google import GoogleModel + + model_config = ModelConfig(provider="gemini", model="gemini-2.0-flash-exp") + result = get_model(model_config) + assert isinstance(result, GoogleModel) + + +@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed") +def test_get_model_gemini_with_thinking(): + """Test get_model configures thinking for Gemini.""" + from pydantic_ai.models.google import GoogleModel + + model_config = ModelConfig( + provider="gemini", model="gemini-2.0-flash-thinking-exp", enable_thinking=True + ) + result = get_model(model_config) + assert isinstance(result, GoogleModel) + + +@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed") +def test_get_model_groq(): + """Test get_model returns GroqModel for Groq.""" + from pydantic_ai.models.groq import GroqModel + + model_config = ModelConfig(provider="groq", model="llama-3.3-70b-versatile") + result = get_model(model_config) + assert isinstance(result, GroqModel) + + +@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed") +def test_get_model_groq_with_thinking(): + """Test get_model configures thinking format for Groq.""" + from pydantic_ai.models.groq import GroqModel + + model_config = ModelConfig( + provider="groq", model="llama-3.3-70b-versatile", enable_thinking=False + ) + result = get_model(model_config) + assert isinstance(result, GroqModel) + + +@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed") +def test_get_model_bedrock(): + """Test get_model returns BedrockConverseModel for Bedrock.""" + from pydantic_ai.models.bedrock import BedrockConverseModel + + model_config = ModelConfig( + provider="bedrock", model="anthropic.claude-3-5-sonnet-20241022-v2:0" + ) + result = get_model(model_config) + assert isinstance(result, BedrockConverseModel) + + +@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed") +def test_get_model_bedrock_with_thinking(): + """Test get_model configures thinking for Bedrock Claude models.""" + from pydantic_ai.models.bedrock import BedrockConverseModel + + model_config = ModelConfig( + provider="bedrock", + model="anthropic.claude-3-5-sonnet-20241022-v2:0", + enable_thinking=True, + ) + result = get_model(model_config) + assert isinstance(result, BedrockConverseModel) + + +def test_get_model_vllm(): + """Test get_model returns OpenAIChatModel for vLLM.""" + model_config = ModelConfig(provider="vllm", model="Qwen/Qwen3-4B") + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_vllm_with_thinking(): + """Test get_model configures thinking for gpt-oss on vLLM.""" + model_config = ModelConfig(provider="vllm", model="gpt-oss", enable_thinking=False) + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) + + +def test_get_model_unknown_provider(): + """Test get_model returns string format for unknown providers.""" + model_config = ModelConfig(provider="mistral", model="mistral-large-latest") + result = get_model(model_config) + assert isinstance(result, str) + assert result == "mistral:mistral-large-latest" + + +def test_get_model_with_all_settings(): + """Test get_model applies all settings together.""" + model_config = ModelConfig( + provider="openai", + model="gpt-4o", + enable_thinking=False, + temperature=0.7, + max_tokens=500, + ) + result = get_model(model_config) + assert isinstance(result, OpenAIChatModel) From 41ea67596450f8552790d7a1e838a58e64c85fee Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 25 Nov 2025 12:13:39 +0200 Subject: [PATCH 2/7] Update and restructure docs --- docs/cli.md | 4 +- docs/config-index.md | 195 +++++++++ docs/configuration.md | 829 -------------------------------------- docs/index.md | 2 +- docs/installation.md | 4 +- docs/processing.md | 250 ++++++++++++ docs/providers.md | 343 ++++++++++++++++ docs/python.md | 2 +- docs/qa-research.md | 71 ++++ docs/remote-processing.md | 2 +- docs/server.md | 2 +- docs/storage.md | 103 +++++ docs/tutorial.md | 4 +- mkdocs.yml | 7 +- 14 files changed, 978 insertions(+), 840 deletions(-) create mode 100644 docs/config-index.md delete mode 100644 docs/configuration.md create mode 100644 docs/processing.md create mode 100644 docs/providers.md create mode 100644 docs/qa-research.md create mode 100644 docs/storage.md diff --git a/docs/cli.md b/docs/cli.md index ec7a164b..a166de81 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -69,7 +69,7 @@ haiku-rag add-src /path/to/documents/ ``` !!! note - When adding a directory, the same content filters configured for [file monitoring](configuration.md#filtering-monitored-files) are applied. This means `ignore_patterns` and `include_patterns` from your configuration will be used to filter which files are added. + When adding a directory, the same content filters configured for [file monitoring](processing.md#filtering-monitored-files) are applied. This means `ignore_patterns` and `include_patterns` from your configuration will be used to filter which files are added. !!! note As you add documents to `haiku.rag` the database keeps growing. By default, LanceDB supports versioning @@ -158,7 +158,7 @@ haiku-rag research "How does haiku.rag organize and query documents?" --verbose Flags: - `--verbose`: Show planning, searching previews, evaluation summary, and stop reason -Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration.md) under the `research` section. +Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](config-index.md) under the `research` section. When `--verbose` is set, the CLI consumes the research graph's AG-UI event stream, displaying step events and activity snapshots as agents progress through planning, search, evaluation, and synthesis. Without `--verbose`, only the final research report is displayed. diff --git a/docs/config-index.md b/docs/config-index.md new file mode 100644 index 00000000..724a50d3 --- /dev/null +++ b/docs/config-index.md @@ -0,0 +1,195 @@ +# Configuration + +Configuration is done through YAML configuration files. + +!!! note + If you create a db with certain settings and later change them, `haiku.rag` will detect incompatibilities (for example, if you change embedding provider) and will exit. You can **rebuild** the database to apply the new settings, see [Rebuild Database](./cli.md#rebuild-database). + +## Getting Started + +Generate a configuration file with defaults: + +```bash +haiku-rag init-config +``` + +This creates a `haiku.rag.yaml` file in your current directory with all available settings. + +## Configuration File Locations + +`haiku.rag` searches for configuration files in this order: + +1. Path specified via `--config` flag: `haiku-rag --config /path/to/config.yaml ` +2. `./haiku.rag.yaml` (current directory) +3. Platform-specific user directory: + - **Linux**: `~/.local/share/haiku.rag/haiku.rag.yaml` + - **macOS**: `~/Library/Application Support/haiku.rag/haiku.rag.yaml` + - **Windows**: `C:/Users//AppData/Roaming/haiku.rag/haiku.rag.yaml` + +## Minimal Configuration + +A minimal configuration file with defaults: + +```yaml +# haiku.rag.yaml +environment: production + +embeddings: + model: + provider: ollama + model: qwen3-embedding:4b + vector_dim: 2560 + +qa: + model: + provider: ollama + model: gpt-oss + enable_thinking: false +``` + +## Complete Configuration Example + +```yaml +# haiku.rag.yaml +environment: production + +storage: + data_dir: "" # Empty = use default platform location + vacuum_retention_seconds: 86400 + +monitor: + directories: + - /path/to/documents + - /another/path + ignore_patterns: [] # Gitignore-style patterns to exclude + include_patterns: [] # Gitignore-style patterns to include + +lancedb: + uri: "" # Empty for local, or db://, s3://, az://, gs:// + api_key: "" + region: "" + +embeddings: + model: + provider: ollama + model: qwen3-embedding:4b + vector_dim: 2560 + +reranking: + model: + provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm + model: "" + +qa: + model: + provider: ollama + model: gpt-oss + enable_thinking: false + max_sub_questions: 3 + max_iterations: 2 + max_concurrency: 1 + +research: + model: + provider: "" # Empty to use qa settings + model: "" + enable_thinking: true + max_iterations: 3 + confidence_threshold: 0.8 + max_concurrency: 1 + +search: + vector_index_metric: cosine # cosine, l2, or dot + vector_refine_factor: 30 + +agui: + host: "0.0.0.0" + port: 8000 + cors_origins: ["*"] + cors_credentials: true + cors_methods: ["GET", "POST", "OPTIONS"] + cors_headers: ["*"] + +processing: + converter: docling-local # docling-local or docling-serve + chunker: docling-local # docling-local or docling-serve + chunker_type: hybrid # hybrid or hierarchical + chunk_size: 256 + context_chunk_radius: 0 + chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" + chunking_merge_peers: true + chunking_use_markdown_tables: false + markdown_preprocessor: "" + conversion_options: + do_ocr: true + force_ocr: false + ocr_lang: [] + do_table_structure: true + table_mode: accurate + table_cell_matching: true + images_scale: 2.0 + +providers: + ollama: + base_url: http://localhost:11434 + + vllm: + embeddings_base_url: "" + rerank_base_url: "" + qa_base_url: "" + research_base_url: "" + + docling_serve: + base_url: http://localhost:5001 + api_key: "" + timeout: 300 +``` + +## Programmatic Configuration + +When using haiku.rag as a Python library, you can pass configuration directly to the `HaikuRAG` client: + +```python +from haiku.rag.config import AppConfig +from haiku.rag.config.models import ModelConfig, QAConfig, EmbeddingsConfig +from haiku.rag.client import HaikuRAG + +# Create custom configuration +custom_config = AppConfig( + qa=QAConfig( + model=ModelConfig( + provider="openai", + model="gpt-4o", + temperature=0.7 + ) + ), + embeddings=EmbeddingsConfig( + model=ModelConfig( + provider="ollama", + model="qwen3-embedding:4b" + ), + vector_dim=2560 + ), + processing={"chunk_size": 512} +) + +# 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 multiple clients with different configurations + +## Configuration Topics + +For detailed configuration of specific topics, see: + +- **[Providers](providers.md)** - Model settings and provider-specific configuration (embeddings, QA, reranking) +- **[QA and Research](qa-research.md)** - Question answering and research workflow configuration +- **[Storage](storage.md)** - Database, remote storage, and vector indexing +- **[Document Processing](processing.md)** - Document conversion, chunking, and file monitoring diff --git a/docs/configuration.md b/docs/configuration.md deleted file mode 100644 index 96d2041b..00000000 --- a/docs/configuration.md +++ /dev/null @@ -1,829 +0,0 @@ -# Configuration - -Configuration is done through YAML configuration files. - -!!! note - If you create a db with certain settings and later change them, `haiku.rag` will detect incompatibilities (for example, if you change embedding provider) and will exit. You can **rebuild** the database to apply the new settings, see [Rebuild Database](./cli.md#rebuild-database). - -## Getting Started - -Generate a configuration file with defaults: - -```bash -haiku-rag init-config -``` - -This creates a `haiku.rag.yaml` file in your current directory with all available settings. - -## Configuration File Locations - -`haiku.rag` searches for configuration files in this order: - -1. Path specified via `--config` flag: `haiku-rag --config /path/to/config.yaml ` -2. `./haiku.rag.yaml` (current directory) -3. Platform-specific user directory: - - **Linux**: `~/.local/share/haiku.rag/haiku.rag.yaml` - - **macOS**: `~/Library/Application Support/haiku.rag/haiku.rag.yaml` - - **Windows**: `C:/Users//AppData/Roaming/haiku.rag/haiku.rag.yaml` - -## Minimal Configuration - -A minimal configuration file with defaults: - -```yaml -# haiku.rag.yaml -environment: production - -embeddings: - provider: ollama - model: qwen3-embedding:4b - vector_dim: 2560 - -qa: - provider: ollama - model: gpt-oss -``` - -## Complete Configuration Example - -```yaml -# haiku.rag.yaml -environment: production - -storage: - data_dir: "" # Empty = use default platform location - vacuum_retention_seconds: 86400 - -monitor: - directories: - - /path/to/documents - - /another/path - ignore_patterns: [] # Gitignore-style patterns to exclude - include_patterns: [] # Gitignore-style patterns to include - -lancedb: - uri: "" # Empty for local, or db://, s3://, az://, gs:// - api_key: "" - region: "" - -embeddings: - provider: ollama - model: qwen3-embedding:4b - vector_dim: 2560 - -reranking: - provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm - model: "" - -qa: - provider: ollama - model: gpt-oss - -research: - provider: "" # Empty to use qa settings - model: "" - max_iterations: 3 - confidence_threshold: 0.8 - max_concurrency: 1 - -search: - vector_index_metric: cosine # cosine, l2, or dot - vector_refine_factor: 30 - -agui: - host: "0.0.0.0" - port: 8000 - cors_origins: ["*"] - cors_credentials: true - cors_methods: ["GET", "POST", "OPTIONS"] - cors_headers: ["*"] - -processing: - converter: docling-local # docling-local or docling-serve - chunker: docling-local # docling-local or docling-serve - chunker_type: hybrid # hybrid or hierarchical - chunk_size: 256 - context_chunk_radius: 0 - chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" - chunking_merge_peers: true - chunking_use_markdown_tables: false - markdown_preprocessor: "" - conversion_options: - do_ocr: true - force_ocr: false - ocr_lang: [] - do_table_structure: true - table_mode: accurate - table_cell_matching: true - images_scale: 2.0 - -providers: - ollama: - base_url: http://localhost:11434 - - vllm: - embeddings_base_url: "" - rerank_base_url: "" - qa_base_url: "" - research_base_url: "" - - docling_serve: - base_url: http://localhost:5001 - api_key: "" - timeout: 300 -``` - -## Programmatic Configuration - -When using haiku.rag as a Python library, you can pass configuration directly to the `HaikuRAG` client: - -```python -from haiku.rag.config import AppConfig -from haiku.rag.client import HaikuRAG - -# Create custom configuration -custom_config = AppConfig( - qa={"provider": "openai", "model": "gpt-4o"}, - embeddings={"provider": "ollama", "model": "qwen3-embedding:4b", "vector_dim": 2560}, - processing={"chunk_size": 512} -) - -# 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 multiple clients with different configurations - -## File Monitoring - -Set directories to monitor for automatic indexing: - -```yaml -monitor: - directories: - - /path/to/documents - - /another_path/to/documents -``` - -### Filtering Monitored Files - -Use gitignore-style patterns to control which files are monitored: - -```yaml -monitor: - directories: - - /path/to/documents - - # Exclude specific files or directories - ignore_patterns: - - "*draft*" # Ignore files with "draft" in the name - - "temp/" # Ignore temp directory - - "**/archive/**" # Ignore all archive directories - - "*.backup" # Ignore backup files - - # Only include specific files (whitelist mode) - include_patterns: - - "*.md" # Only markdown files - - "*.pdf" # Only PDF files - - "**/docs/**" # Only files in docs directories -``` - -**How patterns work:** - -1. **Extension filtering** - Only supported file types are considered -2. **Include patterns** - If specified, only matching files are included (whitelist) -3. **Ignore patterns** - Matching files are excluded (blacklist) -4. **Combining both** - Include patterns are applied first, then ignore patterns - -**Common patterns:** - -```yaml -# Only monitor markdown documentation, but ignore drafts -monitor: - include_patterns: - - "*.md" - ignore_patterns: - - "*draft*" - - "*WIP*" - -# Monitor all supported files except in specific directories -monitor: - ignore_patterns: - - "node_modules/" - - ".git/" - - "**/test/**" - - "**/temp/**" -``` - -Patterns follow [gitignore syntax](https://git-scm.com/docs/gitignore#_pattern_format): - -- `*` matches anything except `/` -- `**` matches zero or more directories -- `?` matches any single character -- `[abc]` matches any character in the set - -## Document Processing - -Configure how documents are converted and chunked: - -```yaml -processing: - # Chunking configuration - chunk_size: 256 # Maximum tokens per chunk - context_chunk_radius: 0 # Context radius for chunk expansion - markdown_preprocessor: "" # Optional preprocessor script - - # Converter selection - converter: docling-local # docling-local or docling-serve - - # Chunker selection and configuration - chunker: docling-local # docling-local or docling-serve - chunker_type: hybrid # hybrid or hierarchical - chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" # HuggingFace model for tokenization - chunking_merge_peers: true # Merge undersized successive chunks - chunking_use_markdown_tables: false # Use markdown tables vs narrative format - - # Conversion options (works with both local and remote converters) - conversion_options: - # OCR settings - do_ocr: true # Enable OCR for bitmap content - force_ocr: false # Replace existing text with OCR - ocr_lang: [] # OCR languages (e.g., ["en", "fr", "de"]) - - # Table extraction - do_table_structure: true # Extract table structure - table_mode: accurate # fast or accurate - table_cell_matching: true # Match table cells back to PDF cells - - # Image settings - images_scale: 2.0 # Image scale factor -``` - -### Conversion Options - -The `conversion_options` section allows fine-grained control over document conversion. These options work with both `docling-local` and `docling-serve` converters. - -#### OCR Settings - -```yaml -conversion_options: - do_ocr: true # Enable OCR for bitmap/scanned content - force_ocr: false # Replace all text with OCR output - ocr_lang: [] # List of OCR languages, e.g., ["en", "fr", "de"] -``` - -- **do_ocr**: When `true`, applies OCR to images and scanned pages. Disable for faster processing if documents contain only native text. -- **force_ocr**: When `true`, replaces existing text layers with OCR output. Useful for documents with poor text extraction. -- **ocr_lang**: List of language codes for OCR. Empty list uses default language detection. Examples: `["en"]`, `["en", "fr", "de"]`. - -#### Table Extraction - -```yaml -conversion_options: - do_table_structure: true # Extract structured table data - table_mode: accurate # fast or accurate - table_cell_matching: true # Match cells back to PDF -``` - -- **do_table_structure**: When `true`, extracts table structure. Disable for faster processing if tables aren't important. -- **table_mode**: - - `accurate`: Better table structure recognition (slower) - - `fast`: Faster processing with simpler table detection -- **table_cell_matching**: When `true`, matches detected table cells back to PDF cells. Disable if tables have merged cells across columns. - -#### Image Settings - -```yaml -conversion_options: - images_scale: 2.0 # Image resolution scale factor -``` - -- **images_scale**: Scale factor for extracted images. Higher values = better quality but larger size. Typical range: 1.0-3.0. - -### Local vs Remote Processing - -**Local processing** (default): - -- Uses `docling` library locally -- No external dependencies -- Good for development and small workloads - -**Remote processing** (docling-serve): - -- Offloads processing to docling-serve API -- Better for heavy workloads and production -- Requires docling-serve instance (see [Remote processing setup](remote-processing.md)) - -To use remote processing: - -```yaml -processing: - converter: docling-serve - chunker: docling-serve - -providers: - docling_serve: - base_url: http://localhost:5001 - api_key: "your-api-key" # Optional - timeout: 300 # Request timeout in seconds -``` - -Conversion options work identically for both local and remote processing. - -### Chunking Strategies - -**Hybrid chunking** (default): -- Structure-aware chunking -- Respects document boundaries -- Best for most use cases - -**Hierarchical chunking**: -- Creates hierarchical chunk structure -- Preserves document hierarchy -- Useful for complex documents - -### Table Serialization - -Control how tables are represented in chunks: - -```yaml -processing: - chunking_use_markdown_tables: false # Default: narrative format -``` - -- `false`: Tables as narrative text ("Value A, Column 2 = Value B") -- `true`: Tables as markdown (preserves table structure) - -## Embedding Providers - -If you use Ollama, you can use any pulled model that supports embeddings. - -### Ollama (Default) - -```yaml -embeddings: - provider: ollama - model: mxbai-embed-large - vector_dim: 1024 -``` - -The Ollama base URL can be configured in your config file or via environment variable: - -```yaml -providers: - ollama: - base_url: http://localhost:11434 -``` - -Or via environment variable: - -```bash -export OLLAMA_BASE_URL=http://localhost:11434 -``` - -If not configured, it defaults to `http://localhost:11434`. - -!!! note - You can use a `.env` file in your project directory to set environment variables like `OLLAMA_BASE_URL` and API keys (e.g., `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`). These will be automatically loaded when running `haiku-rag` commands. - -### VoyageAI - -If you installed `haiku.rag` (full package), VoyageAI is already included. If you installed `haiku.rag-slim`, install with VoyageAI extras: - -```bash -uv pip install haiku.rag-slim[voyageai] -``` - -```yaml -embeddings: - provider: voyageai - model: voyage-3.5 - vector_dim: 1024 -``` - -Set your API key via environment variable: - -```bash -export VOYAGE_API_KEY=your-api-key -``` - -### OpenAI - -OpenAI embeddings are included in the default installation: - -```yaml -embeddings: - provider: openai - model: text-embedding-3-small # or text-embedding-3-large - vector_dim: 1536 -``` - -Set your API key via environment variable: - -```bash -export OPENAI_API_KEY=your-api-key -``` - -### vLLM - -For high-performance local inference, you can use vLLM to serve embedding models with OpenAI-compatible APIs: - -```yaml -embeddings: - provider: vllm - model: mixedbread-ai/mxbai-embed-large-v1 - vector_dim: 512 - -providers: - vllm: - embeddings_base_url: http://localhost:8000 -``` - -**Note:** You need to run a vLLM server separately with an embedding model loaded. - -## Question Answering Providers - -Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used. - -### Ollama (Default) - -```yaml -qa: - provider: ollama - model: gpt-oss -``` - -The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`: - -```bash -export OLLAMA_BASE_URL=http://localhost:11434 -``` - -Or in your config file: - -```yaml -providers: - ollama: - base_url: http://localhost:11434 -``` - -### OpenAI - -OpenAI QA is included in the default installation: - -```yaml -qa: - provider: openai - model: gpt-4o-mini # or gpt-4, gpt-3.5-turbo, etc. -``` - -Set your API key via environment variable: - -```bash -export OPENAI_API_KEY=your-api-key -``` - -### Anthropic - -Anthropic QA is included in the default installation: - -```yaml -qa: - provider: anthropic - model: claude-3-5-haiku-20241022 # or claude-3-5-sonnet-20241022, etc. -``` - -Set your API key via environment variable: - -```bash -export ANTHROPIC_API_KEY=your-api-key -``` - -### vLLM - -For high-performance local inference: - -```yaml -qa: - provider: vllm - model: Qwen/Qwen3-4B # Any model with tool support in vLLM - -providers: - vllm: - qa_base_url: http://localhost:8002 -``` - -**Note:** You need to run a vLLM server separately with a model that supports tool calling loaded. Consult the specific model's documentation for proper vLLM serving configuration. - -### Other Providers - -Any provider supported by Pydantic AI can be used. Examples: - -```yaml -# Google Gemini -qa: - provider: gemini - model: gemini-1.5-flash - -# Groq -qa: - provider: groq - model: llama-3.3-70b-versatile - -# Mistral -qa: - provider: mistral - model: mistral-small-latest -``` - -See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the complete list of supported providers and models. - -## Reranking - -Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (10x the requested limit) and then reranks them to return the most relevant results. - -Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below. - -### MixedBread AI - -If you installed `haiku.rag` (full package), MxBAI is already included. If you installed `haiku.rag-slim`, add the mxbai extra: - -```bash -uv pip install haiku.rag-slim[mxbai] -``` - -Then configure: - -```yaml -reranking: - provider: mxbai - model: mixedbread-ai/mxbai-rerank-base-v2 -``` - -### Cohere - -If you installed `haiku.rag` (full package), Cohere is already included. If you installed `haiku.rag-slim`, add the cohere extra: - -```bash -uv pip install haiku.rag-slim[cohere] -``` - -Then configure: - -```yaml -reranking: - provider: cohere - model: rerank-v3.5 -``` - -Set your API key via environment variable: - -```bash -export CO_API_KEY=your-api-key -``` - -### Zero Entropy - -If you installed `haiku.rag` (full package), Zero Entropy is already included. If you installed `haiku.rag-slim`, add the zeroentropy extra: - -```bash -uv pip install haiku.rag-slim[zeroentropy] -``` - -Then configure: - -```yaml -reranking: - provider: zeroentropy - model: zerank-1 # Currently the only available model -``` - -Set your API key via environment variable: - -```bash -export ZEROENTROPY_API_KEY=your-api-key -``` - -### vLLM - -For high-performance local reranking using dedicated reranking models: - -```yaml -reranking: - provider: vllm - model: mixedbread-ai/mxbai-rerank-base-v2 - -providers: - vllm: - rerank_base_url: http://localhost:8001 -``` - -**Note:** vLLM reranking uses the `/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded. Consult the specific model's documentation for proper vLLM serving configuration. - -## Research Configuration - -Configure the multi-agent research workflow: - -```yaml -research: - provider: "" # Empty to use qa settings - model: "" # Empty to use qa model - max_iterations: 3 # Maximum search/evaluate cycles - confidence_threshold: 0.8 # Stop when confidence meets/exceeds this - max_concurrency: 1 # Sub-questions searched in parallel per iteration -``` - -- **provider/model**: LLM provider and model for research. Leave empty to use the same settings as `qa`. -- **max_iterations**: Maximum number of search/evaluate cycles before stopping (default: 3) -- **confidence_threshold**: Stop research when evaluation confidence score meets or exceeds this threshold (default: 0.8) -- **max_concurrency**: Number of sub-questions to search in parallel during each iteration (default: 1) - -The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations. - -## AG-UI Server Configuration - -Configure the AG-UI HTTP server for streaming graph execution events: - -```yaml -agui: - host: "0.0.0.0" - port: 8000 - cors_origins: ["*"] - cors_credentials: true - cors_methods: ["GET", "POST", "OPTIONS"] - cors_headers: ["*"] -``` - -Start the AG-UI server with: - -```bash -haiku-rag serve --agui -``` - -The server exposes: -- `GET /health` - Health check endpoint -- `POST /v1/agent/stream` - Research graph streaming endpoint (Server-Sent Events) - -See [Server Mode](server.md) for more details. - -## Other Settings - -### Database and Storage - -By default, `haiku.rag` uses a local LanceDB database: - -```yaml -storage: - data_dir: /path/to/data # Empty = use default platform location -``` - -For remote storage, use the `lancedb` settings with various backends: - -```yaml -# LanceDB Cloud -lancedb: - uri: db://your-database-name - api_key: your-api-key - region: us-west-2 # optional - -# Amazon S3 -lancedb: - uri: s3://my-bucket/my-table -# Use AWS credentials or IAM roles - -# Azure Blob Storage -lancedb: - uri: az://my-container/my-table -# Use Azure credentials - -# Google Cloud Storage -lancedb: - uri: gs://my-bucket/my-table -# Use GCP credentials - -# HDFS -lancedb: - uri: hdfs://namenode:port/path/to/table -``` - -Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `api_key` for LanceDB Cloud. - -**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization is still performed locally. - -#### Database Auto-creation - -haiku.rag intelligently handles database creation based on operation type: - -- **Write operations** (add, add-src, delete, rebuild): Automatically create the database and required tables if they don't exist -- **Read operations** (list, get, search, ask, research): Fail with a clear error if the database doesn't exist - -This prevents the common mistake where a search query accidentally creates an empty database. To initialize your database, simply add your first document using `haiku-rag add` or `haiku-rag add-src`. - -### Vector Indexing - -Configure vector indexing behavior for efficient similarity search: - -```yaml -search: - vector_index_metric: cosine # cosine, l2, or dot - vector_refine_factor: 30 # Re-ranking factor for accuracy -``` - -- **vector_index_metric**: Distance metric for vector similarity: - - `cosine`: Cosine similarity (default, best for most embeddings) - - `l2`: Euclidean distance - - `dot`: Dot product similarity -- **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30 - - **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results - -!!! note - Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets. - -**Index creation:** - -Vector indexes are **not created automatically** during document ingestion to avoid slowing down the process. After you've added documents (at least 256 chunks required), create the index manually: - -```bash -haiku-rag create-index -``` - -This command: -- Checks if you have enough data (minimum 256 chunks) -- Creates an IVF_PQ index for fast approximate nearest neighbor (ANN) search -- Uses LanceDB's automatic parameter calculation based on your dataset size and vector dimensions - -**Re-indexing:** - -Indexes are not automatically updated when you add new documents. After adding a significant amount of new data: - -```bash -haiku-rag create-index # Rebuilds the index with all data -``` - -Searches still work with stale indexes - LanceDB uses the index for old data (fast ANN) and brute-force kNN for new unindexed rows, then combines the results. However, performance degrades as more unindexed data accumulates. - -For datasets with fewer than 256 chunks, searches use brute-force kNN scans (exact nearest neighbors, 100% recall) which work well for small datasets but don't scale beyond a few hundred thousand vectors. - -### Document Processing - -```yaml -processing: - # Chunk size for document processing - chunk_size: 256 - - # Number of adjacent chunks to include before/after retrieved chunks for context - # 0 = no expansion (default), 1 = include 1 chunk before and after, etc. - # When expanded chunks overlap or are adjacent, they are automatically merged - # into single chunks with continuous content to eliminate duplication - context_chunk_radius: 0 - - # Optional dotted path or file path to a callable that preprocesses - # markdown content before chunking - markdown_preprocessor: "" - -storage: - # Vacuum retention threshold (seconds) for automatic cleanup - # When documents are added/updated, old table versions older than this are removed - # Default: 86400 seconds (1 day, safe for concurrent connections) - # Set to 0 for aggressive cleanup (removes all old versions immediately) - vacuum_retention_seconds: 86400 -``` - -#### Markdown Preprocessor - -Optionally preprocess Markdown before chunking by pointing to a callable that receives and returns Markdown text. This is useful for normalizing content, stripping boilerplate, or applying custom transformations before chunk boundaries are computed. - -```yaml -processing: - # A callable path in one of these formats: - # - package.module:func - # - package.module.func - # - /abs/or/relative/path/to/file.py:func - markdown_preprocessor: my_pkg.preprocess:clean_md -``` - -!!! note - - The function signature should be `def clean_md(text: str) -> str` or `async def clean_md(text: str) -> str`. - - If the function raises or returns a non-string, haiku.rag logs a warning and proceeds without preprocessing. - - The preprocessor affects only the chunking pipeline. The stored document content remains unchanged. - -Example implementation: - -```python -# my_pkg/preprocess.py -def clean_md(text: str) -> str: - # strip HTML comments and collapse multiple blank lines - lines = [line for line in text.splitlines() if not line.strip().startswith("