Keep chatbot tool-output truncation within its character budget

.github/workflows/chatbot-unit-tests.yml

@@ -0,0 +1,28 @@
+name: Chatbot utility tests
+
+on:
+  pull_request:
+    paths:
+      - 'services/chatbot/src/chatbot/agent_utils.py'
+      - 'services/chatbot/src/chatbot/config.py'
+      - 'services/chatbot/src/chatbot/tests/**'
+      - 'services/chatbot/requirements-dev.txt'
+      - '.github/workflows/chatbot-unit-tests.yml'
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  utilities:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - run: python -m pip install -r services/chatbot/requirements-dev.txt
+      - run: python -m pytest services/chatbot/src/chatbot/tests/test_agent_utils.py -q
+        env:
+          PYTHONPATH: services/chatbot/src
+      - run: python -m black --check services/chatbot/src/chatbot/agent_utils.py services/chatbot/src/chatbot/tests/test_agent_utils.py

services/chatbot/README.md

@@ -0,0 +1,16 @@
+# Chatbot utility tests
+
+From the repository root:
+
+```sh
+python -m pip install -r services/chatbot/requirements-dev.txt
+PYTHONPATH=services/chatbot/src python -m pytest services/chatbot/src/chatbot/tests/test_agent_utils.py -q
+```
+
+The utility tests stub the LLM and database dependencies; they do not require a
+running chatbot, database or API key.
+
+Tool-output truncation includes JSON escaping and truncation markers in
+`MAX_CONTENT_LENGTH`. It preserves JSON structure when the existing minimum
+retained field lengths fit. Otherwise it falls back to bounded text, reserving
+space for a marker when possible. Content already within the limit is unchanged.

services/chatbot/requirements-dev.txt

@@ -1 +1,2 @@
 black
+pytest

services/chatbot/src/chatbot/agent_utils.py

@@ -40,32 +40,62 @@ def _collect(obj):
     return field_info
 
 
+def _truncate_json_string(value, encoded_budget):
+    """Find the longest fitting prefix, or retain the minimum for later fields."""
+    low = INDIVIDUAL_MIN_LENGTH
+    high = len(value) - 1
+    best = value[:low] + f"... [TRUNCATED: {len(value) - low} chars removed]"
+
+    while low <= high:
+        retained = (low + high) // 2
+        candidate = (
+            value[:retained] + f"... [TRUNCATED: {len(value) - retained} chars removed]"
+        )
+        if len(json.dumps(candidate)) <= encoded_budget:
+            best = candidate
+            low = retained + 1
+        else:
+            high = retained - 1
+
+    return best
+
+
 def truncate_by_length(content, max_length):
     """
     Truncate JSON content by recursively truncating the longest fields until content is under limit.
     Preserves structure and smaller fields with minimum loss of information.
     """
+    if max_length <= 0:
+        return ""
+    if len(content) <= max_length:
+        return content
+
     try:
         data = json.loads(content)
         field_info = sorted(collect_long_strings(data), key=lambda x: x["length"])
 
         cur_length = len(json.dumps(data))
         while field_info and cur_length - max_length > 0:
             longest = field_info.pop()
-            excess = cur_length - max_length
-            new_length = max(INDIVIDUAL_MIN_LENGTH, longest["length"] - excess)
-            cur_length -= longest["length"] - new_length
-            longest["dict"][longest["key"]] = (
-                longest["dict"][longest["key"]][:new_length]
-                + f"... [TRUNCATED: {longest['length'] - new_length} chars removed]"
+            value = longest["dict"][longest["key"]]
+            old_length = len(json.dumps(value))
+            replacement = _truncate_json_string(
+                value, max_length - (cur_length - old_length)
             )
+            new_length = len(json.dumps(replacement))
+            if new_length < old_length:
+                longest["dict"][longest["key"]] = replacement
+                cur_length += new_length - old_length
 
         if cur_length <= max_length:
             return json.dumps(data)
-    except (json.JSONDecodeError, Exception):
+    except (json.JSONDecodeError, RecursionError):
         pass
 
-    return content[:max_length] + "\n... [TRUNCATED]"
+    marker = "\n... [TRUNCATED]"
+    if max_length <= len(marker):
+        return content[:max_length]
+    return content[: max_length - len(marker)] + marker
 
 
 @before_model(state_schema=AgentState)

services/chatbot/src/chatbot/tests/test_agent_utils.py

@@ -1,4 +1,6 @@
-"""Tests for trim_messages_to_token_limit and supporting helpers."""
+"""Tests for tool-output truncation, history trimming and supporting helpers."""
+
+import json
 import sys
 from types import ModuleType
 from unittest.mock import MagicMock, patch
@@ -10,12 +12,18 @@
 # ---------------------------------------------------------------------------
 _STUBS = {}
 for mod_name in [
-    "langchain", "langchain.agents", "langchain.agents.middleware",
+    "langchain",
+    "langchain.agents",
+    "langchain.agents.middleware",
     "langchain.agents.middleware.types",
-    "langchain_core", "langchain_core.messages",
-    "langgraph", "langgraph.runtime",
-    "motor", "motor.motor_asyncio",
-    "langchain_community", "langchain_community.agent_toolkits",
+    "langchain_core",
+    "langchain_core.messages",
+    "langgraph",
+    "langgraph.runtime",
+    "motor",
+    "motor.motor_asyncio",
+    "langchain_community",
+    "langchain_community.agent_toolkits",
     "langchain_community.utilities",
     "pymongo",
 ]:
@@ -26,8 +34,8 @@
 
 # Provide the decorator used by agent_utils at import time
 sys.modules["langchain.agents"].AgentState = dict
-sys.modules["langchain.agents.middleware.types"].before_model = (
-    lambda **kw: (lambda fn: fn)
+sys.modules["langchain.agents.middleware.types"].before_model = lambda **kw: (
+    lambda fn: fn
 )
 sys.modules["langchain_core.messages"].ToolMessage = type("ToolMessage", (), {})
 sys.modules["langgraph.runtime"].Runtime = type("Runtime", (), {})
@@ -51,22 +59,66 @@
     _estimate_tokens,
     _message_content,
     trim_messages_to_token_limit,
+    truncate_by_length,
 )
 
-
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
 
+
 def _make_msg(role, content):
     """Return a plain dict message like those stored in chat history."""
     return {"role": role, "content": content}
 
 
+class TestTruncateByLength:
+    @pytest.mark.parametrize(
+        "value, limit",
+        [("x" * 1000, 500), ("\n" * 1000, 500), ("é" * 1000, 900)],
+        ids=["ascii", "escaped", "unicode"],
+    )
+    def test_json_result_includes_marker_and_escaping_in_budget(self, value, limit):
+        content = json.dumps({"description": value, "id": 7})
+        result = truncate_by_length(content, limit)
+        assert len(result) <= limit
+        data = json.loads(result)
+        assert data["id"] == 7
+        assert "[TRUNCATED:" in data["description"]
+
+    def test_multiple_nested_fields_fit_budget(self):
+        content = json.dumps(
+            {"items": [{"description": "a" * 1000}, {"description": "b" * 1000}]}
+        )
+        result = truncate_by_length(content, 400)
+        assert len(result) <= 400
+        assert len(json.loads(result)["items"]) == 2
+
+    @pytest.mark.parametrize("limit", [0, 1, 10, 15, 50, 100])
+    @pytest.mark.parametrize(
+        "content",
+        ["plain text " * 100, json.dumps([1] * 100)],
+        ids=["text", "json-array"],
+    )
+    def test_fallback_includes_marker_in_budget(self, content, limit):
+        result = truncate_by_length(content, limit)
+        assert len(result) == limit
+        marker = "\n... [TRUNCATED]"
+        if limit > len(marker):
+            assert result == content[: limit - len(marker)] + marker
+        else:
+            assert result == content[:limit]
+
+    @pytest.mark.parametrize("content", ["hello", '{"id":7}', '  {"id": 7}  '])
+    def test_content_that_fits_is_unchanged(self, content):
+        assert truncate_by_length(content, 100) == content
+
+
 # ---------------------------------------------------------------------------
 # _estimate_tokens
 # ---------------------------------------------------------------------------
 
+
 class TestEstimateTokens:
     def test_empty_string(self):
         assert _estimate_tokens("") == 0
@@ -83,6 +135,7 @@ def test_short_string(self):
 # _message_content
 # ---------------------------------------------------------------------------
 
+
 class TestMessageContent:
     def test_dict_message(self):
         assert _message_content({"role": "user", "content": "hello"}) == "hello"
@@ -93,11 +146,13 @@ def test_dict_missing_content(self):
     def test_object_message(self):
         class Msg:
             content = "from object"
+
         assert _message_content(Msg()) == "from object"
 
     def test_object_no_content(self):
         class Msg:
             pass
+
         assert _message_content(Msg()) == ""
 
 
@@ -173,8 +228,10 @@ def test_result_fits_within_token_budget(self):
         """After trimming, estimated tokens must be <= budget."""
         token_budget = MAX_CONTENT_LENGTH // CHARS_PER_TOKEN
         # 20 messages each ~2500 tokens = 50000 tokens, well over 25000 budget
-        msgs = [_make_msg("user" if i % 2 == 0 else "assistant", "x" * 10000)
-                for i in range(20)]
+        msgs = [
+            _make_msg("user" if i % 2 == 0 else "assistant", "x" * 10000)
+            for i in range(20)
+        ]
         result = trim_messages_to_token_limit(msgs)
         result_tokens = sum(_estimate_tokens(m["content"]) for m in result)
         assert result_tokens <= token_budget