diff --git a/mempalace/backends/pgvector.py b/mempalace/backends/pgvector.py index a97f29a34c..6269c12fe1 100644 --- a/mempalace/backends/pgvector.py +++ b/mempalace/backends/pgvector.py @@ -1058,6 +1058,21 @@ def count(self) -> int: return 0 return self._client.count_rows(self._table) + def count_by_metadata_field(self, field: str) -> dict[str, int]: + """Return ``{value: count}`` for a top-level metadata key, aggregated in + SQL. Lets callers (e.g. the status overview) tally by ``wing`` / ``room`` + without streaming every drawer's metadata to the client — O(1) memory and + a single query instead of O(n). ``NULL`` values bucket under ``"unknown"``.""" + self._ensure_open() + if not self._table_exists(): + return {} + sql = ( + f"SELECT metadata->>%s AS k, count(*) AS c " + f"FROM {_quote_identifier(self._table)} GROUP BY 1" + ) + rows = self._client._execute(sql, [field], fetch=True) + return {(r[0] if r[0] is not None else "unknown"): int(r[1]) for r in rows} + def lexical_search(self, *, query: str, n_results: int = 10, where: Optional[dict] = None): _validate_where(where) pushdown = None if _requires_local_filter(where) else where diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 9e2b3c6ab4..dd63b3cc85 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -1073,13 +1073,23 @@ def tool_status(): "backend": _selected_backend_name(), } try: - all_meta = _get_cached_metadata(col) - for m in all_meta: - m = m or {} - w = m.get("wing", "unknown") - r = m.get("room", "unknown") - wings[w] = wings.get(w, 0) + 1 - rooms[r] = rooms.get(r, 0) + 1 + grouper = getattr(col, "count_by_metadata_field", None) + if callable(grouper): + # Efficient path: backends that can aggregate server-side (e.g. + # pgvector via SQL GROUP BY) tally wing/room counts without + # streaming every drawer's metadata to the client. On a large + # palace the client-side fetch below is O(n) memory and can take + # minutes (or OOM); this is a single query. + wings.update(grouper("wing")) + rooms.update(grouper("room")) + else: + all_meta = _get_cached_metadata(col) + for m in all_meta: + m = m or {} + 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 as e: logger.exception("tool_status metadata fetch failed") result["error"] = str(e)