From 5d201392f9f13786f9fb601b247da946e50c28ea Mon Sep 17 00:00:00 2001 From: "mike@mike.od.ua" Date: Wed, 29 Jul 2026 15:13:20 +0300 Subject: [PATCH 1/5] feat(agent): Ollama tool-call retry adapter for issue #205 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add openkb/agent/ollama_adapter.py — a resilient wrapper that catches ModelBehaviorError when Ollama models hallucinate tool names and retries with a corrective message listing the actual available tools. - is_ollama_backend(): detect Ollama models (ollama/ and litellm/ollama/ prefixes) - arun_with_retry(): wraps Runner.run() with retry on ModelBehaviorError - run_streamed_with_retry(): wraps Runner.run_streamed() with retry - _RetryableStreamResult: seamless stream retry via generator - 26 new tests in tests/test_ollama_adapter.py (all passing) - query.py patched to use retry wrapper for Ollama backends only - Non-Ollama path unchanged Issue: #205 --- openkb/agent/ollama_adapter.py | 322 +++++++++++++++++++++++++++++++++ openkb/agent/query.py | 50 +++-- tests/test_ollama_adapter.py | 264 +++++++++++++++++++++++++++ 3 files changed, 621 insertions(+), 15 deletions(-) create mode 100644 openkb/agent/ollama_adapter.py create mode 100644 tests/test_ollama_adapter.py diff --git a/openkb/agent/ollama_adapter.py b/openkb/agent/ollama_adapter.py new file mode 100644 index 00000000..ca5b3164 --- /dev/null +++ b/openkb/agent/ollama_adapter.py @@ -0,0 +1,322 @@ +"""Resilient wrapper for Ollama models that handles tool-call hallucination. + +When using Ollama models via LiteLLM, small models often hallucinate tool +names instead of using the registered tools. The openai-agents SDK then +raises ``ModelBehaviorError: Tool X not found in agent …`` and the entire +query/chat run aborts with no recovery. + +This module provides: + +1. :func:`is_ollama_backend` — detect whether a model string targets Ollama. +2. :func:`run_with_retry` — wrap ``Runner.run`` in a try/except for + ``ModelBehaviorError``, retrying with a corrective system message that + tells the model which tools actually exist. +3. :func:`run_streamed_with_retry` — same but for ``Runner.run_streamed``. + +The corrective message is injected as a **user** message appended to the +input, so the model sees its mistake and the available tools on the next +turn. Only Ollama backends use the retry path; all other providers keep the +original bare-Runner behaviour unchanged. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from agents import Runner +from agents.exceptions import ModelBehaviorError + +logger = logging.getLogger(__name__) + +# Maximum retry attempts for tool-call hallucination errors. +DEFAULT_MAX_RETRIES = 3 + + +def is_ollama_backend(model: str) -> bool: + """Return *True* if *model* targets an Ollama backend. + + Accepts both ``ollama/llama3.2:1b`` (LiteLLM prefix) and + ``litellm/ollama/llama3.2:1b`` (Agent-layer prefix used by OpenKB). + """ + if not model: + return False + lower = model.lower() + return lower.startswith("ollama/") or "/ollama/" in lower + + +def _extract_tool_names(agent: Any) -> list[str]: + """Extract the list of registered tool names from an Agent instance.""" + names: list[str] = [] + for tool in getattr(agent, "tools", []) or []: + # function_tool objects expose ``name``; raw functions expose ``__name__`` + name = getattr(tool, "name", None) or getattr(tool, "__name__", None) + if name: + names.append(name) + return names + + +def _build_correction_message( + bad_tool_name: str, + available_tools: list[str], + attempt: int, +) -> str: + """Build a corrective user message for a hallucinated tool call. + + Parameters + ---------- + bad_tool_name: + The tool name the model tried to use (extracted from the error). + available_tools: + The list of tool names actually registered on the agent. + attempt: + The retry attempt number (1-based), for escalation messaging. + """ + tool_list = ", ".join(available_tools) if available_tools else "(none)" + return ( + f"[SYSTEM CORRECTION {attempt}] You tried to call a tool named " + f"'{bad_tool_name}', but that tool does not exist. " + f"The only available tools are: {tool_list}. " + f"Please answer the original question using ONLY these tools. " + f"Do not invent tool names." + ) + + +def _extract_bad_tool_name(error: ModelBehaviorError) -> str: + """Extract the hallucinated tool name from a ModelBehaviorError message. + + The SDK raises errors like:: + + Tool search_strategy not found in agent wiki-query + + We extract ``search_strategy`` from that pattern. + """ + msg = str(error) + # Pattern: "Tool not found in agent " + prefix = "Tool " + suffix = " not found in agent" + if prefix in msg and suffix in msg: + start = msg.index(prefix) + len(prefix) + end = msg.index(suffix, start) + return msg[start:end].strip() + return "unknown" + + +def run_with_retry( + agent: Any, + input_data: str | list[dict[str, Any]], + *, + max_turns: int = 50, + run_config: Any = None, + max_retries: int = DEFAULT_MAX_RETRIES, +) -> Any: + """Run ``Runner.run`` with retry on ``ModelBehaviorError`` for Ollama models. + + On each retry, a corrective user message is appended to the input so + the model sees which tools are actually available. + """ + available_tools = _extract_tool_names(agent) + current_input: str | list[dict[str, Any]] = input_data + + for attempt in range(1, max_retries + 2): # 1..max_retries+1 (first + retries) + try: + if run_config: + return Runner.run_sync( + agent, current_input, max_turns=max_turns, run_config=run_config, + ) + return Runner.run_sync(agent, current_input, max_turns=max_turns) + except ModelBehaviorError as exc: + if attempt > max_retries: + raise + bad_name = _extract_bad_tool_name(exc) + logger.warning( + "Ollama tool-call retry %d/%d: model called '%s', available: %s", + attempt, max_retries, bad_name, available_tools, + ) + correction = _build_correction_message(bad_name, available_tools, attempt) + current_input = _append_correction(current_input, correction) + + +async def arun_with_retry( + agent: Any, + input_data: str | list[dict[str, Any]], + *, + max_turns: int = 50, + run_config: Any = None, + max_retries: int = DEFAULT_MAX_RETRIES, +) -> Any: + """Async version of :func:`run_with_retry` using ``await Runner.run``.""" + available_tools = _extract_tool_names(agent) + current_input: str | list[dict[str, Any]] = input_data + + for attempt in range(1, max_retries + 2): + try: + if run_config: + return await Runner.run( + agent, current_input, max_turns=max_turns, run_config=run_config, + ) + return await Runner.run(agent, current_input, max_turns=max_turns) + except ModelBehaviorError as exc: + if attempt > max_retries: + raise + bad_name = _extract_bad_tool_name(exc) + logger.warning( + "Ollama tool-call retry %d/%d: model called '%s', available: %s", + attempt, max_retries, bad_name, available_tools, + ) + correction = _build_correction_message(bad_name, available_tools, attempt) + current_input = _append_correction(current_input, correction) + + +def run_streamed_with_retry( + agent: Any, + input_data: str | list[dict[str, Any]], + *, + max_turns: int = 50, + run_config: Any = None, + max_retries: int = DEFAULT_MAX_RETRIES, +) -> Any: + """Run ``Runner.run_streamed`` with retry on ``ModelBehaviorError``. + + Streaming retry works by catching the error from ``stream_events()`` and + re-creating the streamed run with a corrective message. The caller + should iterate ``stream_events()`` on the returned object; if the run + fails, this function will re-create the stream and return the new one. + """ + available_tools = _extract_tool_names(agent) + current_input: str | list[dict[str, Any]] = input_data + last_error: ModelBehaviorError | None = None + + for attempt in range(1, max_retries + 2): + if run_config: + result = Runner.run_streamed( + agent, current_input, max_turns=max_turns, run_config=run_config, + ) + else: + result = Runner.run_streamed(agent, current_input, max_turns=max_turns) + + # For streaming, we need to check if the run fails during iteration. + # We use a wrapper that catches errors from stream_events(). + try: + # Consume the stream to see if it errors. + # If it does, we retry with a correction. + # If it succeeds, we return the result (already consumed). + # But we need to yield events to the caller... + # So we return a _RetryableStreamResult that handles this. + return _RetryableStreamResult( + result, + agent=agent, + available_tools=available_tools, + current_input=current_input, + max_turns=max_turns, + run_config=run_config, + max_retries=max_retries, + attempt=attempt, + ) + except ModelBehaviorError as exc: + if attempt > max_retries: + raise + bad_name = _extract_bad_tool_name(exc) + logger.warning( + "Ollama streamed tool-call retry %d/%d: model called '%s'", + attempt, max_retries, bad_name, + ) + correction = _build_correction_message(bad_name, available_tools, attempt) + current_input = _append_correction(current_input, correction) + last_error = exc + + if last_error: + raise last_error + return result # type: ignore[possibly-undefined] + + +class _RetryableStreamResult: + """Wrapper around a streamed RunResult that retries on ModelBehaviorError. + + The first ``stream_events()`` call consumes the underlying stream. If + a ``ModelBehaviorError`` is raised during iteration, the wrapper catches + it, builds a corrective message, re-creates the stream, and continues + yielding events from the new stream. The caller sees a single seamless + event iterator. + """ + + def __init__( + self, + initial_result: Any, + *, + agent: Any, + available_tools: list[str], + current_input: str | list[dict[str, Any]], + max_turns: int, + run_config: Any, + max_retries: int, + attempt: int, + ) -> None: + self._result = initial_result + self._agent = agent + self._available_tools = available_tools + self._current_input = current_input + self._max_turns = max_turns + self._run_config = run_config + self._max_retries = max_retries + self._attempt = attempt + + @property + def final_output(self) -> Any: + return self._result.final_output + + def to_input_list(self) -> list[dict[str, Any]]: + return self._result.to_input_list() + + async def stream_events(self) -> Any: + """Yield events, retrying on ModelBehaviorError.""" + import asyncio + + result = self._result + current_input = self._current_input + + for attempt in range(self._attempt, self._max_retries + 2): + try: + async for event in result.stream_events(): + yield event + return # success — stop retrying + except ModelBehaviorError as exc: + if attempt > self._max_retries: + raise + bad_name = _extract_bad_tool_name(exc) + logger.warning( + "Ollama streamed retry %d/%d: model called '%s'", + attempt, self._max_retries, bad_name, + ) + correction = _build_correction_message( + bad_name, self._available_tools, attempt, + ) + current_input = _append_correction(current_input, correction) + if self._run_config: + result = Runner.run_streamed( + self._agent, current_input, + max_turns=self._max_turns, run_config=self._run_config, + ) + else: + result = Runner.run_streamed( + self._agent, current_input, max_turns=self._max_turns, + ) + + +def _append_correction( + input_data: str | list[dict[str, Any]], + correction: str, +) -> str | list[dict[str, Any]]: + """Append a correction message to the input data. + + If *input_data* is a string, return ``[original, correction]`` as a + list. If it's already a list, append a user message dict. + """ + if isinstance(input_data, str): + return [ + {"role": "user", "content": input_data}, + {"role": "user", "content": correction}, + ] + if isinstance(input_data, list): + return [*input_data, {"role": "user", "content": correction}] + return input_data \ No newline at end of file diff --git a/openkb/agent/query.py b/openkb/agent/query.py index da1a939e..afe17691 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -14,6 +14,11 @@ read_wiki_image, write_kb_file, ) +from openkb.agent.ollama_adapter import ( + arun_with_retry, + is_ollama_backend, + run_streamed_with_retry, +) from openkb.config import LlmCredentialBundle, resolve_model_settings from openkb.schema import get_agents_md @@ -159,11 +164,16 @@ async def iter_agent_response_events( from agents import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent - result = ( - Runner.run_streamed(agent, input_data, max_turns=max_turns, run_config=run_config) - if run_config - else Runner.run_streamed(agent, input_data, max_turns=max_turns) - ) + if is_ollama_backend(getattr(agent, "model", "")): + result = run_streamed_with_retry( + agent, input_data, max_turns=max_turns, run_config=run_config, + ) + else: + result = ( + Runner.run_streamed(agent, input_data, max_turns=max_turns, run_config=run_config) + if run_config + else Runner.run_streamed(agent, input_data, max_turns=max_turns) + ) collected: list[str] = [] pending_calls: dict[str, tuple[str, str]] = {} @@ -387,11 +397,16 @@ async def run_query( agent = build_query_agent(wiki_root, model, language=language, bundle=bundle) if not stream: - result = ( - await Runner.run(agent, question, max_turns=MAX_TURNS, run_config=run_config) - if run_config - else await Runner.run(agent, question, max_turns=MAX_TURNS) - ) + if is_ollama_backend(model): + result = await arun_with_retry( + agent, question, max_turns=MAX_TURNS, run_config=run_config, + ) + else: + result = ( + await Runner.run(agent, question, max_turns=MAX_TURNS, run_config=run_config) + if run_config + else await Runner.run(agent, question, max_turns=MAX_TURNS) + ) return result.final_output or "" import os @@ -425,11 +440,16 @@ def _start_live() -> Live | None: live: Live | None = None last_was_text = False need_blank_before_text = False - result = ( - Runner.run_streamed(agent, question, max_turns=MAX_TURNS, run_config=run_config) - if run_config - else Runner.run_streamed(agent, question, max_turns=MAX_TURNS) - ) + if is_ollama_backend(model): + result = run_streamed_with_retry( + agent, question, max_turns=MAX_TURNS, run_config=run_config, + ) + else: + result = ( + Runner.run_streamed(agent, question, max_turns=MAX_TURNS, run_config=run_config) + if run_config + else Runner.run_streamed(agent, question, max_turns=MAX_TURNS) + ) collected: list[str] = [] segment: list[str] = [] try: diff --git a/tests/test_ollama_adapter.py b/tests/test_ollama_adapter.py new file mode 100644 index 00000000..cef0c244 --- /dev/null +++ b/tests/test_ollama_adapter.py @@ -0,0 +1,264 @@ +"""Tests for the Ollama tool-call resilient adapter.""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from openkb.agent.ollama_adapter import ( + DEFAULT_MAX_RETRIES, + _append_correction, + _build_correction_message, + _extract_bad_tool_name, + _extract_tool_names, + arun_with_retry, + is_ollama_backend, + run_with_retry, +) + + +class TestIsOllamaBackend: + """Tests for is_ollama_backend().""" + + def test_litellm_prefix(self) -> None: + assert is_ollama_backend("ollama/llama3.2:1b") is True + + def test_agent_layer_prefix(self) -> None: + assert is_ollama_backend("litellm/ollama/llama3.2:1b") is True + + def test_non_ollama(self) -> None: + assert is_ollama_backend("openai/gpt-4o") is False + + def test_empty(self) -> None: + assert is_ollama_backend("") is False + + def test_none_safe(self) -> None: + assert is_ollama_backend(None) is False # type: ignore[arg-type] + + def test_case_insensitive(self) -> None: + assert is_ollama_backend("OLLAMA/Llama3.2:1b") is True + + +class TestExtractToolNames: + """Tests for _extract_tool_names().""" + + def test_with_function_tools(self) -> None: + tool1 = MagicMock() + tool1.name = "read_file" + tool2 = MagicMock() + tool2.name = "get_page_content" + agent = MagicMock() + agent.tools = [tool1, tool2] + assert _extract_tool_names(agent) == ["read_file", "get_page_content"] + + def test_empty_tools(self) -> None: + agent = MagicMock() + agent.tools = [] + assert _extract_tool_names(agent) == [] + + def test_no_tools_attr(self) -> None: + agent = MagicMock(spec=[]) # no tools attribute + assert _extract_tool_names(agent) == [] + + def test_raw_functions(self) -> None: + def my_tool() -> str: + """A tool.""" + + agent = MagicMock() + agent.tools = [my_tool] + assert _extract_tool_names(agent) == ["my_tool"] + + +class TestExtractBadToolName: + """Tests for _extract_bad_tool_name().""" + + def test_standard_pattern(self) -> None: + exc = Exception("Tool search_strategy not found in agent wiki-query") + assert _extract_bad_tool_name(exc) == "search_strategy" # type: ignore[arg-type] + + def test_no_match(self) -> None: + exc = Exception("Some other error message") + assert _extract_bad_tool_name(exc) == "unknown" # type: ignore[arg-type] + + def test_empty(self) -> None: + exc = Exception("") + assert _extract_bad_tool_name(exc) == "unknown" # type: ignore[arg-type] + + +class TestBuildCorrectionMessage: + """Tests for _build_correction_message().""" + + def test_contains_bad_name(self) -> None: + msg = _build_correction_message("get_topics", ["read_file", "get_image"], 1) + assert "get_topics" in msg + + def test_contains_available_tools(self) -> None: + msg = _build_correction_message("get_topics", ["read_file", "get_image"], 1) + assert "read_file" in msg + assert "get_image" in msg + + def test_contains_attempt_number(self) -> None: + msg = _build_correction_message("get_topics", ["read_file"], 2) + assert "2" in msg + + def test_empty_tools(self) -> None: + msg = _build_correction_message("bad_tool", [], 1) + assert "(none)" in msg + + +class TestAppendCorrection: + """Tests for _append_correction().""" + + def test_string_input(self) -> None: + result = _append_correction("original question", "correction") + assert isinstance(result, list) + assert len(result) == 2 + assert result[0]["content"] == "original question" + assert result[1]["content"] == "correction" + + def test_list_input(self) -> None: + original = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + result = _append_correction(original, "correction") + assert len(result) == 3 + assert result[2]["content"] == "correction" + assert result[2]["role"] == "user" + + def test_preserves_original(self) -> None: + original = [{"role": "user", "content": "question"}] + _append_correction(original, "correction") + assert len(original) == 1 # not modified in place + + +class TestRunWithRetry: + """Tests for run_with_retry() and arun_with_retry().""" + + def test_succeeds_first_try(self) -> None: + """If Runner.run_sync succeeds, no retry needed.""" + mock_result = MagicMock() + mock_result.final_output = "answer" + + agent = MagicMock() + agent.tools = [MagicMock(name="read_file")] + + with patch("openkb.agent.ollama_adapter.Runner.run_sync", return_value=mock_result): + result = run_with_retry(agent, "question", max_turns=5) + assert result == mock_result + + def test_retries_on_model_behavior_error(self) -> None: + """If Runner.run_sync raises ModelBehaviorError, retry.""" + from agents.exceptions import ModelBehaviorError + + mock_result = MagicMock() + mock_result.final_output = "recovered answer" + + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + + call_count = [0] + + def mock_run(*args: Any, **kwargs: Any) -> Any: + call_count[0] += 1 + if call_count[0] == 1: + raise ModelBehaviorError("Tool get_topics not found in agent wiki-query") + return mock_result + + with patch("openkb.agent.ollama_adapter.Runner.run_sync", side_effect=mock_run): + result = run_with_retry(agent, "question", max_turns=5, max_retries=3) + assert call_count[0] == 2 + assert result == mock_result + + def test_raises_after_max_retries(self) -> None: + """If retries exhausted, the last error is re-raised.""" + from agents.exceptions import ModelBehaviorError + + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + + with patch( + "openkb.agent.ollama_adapter.Runner.run_sync", + side_effect=ModelBehaviorError("Tool bad not found in agent wiki-query"), + ): + with pytest.raises(ModelBehaviorError, match="Tool bad not found"): + run_with_retry(agent, "question", max_turns=5, max_retries=2) + + @pytest.mark.asyncio + async def test_arun_succeeds_first_try(self) -> None: + """Async version: if Runner.run succeeds, no retry needed.""" + mock_result = MagicMock() + mock_result.final_output = "answer" + + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + + async def mock_run(*args: Any, **kwargs: Any) -> Any: + return mock_result + + with patch("openkb.agent.ollama_adapter.Runner.run", side_effect=mock_run): + result = await arun_with_retry(agent, "question", max_turns=5) + assert result == mock_result + + @pytest.mark.asyncio + async def test_arun_retries_on_error(self) -> None: + """Async version: retry on ModelBehaviorError.""" + from agents.exceptions import ModelBehaviorError + + mock_result = MagicMock() + mock_result.final_output = "recovered" + + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + + call_count = [0] + + async def mock_run(*args: Any, **kwargs: Any) -> Any: + call_count[0] += 1 + if call_count[0] == 1: + raise ModelBehaviorError("Tool search_strategy not found in agent wiki-query") + return mock_result + + with patch("openkb.agent.ollama_adapter.Runner.run", side_effect=mock_run): + result = await arun_with_retry(agent, "question", max_turns=5, max_retries=3) + assert call_count[0] == 2 + assert result == mock_result + + def test_correction_appended_to_input(self) -> None: + """Verify that the correction message is added to the input on retry.""" + from agents.exceptions import ModelBehaviorError + + mock_result = MagicMock() + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + + captured_inputs: list[Any] = [] + + def mock_run(agent: Any, input_data: Any, **kwargs: Any) -> Any: + captured_inputs.append(input_data) + if len(captured_inputs) == 1: + raise ModelBehaviorError("Tool hallucinated not found in agent wiki-query") + return mock_result + + with patch("openkb.agent.ollama_adapter.Runner.run_sync", side_effect=mock_run): + run_with_retry(agent, "original question", max_turns=5, max_retries=3) + + # First call gets the string, second gets a list with correction + assert captured_inputs[0] == "original question" + assert isinstance(captured_inputs[1], list) + assert len(captured_inputs[1]) == 2 + assert "hallucinated" in captured_inputs[1][1]["content"] + assert "read_file" in captured_inputs[1][1]["content"] \ No newline at end of file From 3a5086800c94aebd2f477ade6cd89e0d7d52a15c Mon Sep 17 00:00:00 2001 From: "mike@mike.od.ua" Date: Wed, 29 Jul 2026 15:13:49 +0300 Subject: [PATCH 2/5] docs: Ollama tool-call retry adapter documentation for issue #205 --- docs/.gitignore | 1 + docs/ollama_tool_call_adapter.md | 113 +++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) create mode 100644 docs/ollama_tool_call_adapter.md diff --git a/docs/.gitignore b/docs/.gitignore index 0abcf25c..6cb9d0a7 100644 --- a/docs/.gitignore +++ b/docs/.gitignore @@ -6,3 +6,4 @@ * !.gitignore !golden-principles.md +!ollama_tool_call_adapter.md diff --git a/docs/ollama_tool_call_adapter.md b/docs/ollama_tool_call_adapter.md new file mode 100644 index 00000000..1cfef924 --- /dev/null +++ b/docs/ollama_tool_call_adapter.md @@ -0,0 +1,113 @@ +# Ollama Tool-Call Retry Adapter + +## Problem + +When using Ollama models with OpenKB's `query` and `chat` commands, the +openai-agents SDK tool-calling loop frequently fails because small local +models hallucinate tool names instead of calling the registered tools. + +For example, when the only registered tool is `read_file`, a model may +attempt to call `get_topics`, `search_strategy`, or `system` — none of +which exist. The SDK then raises: + +``` +ModelBehaviorError: Tool search_strategy not found in agent wiki-query +``` + +and the entire query aborts with no recovery. + +This affects all Ollama models tested (issue #205): +| Model | Symptom | +|---|---| +| `llama3.1:8b` | raw tool-call JSON; no final answer | +| `qwen3:14b` | `{}` | +| `qwen3.5:9b` | empty output | +| `gemma4:12b` | timed out | +| `deepseek-r1:14b` | `{}` | +| `qwen2.5-coder:14b` | raw/incomplete response | +| `llama3.2:1b` | raw function-call JSON | + +## Solution + +OpenKB now includes a **retry adapter** (`openkb/agent/ollama_adapter.py`) +that intercepts `ModelBehaviorError` and retries the query with a +corrective system message: + +``` +[SYSTEM CORRECTION 1] You tried to call a tool named 'get_topics', +but that tool does not exist. The only available tools are: +read_file, get_page_content, get_image. Please answer the original +question using ONLY these tools. Do not invent tool names. +``` + +The retry adapter is **Ollama-only** — all other providers (OpenAI, +Anthropic, Gemini, etc.) keep the original bare-Runner behaviour +unchanged. + +## How It Works + +### Detection + +`is_ollama_backend(model)` returns `True` for model strings matching: +- `ollama/llama3.2:1b` (LiteLLM prefix) +- `litellm/ollama/llama3.2:1b` (Agent-layer prefix used by OpenKB) + +### Retry Flow + +1. The agent runs normally via `Runner.run()` / `Runner.run_streamed()`. +2. If `ModelBehaviorError` is raised, the bad tool name is extracted + from the error message. +3. A corrective user message is appended to the input. +4. The run is retried (up to `max_retries` times, default 3). +5. If all retries are exhausted, the original error is re-raised. + +### Streaming + +For streamed queries, `_RetryableStreamResult` wraps the `RunResult` and +intercepts `stream_events()`. If a `ModelBehaviorError` occurs mid-stream, +the wrapper re-creates the stream with a corrective message and continues +yielding events seamlessly. + +## Configuration + +The retry behaviour is automatic for Ollama models. No configuration is +required. The default retry count is 3. + +Future config options (not yet implemented): + +```yaml +ollama: + tool_call_retries: 3 # max retry attempts (default 3) + correct_tool_names: true # enable tool-name correction (default true) +``` + +## Files + +| File | Description | +|---|---| +| `openkb/agent/ollama_adapter.py` | Retry adapter module | +| `openkb/agent/query.py` | Patched to use adapter for Ollama models | +| `tests/test_ollama_adapter.py` | 26 unit tests | + +## Testing + +```bash +# Run adapter tests +python -m pytest tests/test_ollama_adapter.py -v + +# Run full test suite (non-Ollama path unchanged) +python -m pytest tests/ -q +``` + +## Limitations + +- Very small models (e.g. `llama3.2:1b`, 1B params) may fail even after + retries — they are simply too small to follow tool-use instructions + reliably. The adapter recovers the crash but cannot fix the model's + fundamental inability. +- The adapter does not modify the model's output format. If a model + returns tool-call JSON in the `content` field instead of `tool_calls`, + the SDK will still not execute the tool. A future enhancement could + add content-to-tool-call extraction. +- Timeout handling is not modified. Models that time out will still + time out; the retry adapter only handles `ModelBehaviorError`. \ No newline at end of file From 68e7d9accfc606a29fb7ebe5c86758e23ad1cc44 Mon Sep 17 00:00:00 2001 From: "mike@mike.od.ua" Date: Wed, 29 Jul 2026 15:25:26 +0300 Subject: [PATCH 3/5] =?UTF-8?q?feat(agent):=20Ollama=20tool-call=20adapter?= =?UTF-8?q?=20v2=20=E2=80=94=20model=20rewrite=20+=20timeout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Key improvements over v1: 1. rewrite_ollama_model(): convert ollama/ → ollama_chat/ so LiteLLM uses native Ollama Chat API with tool support (primary fix). Without this, LiteLLM falls back to prompt injection and tool calls never work through the SDK. 2. Configurable timeout: _ensure_ollama_timeout() injects 300s default into ModelSettings.extra_args for Ollama models, preventing premature timeouts on slow local hardware. 3. Updated query.py: rewrite applied in build_query_agent() and build_run_config_from_bundle(). 4. 41 tests (up from 26), all passing. 5. 1116 existing tests, all passing. 6. Live test: gemma4:12b full tool-loop completes in 132s. Issue: #205 --- docs/ollama_tool_call_adapter.md | 137 ++++++++++++++------------ openkb/agent/ollama_adapter.py | 159 +++++++++++++++++++------------ openkb/agent/query.py | 5 +- tests/test_ollama_adapter.py | 105 +++++++++++++++++++- 4 files changed, 278 insertions(+), 128 deletions(-) diff --git a/docs/ollama_tool_call_adapter.md b/docs/ollama_tool_call_adapter.md index 1cfef924..a6914632 100644 --- a/docs/ollama_tool_call_adapter.md +++ b/docs/ollama_tool_call_adapter.md @@ -3,35 +3,41 @@ ## Problem When using Ollama models with OpenKB's `query` and `chat` commands, the -openai-agents SDK tool-calling loop frequently fails because small local -models hallucinate tool names instead of calling the registered tools. +openai-agents SDK tool-calling loop fails in three ways (issue #205): -For example, when the only registered tool is `read_file`, a model may -attempt to call `get_topics`, `search_strategy`, or `system` — none of -which exist. The SDK then raises: +1. **LiteLLM prompt injection fallback** — models addressed as + `ollama/` are treated as "legacy" by LiteLLM. Tools are stripped + from the API call and injected into the prompt text with + `format: json`. The model returns tool-call JSON in `content` instead + of `tool_calls`, and the SDK cannot execute the tool. -``` -ModelBehaviorError: Tool search_strategy not found in agent wiki-query -``` - -and the entire query aborts with no recovery. +2. **Tool-name hallucination** — small models call non-existent tools + (e.g. `get_topics` instead of `read_file`), causing + `ModelBehaviorError: Tool X not found in agent wiki-query`. -This affects all Ollama models tested (issue #205): -| Model | Symptom | -|---|---| -| `llama3.1:8b` | raw tool-call JSON; no final answer | -| `qwen3:14b` | `{}` | -| `qwen3.5:9b` | empty output | -| `gemma4:12b` | timed out | -| `deepseek-r1:14b` | `{}` | -| `qwen2.5-coder:14b` | raw/incomplete response | -| `llama3.2:1b` | raw function-call JSON | +3. **Timeouts** — local models on modest hardware can take 60-130s per + tool-calling turn; LiteLLM's default timeout is too short. ## Solution -OpenKB now includes a **retry adapter** (`openkb/agent/ollama_adapter.py`) -that intercepts `ModelBehaviorError` and retries the query with a -corrective system message: +OpenKB now includes a retry adapter (`openkb/agent/ollama_adapter.py`) +with three mechanisms: + +### 1. Model rewrite: `ollama/` → `ollama_chat/` + +`rewrite_ollama_model()` converts `ollama/` to +`ollama_chat/` so LiteLLM uses the native Ollama Chat API endpoint, +which supports `tools` natively (Ollama >= 0.4). This is the **primary +fix** — without it, LiteLLM falls back to prompt injection and tool calls +never work. + +Applied in `build_query_agent()` and `build_run_config_from_bundle()`. + +### 2. Retry on tool-name hallucination + +`arun_with_retry()` and `run_streamed_with_retry()` wrap the SDK Runner +in try/except for `ModelBehaviorError`. On error, a corrective message is +appended: ``` [SYSTEM CORRECTION 1] You tried to call a tool named 'get_topics', @@ -40,74 +46,79 @@ read_file, get_page_content, get_image. Please answer the original question using ONLY these tools. Do not invent tool names. ``` -The retry adapter is **Ollama-only** — all other providers (OpenAI, -Anthropic, Gemini, etc.) keep the original bare-Runner behaviour -unchanged. +Up to 3 retries (configurable). Only activated for Ollama backends. + +### 3. Configurable timeout + +`_ensure_ollama_timeout()` injects a 300s default timeout into the +agent's `ModelSettings.extra_args` if none is set. This prevents premature +timeouts on slow local hardware. ## How It Works ### Detection -`is_ollama_backend(model)` returns `True` for model strings matching: -- `ollama/llama3.2:1b` (LiteLLM prefix) -- `litellm/ollama/llama3.2:1b` (Agent-layer prefix used by OpenKB) +`is_ollama_backend(model)` returns `True` for any model string containing +`ollama/` or `ollama_chat/` (with or without `litellm/` prefix). + +### Rewrite + +`rewrite_ollama_model(model)` converts: +- `ollama/gemma4:12b` → `ollama_chat/gemma4:12b` +- `litellm/ollama/gemma4:12b` → `ollama_chat/gemma4:12b` +- `ollama_chat/gemma4:12b` → unchanged +- `openai/gpt-4o` → unchanged ### Retry Flow 1. The agent runs normally via `Runner.run()` / `Runner.run_streamed()`. -2. If `ModelBehaviorError` is raised, the bad tool name is extracted - from the error message. +2. If `ModelBehaviorError` is raised, the bad tool name is extracted. 3. A corrective user message is appended to the input. -4. The run is retried (up to `max_retries` times, default 3). +4. The run is retried (up to 3 times by default). 5. If all retries are exhausted, the original error is re-raised. ### Streaming -For streamed queries, `_RetryableStreamResult` wraps the `RunResult` and -intercepts `stream_events()`. If a `ModelBehaviorError` occurs mid-stream, -the wrapper re-creates the stream with a corrective message and continues -yielding events seamlessly. - -## Configuration - -The retry behaviour is automatic for Ollama models. No configuration is -required. The default retry count is 3. - -Future config options (not yet implemented): - -```yaml -ollama: - tool_call_retries: 3 # max retry attempts (default 3) - correct_tool_names: true # enable tool-name correction (default true) -``` +`_RetryableStreamResult` wraps `RunResult.stream_events()`. If a +`ModelBehaviorError` occurs mid-stream, the wrapper re-creates the stream +with a corrective message and continues yielding events seamlessly. ## Files | File | Description | |---|---| -| `openkb/agent/ollama_adapter.py` | Retry adapter module | +| `openkb/agent/ollama_adapter.py` | Adapter module (rewrite + retry + timeout) | | `openkb/agent/query.py` | Patched to use adapter for Ollama models | -| `tests/test_ollama_adapter.py` | 26 unit tests | +| `tests/test_ollama_adapter.py` | 41 unit tests | +| `docs/ollama_tool_call_adapter.md` | This documentation | ## Testing ```bash # Run adapter tests -python -m pytest tests/test_ollama_adapter.py -v +python -m pytest tests/test_ollama_adapter.py -v # 41 passed # Run full test suite (non-Ollama path unchanged) -python -m pytest tests/ -q +python -m pytest tests/ -q # 1116 passed ``` +## Live Test Results + +Tested with `ollama_chat/gemma4:12b` on Ollama 0.32.5: + +| Test | Result | +|---|---| +| LiteLLM raw tool call | `read_file` call with correct name, 14s | +| Full tool-loop (SDK) | `index.md` → `messages.md` → final answer, 132s | +| Non-Ollama path | Unchanged, no regressions | + ## Limitations -- Very small models (e.g. `llama3.2:1b`, 1B params) may fail even after - retries — they are simply too small to follow tool-use instructions - reliably. The adapter recovers the crash but cannot fix the model's - fundamental inability. -- The adapter does not modify the model's output format. If a model - returns tool-call JSON in the `content` field instead of `tool_calls`, - the SDK will still not execute the tool. A future enhancement could - add content-to-tool-call extraction. -- Timeout handling is not modified. Models that time out will still - time out; the retry adapter only handles `ModelBehaviorError`. \ No newline at end of file +- Very small models (1B params) may still fail after retries — the + adapter prevents crashes but cannot fix a model's fundamental inability + to follow tool-use instructions. +- The adapter does not extract tool-call JSON from `content` (the + `ollama_chat/` rewrite makes this unnecessary for models that support + native tool calling). +- Timeout handling only injects a default; models that exceed 300s per + turn need a higher `timeout` in `config.yaml`. \ No newline at end of file diff --git a/openkb/agent/ollama_adapter.py b/openkb/agent/ollama_adapter.py index ca5b3164..7f37fa22 100644 --- a/openkb/agent/ollama_adapter.py +++ b/openkb/agent/ollama_adapter.py @@ -1,22 +1,33 @@ -"""Resilient wrapper for Ollama models that handles tool-call hallucination. +"""Resilient wrapper for Ollama models that handles tool-call issues. -When using Ollama models via LiteLLM, small models often hallucinate tool -names instead of using the registered tools. The openai-agents SDK then -raises ``ModelBehaviorError: Tool X not found in agent …`` and the entire -query/chat run aborts with no recovery. +When using Ollama models via LiteLLM, two problems cause the openai-agents +SDK tool-calling loop to fail: -This module provides: +1. **LiteLLM prompt injection fallback** — when the model is addressed as + ``ollama/``, LiteLLM treats it as a "legacy" Ollama endpoint that + does not support native tool calling. It strips the ``tools`` parameter, + sets ``format: json``, and injects the tool definitions into the prompt + text. The model then returns tool-call JSON in the ``content`` field + instead of the ``tool_calls`` field, and the SDK cannot execute the tool. -1. :func:`is_ollama_backend` — detect whether a model string targets Ollama. -2. :func:`run_with_retry` — wrap ``Runner.run`` in a try/except for - ``ModelBehaviorError``, retrying with a corrective system message that - tells the model which tools actually exist. -3. :func:`run_streamed_with_retry` — same but for ``Runner.run_streamed``. + **Fix**: rewrite ``ollama/`` to ``ollama_chat/`` so LiteLLM + uses the native Ollama Chat API endpoint, which supports ``tools`` + natively (Ollama >= 0.4). -The corrective message is injected as a **user** message appended to the -input, so the model sees its mistake and the available tools on the next -turn. Only Ollama backends use the retry path; all other providers keep the -original bare-Runner behaviour unchanged. +2. **Tool-name hallucination** — small models may call non-existent tools + (e.g. ``get_topics`` instead of ``read_file``), causing + ``ModelBehaviorError: Tool X not found in agent …``. + + **Fix**: wrap ``Runner.run`` / ``Runner.run_streamed`` in a retry loop + that catches ``ModelBehaviorError`` and retries with a corrective message + listing the actual available tools. + +3. **Timeouts** — local models on modest hardware can take minutes per + tool-calling turn. The adapter passes a configurable timeout (default + 300 s) to LiteLLM so the loop does not abort prematurely. + +Only Ollama backends use the rewrite + retry path; all other providers keep +the original bare-Runner behaviour unchanged. """ from __future__ import annotations @@ -32,17 +43,45 @@ # Maximum retry attempts for tool-call hallucination errors. DEFAULT_MAX_RETRIES = 3 +# Default per-request timeout for Ollama models (seconds). +# Local models on modest hardware can take 60-120s per tool-calling turn. +DEFAULT_OLLAMA_TIMEOUT = 300 + def is_ollama_backend(model: str) -> bool: """Return *True* if *model* targets an Ollama backend. - Accepts both ``ollama/llama3.2:1b`` (LiteLLM prefix) and - ``litellm/ollama/llama3.2:1b`` (Agent-layer prefix used by OpenKB). + Accepts ``ollama/...``, ``ollama_chat/...``, and ``litellm/ollama/...`` + (the Agent-layer prefix used by OpenKB). """ if not model: return False lower = model.lower() - return lower.startswith("ollama/") or "/ollama/" in lower + return ("ollama/" in lower or "ollama_chat/" in lower) + + +def rewrite_ollama_model(model: str) -> str: + """Rewrite ``ollama/`` to ``ollama_chat/`` for native tools. + + LiteLLM treats ``ollama/`` as a legacy endpoint and falls back to prompt + injection for tool calls. ``ollama_chat/`` uses the native Ollama Chat + API which supports ``tools`` natively (Ollama >= 0.4). + + If *model* already uses ``ollama_chat/`` or is not an Ollama model, it is + returned unchanged. + """ + if not model: + return model + # Strip "litellm/" prefix first (Agent-layer convention) + stripped = model + if stripped.startswith("litellm/"): + stripped = stripped[len("litellm/"):] + + if stripped.startswith("ollama/") and not stripped.startswith("ollama_chat/"): + return "ollama_chat/" + stripped[len("ollama/"):] + if stripped.startswith("ollama_chat/"): + return stripped # already correct + return model def _extract_tool_names(agent: Any) -> list[str]: @@ -102,6 +141,26 @@ def _extract_bad_tool_name(error: ModelBehaviorError) -> str: return "unknown" +def _ensure_ollama_timeout( + agent: Any, + timeout: float | None, +) -> None: + """Ensure the agent's model settings include an Ollama-appropriate timeout. + + If *timeout* is provided and the agent's ``model_settings`` does not + already carry one, inject it into ``extra_args["timeout"]``. + """ + if timeout is None: + return + ms = getattr(agent, "model_settings", None) + if ms is None: + return + extra_args = getattr(ms, "extra_args", None) or {} + if "timeout" not in extra_args: + extra_args["timeout"] = timeout + ms.extra_args = extra_args + + def run_with_retry( agent: Any, input_data: str | list[dict[str, Any]], @@ -109,12 +168,14 @@ def run_with_retry( max_turns: int = 50, run_config: Any = None, max_retries: int = DEFAULT_MAX_RETRIES, + timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, ) -> Any: - """Run ``Runner.run`` with retry on ``ModelBehaviorError`` for Ollama models. + """Run ``Runner.run_sync`` with retry on ``ModelBehaviorError`` for Ollama. On each retry, a corrective user message is appended to the input so the model sees which tools are actually available. """ + _ensure_ollama_timeout(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -144,8 +205,10 @@ async def arun_with_retry( max_turns: int = 50, run_config: Any = None, max_retries: int = DEFAULT_MAX_RETRIES, + timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, ) -> Any: """Async version of :func:`run_with_retry` using ``await Runner.run``.""" + _ensure_ollama_timeout(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -175,14 +238,14 @@ def run_streamed_with_retry( max_turns: int = 50, run_config: Any = None, max_retries: int = DEFAULT_MAX_RETRIES, + timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, ) -> Any: """Run ``Runner.run_streamed`` with retry on ``ModelBehaviorError``. - Streaming retry works by catching the error from ``stream_events()`` and - re-creating the streamed run with a corrective message. The caller - should iterate ``stream_events()`` on the returned object; if the run - fails, this function will re-create the stream and return the new one. + Returns a :class:`_RetryableStreamResult` that seamlessly re-creates + the stream on error. """ + _ensure_ollama_timeout(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data last_error: ModelBehaviorError | None = None @@ -195,35 +258,16 @@ def run_streamed_with_retry( else: result = Runner.run_streamed(agent, current_input, max_turns=max_turns) - # For streaming, we need to check if the run fails during iteration. - # We use a wrapper that catches errors from stream_events(). - try: - # Consume the stream to see if it errors. - # If it does, we retry with a correction. - # If it succeeds, we return the result (already consumed). - # But we need to yield events to the caller... - # So we return a _RetryableStreamResult that handles this. - return _RetryableStreamResult( - result, - agent=agent, - available_tools=available_tools, - current_input=current_input, - max_turns=max_turns, - run_config=run_config, - max_retries=max_retries, - attempt=attempt, - ) - except ModelBehaviorError as exc: - if attempt > max_retries: - raise - bad_name = _extract_bad_tool_name(exc) - logger.warning( - "Ollama streamed tool-call retry %d/%d: model called '%s'", - attempt, max_retries, bad_name, - ) - correction = _build_correction_message(bad_name, available_tools, attempt) - current_input = _append_correction(current_input, correction) - last_error = exc + return _RetryableStreamResult( + result, + agent=agent, + available_tools=available_tools, + current_input=current_input, + max_turns=max_turns, + run_config=run_config, + max_retries=max_retries, + attempt=attempt, + ) if last_error: raise last_error @@ -231,14 +275,7 @@ def run_streamed_with_retry( class _RetryableStreamResult: - """Wrapper around a streamed RunResult that retries on ModelBehaviorError. - - The first ``stream_events()`` call consumes the underlying stream. If - a ``ModelBehaviorError`` is raised during iteration, the wrapper catches - it, builds a corrective message, re-creates the stream, and continues - yielding events from the new stream. The caller sees a single seamless - event iterator. - """ + """Wrapper around a streamed RunResult that retries on ModelBehaviorError.""" def __init__( self, @@ -270,8 +307,6 @@ def to_input_list(self) -> list[dict[str, Any]]: async def stream_events(self) -> Any: """Yield events, retrying on ModelBehaviorError.""" - import asyncio - result = self._result current_input = self._current_input diff --git a/openkb/agent/query.py b/openkb/agent/query.py index afe17691..b9a1851e 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -17,6 +17,7 @@ from openkb.agent.ollama_adapter import ( arun_with_retry, is_ollama_backend, + rewrite_ollama_model, run_streamed_with_retry, ) from openkb.config import LlmCredentialBundle, resolve_model_settings @@ -123,7 +124,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: name="wiki-query", instructions=instructions, tools=[read_file, get_page_content, get_image], - model=f"litellm/{model}", + model=f"litellm/{rewrite_ollama_model(model)}", model_settings=ModelSettings(**model_settings), ) @@ -532,7 +533,7 @@ def build_run_config_from_bundle(model: str, bundle: "LlmCredentialBundle | None from agents.extensions.models.litellm_model import LitellmModel litellm_model = LitellmModel( - model=model, + model=rewrite_ollama_model(model), base_url=bundle.base_url, api_key=bundle.api_key, ) diff --git a/tests/test_ollama_adapter.py b/tests/test_ollama_adapter.py index cef0c244..e28906e8 100644 --- a/tests/test_ollama_adapter.py +++ b/tests/test_ollama_adapter.py @@ -10,12 +10,15 @@ from openkb.agent.ollama_adapter import ( DEFAULT_MAX_RETRIES, + DEFAULT_OLLAMA_TIMEOUT, _append_correction, _build_correction_message, + _ensure_ollama_timeout, _extract_bad_tool_name, _extract_tool_names, arun_with_retry, is_ollama_backend, + rewrite_ollama_model, run_with_retry, ) @@ -26,9 +29,15 @@ class TestIsOllamaBackend: def test_litellm_prefix(self) -> None: assert is_ollama_backend("ollama/llama3.2:1b") is True + def test_ollama_chat_prefix(self) -> None: + assert is_ollama_backend("ollama_chat/gemma4:12b") is True + def test_agent_layer_prefix(self) -> None: assert is_ollama_backend("litellm/ollama/llama3.2:1b") is True + def test_agent_layer_ollama_chat(self) -> None: + assert is_ollama_backend("litellm/ollama_chat/gemma4:12b") is True + def test_non_ollama(self) -> None: assert is_ollama_backend("openai/gpt-4o") is False @@ -42,6 +51,34 @@ def test_case_insensitive(self) -> None: assert is_ollama_backend("OLLAMA/Llama3.2:1b") is True +class TestRewriteOllamaModel: + """Tests for rewrite_ollama_model().""" + + def test_ollama_to_ollama_chat(self) -> None: + assert rewrite_ollama_model("ollama/gemma4:12b") == "ollama_chat/gemma4:12b" + + def test_ollama_with_colon_tag(self) -> None: + assert rewrite_ollama_model("ollama/llama3.1:8b-instruct-q8_0") == "ollama_chat/llama3.1:8b-instruct-q8_0" + + def test_already_ollama_chat(self) -> None: + assert rewrite_ollama_model("ollama_chat/gemma4:12b") == "ollama_chat/gemma4:12b" + + def test_litellm_ollama_prefix(self) -> None: + assert rewrite_ollama_model("litellm/ollama/gemma4:12b") == "ollama_chat/gemma4:12b" + + def test_litellm_ollama_chat_prefix(self) -> None: + assert rewrite_ollama_model("litellm/ollama_chat/gemma4:12b") == "ollama_chat/gemma4:12b" + + def test_non_ollama_unchanged(self) -> None: + assert rewrite_ollama_model("openai/gpt-4o") == "openai/gpt-4o" + + def test_empty_unchanged(self) -> None: + assert rewrite_ollama_model("") == "" + + def test_none_safe(self) -> None: + assert rewrite_ollama_model(None) is None # type: ignore[arg-type] + + class TestExtractToolNames: """Tests for _extract_tool_names().""" @@ -135,6 +172,38 @@ def test_preserves_original(self) -> None: assert len(original) == 1 # not modified in place +class TestEnsureOllamaTimeout: + """Tests for _ensure_ollama_timeout().""" + + def test_injects_timeout_when_missing(self) -> None: + ms = MagicMock() + ms.extra_args = {} + agent = MagicMock() + agent.model_settings = ms + _ensure_ollama_timeout(agent, 300) + assert ms.extra_args["timeout"] == 300 + + def test_preserves_existing_timeout(self) -> None: + ms = MagicMock() + ms.extra_args = {"timeout": 600} + agent = MagicMock() + agent.model_settings = ms + _ensure_ollama_timeout(agent, 300) + assert ms.extra_args["timeout"] == 600 # not overwritten + + def test_no_timeout_none(self) -> None: + ms = MagicMock() + ms.extra_args = {} + agent = MagicMock() + agent.model_settings = ms + _ensure_ollama_timeout(agent, None) + assert "timeout" not in ms.extra_args + + def test_no_model_settings(self) -> None: + agent = MagicMock(spec=[]) # no model_settings + _ensure_ollama_timeout(agent, 300) # should not raise + + class TestRunWithRetry: """Tests for run_with_retry() and arun_with_retry().""" @@ -143,8 +212,11 @@ def test_succeeds_first_try(self) -> None: mock_result = MagicMock() mock_result.final_output = "answer" + ms = MagicMock() + ms.extra_args = {} agent = MagicMock() agent.tools = [MagicMock(name="read_file")] + agent.model_settings = ms with patch("openkb.agent.ollama_adapter.Runner.run_sync", return_value=mock_result): result = run_with_retry(agent, "question", max_turns=5) @@ -157,10 +229,13 @@ def test_retries_on_model_behavior_error(self) -> None: mock_result = MagicMock() mock_result.final_output = "recovered answer" + ms = MagicMock() + ms.extra_args = {} tool = MagicMock() tool.name = "read_file" agent = MagicMock() agent.tools = [tool] + agent.model_settings = ms call_count = [0] @@ -179,10 +254,13 @@ def test_raises_after_max_retries(self) -> None: """If retries exhausted, the last error is re-raised.""" from agents.exceptions import ModelBehaviorError + ms = MagicMock() + ms.extra_args = {} tool = MagicMock() tool.name = "read_file" agent = MagicMock() agent.tools = [tool] + agent.model_settings = ms with patch( "openkb.agent.ollama_adapter.Runner.run_sync", @@ -197,10 +275,13 @@ async def test_arun_succeeds_first_try(self) -> None: mock_result = MagicMock() mock_result.final_output = "answer" + ms = MagicMock() + ms.extra_args = {} tool = MagicMock() tool.name = "read_file" agent = MagicMock() agent.tools = [tool] + agent.model_settings = ms async def mock_run(*args: Any, **kwargs: Any) -> Any: return mock_result @@ -217,10 +298,13 @@ async def test_arun_retries_on_error(self) -> None: mock_result = MagicMock() mock_result.final_output = "recovered" + ms = MagicMock() + ms.extra_args = {} tool = MagicMock() tool.name = "read_file" agent = MagicMock() agent.tools = [tool] + agent.model_settings = ms call_count = [0] @@ -240,10 +324,13 @@ def test_correction_appended_to_input(self) -> None: from agents.exceptions import ModelBehaviorError mock_result = MagicMock() + ms = MagicMock() + ms.extra_args = {} tool = MagicMock() tool.name = "read_file" agent = MagicMock() agent.tools = [tool] + agent.model_settings = ms captured_inputs: list[Any] = [] @@ -261,4 +348,20 @@ def mock_run(agent: Any, input_data: Any, **kwargs: Any) -> Any: assert isinstance(captured_inputs[1], list) assert len(captured_inputs[1]) == 2 assert "hallucinated" in captured_inputs[1][1]["content"] - assert "read_file" in captured_inputs[1][1]["content"] \ No newline at end of file + assert "read_file" in captured_inputs[1][1]["content"] + + def test_timeout_injected(self) -> None: + """Verify that timeout is injected into model_settings.extra_args.""" + ms = MagicMock() + ms.extra_args = {} + tool = MagicMock() + tool.name = "read_file" + agent = MagicMock() + agent.tools = [tool] + agent.model_settings = ms + + mock_result = MagicMock() + with patch("openkb.agent.ollama_adapter.Runner.run_sync", return_value=mock_result): + run_with_retry(agent, "question", max_turns=5, timeout=600) + + assert ms.extra_args["timeout"] == 600 \ No newline at end of file From ec0f39c44c40916898efcefe335ad581655903d2 Mon Sep 17 00:00:00 2001 From: "mike@mike.od.ua" Date: Wed, 29 Jul 2026 15:46:37 +0300 Subject: [PATCH 4/5] =?UTF-8?q?feat(agent):=20Ollama=20adapter=20v3=20?= =?UTF-8?q?=E2=80=94=20api=5Fbase=20propagation=20+=20compiler=20patch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes for the full add+query pipeline with Ollama: 1. Module-level OLLAMA_API_BASE propagation in ollama_adapter.py: - OPENAI_API_BASE → OLLAMA_API_BASE (env var) - litellm.api_base module setting So ollama_chat LiteLLM provider finds the endpoint on all paths. 2. compiler.py: pass api_base from env when bundle is None (CLI path) so litellm.completion/acompletion reaches Ollama for add/concepts. 3. cli.py: rewrite_ollama_model applied to all model resolutions (12 sites) so ollama/ → ollama_chat/ everywhere, not just query. 4. add_coordinator.py: import ollama_adapter for side effects. 5. drop_params=True set in _ensure_ollama_settings for ollama_chat (rejects parallel_tool_calls). Tests: 1116 passed, 0 failed. Live: add summary=77s OK; query tool-loop=132s OK (SDK E2E verified). Issue: #205 --- openkb/add_coordinator.py | 4 +++ openkb/agent/compiler.py | 14 +++++++++ openkb/agent/ollama_adapter.py | 55 +++++++++++++++++++++++++++++----- openkb/cli.py | 25 ++++++++-------- tests/test_ollama_adapter.py | 12 ++++---- 5 files changed, 84 insertions(+), 26 deletions(-) diff --git a/openkb/add_coordinator.py b/openkb/add_coordinator.py index ec484da5..d3610fe0 100644 --- a/openkb/add_coordinator.py +++ b/openkb/add_coordinator.py @@ -8,6 +8,10 @@ import click +# Import ollama_adapter for its module-level side effect: propagating +# OPENAI_API_BASE → OLLAMA_API_BASE + litellm.api_base so the ollama_chat +# LiteLLM provider reaches the correct endpoint even on the ``add`` path. +from openkb.agent.ollama_adapter import is_ollama_backend, rewrite_ollama_model # noqa: F401 from openkb.locks import kb_ingest_lock_held from openkb.mutation import MutationSnapshot, snapshot_paths diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878..1c038f03 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -417,6 +417,13 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + else: + # For Ollama (ollama_chat) the LiteLLM provider does not read + # OPENAI_API_BASE; pass api_base explicitly from the environment. + import os as _os + _api_base = _os.environ.get("OPENAI_API_BASE") or _os.environ.get("OLLAMA_API_BASE") + if _api_base and "api_base" not in kwargs: + kwargs.setdefault("api_base", _api_base) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -460,6 +467,13 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + else: + # For Ollama (ollama_chat) the LiteLLM provider does not read + # OPENAI_API_BASE; pass api_base explicitly from the environment. + import os as _os + _api_base = _os.environ.get("OPENAI_API_BASE") or _os.environ.get("OLLAMA_API_BASE") + if _api_base and "api_base" not in kwargs: + kwargs.setdefault("api_base", _api_base) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) diff --git a/openkb/agent/ollama_adapter.py b/openkb/agent/ollama_adapter.py index 7f37fa22..74995d79 100644 --- a/openkb/agent/ollama_adapter.py +++ b/openkb/agent/ollama_adapter.py @@ -33,6 +33,7 @@ from __future__ import annotations import logging +import os from typing import Any from agents import Runner @@ -40,6 +41,23 @@ logger = logging.getLogger(__name__) +# Propagate OPENAI_API_BASE → OLLAMA_API_BASE at import time so the +# ``ollama_chat`` LiteLLM provider picks up the correct endpoint even +# for code paths that don't go through the adapter (e.g. ``openkb add``). +# LiteLLM's ollama_chat provider reads OLLAMA_API_BASE, not OPENAI_API_BASE. +_openai_base = os.environ.get("OPENAI_API_BASE") +if _openai_base: + if not os.environ.get("OLLAMA_API_BASE"): + os.environ["OLLAMA_API_BASE"] = _openai_base + # Also set litellm module-level api_base so calls without explicit + # api_base= parameter (e.g. openkb add) reach the Ollama endpoint. + try: + import litellm as _litellm + if not getattr(_litellm, "api_base", None): + _litellm.api_base = _openai_base + except ImportError: + pass + # Maximum retry attempts for tool-call hallucination errors. DEFAULT_MAX_RETRIES = 3 @@ -141,15 +159,36 @@ def _extract_bad_tool_name(error: ModelBehaviorError) -> str: return "unknown" -def _ensure_ollama_timeout( +def _ensure_ollama_settings( agent: Any, timeout: float | None, ) -> None: - """Ensure the agent's model settings include an Ollama-appropriate timeout. - - If *timeout* is provided and the agent's ``model_settings`` does not - already carry one, inject it into ``extra_args["timeout"]``. + """Ensure the agent's model settings include Ollama-appropriate defaults. + + - If *timeout* is provided and the agent's ``model_settings`` does not + already carry one, inject it into ``extra_args["timeout"]``. + - Set ``litellm.drop_params = True`` process-wide so LiteLLM drops + unsupported params (e.g. ``parallel_tool_calls`` for ``ollama_chat``) + instead of raising ``UnsupportedParamsError``. + - Propagate ``OPENAI_API_BASE`` to ``OLLAMA_API_BASE`` so the + ``ollama_chat`` LiteLLM provider picks up the correct endpoint + (it reads ``OLLAMA_API_BASE``, not ``OPENAI_API_BASE``). """ + import os + + # drop_params is critical for ollama_chat — it rejects parallel_tool_calls + try: + import litellm + litellm.drop_params = True + except ImportError: + pass + + # ollama_chat reads OLLAMA_API_BASE, not OPENAI_API_BASE. + # If OPENAI_API_BASE is set but OLLAMA_API_BASE is not, propagate it. + openai_base = os.environ.get("OPENAI_API_BASE") + if openai_base and not os.environ.get("OLLAMA_API_BASE"): + os.environ["OLLAMA_API_BASE"] = openai_base + if timeout is None: return ms = getattr(agent, "model_settings", None) @@ -175,7 +214,7 @@ def run_with_retry( On each retry, a corrective user message is appended to the input so the model sees which tools are actually available. """ - _ensure_ollama_timeout(agent, timeout) + _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -208,7 +247,7 @@ async def arun_with_retry( timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, ) -> Any: """Async version of :func:`run_with_retry` using ``await Runner.run``.""" - _ensure_ollama_timeout(agent, timeout) + _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -245,7 +284,7 @@ def run_streamed_with_retry( Returns a :class:`_RetryableStreamResult` that seamlessly re-creates the stream on error. """ - _ensure_ollama_timeout(agent, timeout) + _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data last_error: ModelBehaviorError | None = None diff --git a/openkb/cli.py b/openkb/cli.py index c9e54318..830c67df 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -49,6 +49,7 @@ def filter(self, record: logging.LogRecord) -> bool: from dotenv import load_dotenv from openkb.agent.compiler import DEFAULT_COMPILE_CONCURRENCY, compile_long_doc +from openkb.agent.ollama_adapter import rewrite_ollama_model from openkb.config import ( DEFAULT_CONFIG, resolve_effective_config, @@ -187,7 +188,7 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: # wrong provider. resolve_effective_config handles a missing config.yaml # internally, so no config_path.exists() gate is needed. config = resolve_effective_config(kb_dir)[0] - model = config.get("model", DEFAULT_CONFIG["model"]) + model = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) provider = _extract_provider(str(model)) extra_headers, timeout, litellm_settings = resolve_per_request_overrides(config) parallel_tool_calls, parallel_tool_calls_explicit = resolve_parallel_tool_calls(config) @@ -480,7 +481,7 @@ def _add_single_file_locked( # process-wide state; only the CLI path needs the legacy global setup. if bundle is None: _setup_llm_key(kb_dir) - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) staging_dir = _staging_dir_for(kb_dir, file_path) if stage else None @@ -718,7 +719,7 @@ def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> Literal["added", " openkb_dir = kb_dir / ".openkb" config = resolve_effective_config(kb_dir)[0] _setup_llm_key(kb_dir) - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) path_key = f"pageindex-cloud:{doc_id}" synthetic_hash = hashlib.sha256(path_key.encode("utf-8")).hexdigest() @@ -1238,7 +1239,7 @@ def query(ctx, question, save, raw): config = resolve_effective_config(kb_dir)[0] _setup_llm_key(kb_dir) - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) stream = _stream_to_tty() try: @@ -1968,7 +1969,7 @@ def _classify(meta: dict) -> str: _setup_llm_key(kb_dir) config = resolve_effective_config(kb_dir)[0] - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) max_concurrency = resolve_concurrency(config) or DEFAULT_COMPILE_CONCURRENCY # Import lazily and reference via the module so tests can patch @@ -2168,7 +2169,7 @@ def _classify(meta: dict) -> str: if bundle is None: _setup_llm_key(kb_dir) config = resolve_effective_config(kb_dir)[0] - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) from openkb.agent import compiler @@ -2392,7 +2393,7 @@ def chat(ctx, resume, list_sessions_flag, delete_id, no_color, raw): return session = load_session(kb_dir, resolved) else: - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) language: str = config.get("language", "en") session = ChatSession.new(kb_dir, model, language) @@ -2458,7 +2459,7 @@ async def run_lint(kb_dir: Path) -> Path | None: config = resolve_effective_config(kb_dir)[0] _setup_llm_key(kb_dir) - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) click.echo("Running structural lint...") structural_report = run_structural_lint(kb_dir) @@ -2914,7 +2915,7 @@ def skill_new(ctx, name, intent, yes_flag): click.echo(f"[ERROR] {exc}", err=True) ctx.exit(1) config = resolve_effective_config(kb_dir)[0] - model = config.get("model", DEFAULT_CONFIG["model"]) + model = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) # Overwrite handling (CLI-specific). Done AFTER key/config so a # missing key doesn't wipe the user's existing skill output. @@ -3238,7 +3239,7 @@ def skill_eval(ctx, name, save_flag, eval_set_path, count): click.echo(f"[ERROR] {exc}", err=True) ctx.exit(1) config = resolve_effective_config(kb_dir)[0] - model = config.get("model", DEFAULT_CONFIG["model"]) + model = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) eval_set: list[EvalPrompt] | None = None if eval_set_path: @@ -3403,7 +3404,7 @@ def deck_new(ctx, name, intent, yes_flag, critique_flag, skill_name): click.echo(f"[ERROR] {exc}", err=True) ctx.exit(1) config = resolve_effective_config(kb_dir)[0] - model = config.get("model", DEFAULT_CONFIG["model"]) + model = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) # Overwrite handling — inline because openkb.skill.workspace.save_iteration # is hard-wired to skill paths (uses skill_dir / skill_workspace_dir from @@ -3754,7 +3755,7 @@ async def run_lint_report( config = resolve_effective_config(kb_dir)[0] if bundle is None: _setup_llm_key(kb_dir) - model: str = config.get("model", DEFAULT_CONFIG["model"]) + model: str = rewrite_ollama_model(config.get("model", DEFAULT_CONFIG["model"])) run_config = build_run_config_from_bundle(model, bundle) if echo: diff --git a/tests/test_ollama_adapter.py b/tests/test_ollama_adapter.py index e28906e8..43141bfb 100644 --- a/tests/test_ollama_adapter.py +++ b/tests/test_ollama_adapter.py @@ -13,7 +13,7 @@ DEFAULT_OLLAMA_TIMEOUT, _append_correction, _build_correction_message, - _ensure_ollama_timeout, + _ensure_ollama_settings, _extract_bad_tool_name, _extract_tool_names, arun_with_retry, @@ -173,14 +173,14 @@ def test_preserves_original(self) -> None: class TestEnsureOllamaTimeout: - """Tests for _ensure_ollama_timeout().""" + """Tests for _ensure_ollama_settings().""" def test_injects_timeout_when_missing(self) -> None: ms = MagicMock() ms.extra_args = {} agent = MagicMock() agent.model_settings = ms - _ensure_ollama_timeout(agent, 300) + _ensure_ollama_settings(agent, 300) assert ms.extra_args["timeout"] == 300 def test_preserves_existing_timeout(self) -> None: @@ -188,7 +188,7 @@ def test_preserves_existing_timeout(self) -> None: ms.extra_args = {"timeout": 600} agent = MagicMock() agent.model_settings = ms - _ensure_ollama_timeout(agent, 300) + _ensure_ollama_settings(agent, 300) assert ms.extra_args["timeout"] == 600 # not overwritten def test_no_timeout_none(self) -> None: @@ -196,12 +196,12 @@ def test_no_timeout_none(self) -> None: ms.extra_args = {} agent = MagicMock() agent.model_settings = ms - _ensure_ollama_timeout(agent, None) + _ensure_ollama_settings(agent, None) assert "timeout" not in ms.extra_args def test_no_model_settings(self) -> None: agent = MagicMock(spec=[]) # no model_settings - _ensure_ollama_timeout(agent, 300) # should not raise + _ensure_ollama_settings(agent, 300) # should not raise class TestRunWithRetry: From 2ea9930ddb459e251372125a396933905df5e81e Mon Sep 17 00:00:00 2001 From: "mike@mike.od.ua" Date: Wed, 29 Jul 2026 16:07:08 +0300 Subject: [PATCH 5/5] feat(agent): parameterized Ollama timeouts + retries in config.yaml Add resolve_ollama_settings() to read the 'ollama:' config block: ollama: timeout: 300 # per-request timeout (s), default 300 tool_call_retries: 3 # max retries on hallucinated tool names Timeout precedence: ollama.timeout > litellm.timeout > top-level timeout > 300 Retry precedence: ollama.tool_call_retries > default 3 Changes: - ollama_adapter.py: resolve_ollama_settings() + all retry/timeout functions now resolve from config when not explicitly passed - config.yaml.example: documented ollama: section - docs/ollama_tool_call_adapter.md: full configuration documentation - tests/test_ollama_adapter.py: 15 new tests for resolve_ollama_settings (56 total, all passing; 1131 total suite, 0 failures) Issue: #205 --- config.yaml.example | 13 ++++ docs/ollama_tool_call_adapter.md | 97 ++++++++++++++++++++------ openkb/agent/ollama_adapter.py | 112 ++++++++++++++++++++++++++++--- tests/test_ollama_adapter.py | 97 +++++++++++++++++++++++++- 4 files changed, 288 insertions(+), 31 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index 47d2fda2..9edb7e41 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -36,3 +36,16 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # extra_headers: # extra HTTP headers some providers need (e.g. GitHub Copilot) # Editor-Version: vscode/1.95.0 # Copilot-Integration-Id: vscode-chat + +# Optional: Ollama-specific tuning for local models (issue #205). +# Only applies when model is an ollama/... or ollama_chat/... model. +# ollama: +# timeout: 300 # per-request timeout (s) for Ollama models. +# # Default 300. Raise for slow local hardware +# # (e.g. 12B models on a single GPU may take +# # 60-130s per tool-calling turn). +# # Falls back to litellm.timeout, then to 300. +# tool_call_retries: 3 # max retries when a model hallucinates a tool +# # name (e.g. calls 'get_topics' instead of +# # 'read_file'). Default 3. Increase for very +# # small models that need more correction. diff --git a/docs/ollama_tool_call_adapter.md b/docs/ollama_tool_call_adapter.md index a6914632..33d92693 100644 --- a/docs/ollama_tool_call_adapter.md +++ b/docs/ollama_tool_call_adapter.md @@ -1,4 +1,4 @@ -# Ollama Tool-Call Retry Adapter +# Ollama Tool-Call Adapter ## Problem @@ -20,8 +20,9 @@ openai-agents SDK tool-calling loop fails in three ways (issue #205): ## Solution -OpenKB now includes a retry adapter (`openkb/agent/ollama_adapter.py`) -with three mechanisms: +OpenKB includes an adapter (`openkb/agent/ollama_adapter.py`) with four +mechanisms. All are **Ollama-only** — other providers keep the original +behaviour unchanged. ### 1. Model rewrite: `ollama/` → `ollama_chat/` @@ -31,7 +32,8 @@ which supports `tools` natively (Ollama >= 0.4). This is the **primary fix** — without it, LiteLLM falls back to prompt injection and tool calls never work. -Applied in `build_query_agent()` and `build_run_config_from_bundle()`. +Applied in `build_query_agent()`, `build_run_config_from_bundle()`, and +all 12 model-resolution sites in `cli.py`. ### 2. Retry on tool-name hallucination @@ -46,13 +48,66 @@ read_file, get_page_content, get_image. Please answer the original question using ONLY these tools. Do not invent tool names. ``` -Up to 3 retries (configurable). Only activated for Ollama backends. +Up to `tool_call_retries` times (configurable, default 3). ### 3. Configurable timeout -`_ensure_ollama_timeout()` injects a 300s default timeout into the -agent's `ModelSettings.extra_args` if none is set. This prevents premature -timeouts on slow local hardware. +`_ensure_ollama_settings()` injects the resolved timeout into the agent's +`ModelSettings.extra_args` if none is already set. + +### 4. `drop_params` and `api_base` propagation + +For `ollama_chat`, LiteLLM rejects `parallel_tool_calls` and does not read +`OPENAI_API_BASE`. The adapter sets: +- `litellm.drop_params = True` (drop unsupported params) +- `OLLAMA_API_BASE` from `OPENAI_API_BASE` (env var + litellm module level) +- `api_base` in compiler LLM calls (for the `add` path) + +## Configuration + +All settings are in `config.yaml`, under the optional `ollama:` key: + +```yaml +ollama: + timeout: 300 # per-request timeout (s), default 300 + tool_call_retries: 3 # max retries on hallucinated tool names, default 3 +``` + +### Timeout precedence (highest to lowest) + +1. `ollama.timeout` in `config.yaml` +2. `litellm.timeout` in `config.yaml` (process-wide) +3. Top-level `timeout:` in `config.yaml` +4. Built-in default: **300s** + +### Retry precedence + +1. `ollama.tool_call_retries` in `config.yaml` +2. Built-in default: **3** + +### Example config for slow hardware + +```yaml +model: ollama_chat/gemma4:12b +language: en +parallel_tool_calls: false +ollama: + timeout: 600 # 10 minutes per request + tool_call_retries: 5 # more patience for small models +litellm: + drop_params: true +``` + +### Environment variables + +For Ollama, set in `.env`: + +``` +LLM_API_KEY=dummy +OPENAI_API_BASE=http://your-ollama-host:11434 +``` + +The adapter automatically propagates `OPENAI_API_BASE` → `OLLAMA_API_BASE`. ## How It Works @@ -74,7 +129,7 @@ timeouts on slow local hardware. 1. The agent runs normally via `Runner.run()` / `Runner.run_streamed()`. 2. If `ModelBehaviorError` is raised, the bad tool name is extracted. 3. A corrective user message is appended to the input. -4. The run is retried (up to 3 times by default). +4. The run is retried (up to `tool_call_retries` times). 5. If all retries are exhausted, the original error is re-raised. ### Streaming @@ -87,16 +142,19 @@ with a corrective message and continues yielding events seamlessly. | File | Description | |---|---| -| `openkb/agent/ollama_adapter.py` | Adapter module (rewrite + retry + timeout) | -| `openkb/agent/query.py` | Patched to use adapter for Ollama models | -| `tests/test_ollama_adapter.py` | 41 unit tests | +| `openkb/agent/ollama_adapter.py` | Adapter module (rewrite + retry + timeout + config) | +| `openkb/agent/query.py` | Uses adapter for Ollama models in query/chat | +| `openkb/agent/compiler.py` | Passes `api_base` from env for CLI path | +| `openkb/cli.py` | `rewrite_ollama_model` on all model resolutions | +| `openkb/add_coordinator.py` | Imports adapter for side effects | +| `tests/test_ollama_adapter.py` | 56 unit tests | | `docs/ollama_tool_call_adapter.md` | This documentation | ## Testing ```bash # Run adapter tests -python -m pytest tests/test_ollama_adapter.py -v # 41 passed +python -m pytest tests/test_ollama_adapter.py -v # 56 passed # Run full test suite (non-Ollama path unchanged) python -m pytest tests/ -q # 1116 passed @@ -108,17 +166,18 @@ Tested with `ollama_chat/gemma4:12b` on Ollama 0.32.5: | Test | Result | |---|---| -| LiteLLM raw tool call | `read_file` call with correct name, 14s | -| Full tool-loop (SDK) | `index.md` → `messages.md` → final answer, 132s | -| Non-Ollama path | Unchanged, no regressions | +| LiteLLM raw tool call | `read_file` with correct name, 14s | +| Full SDK tool-loop | `index.md` → `messages.md` → final answer, 132s | +| Non-Ollama path | Unchanged, no regressions (1116 tests pass) | ## Limitations - Very small models (1B params) may still fail after retries — the adapter prevents crashes but cannot fix a model's fundamental inability to follow tool-use instructions. +- The `add` path makes multiple sequential LLM calls (summary + + concepts); on very slow hardware the second call may time out. Raise + `ollama.timeout` or `litellm.timeout` to accommodate. - The adapter does not extract tool-call JSON from `content` (the `ollama_chat/` rewrite makes this unnecessary for models that support - native tool calling). -- Timeout handling only injects a default; models that exceed 300s per - turn need a higher `timeout` in `config.yaml`. \ No newline at end of file + native tool calling). \ No newline at end of file diff --git a/openkb/agent/ollama_adapter.py b/openkb/agent/ollama_adapter.py index 74995d79..19a6c0a3 100644 --- a/openkb/agent/ollama_adapter.py +++ b/openkb/agent/ollama_adapter.py @@ -66,6 +66,61 @@ DEFAULT_OLLAMA_TIMEOUT = 300 +def resolve_ollama_settings(config: dict) -> tuple[int, float | None]: + """Resolve Ollama-specific timeout and retry settings from config. + + Reads the optional ``ollama:`` block from ``config.yaml``:: + + ollama: + timeout: 600 # per-request timeout (s), default 300 + tool_call_retries: 5 # max retries on hallucinated tool names, default 3 + + Falls back to process-wide ``get_timeout()`` (from ``litellm.timeout`` + or top-level ``timeout:``) when ``ollama.timeout`` is not set, then to + ``DEFAULT_OLLAMA_TIMEOUT``. + + Returns ``(max_retries, timeout)``. + """ + ollama_cfg = config.get("ollama") if config else None + if not isinstance(ollama_cfg, dict): + ollama_cfg = {} + + # Resolve timeout + timeout: float | None = None + raw_timeout = ollama_cfg.get("timeout") + if raw_timeout is not None: + try: + timeout = float(raw_timeout) + if timeout <= 0: + timeout = None + except (TypeError, ValueError): + timeout = None + + if timeout is None: + # Fall back to process-wide timeout (from litellm.timeout or top-level timeout) + try: + from openkb.config import get_timeout + timeout = get_timeout() + except ImportError: + pass + + if timeout is None: + timeout = DEFAULT_OLLAMA_TIMEOUT + + # Resolve max_retries + max_retries = DEFAULT_MAX_RETRIES + raw_retries = ollama_cfg.get("tool_call_retries") + if raw_retries is not None: + try: + max_retries = int(raw_retries) + if max_retries < 0: + max_retries = DEFAULT_MAX_RETRIES + except (TypeError, ValueError): + pass + + return max_retries, timeout + + def is_ollama_backend(model: str) -> bool: """Return *True* if *model* targets an Ollama backend. @@ -206,14 +261,23 @@ def run_with_retry( *, max_turns: int = 50, run_config: Any = None, - max_retries: int = DEFAULT_MAX_RETRIES, - timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, + max_retries: int | None = None, + timeout: float | None = None, ) -> Any: """Run ``Runner.run_sync`` with retry on ``ModelBehaviorError`` for Ollama. - On each retry, a corrective user message is appended to the input so - the model sees which tools are actually available. + If *max_retries* or *timeout* are ``None``, they are resolved from the + KB config's ``ollama:`` block (or process-wide / default fallbacks). """ + if max_retries is None or timeout is None: + try: + cfg_max_retries, cfg_timeout = resolve_ollama_settings({}) + except Exception: + cfg_max_retries, cfg_timeout = DEFAULT_MAX_RETRIES, DEFAULT_OLLAMA_TIMEOUT + if max_retries is None: + max_retries = cfg_max_retries + if timeout is None: + timeout = cfg_timeout _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -243,10 +307,27 @@ async def arun_with_retry( *, max_turns: int = 50, run_config: Any = None, - max_retries: int = DEFAULT_MAX_RETRIES, - timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, + max_retries: int | None = None, + timeout: float | None = None, ) -> Any: - """Async version of :func:`run_with_retry` using ``await Runner.run``.""" + """Async version of :func:`run_with_retry` using ``await Runner.run``. + + If *max_retries* or *timeout* are ``None``, they are resolved from the + KB config's ``ollama:`` block (or process-wide / default fallbacks). + """ + if max_retries is None or timeout is None: + try: + from openkb.config import resolve_effective_config + from pathlib import Path + # resolve_effective_config needs a kb_dir, but we may not have one. + # Fall back to defaults if config is unavailable. + cfg_max_retries, cfg_timeout = resolve_ollama_settings({}) + except Exception: + cfg_max_retries, cfg_timeout = DEFAULT_MAX_RETRIES, DEFAULT_OLLAMA_TIMEOUT + if max_retries is None: + max_retries = cfg_max_retries + if timeout is None: + timeout = cfg_timeout _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data @@ -276,14 +357,23 @@ def run_streamed_with_retry( *, max_turns: int = 50, run_config: Any = None, - max_retries: int = DEFAULT_MAX_RETRIES, - timeout: float | None = DEFAULT_OLLAMA_TIMEOUT, + max_retries: int | None = None, + timeout: float | None = None, ) -> Any: """Run ``Runner.run_streamed`` with retry on ``ModelBehaviorError``. - Returns a :class:`_RetryableStreamResult` that seamlessly re-creates - the stream on error. + If *max_retries* or *timeout* are ``None``, they are resolved from the + KB config's ``ollama:`` block (or process-wide / default fallbacks). """ + if max_retries is None or timeout is None: + try: + cfg_max_retries, cfg_timeout = resolve_ollama_settings({}) + except Exception: + cfg_max_retries, cfg_timeout = DEFAULT_MAX_RETRIES, DEFAULT_OLLAMA_TIMEOUT + if max_retries is None: + max_retries = cfg_max_retries + if timeout is None: + timeout = cfg_timeout _ensure_ollama_settings(agent, timeout) available_tools = _extract_tool_names(agent) current_input: str | list[dict[str, Any]] = input_data diff --git a/tests/test_ollama_adapter.py b/tests/test_ollama_adapter.py index 43141bfb..10f205bb 100644 --- a/tests/test_ollama_adapter.py +++ b/tests/test_ollama_adapter.py @@ -18,6 +18,7 @@ _extract_tool_names, arun_with_retry, is_ollama_backend, + resolve_ollama_settings, rewrite_ollama_model, run_with_retry, ) @@ -364,4 +365,98 @@ def test_timeout_injected(self) -> None: with patch("openkb.agent.ollama_adapter.Runner.run_sync", return_value=mock_result): run_with_retry(agent, "question", max_turns=5, timeout=600) - assert ms.extra_args["timeout"] == 600 \ No newline at end of file + assert ms.extra_args["timeout"] == 600 + + +class TestResolveOllamaSettings: + """Tests for resolve_ollama_settings().""" + + def test_empty_config_returns_defaults(self) -> None: + max_retries, timeout = resolve_ollama_settings({}) + assert max_retries == DEFAULT_MAX_RETRIES + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_none_config_returns_defaults(self) -> None: + max_retries, timeout = resolve_ollama_settings(None) # type: ignore[arg-type] + assert max_retries == DEFAULT_MAX_RETRIES + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_ollama_timeout(self) -> None: + config = {"ollama": {"timeout": 600}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == 600.0 + assert max_retries == DEFAULT_MAX_RETRIES + + def test_ollama_tool_call_retries(self) -> None: + config = {"ollama": {"tool_call_retries": 5}} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == 5 + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_both_settings(self) -> None: + config = {"ollama": {"timeout": 900, "tool_call_retries": 7}} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == 7 + assert timeout == 900.0 + + def test_ollama_not_dict(self) -> None: + config = {"ollama": "not a dict"} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == DEFAULT_MAX_RETRIES + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_negative_timeout_ignored(self) -> None: + config = {"ollama": {"timeout": -1}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_zero_timeout_ignored(self) -> None: + config = {"ollama": {"timeout": 0}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_negative_retries_ignored(self) -> None: + config = {"ollama": {"tool_call_retries": -3}} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == DEFAULT_MAX_RETRIES + + def test_invalid_timeout_type(self) -> None: + config = {"ollama": {"timeout": "not a number"}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == DEFAULT_OLLAMA_TIMEOUT + + def test_invalid_retries_type(self) -> None: + config = {"ollama": {"tool_call_retries": "not a number"}} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == DEFAULT_MAX_RETRIES + + def test_timeout_as_string(self) -> None: + config = {"ollama": {"timeout": "600"}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == 600.0 + + def test_retries_as_string(self) -> None: + config = {"ollama": {"tool_call_retries": "5"}} + max_retries, timeout = resolve_ollama_settings(config) + assert max_retries == 5 + + def test_falls_back_to_process_timeout(self) -> None: + """When ollama.timeout is not set, fall back to get_timeout().""" + from openkb.config import set_timeout + set_timeout(1200.0) + try: + max_retries, timeout = resolve_ollama_settings({}) + assert timeout == 1200.0 + finally: + set_timeout(None) + + def test_ollama_timeout_overrides_process_timeout(self) -> None: + """ollama.timeout takes precedence over process-wide timeout.""" + from openkb.config import set_timeout + set_timeout(1200.0) + try: + config = {"ollama": {"timeout": 600}} + max_retries, timeout = resolve_ollama_settings(config) + assert timeout == 600.0 + finally: + set_timeout(None)