Message truncator hook to manage tokens (#332)

Author: keyurdoshi03

deploy/docker/docker-compose.yml

@@ -173,6 +173,7 @@ services:
       - API_PASSWORD=Admin!123
       - OPENAPI_SPEC=/app/resources/crapi-openapi-spec.json
       - DEFAULT_MODEL=gpt-4o-mini
+      - MAX_CONTENT_LENGTH=50000
       - CHROMA_HOST=chromadb
       - CHROMA_PORT=8000
       # - CHATBOT_OPENAI_API_KEY=

deploy/helm/templates/chatbot/config.yaml

@@ -22,6 +22,7 @@ data:
     MONGO_DB_NAME: {{ .Values.mongodb.config.mongoDbName }}
     CHATBOT_OPENAI_API_KEY: {{ .Values.openAIApiKey }}
     DEFAULT_MODEL: {{ .Values.chatbot.config.defaultModel | quote }}
+    MAX_CONTENT_LENGTH: {{ .Values.chatbot.config.maxContentLength | quote }}
     CHROMA_HOST: {{ .Values.chromadb.service.name }}
     CHROMA_PORT: {{ .Values.chromadb.port | quote }}
     API_USER: {{ .Values.chatbot.config.apiUser | quote }}

deploy/helm/values.yaml

@@ -153,6 +153,7 @@ chatbot:
     mongoDbDriver: mongodb
     secretKey: crapi
     defaultModel: gpt-4o-mini
+    maxContentLength: 50000
     chromaPersistDirectory: /app/vectorstore
     apiUser: [email protected]
     apiPassword: Admin!123

services/chatbot/.env

@@ -2,6 +2,7 @@ export TLS_ENABLED=false
 export SERVER_PORT=5002
 export WEB_SERVICE=localhost:8888
 export IDENTITY_SERVICE=localhost:8080
+export DB_NAME=crapi
 export DB_USER=admin
 export DB_PASSWORD=crapisecretpassword
 export DB_HOST=localhost
@@ -15,6 +16,7 @@ export [email protected]
 export API_PASSWORD=Admin!123
 export OPENAPI_SPEC=src/resources/crapi-openapi-spec.json
 export DEFAULT_MODEL=gpt-4o-mini
+export MAX_CONTENT_LENGTH=50000
 export CHROMA_HOST=localhost
 export CHROMA_PORT=8000
 export CHATBOT_OPENAI_API_KEY=

services/chatbot/src/chatbot/agent_utils.py

@@ -0,0 +1,73 @@
+import json
+from langchain_core.messages import ToolMessage
+from .config import Config
+
+INDIVIDUAL_MIN_LENGTH = 100
+
+def collect_long_strings(obj):
+    field_info = []
+    def _collect(obj):
+        if isinstance(obj, dict):
+            for key, value in obj.items():
+                if isinstance(value, str) and len(value) > INDIVIDUAL_MIN_LENGTH:
+                    field_info.append({
+                        'length': len(value),
+                        'dict': obj,
+                        'key': key,
+                    })
+                elif isinstance(value, (dict, list)):
+                    _collect(value)
+        elif isinstance(obj, list):
+            for item in obj:
+                if isinstance(item, (dict, list)):
+                    _collect(item)
+    
+    _collect(obj)
+    return field_info
+
+
+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.
+    """
+    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]"
+            )
+        
+        if cur_length <= max_length:
+            return json.dumps(data)
+    except (json.JSONDecodeError, Exception):
+        pass
+
+    return content[:max_length] + "\n... [TRUNCATED]"
+
+
+def truncate_tool_messages(state):
+    """
+    Modify large tool messages to prevent exceeding model's token limits.
+    Truncate to a length such that it keeps messages within your token limit.
+    """
+    messages = state.get("messages", [])
+    modified_messages = []
+    
+    for i,msg in enumerate(messages):
+        if isinstance(msg, ToolMessage) and len(msg.content) > Config.MAX_CONTENT_LENGTH:
+            truncated_msg = msg.model_copy(update={
+                'content': truncate_by_length(msg.content, Config.MAX_CONTENT_LENGTH)
+            })
+            modified_messages.append(truncated_msg)
+        else:
+            modified_messages.append(msg)
+    return {"messages": modified_messages}

services/chatbot/src/chatbot/chat_service.py

@@ -1,6 +1,6 @@
 from uuid import uuid4
 from langgraph.graph.message import Messages
-from .retrieverutils import add_to_chroma_collection
+from .retriever_utils import add_to_chroma_collection
 from .extensions import db
 from .langgraph_agent import execute_langgraph_agent
 
@@ -30,9 +30,10 @@ async def process_user_message(session_id, user_message, api_key, model_name, us
     response = await execute_langgraph_agent(
         api_key, model_name, history, user_jwt, session_id
     )
+    print("Session ID", session_id)
+    print("Messages", history)
     print("Response", response)
     reply: Messages = response.get("messages", [{}])[-1]
-    print("Reply", reply.content)
     response_message_id = uuid4().int & (1 << 63) - 1
     history.append(
         {"id": response_message_id, "role": "assistant", "content": reply.content}

services/chatbot/src/chatbot/config.py

@@ -11,5 +11,6 @@ class Config:
     SECRET_KEY = os.getenv("SECRET_KEY", "super-secret")
     MONGO_URI = MONGO_CONNECTION_URI
     DEFAULT_MODEL_NAME = os.getenv("DEFAULT_MODEL", "gpt-4o-mini")
+    MAX_CONTENT_LENGTH = int(os.getenv("MAX_CONTENT_LENGTH", 50000))
     CHROMA_HOST = CHROMA_HOST
     CHROMA_PORT = CHROMA_PORT

services/chatbot/src/chatbot/langgraph_agent.py

@@ -1,28 +1,12 @@
-import os
 import textwrap
-from typing import Annotated, Sequence, TypedDict
-
-from langchain.agents.agent_toolkits import create_retriever_tool
-from langchain.chains import LLMChain, RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain.schema import BaseMessage
-from langchain.tools import Tool
 from langchain_community.agent_toolkits import SQLDatabaseToolkit
-from langchain_community.agent_toolkits.sql.base import create_sql_agent
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_community.embeddings import OpenAIEmbeddings
-from langchain_community.vectorstores import FAISS  # or Chroma, Weaviate, etc.
 from langchain_openai import ChatOpenAI
-from langgraph.graph import MessageGraph, StateGraph
-from langgraph.graph.message import add_messages
 from langgraph.prebuilt import create_react_agent
-from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE, Settings
-from .retrieverutils import get_retriever_tool
+from .retriever_utils import get_retriever_tool
 
 from .extensions import postgresdb
-from .config import Config
 from .mcp_client import get_mcp_client
-import chromadb
+from .agent_utils import truncate_tool_messages
 
 
 async def build_langgraph_agent(api_key, model_name, user_jwt):
@@ -71,16 +55,13 @@ async def build_langgraph_agent(api_key, model_name, user_jwt):
     tools = mcp_tools + db_tools
     retriever_tool = get_retriever_tool(api_key)
     tools.append(retriever_tool)
-    agent_node = create_react_agent(model=llm, tools=tools, prompt=system_prompt)
+    agent_node = create_react_agent(model=llm, tools=tools, prompt=system_prompt, pre_model_hook=truncate_tool_messages)
     return agent_node
 
 
 async def execute_langgraph_agent(
     api_key, model_name, messages, user_jwt, session_id=None
 ):
     agent = await build_langgraph_agent(api_key, model_name, user_jwt)
-    print("messages", messages)
-    print("Session ID", session_id)
     response = await agent.ainvoke({"messages": messages})
-    print("Response", response)
     return response

…es/chatbot/src/chatbot/retrieverutils.py …s/chatbot/src/chatbot/retriever_utils.pyservices/chatbot/src/chatbot/retrieverutils.py renamed to services/chatbot/src/chatbot/retriever_utils.py services/chatbot/src/chatbot/retrieverutils.py renamed to services/chatbot/src/chatbot/retriever_utils.py

@@ -1,28 +1,7 @@
-import os
-import textwrap
-from typing import Annotated, Sequence, TypedDict
-
 from langchain.agents.agent_toolkits import create_retriever_tool
-from langchain.chains import LLMChain, RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain.schema import BaseMessage
-from langchain.tools import Tool
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
-from langchain_community.agent_toolkits.sql.base import create_sql_agent
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_community.vectorstores import FAISS  # or Chroma, Weaviate, etc.
-from langchain_openai import ChatOpenAI
-from langgraph.graph import MessageGraph, StateGraph
-from langgraph.graph.message import add_messages
-from langgraph.prebuilt import create_react_agent
-from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE
 import chromadb
 from langchain_chroma import Chroma as ChromaClient
-
-from .extensions import postgresdb
 from .config import Config
-from .mcp_client import get_mcp_client
-
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_core.documents import Document