fix/over-counts tool outputs and blocks legitimate conversations

apps/chatbots/views.py

@@ -45,7 +45,7 @@
 from apps.experiments.views.experiment import (
     start_session_public,
 )
-from apps.experiments.views.utils import get_channels_context, get_max_char_limit
+from apps.experiments.views.utils import get_channels_context
 from apps.filters.models import FilterSet
 from apps.generics import actions
 from apps.generics.help import render_help_with_link
@@ -631,7 +631,6 @@ def chatbot_chat_session(request, team_slug: str, experiment_id: int, version_nu
         "experiment_name": experiment_version.name,
         "experiment_version": experiment_version,
         "experiment_version_number": experiment_version.version_number,
-        "max_char_limit": get_max_char_limit(experiment_version),
     }
     return TemplateResponse(
         request,
@@ -769,7 +768,6 @@ def _chatbot_chat_ui(request, embedded=False):
         "chatbot_name": chatbot_version.name,
         "experiment_version": chatbot_version,
         "experiment_version_number": chatbot_version.version_number,
-        "max_char_limit": get_max_char_limit(chatbot_version),
     }
     return TemplateResponse(
         request,

apps/experiments/views/experiment.py

@@ -79,7 +79,6 @@
     async_export_chat,
     get_response_for_webchat_task,
 )
-from apps.experiments.views.utils import get_max_char_limit
 from apps.files.models import File
 from apps.service_providers.llm_service.default_models import get_default_translation_models_by_provider
 from apps.service_providers.models import LlmProvider, LlmProviderModel
@@ -167,7 +166,6 @@ def _experiment_session_message(request, version_number: int, embedded=False):
     version_specific_vars = {
         "assistant": experiment_version.get_assistant(),
         "experiment_version_number": experiment_version.version_number,
-        "max_char_limit": get_max_char_limit(experiment_version),
     }
     return TemplateResponse(
         request,

apps/experiments/views/utils.py

@@ -1,16 +1,8 @@
 from apps.channels.models import ChannelPlatform, ExperimentChannel
-from apps.pipelines.utils import compute_max_char_limit
 
 
 def get_channels_context(experiment) -> tuple[list[ExperimentChannel], dict[ChannelPlatform, bool]]:
     channels = experiment.experimentchannel_set.exclude(platform__in=ChannelPlatform.team_global_platforms()).all()
     used_platforms = {channel.platform_enum for channel in channels}
     available_platforms = ChannelPlatform.for_dropdown(used_platforms, experiment.team)
     return channels, available_platforms
-
-
-def get_max_char_limit(experiment_version) -> int | None:
-    pipeline = experiment_version.pipeline
-    if not pipeline:
-        return None
-    return compute_max_char_limit(pipeline)

apps/pipelines/nodes/history_middleware.py

@@ -5,7 +5,7 @@
 
 from langchain.agents.middleware import AgentMiddleware
 from langchain.agents.middleware.summarization import SummarizationMiddleware
-from langchain_core.messages import BaseMessage, RemoveMessage
+from langchain_core.messages import BaseMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage
 from langchain_core.messages.utils import count_tokens_approximately
 from langgraph.graph.message import (
     REMOVE_ALL_MESSAGES,
@@ -14,6 +14,15 @@
 from apps.chat.conversation import COMPRESSION_MARKER
 from apps.pipelines.exceptions import MessageTooLargeError
 
+
+def get_token_counter(model=None):
+    """Returns a token-counting callable for the given model, falling back to
+    langchain's approximate counter when no model is provided."""
+    if model is None:
+        return count_tokens_approximately
+    return model.get_num_tokens_from_messages
+
+
 if TYPE_CHECKING:
     from apps.pipelines.nodes.nodes import PipelineNode
 
@@ -134,22 +143,42 @@ def persist_summary(self, messages: list[BaseMessage]):
 
 
 class MessageSizeValidationMiddleware(AgentMiddleware):
-    """Raises MessageTooLargeError before the model call if the full context exceeds the token budget.
+    """Raises MessageTooLargeError if the current user message or accumulated history
+    exceeds the model's token budget.
 
-    Runs after history compression and includes all messages in state
-    (history + current user input + any attachments).
+    Two checks run on each call: the last HumanMessage is validated individually, and
+    the total conversation history (HumanMessages + AIMessages) is validated together.
+    ToolMessages and SystemMessages are excluded from the total — tool responses are
+    transient and the token_limit already has system-prompt tokens subtracted.
     """
 
-    def __init__(self, token_limit: int):
+    def __init__(self, token_limit: int | None, token_counter):
         self._token_limit = token_limit
+        self._token_counter = token_counter
 
     def before_model(self, state, runtime):
-        if not self._token_limit:
+        if self._token_limit is None:
+            return None
+        all_messages = state["messages"]
+        human_messages = [m for m in all_messages if isinstance(m, HumanMessage)]
+        if not human_messages:
             return None
-        token_count = count_tokens_approximately(state["messages"])
+
+        # Validate the current user message individually.
+        token_count = self._token_counter([human_messages[-1]])
         if token_count > self._token_limit:
             raise MessageTooLargeError(
                 f"Your message is too large for this model. "
                 f"It uses approximately {token_count} tokens, but only {self._token_limit} tokens are available."
             )
+
+        # Validate total history, excluding tool responses and the system prompt.
+        conversation_messages = [m for m in all_messages if not isinstance(m, (ToolMessage, SystemMessage))]
+        total_token_count = self._token_counter(conversation_messages)
+        if total_token_count > self._token_limit:
+            raise MessageTooLargeError(
+                f"The conversation history is too large for this model. "
+                f"The total context uses approximately {total_token_count} tokens, but only "
+                f"{self._token_limit} tokens are available. Please start a new conversation."
+            )
         return None

apps/pipelines/nodes/llm_node.py

@@ -6,7 +6,6 @@
 from langchain.agents import create_agent
 from langchain.agents.middleware import AgentState
 from langchain_core.messages import AIMessage
-from langchain_core.messages.utils import count_tokens_approximately
 from langchain_core.tools import BaseTool
 
 from apps.chat.agent.tools import SearchCollectionByIdTool, SearchIndexTool, SearchToolConfig, get_node_tools
@@ -15,7 +14,6 @@
 from apps.files.models import File
 from apps.pipelines.nodes.base import PipelineNode, PipelineState
 from apps.pipelines.nodes.helpers import get_system_message
-from apps.pipelines.nodes.history_middleware import MessageSizeValidationMiddleware
 from apps.pipelines.nodes.tool_callbacks import ToolCallbacks
 from apps.service_providers.llm_service.datamodels import LlmChatResponse
 from apps.service_providers.llm_service.prompt_context import PromptTemplateContext
@@ -102,31 +100,19 @@ def build_node_agent(
     tools = _get_configured_tools(node, session=session, tool_callbacks=tool_callbacks)
     system_message = get_system_message(prompt_template=node.prompt, prompt_context=prompt_context)
 
-    middleware = []
-    if history_middleware := node.build_history_middleware(system_message=system_message):
-        middleware.append(history_middleware)
-    # MessageSizeValidationMiddleware temporarily disabled — over-counts tool outputs and blocks
-    # legitimate conversations. Re-enable after switching to a tool-aware token check.
+    model = node.get_chat_model()
+    middleware = node.build_history_middleware(system_message=system_message, model=model)
 
     return create_agent(
         # TODO: I think this will fail with google builtin tools
-        model=node.get_chat_model(),
+        model=model,
         tools=tools,
         system_prompt=system_message,
         middleware=middleware,
         state_schema=StateSchema,
     )
 
 
-def _build_size_validation_middleware(node: PipelineNode, system_message) -> MessageSizeValidationMiddleware | None:
-    max_token_limit = node.repo.get_llm_provider_model(node.llm_provider_model_id).max_token_limit
-    if not max_token_limit:
-        return None
-    system_tokens = count_tokens_approximately([system_message])
-    effective_limit = max(max_token_limit - system_tokens, 0)
-    return MessageSizeValidationMiddleware(token_limit=effective_limit)
-
-
 def _process_files(node: PipelineNode, cited_files: set[File], generated_files: set[File]) -> dict:
     """`cited_files` is a list of files that are cited in the response whereas generated files are those generated
     by the LLM

apps/pipelines/nodes/mixins.py

@@ -28,10 +28,11 @@
     Widgets,
 )
 from apps.pipelines.nodes.history_middleware import (
-    BaseNodeHistoryMiddleware,
     MaxHistoryLengthHistoryMiddleware,
+    MessageSizeValidationMiddleware,
     SummarizeHistoryMiddleware,
     TruncateTokensHistoryMiddleware,
+    get_token_counter,
 )
 from apps.pipelines.repository import ORMRepository, RepositoryLookupError
 from apps.service_providers.exceptions import ServiceProviderConfigError
@@ -222,34 +223,44 @@ def store_compression_checkpoint(self, compression_marker: str, checkpoint_messa
             history_mode=self.get_history_mode(),
         )
 
-    def build_history_middleware(self, system_message: BaseMessage) -> BaseNodeHistoryMiddleware | None:
-        """Construct the history compression middleware configured for this node."""
-        if self.history_is_disabled:
-            return None
-
-        history_mode = self.get_history_mode()
-
-        compressor_kwargs = {
-            "node": self,
-        }
-        if history_mode == PipelineChatHistoryModes.MAX_HISTORY_LENGTH:
-            return MaxHistoryLengthHistoryMiddleware(max_history_length=self.max_history_length, **compressor_kwargs)
-
-        specified_token_limit = (
-            self.user_max_token_limit
-            if self.user_max_token_limit is not None
-            else self.repo.get_llm_provider_model(self.llm_provider_model_id).max_token_limit
-        )
-
-        # Reserve space for the system message so trigger/keep thresholds reflect usable context
-        system_message_tokens = count_tokens_approximately([system_message])
-        token_limit = max(specified_token_limit - system_message_tokens, 100)
-
-        if history_mode == PipelineChatHistoryModes.SUMMARIZE:
-            return SummarizeHistoryMiddleware(token_limit=token_limit, **compressor_kwargs)
-
-        if history_mode == PipelineChatHistoryModes.TRUNCATE_TOKENS:
-            return TruncateTokensHistoryMiddleware(token_limit=token_limit, **compressor_kwargs)
+    def build_history_middleware(self, system_message: BaseMessage, model=None) -> list:
+        """Returns compression and size-validation middleware for this node."""
+        middleware = []
+        llm_provider_model = self.repo.get_llm_provider_model(self.llm_provider_model_id)
+        max_token_limit = llm_provider_model.max_token_limit
+        try:
+            token_counter = get_token_counter(model)
+            system_message_tokens = token_counter([system_message])
+        except Exception:
+            # Azure deployment names are often unrecognised by tiktoken.
+            token_counter = count_tokens_approximately
+            system_message_tokens = token_counter([system_message])
+
+        if not self.history_is_disabled:
+            history_mode = self.get_history_mode()
+            compressor_kwargs = {"node": self}
+
+            if history_mode == PipelineChatHistoryModes.MAX_HISTORY_LENGTH:
+                middleware.append(
+                    MaxHistoryLengthHistoryMiddleware(max_history_length=self.max_history_length, **compressor_kwargs)
+                )
+            else:
+                # Reserve space for the system message so trigger/keep thresholds reflect usable context
+                specified_token_limit = (
+                    self.user_max_token_limit if self.user_max_token_limit is not None else max_token_limit
+                )
+                token_limit = max(specified_token_limit - system_message_tokens, 100)
+
+                if history_mode == PipelineChatHistoryModes.SUMMARIZE:
+                    middleware.append(SummarizeHistoryMiddleware(token_limit=token_limit, **compressor_kwargs))
+                elif history_mode == PipelineChatHistoryModes.TRUNCATE_TOKENS:
+                    middleware.append(TruncateTokensHistoryMiddleware(token_limit=token_limit, **compressor_kwargs))
+
+        if max_token_limit:
+            effective_limit = max(max_token_limit - system_message_tokens, 0)
+            middleware.append(MessageSizeValidationMiddleware(token_limit=effective_limit, token_counter=token_counter))
+
+        return middleware
 
     def save_history(self, human_message: str, ai_message: str):
         if self.history_is_disabled: