diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 7d263a6..886f122 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -96,16 +96,22 @@ def tool_status(): count = col.count() wings = {} rooms = {} - try: - all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"] - for m in all_meta: - w = m.get("wing", "unknown") - r = m.get("room", "unknown") - wings[w] = wings.get(w, 0) + 1 - rooms[r] = rooms.get(r, 0) + 1 - except Exception: - pass - return { + batch_size = 5000 + offset = 0 + error_info = None + while offset < count: + try: + batch = col.get(include=["metadatas"], limit=batch_size, offset=offset) + for m in batch["metadatas"]: + w = m.get("wing", "unknown") + r = m.get("room", "unknown") + wings[w] = wings.get(w, 0) + 1 + rooms[r] = rooms.get(r, 0) + 1 + offset += batch_size + except Exception as e: + error_info = f"Partial result, failed at offset {offset}: {str(e)}" + break + result = { "total_drawers": count, "wings": wings, "rooms": rooms, @@ -113,6 +119,10 @@ def tool_status(): "protocol": PALACE_PROTOCOL, "aaak_dialect": AAAK_SPEC, } + if error_info: + result["error"] = error_info + result["partial"] = True + return result # ── AAAK Dialect Spec ───────────────────────────────────────────────────────── @@ -153,13 +163,21 @@ def tool_list_wings(): if not col: return _no_palace() wings = {} + batch_size = 5000 + offset = 0 try: - all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"] - for m in all_meta: - w = m.get("wing", "unknown") - wings[w] = wings.get(w, 0) + 1 - except Exception: - pass + total = col.count() + except Exception as e: + return {"wings": {}, "error": str(e)} + while offset < total: + try: + batch = col.get(include=["metadatas"], limit=batch_size, offset=offset) + for m in batch["metadatas"]: + w = m.get("wing", "unknown") + wings[w] = wings.get(w, 0) + 1 + offset += batch_size + except Exception as e: + return {"wings": wings, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True} return {"wings": wings} @@ -168,16 +186,25 @@ def tool_list_rooms(wing: str = None): if not col: return _no_palace() rooms = {} + batch_size = 5000 + offset = 0 + where = {"wing": wing} if wing else None try: - kwargs = {"include": ["metadatas"], "limit": 10000} - if wing: - kwargs["where"] = {"wing": wing} - all_meta = col.get(**kwargs)["metadatas"] - for m in all_meta: - r = m.get("room", "unknown") - rooms[r] = rooms.get(r, 0) + 1 - except Exception: - pass + total = col.count() + except Exception as e: + return {"wing": wing or "all", "rooms": {}, "error": str(e)} + while offset < total: + try: + kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset} + if where: + kwargs["where"] = where + batch = col.get(**kwargs) + for m in batch["metadatas"]: + r = m.get("room", "unknown") + rooms[r] = rooms.get(r, 0) + 1 + offset += batch_size + except Exception as e: + return {"wing": wing or "all", "rooms": rooms, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True} return {"wing": wing or "all", "rooms": rooms} @@ -186,16 +213,24 @@ def tool_get_taxonomy(): if not col: return _no_palace() taxonomy = {} + batch_size = 5000 + offset = 0 try: - all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"] - for m in all_meta: - w = m.get("wing", "unknown") - r = m.get("room", "unknown") - if w not in taxonomy: - taxonomy[w] = {} - taxonomy[w][r] = taxonomy[w].get(r, 0) + 1 - except Exception: - pass + total = col.count() + except Exception as e: + return {"taxonomy": {}, "error": str(e)} + while offset < total: + try: + batch = col.get(include=["metadatas"], limit=batch_size, offset=offset) + for m in batch["metadatas"]: + w = m.get("wing", "unknown") + r = m.get("room", "unknown") + if w not in taxonomy: + taxonomy[w] = {} + taxonomy[w][r] = taxonomy[w].get(r, 0) + 1 + offset += batch_size + except Exception as e: + return {"taxonomy": taxonomy, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True} return {"taxonomy": taxonomy}