diff --git a/CHANGELOG.md b/CHANGELOG.md index ecfa0e71..4c8577c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,9 +3,10 @@ ### Changed -- Bump `pydantic-ai` to 2.18.0. -- `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; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`. +- 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 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 7131364b..96b32c73 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 18a3cef9..de250785 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -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( diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 0f0f0385..f12062d6 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -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( diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index acdf6425..4e8bf8e5 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -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." ) diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index cce2e7e0..601082d1 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -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( diff --git a/tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_return_cap_message.yaml b/tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_fail_the_tool.yaml similarity index 100% rename from tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_return_cap_message.yaml rename to tests/cassettes/test_search_tools/TestSearchMaxSearches.test_searches_beyond_limit_fail_the_tool.yaml diff --git a/tests/tools/test_document.py b/tests/tools/test_document.py index ab935cf4..967a1bf8 100644 --- a/tests/tools/test_document.py +++ b/tests/tools/test_document.py @@ -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 diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 2cd8499e..60bf3404 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -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