feat(serve): honor a token_budget on get_neighbors/get_community

get_neighbors and get_community render every edge/member line unbounded — on a
god node or a large community that floods an MCP client's context window with
100KB+ of text in one tool result. query_graph already solves this with the
~3-chars/token cut in _subgraph_to_text; this applies the same budget rule to
the two line-list tools via a shared helper (_cut_lines_to_budget): cut at a
line boundary, report how many lines were dropped, and point at the narrowing
path (relation_filter / get_node). Default 2000 like query_graph; output under
budget is byte-identical to today.
This commit is contained in:
B
2026-07-21 21:36:30 +01:00
committed by safishamsi
parent a59abaf99d
commit fef9dbb5e7
+33 -3
View File
@@ -908,6 +908,26 @@ def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_bu
return output
def _cut_lines_to_budget(lines: list[str], token_budget: int, narrow_hint: str) -> str:
"""Render pre-built lines under the same ~3-chars/token budget rule as
_subgraph_to_text; over-budget output is cut at a line boundary with a count and a
narrowing hint instead of flooding the caller's context window."""
output = "\n".join(lines)
char_budget = token_budget * 3
if len(output) <= char_budget:
return output
cut_at = output[:char_budget].rfind("\n")
cut_at = cut_at if cut_at > 0 else char_budget
kept = output[:cut_at]
cut_count = len(lines) - kept.count("\n") - 1
return (
kept
+ f"\n... (truncated — {cut_count} more lines cut by ~{token_budget}-token budget. "
+ narrow_hint
+ ")"
)
def _query_graph_text(
G: nx.Graph,
question: str,
@@ -1189,6 +1209,7 @@ def _build_server(graph_path: str):
"properties": {
"label": {"type": "string"},
"relation_filter": {"type": "string", "description": "Optional: filter by relation type"},
"token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"},
},
"required": ["label"],
},
@@ -1198,7 +1219,10 @@ def _build_server(graph_path: str):
description="Get all nodes in a community by community ID.",
inputSchema={
"type": "object",
"properties": {"community_id": {"type": "integer", "description": "Community ID (0-indexed by size)"}},
"properties": {
"community_id": {"type": "integer", "description": "Community ID (0-indexed by size)"},
"token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"},
},
"required": ["community_id"],
},
),
@@ -1367,7 +1391,10 @@ def _build_server(graph_path: str):
f" <-- {sanitize_label(G.nodes[nb].get('label', nb))} "
f"[{sanitize_label(str(rel))}] [{sanitize_label(str(d.get('confidence', '')))}]{_edge_at(d)}"
)
return "\n".join(lines)
budget = int(arguments.get("token_budget", 2000))
return _cut_lines_to_budget(
lines, budget, "Narrow with relation_filter or use get_node for a specific symbol"
)
def _tool_get_community(arguments: dict) -> str:
cid = int(arguments["community_id"])
@@ -1383,7 +1410,10 @@ def _build_server(graph_path: str):
f" {sanitize_label(d.get('label', n))} "
f"[{sanitize_label(str(d.get('source_file', '')))}]"
)
return "\n".join(lines)
budget = int(arguments.get("token_budget", 2000))
return _cut_lines_to_budget(
lines, budget, "Raise token_budget or use get_node for specific members"
)
def _tool_god_nodes(arguments: dict) -> str:
from graphify.analyze import god_nodes as _god_nodes