From 36c2b2746b4cd55fe06136efe1c93bb3e0e35cc7 Mon Sep 17 00:00:00 2001 From: 0x5t4l1n Date: Tue, 27 Oct 2026 10:25:16 +0530 Subject: [PATCH] Include JSON escaping and markers in chatbot truncation budgets --- .github/workflows/chatbot-unit-tests.yml | 28 +++++++ services/chatbot/README.md | 16 ++++ services/chatbot/requirements-dev.txt | 1 + services/chatbot/src/chatbot/agent_utils.py | 46 +++++++++-- .../src/chatbot/tests/test_agent_utils.py | 79 ++++++++++++++++--- 5 files changed, 151 insertions(+), 19 deletions(-) create mode 100644 .github/workflows/chatbot-unit-tests.yml create mode 100644 services/chatbot/README.md diff --git a/.github/workflows/chatbot-unit-tests.yml b/.github/workflows/chatbot-unit-tests.yml new file mode 100644 index 000000000..109615646 --- /dev/null +++ b/.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 diff --git a/services/chatbot/README.md b/services/chatbot/README.md new file mode 100644 index 000000000..135d420bc --- /dev/null +++ b/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. diff --git a/services/chatbot/requirements-dev.txt b/services/chatbot/requirements-dev.txt index 7e66a17d4..6e9e9ddd2 100644 --- a/services/chatbot/requirements-dev.txt +++ b/services/chatbot/requirements-dev.txt @@ -1 +1,2 @@ black +pytest diff --git a/services/chatbot/src/chatbot/agent_utils.py b/services/chatbot/src/chatbot/agent_utils.py index 16cd57192..090f88ef3 100644 --- a/services/chatbot/src/chatbot/agent_utils.py +++ b/services/chatbot/src/chatbot/agent_utils.py @@ -40,11 +40,36 @@ 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"]) @@ -52,20 +77,25 @@ def truncate_by_length(content, max_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) diff --git a/services/chatbot/src/chatbot/tests/test_agent_utils.py b/services/chatbot/src/chatbot/tests/test_agent_utils.py index 790234116..d1b73faff 100644 --- a/services/chatbot/src/chatbot/tests/test_agent_utils.py +++ b/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