fix list entities by type bug

gradio_mcp_space.py

@@ -433,22 +433,24 @@ Edge Relations:
     except Exception as e:
         return f"Error: {str(e)}"
 
-
 @observe(as_type="tool")
 def list_nodes_by_type(node_type: str, limit: int = 20, page: int = 1) -> str:
     """
     List nodes of a specific type in the knowledge graph.
-
+    
+    For entities, use entity_type (e.g., 'class', 'function', 'method').
+    For other nodes, use node_type (e.g., 'file', 'chunk', 'directory').
+    
     Args:
         node_type: The type of nodes to list (e.g., 'function', 'class', 'file')
         limit: Maximum number of nodes to return (default: 20)
-
+        page: Page number for pagination (default: 1)
     Returns:
         str: A formatted string with matching nodes
     """
     if knowledge_graph is None:
         return "Error: Knowledge graph not initialized"
-
+    
     try:
         # Convert limit/page to int if they're strings (MCP/Gradio may pass strings)
         if isinstance(limit, str):
@@ -456,52 +458,77 @@ def list_nodes_by_type(node_type: str, limit: int = 20, page: int = 1) -> str:
                 limit = int(limit)
             except ValueError:
                 return f"Error: 'limit' must be an integer, got '{limit}'"
-
+        
         if isinstance(page, str):
             try:
                 page = int(page)
             except ValueError:
                 return f"Error: 'page' must be an integer, got '{page}'"
-
+        
         if limit <= 0:
             return "Error: limit must be a positive integer"
         if page < 1:
             return "Error: 'page' must be a positive integer (1 or greater)"
-
+        
         g = knowledge_graph.graph
-        matching_nodes = [
-            {
-                "id": node_id,
-                "name": getattr(data['data'], 'name', 'Unknown')
-            }
-            for node_id, data in g.nodes(data=True)
-            if getattr(data['data'], 'node_type', None) == node_type
-        ]
-
+        matching_nodes = []
+        
+        for node_id, data in g.nodes(data=True):
+            node = data['data']
+            current_node_type = getattr(node, 'node_type', None)
+            node_name = getattr(node, 'name', 'Unknown')
+            
+            # For entity nodes, check entity_type instead of node_type
+            if current_node_type == 'entity':
+                entity_type = getattr(node, 'entity_type', '')
+                
+                # Fallback: if entity_type is empty, check the entities dictionary
+                if not entity_type and node_id in knowledge_graph.entities:
+                    entity_types = knowledge_graph.entities[node_id].get('type', [])
+                    entity_type = entity_types[0] if entity_types else ''
+                
+                if entity_type and entity_type.lower() == node_type.lower():
+                    matching_nodes.append({
+                        "id": node_id,
+                        "name": node_name,
+                        "type": f"entity ({entity_type})"
+                    })
+            # For other nodes, check node_type directly
+            elif current_node_type == node_type:
+                matching_nodes.append({
+                    "id": node_id,
+                    "name": node_name,
+                    "type": current_node_type
+                })
+        
+        # Sort by name for consistent ordering
+        matching_nodes.sort(key=lambda x: x['name'].lower())
+        
         total = len(matching_nodes)
         if total == 0:
             return f"No nodes found of type '{node_type}'."
-
+        
         # Pagination
         total_pages = (total + limit - 1) // limit
         if page > total_pages:
             return f"Error: Page {page} does not exist. Total pages: {total_pages} (with {total} nodes at {limit} per page)"
-
+        
         start_idx = (page - 1) * limit
         end_idx = start_idx + limit
         page_slice = matching_nodes[start_idx:end_idx]
-
+        
         result = f"Nodes of type '{node_type}' (Page {page}/{total_pages}, {total} total):\n"
         result += "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n\n"
-
+        
         for i, node in enumerate(page_slice, start=start_idx + 1):
             result += f"{i}. {node['name']}\n"
-            result += f"   ID: {node['id']}\n\n"
-
+            result += f"   ID: {node['id']}\n"
+            result += f"   Type: {node['type']}\n\n"
+        
         # Pagination hint
         if page < total_pages:
             result += f"Use page={page + 1} to see the next page\n"
-
+        
         return result
     except Exception as e:
         return f"Error: {str(e)}"
@@ -2043,7 +2070,7 @@ def create_gradio_app():
             with gr.Row():
                 with gr.Column():
                     node_type_input = gr.Dropdown(
-                        choices=["file", "directory", "chunk", "entity", "function", "class", "method"],
+                        choices=["file", "directory", "chunk", "function", "class", "method"],
                         label="Node Type"
                     )
                     type_limit = gr.Slider(1, 100, value=20, step=1, label="Max Results")
@@ -2059,7 +2086,7 @@ def create_gradio_app():
             with gr.Row():
                 with gr.Column():
                     search_type = gr.Dropdown(
-                        choices=["file", "directory", "chunk", "entity", "function", "class", "method"],
+                        choices=["file", "directory", "chunk", "function", "class", "method"],
                         label="Node Type"
                     )
                     search_name = gr.Textbox(label="Name Contains", placeholder="Enter partial name...")