From ba57dbe5003593ec97b852465f194833e61af4f1 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 10:58:51 +0100 Subject: [PATCH 1/8] feat: add new `MaxStepsCallback` to stop the execution after a maxmum number of steps --- src/askui/__init__.py | 4 ++ src/askui/models/exceptions.py | 19 +++++++++ src/askui/models/shared/max_steps_callback.py | 40 +++++++++++++++++++ tests/unit/models/shared/__init__.py | 0 .../models/shared/test_max_steps_callback.py | 33 +++++++++++++++ 5 files changed, 96 insertions(+) create mode 100644 src/askui/models/shared/max_steps_callback.py create mode 100644 tests/unit/models/shared/__init__.py create mode 100644 tests/unit/models/shared/test_max_steps_callback.py diff --git a/src/askui/__init__.py b/src/askui/__init__.py index 29e24efb..63a1effb 100644 --- a/src/askui/__init__.py +++ b/src/askui/__init__.py @@ -30,7 +30,9 @@ ToolUseBlockParam, UrlImageSourceParam, ) +from .models.exceptions import MaxStepsReachedError from .models.shared.conversation_callback import ConversationCallback +from .models.shared.max_steps_callback import MaxStepsCallback from .models.shared.settings import ( DEFAULT_GET_RESOLUTION, DEFAULT_LOCATE_RESOLUTION, @@ -81,6 +83,8 @@ "ConfigurableRetry", "ContentBlockParam", "ConversationCallback", + "MaxStepsCallback", + "MaxStepsReachedError", "DEFAULT_GET_RESOLUTION", "DEFAULT_LOCATE_RESOLUTION", "GetSettings", diff --git a/src/askui/models/exceptions.py b/src/askui/models/exceptions.py index c69a9401..34bf72cb 100644 --- a/src/askui/models/exceptions.py +++ b/src/askui/models/exceptions.py @@ -147,6 +147,25 @@ def __init__(self, max_tokens: int, message: str | None = None): super().__init__(error_msg) +class MaxStepsReachedError(AutomationError): + """Exception raised when the agent reaches the maximum number of steps. + + Args: + max_steps (int): The maximum step limit that was reached. + message (str, optional): Custom error message. If not provided, a default + message will be generated. + """ + + def __init__(self, max_steps: int, message: str | None = None): + self.max_steps = max_steps + error_msg = ( + f"Agent stopped due to reaching maximum step limit of {max_steps} steps" + if message is None + else message + ) + super().__init__(error_msg) + + class ModelRefusalError(AutomationError): """Exception raised when the model refuses to process the request. diff --git a/src/askui/models/shared/max_steps_callback.py b/src/askui/models/shared/max_steps_callback.py new file mode 100644 index 00000000..f954a0a4 --- /dev/null +++ b/src/askui/models/shared/max_steps_callback.py @@ -0,0 +1,40 @@ +"""Callback for terminating the agentic loop after a maximum number of steps.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from typing_extensions import override + +from askui.models.exceptions import MaxStepsReachedError +from askui.models.shared.conversation_callback import ConversationCallback + +if TYPE_CHECKING: + from askui.models.shared.conversation import Conversation + + +class MaxStepsCallback(ConversationCallback): + """Terminates the agentic loop after a maximum number of steps. + + Args: + max_steps (int): The maximum number of steps before the loop is terminated. + + Raises: + MaxStepsReachedError: When the step limit is reached. + + Example: + ```python + from askui import ComputerAgent, MaxStepsCallback + + with ComputerAgent(callbacks=[MaxStepsCallback(max_steps=10)]) as agent: + agent.act("Open the settings menu") + ``` + """ + + def __init__(self, max_steps: int) -> None: + self._max_steps = max_steps + + @override + def on_step_start(self, conversation: Conversation, step_index: int) -> None: + if step_index >= self._max_steps: + raise MaxStepsReachedError(self._max_steps) diff --git a/tests/unit/models/shared/__init__.py b/tests/unit/models/shared/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/models/shared/test_max_steps_callback.py b/tests/unit/models/shared/test_max_steps_callback.py new file mode 100644 index 00000000..f6b82c53 --- /dev/null +++ b/tests/unit/models/shared/test_max_steps_callback.py @@ -0,0 +1,33 @@ +from unittest.mock import MagicMock + +import pytest + +from askui.models.exceptions import MaxStepsReachedError +from askui.models.shared.max_steps_callback import MaxStepsCallback + + +class TestMaxStepsCallback: + def test_raises_when_step_index_equals_max_steps(self) -> None: + callback = MaxStepsCallback(max_steps=3) + conversation = MagicMock() + with pytest.raises(MaxStepsReachedError, match="3 steps"): + callback.on_step_start(conversation, step_index=3) + + def test_raises_when_step_index_exceeds_max_steps(self) -> None: + callback = MaxStepsCallback(max_steps=3) + conversation = MagicMock() + with pytest.raises(MaxStepsReachedError): + callback.on_step_start(conversation, step_index=5) + + def test_does_not_raise_when_under_limit(self) -> None: + callback = MaxStepsCallback(max_steps=3) + conversation = MagicMock() + for i in range(3): + callback.on_step_start(conversation, step_index=i) + + def test_max_steps_of_one(self) -> None: + callback = MaxStepsCallback(max_steps=1) + conversation = MagicMock() + callback.on_step_start(conversation, step_index=0) + with pytest.raises(MaxStepsReachedError): + callback.on_step_start(conversation, step_index=1) From d2dbb6cd9ae4ea50b7a0db4b5af77e7d9d54a9bc Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:04:40 +0100 Subject: [PATCH 2/8] chore: move callbacks to a new directory so that they can be imported with `from askui.callbacks import ...` --- src/askui/callbacks/__init__.py | 9 +++++++++ .../shared => callbacks}/conversation_callback.py | 0 .../{models/shared => callbacks}/max_steps_callback.py | 2 +- .../shared => callbacks}/usage_tracking_callback.py | 2 +- 4 files changed, 11 insertions(+), 2 deletions(-) create mode 100644 src/askui/callbacks/__init__.py rename src/askui/{models/shared => callbacks}/conversation_callback.py (100%) rename src/askui/{models/shared => callbacks}/max_steps_callback.py (93%) rename src/askui/{models/shared => callbacks}/usage_tracking_callback.py (97%) diff --git a/src/askui/callbacks/__init__.py b/src/askui/callbacks/__init__.py new file mode 100644 index 00000000..09eb5f4e --- /dev/null +++ b/src/askui/callbacks/__init__.py @@ -0,0 +1,9 @@ +from .conversation_callback import ConversationCallback +from .max_steps_callback import MaxStepsCallback +from .usage_tracking_callback import UsageTrackingCallback + +__all__ = [ + "ConversationCallback", + "MaxStepsCallback", + "UsageTrackingCallback", +] diff --git a/src/askui/models/shared/conversation_callback.py b/src/askui/callbacks/conversation_callback.py similarity index 100% rename from src/askui/models/shared/conversation_callback.py rename to src/askui/callbacks/conversation_callback.py diff --git a/src/askui/models/shared/max_steps_callback.py b/src/askui/callbacks/max_steps_callback.py similarity index 93% rename from src/askui/models/shared/max_steps_callback.py rename to src/askui/callbacks/max_steps_callback.py index f954a0a4..2568df83 100644 --- a/src/askui/models/shared/max_steps_callback.py +++ b/src/askui/callbacks/max_steps_callback.py @@ -6,8 +6,8 @@ from typing_extensions import override +from askui.callbacks.conversation_callback import ConversationCallback from askui.models.exceptions import MaxStepsReachedError -from askui.models.shared.conversation_callback import ConversationCallback if TYPE_CHECKING: from askui.models.shared.conversation import Conversation diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/callbacks/usage_tracking_callback.py similarity index 97% rename from src/askui/models/shared/usage_tracking_callback.py rename to src/askui/callbacks/usage_tracking_callback.py index 3bc58206..0aa199fe 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/src/askui/callbacks/usage_tracking_callback.py @@ -7,8 +7,8 @@ from opentelemetry import trace from typing_extensions import override +from askui.callbacks.conversation_callback import ConversationCallback from askui.models.shared.agent_message_param import UsageParam -from askui.models.shared.conversation_callback import ConversationCallback from askui.reporting import NULL_REPORTER, Reporter if TYPE_CHECKING: From 645501031d6190ee91046cb20b24ad2eae139886 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:05:13 +0100 Subject: [PATCH 3/8] fix: callback imports as they have been moved --- docs/11_callbacks.md | 15 +++++++++++++++ src/askui/__init__.py | 3 +-- src/askui/agent_base.py | 3 +-- src/askui/android_agent.py | 2 +- src/askui/computer_agent.py | 2 +- src/askui/models/shared/conversation.py | 2 +- .../unit/models/shared/test_max_steps_callback.py | 2 +- 7 files changed, 21 insertions(+), 8 deletions(-) diff --git a/docs/11_callbacks.md b/docs/11_callbacks.md index d9b354fa..8bf23e2d 100644 --- a/docs/11_callbacks.md +++ b/docs/11_callbacks.md @@ -2,6 +2,8 @@ Callbacks provide hooks into the agent's conversation lifecycle, similar to PyTorch Lightning's callback system. Use them for logging, monitoring, custom metrics, or extending agent behavior. +All callbacks live in `askui.callbacks` and can be imported from there. + ## Usage Subclass `ConversationCallback` and override the hooks you need: @@ -71,6 +73,19 @@ with ComputerAgent(callbacks=[TimingCallback()]) as agent: agent.act("Search for documents") ``` +## Built-in Callbacks + +### `MaxStepsCallback` + +Terminates the agentic loop after a maximum number of steps. Raises `MaxStepsReachedError` when the limit is reached. + +```python +from askui import ComputerAgent, MaxStepsCallback + +with ComputerAgent(callbacks=[MaxStepsCallback(max_steps=10)]) as agent: + agent.act("Open the settings menu") +``` + ## Multiple Callbacks Pass multiple callbacks to combine behaviors: diff --git a/src/askui/__init__.py b/src/askui/__init__.py index 63a1effb..526562e1 100644 --- a/src/askui/__init__.py +++ b/src/askui/__init__.py @@ -9,6 +9,7 @@ from .agent_base import Agent from .agent_settings import AgentSettings +from .callbacks import ConversationCallback, MaxStepsCallback from .computer_agent import ComputerAgent, VisionAgent from .locators import Locator from .models import ( @@ -31,8 +32,6 @@ UrlImageSourceParam, ) from .models.exceptions import MaxStepsReachedError -from .models.shared.conversation_callback import ConversationCallback -from .models.shared.max_steps_callback import MaxStepsCallback from .models.shared.settings import ( DEFAULT_GET_RESOLUTION, DEFAULT_LOCATE_RESOLUTION, diff --git a/src/askui/agent_base.py b/src/askui/agent_base.py index 3f72d84c..cf53ef21 100644 --- a/src/askui/agent_base.py +++ b/src/askui/agent_base.py @@ -10,11 +10,11 @@ from typing_extensions import Self from askui.agent_settings import AgentSettings +from askui.callbacks import ConversationCallback, UsageTrackingCallback from askui.container import telemetry from askui.locators.locators import Locator from askui.models.shared.agent_message_param import MessageParam from askui.models.shared.conversation import Conversation, Speakers -from askui.models.shared.conversation_callback import ConversationCallback from askui.models.shared.settings import ( ActSettings, CacheWritingSettings, @@ -23,7 +23,6 @@ LocateSettings, ) from askui.models.shared.tools import Tool, ToolCollection -from askui.models.shared.usage_tracking_callback import UsageTrackingCallback from askui.prompts.act_prompts import CACHE_USE_PROMPT, create_default_prompt from askui.tools.agent_os import AgentOs from askui.tools.android.agent_os import AndroidAgentOs diff --git a/src/askui/android_agent.py b/src/askui/android_agent.py index f9779821..b4fb0182 100644 --- a/src/askui/android_agent.py +++ b/src/askui/android_agent.py @@ -6,10 +6,10 @@ from askui.agent_base import Agent from askui.agent_settings import AgentSettings +from askui.callbacks import ConversationCallback from askui.container import telemetry from askui.locators.locators import Locator from askui.models.models import Point -from askui.models.shared.conversation_callback import ConversationCallback from askui.models.shared.settings import ActSettings, MessageSettings from askui.models.shared.tools import Tool from askui.prompts.act_prompts import create_android_agent_prompt diff --git a/src/askui/computer_agent.py b/src/askui/computer_agent.py index a71ef63c..d35a97a7 100644 --- a/src/askui/computer_agent.py +++ b/src/askui/computer_agent.py @@ -6,10 +6,10 @@ from askui.agent_base import Agent from askui.agent_settings import AgentSettings +from askui.callbacks import ConversationCallback from askui.container import telemetry from askui.locators.locators import Locator from askui.models.models import Point -from askui.models.shared.conversation_callback import ConversationCallback from askui.models.shared.settings import ActSettings, LocateSettings, MessageSettings from askui.models.shared.tools import Tool from askui.prompts.act_prompts import ( diff --git a/src/askui/models/shared/conversation.py b/src/askui/models/shared/conversation.py index 167df491..da0bde92 100644 --- a/src/askui/models/shared/conversation.py +++ b/src/askui/models/shared/conversation.py @@ -22,7 +22,7 @@ from askui.tools.switch_speaker_tool import SwitchSpeakerTool if TYPE_CHECKING: - from askui.models.shared.conversation_callback import ConversationCallback + from askui.callbacks import ConversationCallback from askui.utils.caching.cache_manager import CacheManager logger = logging.getLogger(__name__) diff --git a/tests/unit/models/shared/test_max_steps_callback.py b/tests/unit/models/shared/test_max_steps_callback.py index f6b82c53..70835c1d 100644 --- a/tests/unit/models/shared/test_max_steps_callback.py +++ b/tests/unit/models/shared/test_max_steps_callback.py @@ -2,8 +2,8 @@ import pytest +from askui.callbacks import MaxStepsCallback from askui.models.exceptions import MaxStepsReachedError -from askui.models.shared.max_steps_callback import MaxStepsCallback class TestMaxStepsCallback: From 057810ba2983bd8fea711f904d8957ef273bd146 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 15:45:36 +0100 Subject: [PATCH 4/8] chore: remove max_steps callback --- docs/11_callbacks.md | 11 +---- src/askui/__init__.py | 5 +-- src/askui/callbacks/__init__.py | 2 - src/askui/callbacks/max_steps_callback.py | 40 ------------------- src/askui/models/exceptions.py | 19 --------- tests/unit/models/__init__.py | 0 tests/unit/models/shared/__init__.py | 0 .../models/shared/test_max_steps_callback.py | 33 --------------- 8 files changed, 2 insertions(+), 108 deletions(-) delete mode 100644 src/askui/callbacks/max_steps_callback.py delete mode 100644 tests/unit/models/__init__.py delete mode 100644 tests/unit/models/shared/__init__.py delete mode 100644 tests/unit/models/shared/test_max_steps_callback.py diff --git a/docs/11_callbacks.md b/docs/11_callbacks.md index 8bf23e2d..ce5c8a6f 100644 --- a/docs/11_callbacks.md +++ b/docs/11_callbacks.md @@ -75,16 +75,7 @@ with ComputerAgent(callbacks=[TimingCallback()]) as agent: ## Built-in Callbacks -### `MaxStepsCallback` - -Terminates the agentic loop after a maximum number of steps. Raises `MaxStepsReachedError` when the limit is reached. - -```python -from askui import ComputerAgent, MaxStepsCallback - -with ComputerAgent(callbacks=[MaxStepsCallback(max_steps=10)]) as agent: - agent.act("Open the settings menu") -``` +(we will add built-in callbacks at a later stage) ## Multiple Callbacks diff --git a/src/askui/__init__.py b/src/askui/__init__.py index 526562e1..0d04f17e 100644 --- a/src/askui/__init__.py +++ b/src/askui/__init__.py @@ -9,7 +9,7 @@ from .agent_base import Agent from .agent_settings import AgentSettings -from .callbacks import ConversationCallback, MaxStepsCallback +from .callbacks import ConversationCallback from .computer_agent import ComputerAgent, VisionAgent from .locators import Locator from .models import ( @@ -31,7 +31,6 @@ ToolUseBlockParam, UrlImageSourceParam, ) -from .models.exceptions import MaxStepsReachedError from .models.shared.settings import ( DEFAULT_GET_RESOLUTION, DEFAULT_LOCATE_RESOLUTION, @@ -82,8 +81,6 @@ "ConfigurableRetry", "ContentBlockParam", "ConversationCallback", - "MaxStepsCallback", - "MaxStepsReachedError", "DEFAULT_GET_RESOLUTION", "DEFAULT_LOCATE_RESOLUTION", "GetSettings", diff --git a/src/askui/callbacks/__init__.py b/src/askui/callbacks/__init__.py index 09eb5f4e..29eb7029 100644 --- a/src/askui/callbacks/__init__.py +++ b/src/askui/callbacks/__init__.py @@ -1,9 +1,7 @@ from .conversation_callback import ConversationCallback -from .max_steps_callback import MaxStepsCallback from .usage_tracking_callback import UsageTrackingCallback __all__ = [ "ConversationCallback", - "MaxStepsCallback", "UsageTrackingCallback", ] diff --git a/src/askui/callbacks/max_steps_callback.py b/src/askui/callbacks/max_steps_callback.py deleted file mode 100644 index 2568df83..00000000 --- a/src/askui/callbacks/max_steps_callback.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Callback for terminating the agentic loop after a maximum number of steps.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from typing_extensions import override - -from askui.callbacks.conversation_callback import ConversationCallback -from askui.models.exceptions import MaxStepsReachedError - -if TYPE_CHECKING: - from askui.models.shared.conversation import Conversation - - -class MaxStepsCallback(ConversationCallback): - """Terminates the agentic loop after a maximum number of steps. - - Args: - max_steps (int): The maximum number of steps before the loop is terminated. - - Raises: - MaxStepsReachedError: When the step limit is reached. - - Example: - ```python - from askui import ComputerAgent, MaxStepsCallback - - with ComputerAgent(callbacks=[MaxStepsCallback(max_steps=10)]) as agent: - agent.act("Open the settings menu") - ``` - """ - - def __init__(self, max_steps: int) -> None: - self._max_steps = max_steps - - @override - def on_step_start(self, conversation: Conversation, step_index: int) -> None: - if step_index >= self._max_steps: - raise MaxStepsReachedError(self._max_steps) diff --git a/src/askui/models/exceptions.py b/src/askui/models/exceptions.py index 34bf72cb..c69a9401 100644 --- a/src/askui/models/exceptions.py +++ b/src/askui/models/exceptions.py @@ -147,25 +147,6 @@ def __init__(self, max_tokens: int, message: str | None = None): super().__init__(error_msg) -class MaxStepsReachedError(AutomationError): - """Exception raised when the agent reaches the maximum number of steps. - - Args: - max_steps (int): The maximum step limit that was reached. - message (str, optional): Custom error message. If not provided, a default - message will be generated. - """ - - def __init__(self, max_steps: int, message: str | None = None): - self.max_steps = max_steps - error_msg = ( - f"Agent stopped due to reaching maximum step limit of {max_steps} steps" - if message is None - else message - ) - super().__init__(error_msg) - - class ModelRefusalError(AutomationError): """Exception raised when the model refuses to process the request. diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/unit/models/shared/__init__.py b/tests/unit/models/shared/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/unit/models/shared/test_max_steps_callback.py b/tests/unit/models/shared/test_max_steps_callback.py deleted file mode 100644 index 70835c1d..00000000 --- a/tests/unit/models/shared/test_max_steps_callback.py +++ /dev/null @@ -1,33 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -from askui.callbacks import MaxStepsCallback -from askui.models.exceptions import MaxStepsReachedError - - -class TestMaxStepsCallback: - def test_raises_when_step_index_equals_max_steps(self) -> None: - callback = MaxStepsCallback(max_steps=3) - conversation = MagicMock() - with pytest.raises(MaxStepsReachedError, match="3 steps"): - callback.on_step_start(conversation, step_index=3) - - def test_raises_when_step_index_exceeds_max_steps(self) -> None: - callback = MaxStepsCallback(max_steps=3) - conversation = MagicMock() - with pytest.raises(MaxStepsReachedError): - callback.on_step_start(conversation, step_index=5) - - def test_does_not_raise_when_under_limit(self) -> None: - callback = MaxStepsCallback(max_steps=3) - conversation = MagicMock() - for i in range(3): - callback.on_step_start(conversation, step_index=i) - - def test_max_steps_of_one(self) -> None: - callback = MaxStepsCallback(max_steps=1) - conversation = MagicMock() - callback.on_step_start(conversation, step_index=0) - with pytest.raises(MaxStepsReachedError): - callback.on_step_start(conversation, step_index=1) From f3183f90c262717de99e922bf266b7e47d592ed9 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 15:56:03 +0100 Subject: [PATCH 5/8] feat: add `max_steps` parameter to ActSettings --- src/askui/models/shared/conversation.py | 13 +++++++++++++ src/askui/models/shared/settings.py | 1 + 2 files changed, 14 insertions(+) diff --git a/src/askui/models/shared/conversation.py b/src/askui/models/shared/conversation.py index da0bde92..cb6521ff 100644 --- a/src/askui/models/shared/conversation.py +++ b/src/askui/models/shared/conversation.py @@ -207,9 +207,22 @@ def _execute_control_loop(self) -> None: self._step_index = 0 continue_execution = True while continue_execution: + if self._is_max_steps_reached(): + break continue_execution = self._execute_step() self._on_control_loop_end() + def _is_max_steps_reached(self) -> bool: + if self.settings.max_steps is None: + return False + if self._step_index >= self.settings.max_steps: + logger.info( + "Reached max_steps limit (%d), stopping conversation", + self.settings.max_steps, + ) + return True + return False + @tracer.start_as_current_span("_teardown_control_loop") def _teardown_control_loop(self) -> None: # Finish recording if cache_manager is active and not executing from cache diff --git a/src/askui/models/shared/settings.py b/src/askui/models/shared/settings.py index 3fc7624a..293d7fb2 100644 --- a/src/askui/models/shared/settings.py +++ b/src/askui/models/shared/settings.py @@ -89,6 +89,7 @@ class ActSettings(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) messages: MessageSettings = Field(default_factory=MessageSettings) + max_steps: int | None = None class GetSettings(BaseModel): From 82aa192f18726603d09c1411691e03fdba355189 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 16:02:31 +0100 Subject: [PATCH 6/8] fix: add missing init to tests subdir --- tests/unit/models/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/models/__init__.py diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py new file mode 100644 index 00000000..e69de29b From b70eb30eb8bc65350cbdbe18944a1810701b4d94 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 16:08:51 +0100 Subject: [PATCH 7/8] fix: broken imports --- src/askui/reporting.py | 2 +- tests/unit/model_providers/test_model_pricing.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/askui/reporting.py b/src/askui/reporting.py index ec601286..2feea108 100644 --- a/src/askui/reporting.py +++ b/src/askui/reporting.py @@ -21,7 +21,7 @@ if TYPE_CHECKING: from PIL import Image - from askui.models.shared.usage_tracking_callback import UsageSummary + from askui.callbacks.usage_tracking_callback import UsageSummary def normalize_to_pil_images( diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index 7fa00c6d..e1b2103b 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/tests/unit/model_providers/test_model_pricing.py @@ -4,11 +4,11 @@ import pytest -from askui.models.shared.agent_message_param import UsageParam -from askui.models.shared.usage_tracking_callback import ( +from askui.callbacks.usage_tracking_callback import ( UsageSummary, UsageTrackingCallback, ) +from askui.models.shared.agent_message_param import UsageParam from askui.utils.model_pricing import ModelPricing From c84fd5591059126f481a447812df605cada6b9a3 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 16:09:20 +0100 Subject: [PATCH 8/8] refactor: raise Exception if max steps is reached i/o logging it --- src/askui/models/shared/conversation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/askui/models/shared/conversation.py b/src/askui/models/shared/conversation.py index cb6521ff..9205d259 100644 --- a/src/askui/models/shared/conversation.py +++ b/src/askui/models/shared/conversation.py @@ -216,11 +216,10 @@ def _is_max_steps_reached(self) -> bool: if self.settings.max_steps is None: return False if self._step_index >= self.settings.max_steps: - logger.info( - "Reached max_steps limit (%d), stopping conversation", - self.settings.max_steps, + msg = ( + f"Reached max_steps limit {self.settings.max_steps}, stopping execution" ) - return True + raise ConversationException(msg) return False @tracer.start_as_current_span("_teardown_control_loop")