pass at cleaning up keys to reduce tokens

Author: Btibert3

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

@@ -62,6 +62,69 @@ async def get_neo4j_schema() -> list[types.TextContent]:
         CALL apoc.meta.schema();
         """
 
+        def clean_schema(schema: dict) -> dict:
+            cleaned = {}
+
+            for key, entry in schema.items():
+
+                new_entry = {
+                    "type": entry["type"]
+                }
+                if "count" in entry:
+                    new_entry["count"] = entry["count"]
+
+                labels = entry.get("labels", [])
+                if labels:
+                    new_entry["labels"] = labels
+
+                props = entry.get("properties", {})
+                clean_props = {}
+                for pname, pinfo in props.items():
+                    cp = {}
+                    if "indexed" in pinfo:
+                        cp["indexed"] = pinfo["indexed"]
+                    if "type" in pinfo:
+                        cp["type"] = pinfo["type"]
+                    if cp:
+                        clean_props[pname] = cp
+                if clean_props:
+                    new_entry["properties"] = clean_props
+
+                if entry.get("relationships"):
+                    rels_out = {}
+                    for rel_name, rel in entry["relationships"].items():
+                        cr = {}
+                        if "direction" in rel:
+                            cr["direction"] = rel["direction"]
+                        # nested labels
+                        rlabels = rel.get("labels", [])
+                        if rlabels:
+                            cr["labels"] = rlabels
+                        # nested properties
+                        rprops = rel.get("properties", {})
+                        clean_rprops = {}
+                        for rpname, rpinfo in rprops.items():
+                            crp = {}
+                            if "indexed" in rpinfo:
+                                crp["indexed"] = rpinfo["indexed"]
+                            if "type" in rpinfo:
+                                crp["type"] = rpinfo["type"]
+                            if crp:
+                                clean_rprops[rpname] = crp
+                        if clean_rprops:
+                            cr["properties"] = clean_rprops
+
+                        if cr:
+                            rels_out[rel_name] = cr
+
+                    if rels_out:
+                        new_entry["relationships"] = rels_out
+
+                cleaned[key] = new_entry
+
+            return cleaned
+
+
         try:
             async with neo4j_driver.session(database=database) as session:
                 results_json_str = await session.execute_read(
@@ -70,7 +133,11 @@ async def get_neo4j_schema() -> list[types.TextContent]:
 
                 logger.debug(f"Read query returned {len(results_json_str)} rows")
 
-                return [types.TextContent(type="text", text=results_json_str)]
+                schema = json.loads(results_json_str)[0].get('value')
+                schema_clean = clean_schema(schema)
+                schema_clean_str = json.dumps(schema_clean)
+
+                return [types.TextContent(type="text", text=schema_clean_str)]
 
         except Exception as e:
             logger.error(f"Database error retrieving schema: {e}")