cypher - Add configurable token limit to cypher read tool (#157)

* Add configurable token limit to cypher read tool

* update changelog, change env var name, update docstrings

* Update README.md

* add unit tests

---------

Co-authored-by: alex <[email protected]>

Author: tomasonjo

servers/mcp-neo4j-cypher/CHANGELOG.md

@@ -7,6 +7,7 @@
 
 ### Added
 * Added Cypher result sanitation function from Neo4j GraphRAG that removes embedding values from the result
+* Add response token limit for read Cypher responses
 
 ## v0.3.1
 

servers/mcp-neo4j-cypher/README.md

@@ -51,6 +51,37 @@ The server supports namespacing to allow multiple Neo4j MCP servers to be used s
 
 This is useful when you need to connect to multiple Neo4j databases or instances from the same session.
 
+### ⚙️ Query Configuration
+
+The server provides configuration options to optimize query performance and manage response sizes:
+
+#### 📏 Token Limits
+
+Control the maximum size of query responses to prevent overwhelming the AI model:
+
+**Command Line:**
+```bash
+mcp-neo4j-cypher --token-limit 4000
+```
+
+**Environment Variable:**
+```bash
+export NEO4J_RESPONSE_TOKEN_LIMIT=4000
+```
+
+**Docker:**
+```bash
+docker run -e NEO4J_RESPONSE_TOKEN_LIMIT=4000 mcp-neo4j-cypher:latest
+```
+
+When a response exceeds the token limit, it will be automatically truncated to fit within the specified limit using `tiktoken`. This ensures:
+
+- **Consistent Performance**: Responses stay within model context limits
+- **Cost Control**: Prevents excessive token usage in AI interactions  
+- **Reliability**: Large datasets don't break the conversation flow
+
+**Note**: Token limits only apply to `read_neo4j_cypher` responses. Schema queries and write operations return summary information and are not affected.
+
 ## 🏗️ Local Development & Deployment
 
 ### 🐳 Local Docker Development
@@ -261,17 +292,18 @@ docker run --rm -p 8000:8000 \
 
 ### 🔧 Environment Variables
 
-| Variable                | Default                                 | Description                                    |
-| ----------------------- | --------------------------------------- | ---------------------------------------------- |
-| `NEO4J_URI`             | `bolt://localhost:7687`                 | Neo4j connection URI                           |
-| `NEO4J_USERNAME`        | `neo4j`                                 | Neo4j username                                 |
-| `NEO4J_PASSWORD`        | `password`                              | Neo4j password                                 |
-| `NEO4J_DATABASE`        | `neo4j`                                 | Neo4j database name                            |
-| `NEO4J_TRANSPORT`       | `stdio` (local), `http` (remote)        | Transport protocol (`stdio`, `http`, or `sse`) |
-| `NEO4J_NAMESPACE`       | _(empty)_                               | Tool namespace prefix                          |
-| `NEO4J_MCP_SERVER_HOST` | `127.0.0.1` (local)                     | Host to bind to                                |
-| `NEO4J_MCP_SERVER_PORT` | `8000`                                  | Port for HTTP/SSE transport                    |
-| `NEO4J_MCP_SERVER_PATH` | `/api/mcp/`                             | Path for accessing MCP server                  |
+| Variable                      | Default                                 | Description                                    |
+| ----------------------------- | --------------------------------------- | ---------------------------------------------- |
+| `NEO4J_URI`                   | `bolt://localhost:7687`                 | Neo4j connection URI                           |
+| `NEO4J_USERNAME`              | `neo4j`                                 | Neo4j username                                 |
+| `NEO4J_PASSWORD`              | `password`                              | Neo4j password                                 |
+| `NEO4J_DATABASE`              | `neo4j`                                 | Neo4j database name                            |
+| `NEO4J_TRANSPORT`             | `stdio` (local), `http` (remote)        | Transport protocol (`stdio`, `http`, or `sse`) |
+| `NEO4J_NAMESPACE`             | _(empty)_                               | Tool namespace prefix                          |
+| `NEO4J_MCP_SERVER_HOST`       | `127.0.0.1` (local)                     | Host to bind to                                |
+| `NEO4J_MCP_SERVER_PORT`       | `8000`                                  | Port for HTTP/SSE transport                    |
+| `NEO4J_MCP_SERVER_PATH`       | `/api/mcp/`                             | Path for accessing MCP server                  |
+| `NEO4J_RESPONSE_TOKEN_LIMIT`  | _(none)_                                | Maximum tokens for read query responses        |
 
 ### 🌐 SSE Transport for Legacy Web Access
 

servers/mcp-neo4j-cypher/manifest.json

@@ -31,7 +31,8 @@
         "NEO4J_NAMESPACE": "${user_config.neo4j_namespace}",
         "NEO4J_MCP_SERVER_HOST": "${user_config.mcp_server_host}",
         "NEO4J_MCP_SERVER_PORT": "${user_config.mcp_server_port}",
-        "NEO4J_MCP_SERVER_PATH": "${user_config.mcp_server_path}"
+        "NEO4J_MCP_SERVER_PATH": "${user_config.mcp_server_path}",
+        "NEO4J_RESPONSE_TOKEN_LIMIT": "${user_config.token_limit}"
       }
     }
   },
@@ -124,6 +125,14 @@
       "default": "/mcp/",
       "required": false,
       "sensitive": false
+    },
+    "token_limit": {
+      "type": "int",
+      "title": "Response token limit",
+      "description": "Optional response token limit for the read tool.",
+      "default": "",
+      "required": false,
+      "sensitive": false
     }
   }
 }

servers/mcp-neo4j-cypher/pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "fastmcp>=2.10.5",
     "neo4j>=5.26.0",
     "pydantic>=2.10.1",
+    "tiktoken>=0.11.0",
 ]
 
 [build-system]

servers/mcp-neo4j-cypher/src/mcp_neo4j_cypher/__init__.py

@@ -21,6 +21,7 @@ def main():
     )
     parser.add_argument("--server-host", default=None, help="Server host")
     parser.add_argument("--server-port", default=None, help="Server port")
+    parser.add_argument("--token-limit", default=None, help="Response token limit")
 
     args = parser.parse_args()
     config = process_config(args)

servers/mcp-neo4j-cypher/src/mcp_neo4j_cypher/server.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import re
-from typing import Any, Literal
+from typing import Any, Literal, Optional
 
 from fastmcp.exceptions import ToolError
 from fastmcp.server import FastMCP
@@ -10,7 +10,7 @@
 from neo4j import AsyncDriver, AsyncGraphDatabase, RoutingControl
 from neo4j.exceptions import ClientError, Neo4jError
 from pydantic import Field
-from .utils import _value_sanitize
+from .utils import _value_sanitize, _truncate_string_to_tokens
 
 logger = logging.getLogger("mcp_neo4j_cypher")
 
@@ -34,7 +34,10 @@ def _is_write_query(query: str) -> bool:
 
 
 def create_mcp_server(
-    neo4j_driver: AsyncDriver, database: str = "neo4j", namespace: str = ""
+    neo4j_driver: AsyncDriver,
+    database: str = "neo4j",
+    namespace: str = "",
+    token_limit: Optional[int] = None,
 ) -> FastMCP:
     mcp: FastMCP = FastMCP(
         "mcp-neo4j-cypher", dependencies=["neo4j", "pydantic"], stateless_http=True
@@ -183,6 +186,10 @@ async def read_neo4j_cypher(
             )
             sanitized_results = [_value_sanitize(el) for el in results]
             results_json_str = json.dumps(sanitized_results, default=str)
+            if token_limit:
+                results_json_str = _truncate_string_to_tokens(
+                    results_json_str, token_limit
+                )
 
             logger.debug(f"Read query returned {len(results_json_str)} rows")
 
@@ -254,6 +261,7 @@ async def main(
     host: str = "127.0.0.1",
     port: int = 8000,
     path: str = "/mcp/",
+    token_limit: Optional[int] = None,
 ) -> None:
     logger.info("Starting MCP neo4j Server")
 
@@ -265,7 +273,7 @@ async def main(
         ),
     )
 
-    mcp = create_mcp_server(neo4j_driver, database, namespace)
+    mcp = create_mcp_server(neo4j_driver, database, namespace, token_limit)
 
     # Run the server with the specified transport
     match transport:

servers/mcp-neo4j-cypher/src/mcp_neo4j_cypher/utils.py

@@ -1,3 +1,5 @@
+import tiktoken
+
 import argparse
 import logging
 import os
@@ -167,9 +169,19 @@ def process_config(args: argparse.Namespace) -> dict[str, Union[str, int, None]]
                 "Info: No server path provided and transport is `stdio`. `server_path` will be None."
             )
             config["path"] = None
+    # parse token limit
+    if args.token_limit is not None:
+        config["token_limit"] = args.token_limit
+    else:
+        if os.getenv("NEO4J_RESPONSE_TOKEN_LIMIT") is not None:
+            config["token_limit"] = int(os.getenv("NEO4J_RESPONSE_TOKEN_LIMIT"))
+        else:
+            logger.info("Info: No token limit provided. No token limit will be used.")
+            config["token_limit"] = None
 
     return config
 
+
 def _value_sanitize(d: Any, list_limit: int = 128) -> Any:
     """
     Sanitize the input dictionary or list.
@@ -222,4 +234,39 @@ def _value_sanitize(d: Any, list_limit: int = 128) -> Any:
         else:
             return None
     else:
-        return d
+        return d
+
+
+def _truncate_string_to_tokens(
+    text: str, token_limit: int, model: str = "gpt-4"
+) -> str:
+    """
+    Truncates the input string to fit within the specified token limit.
+
+    Parameters
+    ----------
+    text : str
+        The input text string.
+    token_limit : int
+        Maximum number of tokens allowed.
+    model : str
+        Model name (affects tokenization). Defaults to "gpt-4".
+
+    Returns
+    -------
+    str
+        The truncated string that fits within the token limit.
+    """
+    # Load encoding for the chosen model
+    encoding = tiktoken.encoding_for_model(model)
+
+    # Encode text into tokens
+    tokens = encoding.encode(text)
+
+    # Truncate tokens if they exceed the limit
+    if len(tokens) > token_limit:
+        tokens = tokens[:token_limit]
+
+    # Decode back into text
+    truncated_text = encoding.decode(tokens)
+    return truncated_text