From 1bb7b6d5bdb71a8e7b0a8edcf6c23e64164587c4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 25 Nov 2025 11:58:57 +0200 Subject: [PATCH] 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)