fix(tools): implement tool output truncation to prevent LLM context overflow

Author: Adar5

chatbot-core/api/config/config.yml

@@ -39,3 +39,6 @@ tool_names:
 cors:
   allowed_origins:
     - "*"
+
+tools:
+  max_tool_output_length: 4000

chatbot-core/api/tools/tools.py

@@ -2,6 +2,8 @@
 Definition of the tools avaialable to the Agent.
 """
 
+import logging
+from functools import wraps
 from typing import Optional
 from types import MappingProxyType
 from api.models.embedding_model import EMBEDDING_MODEL
@@ -12,9 +14,39 @@
     extract_top_chunks
 )
 from api.config.loader import CONFIG
+# 1. Rename the logger to avoid clashing with the other functions
+decorator_logger = logging.getLogger(__name__)
+
+tool_config = CONFIG.get("tools", {})
+MAX_TOOL_OUTPUT_LENGTH = tool_config.get("max_tool_output_length", 4000)
+
+
+def truncate_tool_output(func):
+    """Decorator to prevent tool outputs from crashing the LLM context window."""
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        result = func(*args, **kwargs)
+        if isinstance(result, str) and len(result) > MAX_TOOL_OUTPUT_LENGTH:
+            truncated = result[:MAX_TOOL_OUTPUT_LENGTH]
+            # 2. Use lazy formatting instead of f-strings
+            decorator_logger.warning(
+                "[SECURITY] Tool '%s' output truncated from %d to %d characters.",
+                func.__name__,
+                len(result),
+                MAX_TOOL_OUTPUT_LENGTH
+            )
+            # 3. Break the long string into two lines
+            return (
+                f"{truncated}\n\n...[SYSTEM WARNING: Tool output truncated "
+                "to prevent context overflow.]"
+            )
+        return result
+    return wrapper
+
 
 retrieval_config = CONFIG["retrieval"]
 
+
 def search_plugin_docs(query: str, keywords: str, logger, plugin_name: Optional[str] = None) -> str:
     """
     Search tool for the plugin docs. Exploits both a sparse and dense search, resulting in a 
@@ -24,7 +56,7 @@ def search_plugin_docs(query: str, keywords: str, logger, plugin_name: Optional[
         query (str): The user query.
         keywords (str): Keywords extracted from the user query.
         plugin_name (Optional[str]): The refered plugin name in the query (if available).
-    
+
     Returns:
         str: The result of the research of the plugin search tool.
     """
@@ -55,6 +87,7 @@ def search_plugin_docs(query: str, keywords: str, logger, plugin_name: Optional[
         logger=logger
     )
 
+
 def search_jenkins_docs(query: str, keywords: str, logger) -> str:
     """
     Search tool for the Jenkins docs. Exploits both a sparse and dense search, resulting in a 
@@ -63,7 +96,7 @@ def search_jenkins_docs(query: str, keywords: str, logger) -> str:
     Args:
         query (str): The user query.
         keywords (str): Keywords extracted from the user query.
-    
+
     Returns:
         str: The result of the research of the docs search tool.
     """
@@ -87,6 +120,7 @@ def search_jenkins_docs(query: str, keywords: str, logger) -> str:
         logger=logger
     )
 
+
 def search_stackoverflow_threads(query: str) -> str:
     """
     Stackoverflow Search tool
@@ -95,6 +129,7 @@ def search_stackoverflow_threads(query: str) -> str:
         pass
     return "Nothing relevant"
 
+
 def search_community_threads(query: str, keywords: str, logger) -> str:
     """
     Search tool for the community discourse threads. Exploits both a sparse and 
@@ -104,7 +139,7 @@ def search_community_threads(query: str, keywords: str, logger) -> str:
     Args:
         query (str): The user query.
         keywords (str): Keywords extracted from the user query.
-    
+
     Returns:
         str: The result of the research of the docs search tool.
     """
@@ -129,9 +164,10 @@ def search_community_threads(query: str, keywords: str, logger) -> str:
         semantic_weight=0.7
     )
 
+
 TOOL_REGISTRY = MappingProxyType({
-    "search_plugin_docs": search_plugin_docs,
-    "search_jenkins_docs": search_jenkins_docs,
-    "search_stackoverflow_threads": search_stackoverflow_threads,
-    "search_community_threads": search_community_threads,
+    "search_plugin_docs": truncate_tool_output(search_plugin_docs),
+    "search_jenkins_docs": truncate_tool_output(search_jenkins_docs),
+    "search_stackoverflow_threads": truncate_tool_output(search_stackoverflow_threads),
+    "search_community_threads": truncate_tool_output(search_community_threads),
 })

chatbot-core/tests/unit/chunking/test_tool_overflow.py

@@ -0,0 +1,18 @@
+from api.tools.tools import truncate_tool_output, MAX_TOOL_OUTPUT_LENGTH
+
+
+def test_truncate_tool_output_prevents_overflow():
+    # 1. Create a dummy tool wrapped with our new security decorator
+    @truncate_tool_output
+    def massive_log_generator():
+        return "ERROR: Stack trace line. " * 50000  # Creates a massive string
+
+    # 2. Execute the tool
+    result = massive_log_generator()
+
+    # 3. Assert the string was successfully chopped down
+    # +100 to account for our warning message
+    assert len(result) <= (MAX_TOOL_OUTPUT_LENGTH + 100)
+
+    # 4. Assert our system warning was appended
+    assert "[SYSTEM WARNING: Tool output truncated" in result