Merge pull request #257 from askui/fix/truncation_strategy

fix(truncation strategy): add proper message history management through summarization

Author: philipph-askui

src/askui/agent_base.py

@@ -23,6 +23,7 @@
     LocateSettings,
 )
 from askui.models.shared.tools import Tool, ToolCollection
+from askui.models.shared.truncation_strategies import TruncationStrategy
 from askui.prompts.act_prompts import CACHE_USE_PROMPT, create_default_prompt
 from askui.telemetry.otel import OtelSettings, setup_opentelemetry_tracing
 from askui.tools.agent_os import AgentOs
@@ -59,6 +60,7 @@ def __init__(
         agent_os: AgentOs | AndroidAgentOs | None = None,
         settings: AgentSettings | None = None,
         callbacks: list[ConversationCallback] | None = None,
+        truncation_strategy: TruncationStrategy | None = None,
     ) -> None:
         load_dotenv()
         self._reporter: Reporter = reporter or CompositeReporter(reporters=None)
@@ -87,6 +89,7 @@ def __init__(
             image_qa_provider=self._image_qa_provider,
             detection_provider=self._detection_provider,
             reporter=self._reporter,
+            truncation_strategy=truncation_strategy,
             callbacks=_callbacks,
         )
 

src/askui/android_agent.py

@@ -12,6 +12,7 @@
 from askui.models.models import Point
 from askui.models.shared.settings import ActSettings, MessageSettings
 from askui.models.shared.tools import Tool
+from askui.models.shared.truncation_strategies import TruncationStrategy
 from askui.prompts.act_prompts import create_android_agent_prompt
 from askui.tools.android.agent_os import ANDROID_KEY
 from askui.tools.android.agent_os_facade import AndroidAgentOsFacade
@@ -64,7 +65,15 @@ class AndroidAgent(Agent):
         ```
     """
 
-    @telemetry.record_call(exclude={"reporters", "settings", "act_tools", "callbacks"})
+    @telemetry.record_call(
+        exclude={
+            "reporters",
+            "settings",
+            "act_tools",
+            "callbacks",
+            "truncation_strategy",
+        }
+    )
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def __init__(
         self,
@@ -74,6 +83,7 @@ def __init__(
         retry: Retry | None = None,
         act_tools: list[Tool] | None = None,
         callbacks: list[ConversationCallback] | None = None,
+        truncation_strategy: TruncationStrategy | None = None,
     ) -> None:
         reporter = CompositeReporter(reporters=reporters)
         self.os = PpadbAgentOs(device_identifier=device, reporter=reporter)
@@ -85,6 +95,7 @@ def __init__(
             agent_os=self.os,
             settings=settings,
             callbacks=callbacks,
+            truncation_strategy=truncation_strategy,
         )
         self.act_tool_collection.add_agent_os(self.act_agent_os_facade)
         # Override default act settings with Android-specific settings

src/askui/callbacks/conversation_callback.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING
 
+from askui.models.shared.agent_message_param import UsageParam
+
 if TYPE_CHECKING:
     from askui.models.shared.conversation import Conversation
     from askui.speaker.speaker import SpeakerResult
@@ -123,3 +125,10 @@ def on_tool_execution_end(
             conversation: The conversation instance.
             tool_names: Names of tools that were executed.
         """
+
+    def on_truncation_summarize(self, usage: UsageParam) -> None:
+        """Called when a truncation strategy summarizes message history.
+
+        Args:
+            usage: Token usage from the summarization LLM call.
+        """

src/askui/callbacks/usage_tracking_callback.py

@@ -227,6 +227,11 @@ def on_step_end(
         current_span = trace.get_current_span()
         current_span.set_attributes(step_summary.token_attributes())
 
+    @override
+    def on_truncation_summarize(self, usage: UsageParam) -> None:
+        self._per_conversation_usage.add_usage(usage)
+        self._summary.add_usage(usage)
+
     @override
     def on_conversation_end(self, conversation: Conversation) -> None:
         generated_steps: list[StepUsageSummary] = [

src/askui/computer_agent.py

@@ -12,6 +12,7 @@
 from askui.models.models import Point
 from askui.models.shared.settings import ActSettings, LocateSettings, MessageSettings
 from askui.models.shared.tools import Tool
+from askui.models.shared.truncation_strategies import TruncationStrategy
 from askui.prompts.act_prompts import (
     create_computer_agent_prompt,
 )
@@ -69,7 +70,14 @@ class ComputerAgent(Agent):
     """
 
     @telemetry.record_call(
-        exclude={"reporters", "tools", "settings", "act_tools", "callbacks"}
+        exclude={
+            "reporters",
+            "tools",
+            "settings",
+            "act_tools",
+            "callbacks",
+            "truncation_strategy",
+        }
     )
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def __init__(
@@ -81,6 +89,7 @@ def __init__(
         retry: Retry | None = None,
         act_tools: list[Tool] | None = None,
         callbacks: list[ConversationCallback] | None = None,
+        truncation_strategy: TruncationStrategy | None = None,
     ) -> None:
         reporter = CompositeReporter(reporters=reporters)
         self.tools = tools or AgentToolbox(
@@ -96,6 +105,7 @@ def __init__(
             agent_os=self.tools.os,
             settings=settings,
             callbacks=callbacks,
+            truncation_strategy=truncation_strategy,
         )
         self.act_agent_os_facade: ComputerAgentOsFacade = ComputerAgentOsFacade(
             self.tools.os

src/askui/models/shared/conversation.py

@@ -13,9 +13,8 @@
 from askui.models.shared.settings import ActSettings
 from askui.models.shared.tools import ToolCollection
 from askui.models.shared.truncation_strategies import (
-    SimpleTruncationStrategyFactory,
+    SummarizingTruncationStrategy,
     TruncationStrategy,
-    TruncationStrategyFactory,
 )
 from askui.reporting import NULL_REPORTER, Reporter
 from askui.speaker.speaker import SpeakerResult, Speakers
@@ -55,7 +54,7 @@ class Conversation:
         detection_provider: Detection provider (optional)
         reporter: Reporter for logging messages and actions
         cache_manager: Cache manager for recording/playback (optional)
-        truncation_strategy_factory: Factory for creating truncation strategies
+        truncation_strategy: truncation strategies (optional)
         callbacks: List of callbacks for conversation lifecycle hooks (optional)
     """
 
@@ -67,7 +66,7 @@ def __init__(
         detection_provider: DetectionProvider | None = None,
         reporter: Reporter = NULL_REPORTER,
         cache_manager: "CacheManager | None" = None,
-        truncation_strategy_factory: TruncationStrategyFactory | None = None,
+        truncation_strategy: TruncationStrategy | None = None,
         callbacks: "list[ConversationCallback] | None" = None,
     ) -> None:
         """Initialize conversation with speakers and model providers."""
@@ -90,10 +89,6 @@ def __init__(
         # Infrastructure
         self._reporter = reporter
         self.cache_manager = cache_manager
-        self._truncation_strategy_factory = (
-            truncation_strategy_factory or SimpleTruncationStrategyFactory()
-        )
-        self._truncation_strategy: TruncationStrategy | None = None
         self._callbacks: "list[ConversationCallback]" = callbacks or []
 
         # State for current execution (set in start())
@@ -102,6 +97,22 @@ def __init__(
         self._reporters: list[Reporter] = []
         self._step_index: int = 0
 
+        # Truncation strategy. Conversation-owned dependencies are
+        # auto-injected so users can pass a custom strategy with only
+        # strategy-specific config (e.g. n_messages_to_keep) without
+        # needing access to vlm_provider/reporter/callbacks/conversation
+        # at construction time. ``vlm_provider`` is only injected when
+        # not pre-set, allowing callers to override the summarization
+        # VLM (e.g. with a cheaper model).
+        self._truncation_strategy: TruncationStrategy = (
+            truncation_strategy or SummarizingTruncationStrategy()
+        )
+        if self._truncation_strategy.vlm_provider is None:
+            self._truncation_strategy.vlm_provider = vlm_provider
+        self._truncation_strategy.reporter = reporter
+        self._truncation_strategy.callbacks = self._callbacks
+        self._truncation_strategy.conversation = self
+
         # Track if cache execution was used (to prevent recording during playback)
         self._executed_from_cache: bool = False
 
@@ -180,6 +191,7 @@ def _setup_control_loop(
         reporters: list[Reporter] | None = None,
     ) -> None:
         # Reset state
+        self._truncation_strategy.reset(messages)
         self._executed_from_cache = False
         self.speakers.reset_state()
 
@@ -191,16 +203,6 @@ def _setup_control_loop(
         # Auto-populate speaker descriptions and switch_speaker tool
         self._setup_speaker_handoff()
 
-        # Initialize truncation strategy
-        self._truncation_strategy = (
-            self._truncation_strategy_factory.create_truncation_strategy(
-                tools=self.tools.to_params(),
-                system=self.settings.messages.system,
-                messages=messages,
-                model=self.vlm_provider.model_id,
-            )
-        )
-
     @tracer.start_as_current_span("_execute_control_loop")
     def _execute_control_loop(self) -> None:
         self._on_control_loop_start()
@@ -448,7 +450,9 @@ def get_messages(self) -> list[MessageParam]:
         Returns:
             List of messages in current conversation
         """
-        return self._truncation_strategy.messages if self._truncation_strategy else []
+        return (
+            self._truncation_strategy.full_messages if self._truncation_strategy else []
+        )
 
     def get_truncation_strategy(self) -> TruncationStrategy | None:
         """Get current truncation strategy.