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

Add support for per-model configuration settings including thinking, temperature and max_tokens

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-25 11:58:57 +02:00
parent bfcbbbb91f
commit 1bb7b6d5bd
No known key found for this signature in database
30 changed files with 597 additions and 197 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

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

16
evaluations/evaluations/benchmark.py
View file

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

14
evaluations/evaluations/evaluators/judge.py
View file

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

2
examples/a2a-server/haiku_rag_a2a/a2a/__init__.py
View file

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

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

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

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

@ -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(),
}

48
haiku_rag_slim/haiku/rag/config/models.py
View file

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

20
haiku_rag_slim/haiku/rag/embeddings/__init__.py
View file

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

2
haiku_rag_slim/haiku/rag/embeddings/base.py
View file

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

2
haiku_rag_slim/haiku/rag/graph/common/__init__.py
View file

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

24
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

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

48
haiku_rag_slim/haiku/rag/graph/common/utils.py
View file

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

13
haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py
View file

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

15
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

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

6
haiku_rag_slim/haiku/rag/qa/__init__.py
View file

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

30
haiku_rag_slim/haiku/rag/qa/agent.py
View file

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

12
haiku_rag_slim/haiku/rag/reranking/__init__.py
View file

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

4
haiku_rag_slim/haiku/rag/reranking/base.py
View file

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

3
haiku_rag_slim/haiku/rag/reranking/cohere.py
View file

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

7
haiku_rag_slim/haiku/rag/reranking/mxbai.py
View file

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

3
haiku_rag_slim/haiku/rag/reranking/zeroentropy.py
View file

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

246
haiku_rag_slim/haiku/rag/utils.py
View file

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

10
tests/graph/test_deep_qa.py
View file

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

13
tests/graph/test_research_graph.py
View file

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

1
tests/test_chunk.py
View file

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

8
tests/test_client.py
View file

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

29
tests/test_embedder_config.py
View file

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

17
tests/test_qa.py
View file

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

8
tests/test_reranker.py
View file

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

181
tests/test_utils.py
View file

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