perf(mcp): sqlite fast path for graph_stats to fix large-palace timeouts (#1379)
tool_graph_stats built the whole palace graph via build_graph(), which pages every metadata row (col.get limit/offset) and cold-loads the HNSW index — the remaining overview-tool timeout from #1379 (#1836 fixed status / list_wings / list_rooms / get_taxonomy but deliberately left graph_stats out, as it builds an in-memory graph rather than a flat tally). Add _sqlite_graph_stats(): one GROUP BY room, wing, hall over chroma.sqlite3, reconstructing build_graph's room_data and the same stats (total_rooms, tunnel_rooms, total_edges, rooms_per_wing, top_tunnels) with the same per-drawer filter (room present, != "general", wing present) and edge semantics (C(wings, 2) * halls per multi-wing room). Same _is_chroma_backend() guard + client-path fallback as the #1748 overview tools. Test seeds a real chroma palace mirroring the build_graph parity case in test_palace_graph, with a tripwire on graph_stats proving the fast path runs and that "general"/wing-less drawers are excluded. Idea adapted from #1381's _sqlite_graph_stats.
This commit is contained in:
Igor Lins e Silva
parent
2eda1f9130
commit
477aa362cd
|
|
@ -1014,6 +1014,116 @@ def _sqlite_taxonomy():
|
|||
return total, normalized
|
||||
|
||||
|
||||
def _sqlite_graph_stats():
|
||||
"""Compute ``graph_stats`` from one grouped sqlite read (#1379, graph_stats
|
||||
half; follow-up to #1748).
|
||||
|
||||
``graph_stats`` only needs grouped counts, but the client path builds the
|
||||
whole graph by paging every metadata row (``build_graph`` →
|
||||
``col.get(limit, offset)``) and cold-loads the HNSW index — which times out
|
||||
on six-figure palaces. This reads the same wing/room/hall grouping straight
|
||||
from ``chroma.sqlite3`` and reconstructs the stats.
|
||||
|
||||
Returns the stats dict, or ``None`` to fall back to the client path
|
||||
(non-chroma backend, missing/unbootstrapped palace, sqlite error). The
|
||||
reconstruction mirrors ``palace_graph.build_graph`` /
|
||||
``palace_graph.graph_stats`` exactly: a node is a room with a non-empty
|
||||
wing and a usable room name (the catch-all ``"general"`` is excluded), and
|
||||
edges are the per-hall cross-wing crossings of multi-wing rooms.
|
||||
"""
|
||||
if not _is_chroma_backend():
|
||||
return None
|
||||
import sqlite3 as _sqlite3
|
||||
from collections import Counter, defaultdict
|
||||
|
||||
db_path = os.path.join(_config.palace_path, "chroma.sqlite3")
|
||||
if not os.path.isfile(db_path):
|
||||
return None
|
||||
collection_name = _config.collection_name
|
||||
try:
|
||||
conn = _sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
|
||||
try:
|
||||
conn.execute("PRAGMA busy_timeout = 3000")
|
||||
if (
|
||||
conn.execute(
|
||||
"SELECT 1 FROM collections WHERE name = ?", (collection_name,)
|
||||
).fetchone()
|
||||
is None
|
||||
):
|
||||
return None
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT
|
||||
COALESCE(rm.string_value, CAST(rm.int_value AS TEXT),
|
||||
CAST(rm.float_value AS TEXT), '') AS room,
|
||||
COALESCE(wm.string_value, CAST(wm.int_value AS TEXT),
|
||||
CAST(wm.float_value AS TEXT), '') AS wing,
|
||||
COALESCE(hm.string_value, CAST(hm.int_value AS TEXT),
|
||||
CAST(hm.float_value AS TEXT), '') AS hall,
|
||||
COUNT(*) AS n
|
||||
FROM embeddings e
|
||||
JOIN segments s ON e.segment_id = s.id AND s.scope = 'METADATA'
|
||||
JOIN collections c ON s.collection = c.id
|
||||
LEFT JOIN embedding_metadata rm ON rm.id = e.id AND rm.key = 'room'
|
||||
LEFT JOIN embedding_metadata wm ON wm.id = e.id AND wm.key = 'wing'
|
||||
LEFT JOIN embedding_metadata hm ON hm.id = e.id AND hm.key = 'hall'
|
||||
WHERE c.name = ?
|
||||
GROUP BY room, wing, hall
|
||||
""",
|
||||
(collection_name,),
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
except _sqlite3.Error:
|
||||
logger.debug("sqlite graph_stats fast path failed; falling back", exc_info=True)
|
||||
return None
|
||||
|
||||
# Reconstruct build_graph()'s room_data, applying its per-drawer filter
|
||||
# (`if room and room != "general" and wing`).
|
||||
room_data = defaultdict(lambda: {"wings": set(), "halls": set(), "count": 0})
|
||||
for room, wing, hall, n in rows:
|
||||
if not room or room == "general" or not wing:
|
||||
continue
|
||||
node = room_data[room]
|
||||
node["wings"].add(wing)
|
||||
if hall:
|
||||
node["halls"].add(hall)
|
||||
node["count"] += int(n)
|
||||
|
||||
tunnel_rooms = 0
|
||||
total_edges = 0
|
||||
wing_counts = Counter()
|
||||
for data in room_data.values():
|
||||
n_wings = len(data["wings"])
|
||||
for wing in data["wings"]:
|
||||
wing_counts[wing] += 1
|
||||
if n_wings >= 2:
|
||||
tunnel_rooms += 1
|
||||
# Edges per multi-wing room: one per wing-pair per hall, matching
|
||||
# build_graph's nested wa<wb × hall expansion.
|
||||
total_edges += (n_wings * (n_wings - 1) // 2) * len(data["halls"])
|
||||
|
||||
top_tunnels = [
|
||||
{"room": room, "wings": sorted(data["wings"]), "count": data["count"]}
|
||||
# build_graph's graph_stats slices the top 10 by wing-count first, then
|
||||
# keeps the multi-wing ones. Secondary sort by room name keeps the
|
||||
# fast-path order deterministic (the client path's tie order follows
|
||||
# dict-insertion order, which sqlite grouping can't reproduce).
|
||||
for room, data in sorted(room_data.items(), key=lambda kv: (-len(kv[1]["wings"]), kv[0]))[
|
||||
:10
|
||||
]
|
||||
if len(data["wings"]) >= 2
|
||||
]
|
||||
|
||||
return {
|
||||
"total_rooms": len(room_data),
|
||||
"tunnel_rooms": tunnel_rooms,
|
||||
"total_edges": total_edges,
|
||||
"rooms_per_wing": dict(wing_counts.most_common()),
|
||||
"top_tunnels": top_tunnels,
|
||||
}
|
||||
|
||||
|
||||
def tool_status():
|
||||
# Run the safe sqlite/pickle probe before we touch chromadb. In the
|
||||
# #1222 failure mode, opening the persistent client to call .count()
|
||||
|
|
@ -1355,6 +1465,12 @@ def tool_find_tunnels(wing_a: str = None, wing_b: str = None):
|
|||
|
||||
def tool_graph_stats():
|
||||
"""Palace graph overview: nodes, tunnels, edges, connectivity."""
|
||||
# Fast path: grouped sqlite read instead of paging all metadata and
|
||||
# cold-loading HNSW via build_graph(), which times out on large palaces
|
||||
# (#1379). Falls through to the client path for non-chroma backends.
|
||||
fast = _sqlite_graph_stats()
|
||||
if fast is not None:
|
||||
return fast
|
||||
col = _get_collection()
|
||||
if not col:
|
||||
return _collection_error_or_no_palace()
|
||||
|
|
|
|||
|
|
@ -897,6 +897,49 @@ class TestReadTools:
|
|||
assert status["wings"] == {"unknown": 1}
|
||||
assert status["rooms"] == {"unknown": 1}
|
||||
|
||||
def test_graph_stats_uses_sqlite_fast_path(
|
||||
self, monkeypatch, config, palace_path, collection, kg
|
||||
):
|
||||
"""graph_stats must aggregate from sqlite without paging metadata
|
||||
through build_graph()/HNSW (#1379). Mirrors the build_graph parity
|
||||
case in test_palace_graph; the tripwire on graph_stats fails loudly if
|
||||
the fast path regresses to the slow client build."""
|
||||
collection.add(
|
||||
ids=["d_db_code", "d_db_proj", "d_auth", "d_general", "d_orphan"],
|
||||
documents=[
|
||||
"chromadb setup in the code wing",
|
||||
"chromadb usage in the project wing",
|
||||
"auth and security notes",
|
||||
"a general catch-all drawer",
|
||||
"a drawer with no wing",
|
||||
],
|
||||
metadatas=[
|
||||
{"room": "chromadb", "wing": "wing_code", "hall": "db"},
|
||||
{"room": "chromadb", "wing": "wing_project", "hall": "db"},
|
||||
{"room": "auth", "wing": "wing_code", "hall": "security"},
|
||||
{"room": "general", "wing": "wing_code", "hall": "misc"},
|
||||
{"room": "orphan", "source_file": "loose.txt"},
|
||||
],
|
||||
)
|
||||
_patch_mcp_server(monkeypatch, config, kg)
|
||||
from mempalace import mcp_server
|
||||
|
||||
def _boom(*_a, **_k):
|
||||
raise AssertionError("build_graph client path used instead of sqlite fast path")
|
||||
|
||||
monkeypatch.setattr(mcp_server, "graph_stats", _boom)
|
||||
|
||||
stats = mcp_server.tool_graph_stats()
|
||||
# "general" room and the wing-less drawer are excluded, matching
|
||||
# build_graph's per-drawer filter.
|
||||
assert stats["total_rooms"] == 2
|
||||
assert stats["tunnel_rooms"] == 1
|
||||
assert stats["total_edges"] == 1
|
||||
assert stats["rooms_per_wing"] == {"wing_code": 2, "wing_project": 1}
|
||||
assert stats["top_tunnels"] == [
|
||||
{"room": "chromadb", "wings": ["wing_code", "wing_project"], "count": 2}
|
||||
]
|
||||
|
||||
def test_no_palace_returns_error(self, monkeypatch, config, kg):
|
||||
_patch_mcp_server(monkeypatch, config, kg)
|
||||
from mempalace.mcp_server import tool_status
|
||||
|
|
|
|||
Loading…
Reference in New Issue