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

Merge pull request #519 from ggozad/chore/update-pydantic-ai

Browse source 
Update pydantic-ai to 2.18 and adopt its unified thinking setting and ToolFailed
This commit is contained in:
Yiorgis Gozadinos 2026-07-28 14:32:50 +03:00 committed by GitHub
commit 47015758a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 325 additions and 221 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.claude/skills/debug-evals/SKILL.md
View file

@ -31,7 +31,7 @@ A run is one experiment span; its cases are direct children sharing its
- `attributes->>'name'` — run label (the `--name` arg, or `{dataset}_qa_evaluation` / `{dataset}_retrieval_evaluation`).
- `attributes->>'dataset_name'` — dataset.
- `(attributes->>'assertion_pass_rate')::float` — overall judge pass rate (QA runs).
- `attributes->'logfire.experiment.metadata'->'metadata'` — run config: `target` (`rag-skill`|`analysis-skill`), `qa_model`, `embedder_model`, `chunk_size`, `search_limit`, `rerank_model`, `judge_model`, `skill_model`, etc.
- `attributes->'logfire.experiment.metadata'->'metadata'` — run config: `target` (`rag-capability`|`analysis-capability`), `qa_model`, `embedder_model`, `chunk_size`, `search_limit`, `rerank_model`, `judge_model`, `qa_max_searches`, etc.
- `trace_id` — scopes the whole run.
- Case span: `span_name = 'case: {case_name}'` (scope `pydantic-evals`).
- `message` — `case: <id>`.
@ -39,8 +39,8 @@ A run is one experiment span; its cases are direct children sharing its
- `attributes->'scores'->'cited_map'->>'value'` — citation average precision (0..1).
- `attributes->'scores'->'number_match'->>'value'` — numeric-answer match (datasets that use it).
- `duration` — task time in seconds.
- Inside each case the skill under test emits agent spans (scope `pydantic-ai`):
`execute {task}`, `agent run`, `running tool`, `chat {model}`.
- Inside each case the capability under test emits agent spans (scope `pydantic-ai`):
`execute {task}`, `invoke_agent agent`, `execute_tool {tool_name}`, `chat {model}`.
The service is `evals` regardless of model, so filter on `service_name = 'evals'`
first. `otel_scope_name` separates the layers (`pydantic-evals` for run/case,

8
CHANGELOG.md
View file

@ -1,10 +1,18 @@
# Changelog
## [Unreleased]
### Changed
- Require `pydantic-ai-slim>=2.18,<3`.
- `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096, and `max_tokens` must exceed it on budget-based Claude models; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`; Bedrock Qwen with `enable_thinking: false` no longer sends `reasoning_config`, while Bedrock-served Claude keeps an explicit `thinking: disabled`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`.
- `provider: bedrock` with a proprietary OpenAI model such as `openai.o3-mini-v1:0` raises `UserError`: Bedrock Converse serves only the `gpt-oss` family. Use `provider: bedrock-mantle`.
- Tool failures raise `pydantic_ai.ToolFailed` instead of returning failure text: search and code-execution limits, sandbox execution errors, and `get_document`/`summarize_document` misses.
### Fixed
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
- `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild.
- `evaluations run` opens the database read-only outside the population phase, so an embedder identity differing from the stored one warns instead of aborting the run.
### Removed

2
app/backend/pyproject.toml
View file

@ -6,7 +6,7 @@ requires-python = ">=3.12"
dependencies = [
"starlette>=0.50.0",
"uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.11.0,<3.0.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0",
"python-dotenv>=1.2.1",
"haiku.rag-slim>=0.70.0",
"logfire[pydantic-ai]>=3.17.0",

5
docs/configuration/providers.md
View file

@ -54,7 +54,7 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- **Anthropic**: All Claude models
- **Google**: Gemini models with thinking support
- **Groq**: Models with reasoning capabilities
- **Bedrock**: Claude, OpenAI, and Qwen models
- **Bedrock**: Claude, Qwen, and `gpt-oss` models. Bedrock Converse does not serve the proprietary OpenAI models, so configuring one raises an error. Reach those through `provider: bedrock-mantle`.
- **Ollama**: Models supporting reasoning (gpt-oss, etc.)
- **vLLM**: Models with a pydantic-ai reasoning profile (gpt-oss). Qwen3, Gemma, and similar templates ignore the OpenAI `reasoning_effort` that `enable_thinking` translates to — use [`extra_body`](#raw-provider-pass-through) to drive them.
- **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
@ -63,6 +63,9 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- Enable for QA, complex reasoning, and mathematical problems
- Disable for speed-critical applications, title generation, and simple tasks
!!! note "Anthropic thinking and max_tokens"
Anthropic requires `max_tokens` to exceed the thinking budget, and `enable_thinking: true` requests Pydantic AI's default budget of 10000 tokens. Set `max_tokens` above 10000 on Claude models that use budget-based thinking, or leave it unset on Sonnet 4.6+ and Opus 4.6+, which use adaptive thinking instead of a budget.
!!! note "vLLM-served models without a reasoning profile"
On `provider: openai` with a custom `base_url`, `enable_thinking` only takes effect for models whose pydantic-ai profile advertises reasoning support (o-series, gpt-5, gpt-oss). For other vLLM-served models (Qwen3, Gemma family, …) the field is a silent no-op. Reach the chat template's thinking switch directly via [`extra_body`](#raw-provider-pass-through).

2
evaluations/evaluations/benchmark.py
View file

@ -227,7 +227,7 @@ async def run_retrieval_benchmark(
)
db = spec.db_path(db_path)
async with HaikuRAG(db, config=config) as rag:
async with HaikuRAG(db, config=config, read_only=True) as rag:
async def retrieval_target(question: str) -> list[str]:
chunks = await rag.search(query=question, limit=5)

22
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -5,7 +5,7 @@ from pathlib import Path
from typing import Any, cast
from pydantic import BaseModel
from pydantic_ai import ModelRetry, RunContext
from pydantic_ai import ModelRetry, RunContext, ToolFailed
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import (
InstructionPart,
@ -198,16 +198,21 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
self.outer_state[self.state_namespace] = self.state.model_dump(mode="json")
async def _with_state(self, operation: Any) -> Any:
"""Execute an operation and copy its state back to the host dependencies."""
result = await operation
self._sync_state()
return result
"""Execute an operation and copy its state back to the host dependencies.
A failing tool still syncs, so evidence it gathered before the failure
reaches the host.
"""
try:
return await operation
finally:
self._sync_state()
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
assert self.state is not None
self.search_count += 1
if self.search_count > self.config.qa.max_searches:
return (
raise ToolFailed(
"Search limit reached. Answer the question using "
"the results you already have."
)
@ -227,7 +232,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
async def _cite(self, chunk_ids: list[str]) -> str:
assert self.state is not None
if not chunk_ids:
return "Registered 0 citations (empty chunk_ids)."
raise ModelRetry(
"No citations registered: chunk_ids was empty. Pass the chunk_ids "
"you want to cite, copied verbatim from search results."
)
all_results: list[SearchResult] = []
state = cast(Any, self.state)

10
haiku_rag_slim/haiku/rag/capabilities/analysis.py
View file

@ -4,7 +4,7 @@ from pathlib import Path
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from pydantic_ai import RunContext, ToolFailed
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
@ -74,7 +74,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
assert self.state is not None
self.execute_count += 1
if self.execute_count > self.config.analysis.max_executions:
return (
raise ToolFailed(
"Code-execution limit reached. Give your final answer now from what "
"you already have; do not call analysis_execute_code again."
)
@ -96,9 +96,9 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
success=result.success,
)
)
if result.success:
return result.stdout or "No output."
return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
if not result.success:
raise ToolFailed(f"{result.stderr}\n\nOutput: {result.stdout}")
return result.stdout or "No output."
def get_toolset(self) -> FunctionToolset[Any]:
async def analysis_search(

10
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -1,5 +1,5 @@
from pydantic import BaseModel
from pydantic_ai import Agent, FunctionToolset, RunContext
from pydantic_ai import Agent, FunctionToolset, RunContext, ToolFailed
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
@ -123,14 +123,14 @@ def create_document_toolset(
query: The document title or URI to look up.
Returns:
Document content and metadata, or not found message.
Document content and metadata.
"""
client = ctx.deps.client
doc = await find_document(client, query)
if doc is None:
return f"Document not found: {query}"
raise ToolFailed(f"Document not found: {query}")
return (
f"**{doc.title or 'Untitled'}**\n\n"
@ -147,14 +147,14 @@ def create_document_toolset(
query: The document title or URI to summarize.
Returns:
Generated summary or not found message.
Generated summary.
"""
client = ctx.deps.client
doc = await find_document(client, query)
if doc is None:
return f"Document not found: {query}"
raise ToolFailed(f"Document not found: {query}")
summary_model = get_model(config.qa.model, config)
summary_agent: Agent[None, str] = Agent(

9
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -3,7 +3,7 @@ from collections.abc import Callable
from io import BytesIO
from PIL import Image
from pydantic_ai import FunctionToolset, RunContext
from pydantic_ai import FunctionToolset, RunContext, ToolFailed
from pydantic_ai.messages import BinaryContent, ToolReturn
from haiku.rag.config.models import AppConfig
@ -68,8 +68,9 @@ def create_search_toolset(
tool_name: Name for the search tool. Defaults to "search".
on_results: Optional callback invoked with search results after each search.
Useful for accumulating results externally (e.g., for citation resolution).
max_searches: Maximum number of searches allowed. When exceeded, returns
a message directing the agent to answer with existing results.
max_searches: Maximum number of searches allowed. When exceeded, the
tool fails with a message directing the agent to answer with
existing results.
Returns:
FunctionToolset with a search tool.
@ -99,7 +100,7 @@ def create_search_toolset(
rid = ctx.run_id or ""
search_counts[rid] = search_counts.get(rid, 0) + 1
if max_searches is not None and search_counts[rid] > max_searches:
return (
raise ToolFailed(
"Search limit reached. "
"Answer the question using the results you already have."
)

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

@ -55,30 +55,35 @@ def image_binary_content(data: bytes) -> "BinaryContent":
def apply_common_settings(
settings: Any | None,
settings_class: type[Any],
model_config: Any,
*,
map_thinking: bool = True,
) -> Any | None:
"""Apply common settings (temperature, max_tokens) to model settings.
"""Apply the settings every provider shares onto a model settings dict.
Args:
settings: Existing settings instance or None
settings_class: Settings class to instantiate if needed
model_config: ModelConfig with temperature and max_tokens
map_thinking: Whether to map `enable_thinking` onto the unified
`thinking` setting. The OpenAI-compatible branches opt out and set
`openai_reasoning_effort` themselves, so that models whose profile
advertises thinking without OpenAI reasoning support (Ollama's
deepseek-r1, for one) keep receiving no `reasoning_effort`.
Returns:
Updated settings instance or None if no settings to apply
"""
thinking = model_config.enable_thinking if map_thinking else None
if (
model_config.temperature is None
and model_config.max_tokens is None
and model_config.extra_body is None
and thinking is None
):
return settings
if settings is None:
settings_dict = settings_class()
else:
settings_dict = settings
settings_dict = {} if settings is None else settings
if model_config.temperature is not None:
settings_dict["temperature"] = model_config.temperature
@ -89,6 +94,9 @@ def apply_common_settings(
if model_config.extra_body is not None:
settings_dict["extra_body"] = model_config.extra_body
if thinking is not None:
settings_dict["thinking"] = thinking
return settings_dict
@ -139,7 +147,7 @@ def get_model(
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
model_settings = apply_common_settings(
model_settings, OpenAIChatModelSettings, model_config
model_settings, model_config, map_thinking=False
)
# Ollama's OpenAI-compatible API lives under /v1. Append it if the
@ -173,7 +181,7 @@ def get_model(
)
openai_settings = apply_common_settings(
openai_settings, OpenAIChatModelSettings, model_config
openai_settings, model_config, map_thinking=False
)
# Use model-level base_url if set (for vLLM, LM Studio, etc.)
@ -188,75 +196,42 @@ def get_model(
return OpenAIChatModel(model_name=model, settings=openai_settings)
elif provider == "anthropic":
from anthropic.types.beta import (
BetaThinkingConfigDisabledParam,
BetaThinkingConfigEnabledParam,
)
from anthropic.types.beta import BetaThinkingConfigDisabledParam
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:
thinking_config: BetaThinkingConfigEnabledParam = {
"type": "enabled",
"budget_tokens": 4096,
}
anthropic_settings = AnthropicModelSettings(
anthropic_thinking=thinking_config
)
else:
thinking_disabled: BetaThinkingConfigDisabledParam = {
"type": "disabled"
}
anthropic_settings = AnthropicModelSettings(
anthropic_thinking=thinking_disabled
)
# Unified `thinking=False` omits the request field, which leaves the
# adaptive-thinking models (Sonnet 4.6+, Opus 4.6+) thinking by default.
disable_thinking = model_config.enable_thinking is False
if disable_thinking:
thinking_disabled: BetaThinkingConfigDisabledParam = {"type": "disabled"}
anthropic_settings = AnthropicModelSettings(
anthropic_thinking=thinking_disabled
)
anthropic_settings = apply_common_settings(
anthropic_settings, AnthropicModelSettings, model_config
anthropic_settings, model_config, map_thinking=not disable_thinking
)
return AnthropicModel(model_name=model, settings=anthropic_settings)
elif provider == "gemini":
from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
from pydantic_ai.models.google import GoogleModel
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=apply_common_settings(None, model_config),
)
return GoogleModel(model_name=model, settings=gemini_settings)
elif provider == "groq":
from pydantic_ai.models.groq import GroqModel, GroqModelSettings
from pydantic_ai.models.groq import GroqModel
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=apply_common_settings(None, model_config),
)
return GroqModel(model_name=model, settings=groq_settings)
elif provider == "bedrock":
from pydantic_ai.models.bedrock import (
BedrockConverseModel,
@ -265,41 +240,25 @@ def get_model(
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 "o1" in model or "o3" in model:
# OpenAI reasoning models on Bedrock (o-series only, not gpt-4o)
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
# Same omission as the direct Anthropic branch: unified `thinking=False`
# leaves the adaptive-thinking Claude models thinking. Bedrock ids are
# `[<geo>.]<family>.<model>`, as in `us.anthropic.claude-...`.
disable_claude_thinking = (
model_config.enable_thinking is False and "anthropic." in model
)
if disable_claude_thinking:
bedrock_settings = BedrockModelSettings(
bedrock_additional_model_requests_fields={
"thinking": {"type": "disabled"}
}
)
return BedrockConverseModel(model_name=model, settings=bedrock_settings)
return BedrockConverseModel(
model_name=model,
settings=apply_common_settings(
bedrock_settings, model_config, map_thinking=not disable_claude_thinking
),
)
else:
# For any other provider, use string format and let Pydantic AI handle it

2
haiku_rag_slim/pyproject.toml
View file

@ -29,7 +29,7 @@ dependencies = [
"lancedb==0.34.0",
"pathspec>=1.0.4",
"pydantic>=2.12.5",
"pydantic-ai-slim[openai,logfire,ag-ui]>=2.11.0,<3.0.0",
"pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0",
"pydantic-monty>=0.0.17",
"pypdfium2>=5.0",
"python-dotenv>=1.2.2",

81
tests/capabilities/test_capabilities.py
View file

@ -4,7 +4,7 @@ from typing import Any, cast
from unittest.mock import AsyncMock, patch
import pytest
from pydantic_ai import Agent, RunContext
from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
@ -291,13 +291,11 @@ async def test_search_and_empty_citation_limits(temp_db_path):
capability = create_rag(db_path=temp_db_path, config=config)
capability.state = RAGState()
result = await capability._search("anything", None)
with pytest.raises(ToolFailed, match="Search limit reached"):
await capability._search("anything", None)
assert (
result
== "Search limit reached. Answer the question using the results you already have."
)
assert await capability._cite([]) == "Registered 0 citations (empty chunk_ids)."
with pytest.raises(ModelRetry, match="chunk_ids was empty"):
await capability._cite([])
@pytest.mark.asyncio
@ -372,6 +370,75 @@ async def test_analysis_records_new_sandbox_search_results(temp_db_path):
]
@pytest.mark.asyncio
async def test_failed_tool_reaches_the_model_and_the_run_continues(temp_db_path):
"""A `ToolFailed` tool leaves a failed result in history and answers anyway."""
config = AppConfig()
config.qa.max_searches = 0
calls = 0
def model_function(_messages, _info):
nonlocal calls
calls += 1
if calls == 1:
return ModelResponse(parts=[ToolCallPart("rag_search", {"query": "x"})])
return ModelResponse(parts=[TextPart("answered from what I had")])
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
capabilities=[
create_rag(db_path=temp_db_path, config=config, defer_loading=False)
],
)
result = await agent.run("question", deps=Deps())
assert result.output == "answered from what I had"
failed = [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolReturnPart) and part.outcome == "failed"
]
assert [part.tool_name for part in failed] == ["rag_search"]
assert "Search limit reached" in str(failed[0].content)
@pytest.mark.asyncio
async def test_analysis_execution_limit_fails_the_tool(temp_db_path):
config = AppConfig()
config.analysis.max_executions = 0
capability = create_analysis(db_path=temp_db_path, config=config)
capability.state = AnalysisState()
with pytest.raises(ToolFailed, match="Code-execution limit reached"):
await capability._execute_code("print('done')")
@pytest.mark.asyncio
async def test_analysis_sandbox_failure_records_execution_and_fails_the_tool(
temp_db_path,
):
capability = create_analysis(db_path=temp_db_path, config=AppConfig())
capability.state = AnalysisState()
capability.outer_state = {}
sandbox = AsyncMock()
sandbox.execute.return_value = SandboxResult(
stdout="partial", stderr="NameError: undefined", success=False
)
sandbox._search_results = []
capability.sandbox = cast(Sandbox, sandbox)
with pytest.raises(ToolFailed, match="NameError: undefined"):
await capability._with_state(capability._execute_code("boom"))
entry = capability.state.executions[-1]
assert entry.success is False
assert entry.stderr == "NameError: undefined"
assert capability.outer_state["analysis"]["executions"][-1]["code"] == "boom"
@pytest.mark.asyncio
async def test_native_agent_composition_initializes_host_state(temp_db_path):
capability = create_rag(

0
tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_return_cap_message.yaml → tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_fail_the_tool.yaml
View file

183
tests/test_utils.py
View file

@ -237,6 +237,27 @@ def test_get_model_openai_non_reasoning_model_ignores_thinking():
assert result._settings is None
def test_get_model_vllm_model_without_reasoning_profile_sends_no_thinking():
"""A vLLM-served model with no reasoning profile carries no thinking settings.
Its chat template reads the switch from `chat_template_kwargs`, which only
`extra_body` can reach, and the endpoint rejects `reasoning_effort`.
"""
model_config = ModelConfig(
provider="openai",
name="Qwen/Qwen3-32B",
base_url="http://vllm:8000/v1",
enable_thinking=True,
temperature=0.2,
)
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
assert result._settings is not None
assert "thinking" not in result._settings
assert "openai_reasoning_effort" not in result._settings
def test_get_model_openai_extra_body_forwarded():
"""`extra_body` on ModelConfig is forwarded to ModelSettings.extra_body.
@ -328,27 +349,41 @@ def test_get_model_anthropic():
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
@pytest.mark.parametrize(
"enable_thinking,expected_thinking",
[
(True, {"type": "enabled", "budget_tokens": 4096}),
(False, {"type": "disabled"}),
],
)
def test_get_model_anthropic_with_thinking(enable_thinking, expected_thinking):
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",
name="claude-3-5-sonnet-20241022",
enable_thinking=enable_thinking,
enable_thinking=True,
)
result = get_model(model_config)
assert isinstance(result, AnthropicModel)
assert result.settings is not None
assert result.settings.get("anthropic_thinking") == expected_thinking
assert result.settings.get("thinking") is True
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
def test_get_model_anthropic_thinking_off_disables_adaptive_models():
"""Adaptive-thinking models think by default, so off must be explicit.
The unified `thinking=False` omits the request field, which leaves Sonnet
4.6+ and Opus 4.6+ thinking.
"""
from pydantic_ai.models.anthropic import AnthropicModel
model_config = ModelConfig(
provider="anthropic", name="claude-sonnet-4-6", enable_thinking=False
)
result = get_model(model_config)
assert isinstance(result, AnthropicModel)
assert result.settings is not None
assert result.settings.get("anthropic_thinking") == {"type": "disabled"}
# The explicit disable replaces the unified key rather than joining it.
assert "thinking" not in result.settings
@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
@ -362,15 +397,21 @@ def test_get_model_gemini():
@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
def test_get_model_gemini_with_thinking():
@pytest.mark.parametrize("enable_thinking", [True, False])
def test_get_model_gemini_with_thinking(enable_thinking):
"""Test get_model configures thinking for Gemini."""
from pydantic_ai.models.google import GoogleModel
model_config = ModelConfig(
provider="gemini", name="gemini-2.0-flash-thinking-exp", enable_thinking=True
provider="gemini",
name="gemini-2.0-flash-thinking-exp",
enable_thinking=enable_thinking,
)
result = get_model(model_config)
assert isinstance(result, GoogleModel)
assert result.settings is not None
assert result.settings.get("thinking") == enable_thinking
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
@ -384,11 +425,9 @@ def test_get_model_groq():
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
@pytest.mark.parametrize(
"enable_thinking,expected_format", [(True, "parsed"), (False, "hidden")]
)
def test_get_model_groq_with_thinking(enable_thinking, expected_format):
"""Test get_model configures thinking format for Groq."""
@pytest.mark.parametrize("enable_thinking", [True, False])
def test_get_model_groq_with_thinking(enable_thinking):
"""Test get_model configures thinking for Groq."""
from pydantic_ai.models.groq import GroqModel
model_config = ModelConfig(
@ -400,7 +439,7 @@ def test_get_model_groq_with_thinking(enable_thinking, expected_format):
assert isinstance(result, GroqModel)
assert result.settings is not None
assert result.settings.get("groq_reasoning_format") == expected_format
assert result.settings.get("thinking") == enable_thinking
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@ -417,57 +456,81 @@ def test_get_model_bedrock():
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@pytest.mark.parametrize(
"name,enable_thinking,expected_fields",
"name",
[
(
"anthropic.claude-3-5-sonnet-20241022-v2:0",
True,
{"thinking": {"type": "enabled", "budget_tokens": 4096}},
),
(
"anthropic.claude-3-5-sonnet-20241022-v2:0",
False,
{"thinking": {"type": "disabled"}},
),
("openai.o3-mini-v1:0", True, {"reasoning_effort": "high"}),
("openai.o3-mini-v1:0", False, {"reasoning_effort": "low"}),
("qwen.qwen3-32b-v1:0", True, {"reasoning_config": "high"}),
("qwen.qwen3-32b-v1:0", False, {"reasoning_config": "low"}),
# A family with no reasoning mapping leaves the request fields untouched.
("meta.llama3-70b-instruct-v1:0", True, None),
("meta.llama3-70b-instruct-v1:0", False, None),
],
ids=[
"claude_on",
"claude_off",
"o_series_on",
"o_series_off",
"qwen_on",
"qwen_off",
"unmapped_on",
"unmapped_off",
"anthropic.claude-3-5-sonnet-20241022-v2:0",
"openai.gpt-oss-120b-1:0",
"qwen.qwen3-32b-v1:0",
"meta.llama3-70b-instruct-v1:0",
],
ids=["claude", "gpt_oss", "qwen", "unmapped"],
)
def test_get_model_bedrock_with_thinking(name, enable_thinking, expected_fields):
"""Each Bedrock model family maps thinking onto its own request field."""
def test_get_model_bedrock_with_thinking(name):
"""Every Bedrock family carries the unified thinking setting."""
from pydantic_ai.models.bedrock import BedrockConverseModel
model_config = ModelConfig(
provider="bedrock",
name=name,
enable_thinking=enable_thinking,
enable_thinking=True,
)
result = get_model(model_config)
assert isinstance(result, BedrockConverseModel)
if expected_fields is None:
assert result.settings is None
return
assert result.settings is not None
assert (
result.settings.get("bedrock_additional_model_requests_fields")
== expected_fields
assert result.settings.get("thinking") is True
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@pytest.mark.parametrize(
"name",
[
"anthropic.claude-sonnet-4-6-20260514-v1:0",
"us.anthropic.claude-sonnet-4-6-20260514-v1:0",
],
ids=["plain", "cross_region"],
)
def test_get_model_bedrock_thinking_off_disables_adaptive_claude(name):
"""Bedrock omits the field for adaptive Claude, which leaves it thinking."""
from pydantic_ai.models.bedrock import BedrockConverseModel
model_config = ModelConfig(provider="bedrock", name=name, enable_thinking=False)
result = get_model(model_config)
assert isinstance(result, BedrockConverseModel)
assert result.settings is not None
assert result.settings.get("bedrock_additional_model_requests_fields") == {
"thinking": {"type": "disabled"}
}
# The explicit disable replaces the unified key rather than joining it.
assert "thinking" not in result.settings
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock_thinking_off_leaves_non_claude_families_alone():
"""Only the Anthropic variant takes a `thinking` request field."""
from pydantic_ai.models.bedrock import BedrockConverseModel
model_config = ModelConfig(
provider="bedrock", name="qwen.qwen3-32b-v1:0", enable_thinking=False
)
result = get_model(model_config)
assert isinstance(result, BedrockConverseModel)
assert result.settings is not None
assert "bedrock_additional_model_requests_fields" not in result.settings
assert result.settings.get("thinking") is False
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock_rejects_mantle_only_model():
"""Proprietary OpenAI models are Bedrock Mantle-only, not served by Converse."""
from pydantic_ai.exceptions import UserError
model_config = ModelConfig(provider="bedrock", name="openai.o3-mini-v1:0")
with pytest.raises(UserError):
get_model(model_config)
def test_get_model_unknown_provider():
@ -491,7 +554,7 @@ def test_get_package_versions():
assert "docling_document_schema" in versions
# All should be non-empty strings
for key, value in versions.items():
for value in versions.values():
assert isinstance(value, str)
assert len(value) > 0
@ -504,7 +567,7 @@ def test_apply_common_settings_no_settings():
from haiku.rag.utils import apply_common_settings
mc = ModelConfig(provider="openai", name="gpt-4o")
result = apply_common_settings(None, dict, mc)
result = apply_common_settings(None, mc)
assert result is None
@ -513,7 +576,7 @@ def test_apply_common_settings_temperature():
from haiku.rag.utils import apply_common_settings
mc = ModelConfig(provider="openai", name="gpt-4o", temperature=0.7)
result = apply_common_settings(None, dict, mc)
result = apply_common_settings(None, mc)
assert result is not None
assert result["temperature"] == 0.7
@ -523,7 +586,7 @@ def test_apply_common_settings_max_tokens():
from haiku.rag.utils import apply_common_settings
mc = ModelConfig(provider="openai", name="gpt-4o", max_tokens=500)
result = apply_common_settings(None, dict, mc)
result = apply_common_settings(None, mc)
assert result is not None
assert result["max_tokens"] == 500
@ -534,7 +597,7 @@ def test_apply_common_settings_existing():
mc = ModelConfig(provider="openai", name="gpt-4o", temperature=0.5)
existing = {"some_key": "value"}
result = apply_common_settings(existing, dict, mc)
result = apply_common_settings(existing, mc)
assert result is not None
assert result["temperature"] == 0.5
assert result["some_key"] == "value"

9
tests/tools/test_document.py
View file

@ -2,6 +2,7 @@ from pathlib import Path
from types import SimpleNamespace
import pytest
from pydantic_ai import ToolFailed
from haiku.rag.tools.document import (
DocumentInfo,
@ -130,9 +131,9 @@ class TestDocumentToolExecution:
get_tool = toolset.tools["get_document"]
ctx = make_ctx(doc_client)
result = await get_tool.function(ctx, "nonexistent")
assert "Document not found" in result
with pytest.raises(ToolFailed, match="Document not found: nonexistent"):
await get_tool.function(ctx, "nonexistent")
@pytest.mark.asyncio
async def test_list_documents_with_base_filter(self, doc_client, doc_config):
@ -183,9 +184,9 @@ class TestSummarizeDocumentTool:
summarize_tool = toolset.tools["summarize_document"]
ctx = make_ctx(doc_client)
result = await summarize_tool.function(ctx, "nonexistent document")
assert "Document not found" in result
with pytest.raises(ToolFailed, match="Document not found: nonexistent"):
await summarize_tool.function(ctx, "nonexistent document")
@pytest.mark.vcr()
@pytest.mark.asyncio

32
tests/tools/test_search.py
View file

@ -2,6 +2,7 @@ from pathlib import Path
from types import SimpleNamespace
import pytest
from pydantic_ai import ToolFailed
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
@ -139,26 +140,22 @@ class TestSearchMaxSearches:
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client)
result1 = await search_tool.function(ctx, "Python")
assert "Search limit reached" not in result1
result2 = await search_tool.function(ctx, "JavaScript")
assert "Search limit reached" not in result2
assert await search_tool.function(ctx, "Python")
assert await search_tool.function(ctx, "JavaScript")
@pytest.mark.asyncio
async def test_searches_beyond_limit_return_cap_message(
async def test_searches_beyond_limit_fail_the_tool(
self, search_client, search_config
):
"""Searches beyond max_searches return limit message."""
"""Searches beyond max_searches fail with the limit message."""
toolset = create_search_toolset(search_config, max_searches=1)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client)
result1 = await search_tool.function(ctx, "Python")
assert "Search limit reached" not in result1
assert await search_tool.function(ctx, "Python")
result2 = await search_tool.function(ctx, "JavaScript")
assert "Search limit reached" in result2
with pytest.raises(ToolFailed, match="Search limit reached"):
await search_tool.function(ctx, "JavaScript")
@pytest.mark.asyncio
async def test_counter_resets_across_runs(self, search_client, search_config):
@ -167,15 +164,13 @@ class TestSearchMaxSearches:
search_tool = toolset.tools["search"]
ctx_run1 = make_ctx(search_client, run_id="run-1")
result = await search_tool.function(ctx_run1, "Python")
assert "Search limit reached" not in result
assert await search_tool.function(ctx_run1, "Python")
result2 = await search_tool.function(ctx_run1, "JavaScript")
assert "Search limit reached" in result2
with pytest.raises(ToolFailed, match="Search limit reached"):
await search_tool.function(ctx_run1, "JavaScript")
ctx_run2 = make_ctx(search_client, run_id="run-2")
result3 = await search_tool.function(ctx_run2, "Python")
assert "Search limit reached" not in result3
assert await search_tool.function(ctx_run2, "Python")
@pytest.mark.asyncio
async def test_no_limit_by_default(self, search_client, search_config):
@ -185,8 +180,7 @@ class TestSearchMaxSearches:
ctx = make_ctx(search_client)
for _ in range(5):
result = await search_tool.function(ctx, "Python")
assert "Search limit reached" not in result
assert await search_tool.function(ctx, "Python")
@pytest.fixture

20
uv.lock
View file

@ -1768,7 +1768,7 @@ requires-dist = [
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'vertexai'" },
{ name = "pydantic-ai-slim", extras = ["groq"], marker = "extra == 'groq'" },
{ name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" },
{ name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=2.11.0,<3.0.0" },
{ name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=2.18.0,<3.0.0" },
{ name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" },
{ name = "pydantic-monty", specifier = ">=0.0.17" },
{ name = "pypdfium2", specifier = ">=5.0" },
@ -3826,7 +3826,7 @@ email = [
[[package]]
name = "pydantic-ai-slim"
version = "2.16.0"
version = "2.18.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "genai-prices" },
@ -3837,9 +3837,9 @@ dependencies = [
{ name = "pydantic-graph" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/76/6a/3048579c646f4cea7966889009a1d7eef1365f9126393945d39cefad9e83/pydantic_ai_slim-2.16.0.tar.gz", hash = "sha256:36d17cb12edd72ffc62f9e06cc49ac5f23cb77cf6b665c4b1fd11c00dbe6a852", size = 889343, upload-time = "2026-07-23T02:47:20.696Z" }
sdist = { url = "https://files.pythonhosted.org/packages/4b/39/c3a941027be87f6bc07e50e1e72ae93e66a1b11b838e33ad5d135d2fe2f0/pydantic_ai_slim-2.18.0.tar.gz", hash = "sha256:dfe47a3602049f779223702a684ed8091b111cfae1851236616d6b0eb3b0b077", size = 902113, upload-time = "2026-07-25T01:21:05.614Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2e/95/2c79b9f8e875562bae8141078af27428120e8415d493bb07ea36ca5905f9/pydantic_ai_slim-2.16.0-py3-none-any.whl", hash = "sha256:7cad27fb8f45ce4af4e8da83d7f206a3e1038dbc7535625c0ce5518fe378a55d", size = 1077057, upload-time = "2026-07-23T02:47:12.633Z" },
{ url = "https://files.pythonhosted.org/packages/bf/a0/a0619bddf602f69a754540cb82147f9b53dfa2d093e1fe7051c5e6a8eceb/pydantic_ai_slim-2.18.0-py3-none-any.whl", hash = "sha256:4c4076166a63ad96fe6ed5a517223c4a5aba6c9682a17d8d0ac42839cbc83622", size = 1091565, upload-time = "2026-07-25T01:20:57.323Z" },
]
[package.optional-dependencies]
@ -3949,7 +3949,7 @@ wheels = [
[[package]]
name = "pydantic-evals"
version = "2.16.0"
version = "2.18.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
@ -3959,14 +3959,14 @@ dependencies = [
{ name = "pyyaml" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/52/f5/7c2bf8ce45da52ed70c030e91ce43747317f61bd0615a40e91f1ec515c70/pydantic_evals-2.16.0.tar.gz", hash = "sha256:717e9615c7688650cdc716046f1d8edfbe0bac1748286c0a2744e3689fb518c7", size = 85147, upload-time = "2026-07-23T02:47:22.089Z" }
sdist = { url = "https://files.pythonhosted.org/packages/fd/4a/e2b629f4724a6026c792a5bd7ceb470bf38d48c6a2ebe8cec091ed85485f/pydantic_evals-2.18.0.tar.gz", hash = "sha256:d26ab006290564a5e56394b9033fbd5925e4a172d9e23be5f3c034ae3637eb18", size = 85222, upload-time = "2026-07-25T01:21:07.101Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f0/f6/d7f20f976f81d073856138bd046bf1ce5dc96561391526753c1c9076d33c/pydantic_evals-2.16.0-py3-none-any.whl", hash = "sha256:6705b427ea7c77d7f6b7d152ced6f86728931eceb9c57a8aa07ab8225fb1a934", size = 100439, upload-time = "2026-07-23T02:47:14.523Z" },
{ url = "https://files.pythonhosted.org/packages/1c/8d/be8abc897f9cd9ef83e000bfaa8d9bec0d93fe2f59f219d2c3835d66b033/pydantic_evals-2.18.0-py3-none-any.whl", hash = "sha256:98d20973df83c3ca6d7341ce2f704ae59b53e02698b55df765f9a1cd5ec5ea1d", size = 100524, upload-time = "2026-07-25T01:20:59.284Z" },
]
[[package]]
name = "pydantic-graph"
version = "2.16.0"
version = "2.18.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "httpx" },
@ -3974,9 +3974,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c3/22/6b6426e14607275f6b15f2a0c4530514c48ff01f86c53bbaac4041d09df9/pydantic_graph-2.16.0.tar.gz", hash = "sha256:f71e5c8e78a4ce56bc044861178e506ebd99881717ae0993924c457f5de6230e", size = 43979, upload-time = "2026-07-23T02:47:23.328Z" }
sdist = { url = "https://files.pythonhosted.org/packages/e5/f9/1554e818e6d38bcb3f66ec7da7abe17c28dcd00caa1550b034e175b8841f/pydantic_graph-2.18.0.tar.gz", hash = "sha256:9423defa047b477a561a06eadfd37aaf101a2b690d044adb79de80d7ea203e4e", size = 44017, upload-time = "2026-07-25T01:21:08.152Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/80/c0/362a7fb50562d7b51e3d02d454826e0a7b5652f27e5be8e835ddfaf84da6/pydantic_graph-2.16.0-py3-none-any.whl", hash = "sha256:99c25852c436d4d510d1ecdcb53ce4a0b11aa4d1d9d81472727587f2b5a3bc13", size = 51661, upload-time = "2026-07-23T02:47:15.921Z" },
{ url = "https://files.pythonhosted.org/packages/2f/29/fd05a4a84db9dae16c0417d8bcd6847148439a55329d000ef101d7243ded/pydantic_graph-2.18.0-py3-none-any.whl", hash = "sha256:48df74e3ae12ca44f99890e9b4f7020ae70a5a38d20df02db7cb16b0db3c3711", size = 51662, upload-time = "2026-07-25T01:21:00.764Z" },
]
[[package]]