hanzo-mcp 0.7.7__py3-none-any.whl → 0.8.0__py3-none-any.whl

hanzo_mcp/tools/database/graph_query.py

@@ -2,18 +2,24 @@
 
 import json
 import sqlite3
-from typing import Annotated, Optional, TypedDict, Unpack, final, override, List, Dict, Any
+from typing import (
+    Unpack,
+    Optional,
+    Annotated,
+    TypedDict,
+    final,
+    override,
+)
 from collections import deque
 
-from mcp.server.fastmcp import Context as MCPContext
 from pydantic import Field
+from mcp.server.fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
 from hanzo_mcp.tools.common.permissions import PermissionManager
 from hanzo_mcp.tools.database.database_manager import DatabaseManager
 
-
 Query = Annotated[
     str,
     Field(
@@ -96,7 +102,9 @@ class GraphQueryParams(TypedDict, total=False):
 class GraphQueryTool(BaseTool):
     """Tool for querying the graph database."""
 
-    def __init__(self, permission_manager: PermissionManager, db_manager: DatabaseManager):
+    def __init__(
+        self, permission_manager: PermissionManager, db_manager: DatabaseManager
+    ):
         """Initialize the graph query tool.
 
         Args:
@@ -172,12 +180,22 @@ Examples:
         project_path = params.get("project_path")
 
         # Validate query type
-        valid_queries = ["neighbors", "path", "subgraph", "connected", "ancestors", "descendants"]
+        valid_queries = [
+            "neighbors",
+            "path",
+            "subgraph",
+            "connected",
+            "ancestors",
+            "descendants",
+        ]
         if query not in valid_queries:
             return f"Error: Invalid query '{query}'. Must be one of: {', '.join(valid_queries)}"
 
         # Validate required parameters
-        if query in ["neighbors", "subgraph", "connected", "ancestors", "descendants"] and not node_id:
+        if (
+            query in ["neighbors", "subgraph", "connected", "ancestors", "descendants"]
+            and not node_id
+        ):
             return f"Error: node_id is required for '{query}' query"
 
         if query == "path" and (not node_id or not target_id):
@@ -189,8 +207,9 @@ Examples:
                 project_db = self.db_manager.get_project_db(project_path)
             else:
                 import os
+
                 project_db = self.db_manager.get_project_for_path(os.getcwd())
-                
+
             if not project_db:
                 return "Error: Could not find project database"
 
@@ -206,28 +225,43 @@ Examples:
 
         try:
             if query == "neighbors":
-                return self._query_neighbors(graph_conn, node_id, relationship, node_type, direction)
+                return self._query_neighbors(
+                    graph_conn, node_id, relationship, node_type, direction
+                )
             elif query == "path":
                 return self._query_path(graph_conn, node_id, target_id, relationship)
             elif query == "subgraph":
-                return self._query_subgraph(graph_conn, node_id, depth, relationship, node_type, direction)
+                return self._query_subgraph(
+                    graph_conn, node_id, depth, relationship, node_type, direction
+                )
             elif query == "connected":
-                return self._query_connected(graph_conn, node_id, relationship, node_type, direction)
+                return self._query_connected(
+                    graph_conn, node_id, relationship, node_type, direction
+                )
             elif query == "ancestors":
-                return self._query_ancestors(graph_conn, node_id, depth, relationship, node_type)
+                return self._query_ancestors(
+                    graph_conn, node_id, depth, relationship, node_type
+                )
             elif query == "descendants":
-                return self._query_descendants(graph_conn, node_id, depth, relationship, node_type)
+                return self._query_descendants(
+                    graph_conn, node_id, depth, relationship, node_type
+                )
 
         except Exception as e:
             await tool_ctx.error(f"Failed to execute query: {str(e)}")
             return f"Error executing query: {str(e)}"
 
-    def _query_neighbors(self, conn: sqlite3.Connection, node_id: str, 
-                        relationship: Optional[str], node_type: Optional[str], 
-                        direction: str) -> str:
+    def _query_neighbors(
+        self,
+        conn: sqlite3.Connection,
+        node_id: str,
+        relationship: Optional[str],
+        node_type: Optional[str],
+        direction: str,
+    ) -> str:
         """Get direct neighbors of a node."""
         cursor = conn.cursor()
-        
+
         # Check if node exists
         cursor.execute("SELECT type, properties FROM nodes WHERE id = ?", (node_id,))
         node_info = cursor.fetchone()
@@ -235,76 +269,89 @@ Examples:
             return f"Error: Node '{node_id}' not found"
 
         neighbors = []
-        
+
         # Get outgoing edges
         if direction in ["both", "outgoing"]:
             query = """SELECT e.target, e.relationship, e.weight, n.type, n.properties
                       FROM edges e JOIN nodes n ON e.target = n.id
                       WHERE e.source = ?"""
             params = [node_id]
-            
+
             if relationship:
                 query += " AND e.relationship = ?"
                 params.append(relationship)
             if node_type:
                 query += " AND n.type = ?"
                 params.append(node_type)
-                
+
             cursor.execute(query, params)
             for row in cursor.fetchall():
-                neighbors.append({
-                    "direction": "outgoing",
-                    "node_id": row[0],
-                    "relationship": row[1],
-                    "weight": row[2],
-                    "node_type": row[3],
-                    "properties": json.loads(row[4]) if row[4] else {}
-                })
-        
+                neighbors.append(
+                    {
+                        "direction": "outgoing",
+                        "node_id": row[0],
+                        "relationship": row[1],
+                        "weight": row[2],
+                        "node_type": row[3],
+                        "properties": json.loads(row[4]) if row[4] else {},
+                    }
+                )
+
         # Get incoming edges
         if direction in ["both", "incoming"]:
             query = """SELECT e.source, e.relationship, e.weight, n.type, n.properties
                       FROM edges e JOIN nodes n ON e.source = n.id
                       WHERE e.target = ?"""
             params = [node_id]
-            
+
             if relationship:
                 query += " AND e.relationship = ?"
                 params.append(relationship)
             if node_type:
                 query += " AND n.type = ?"
                 params.append(node_type)
-                
+
             cursor.execute(query, params)
             for row in cursor.fetchall():
-                neighbors.append({
-                    "direction": "incoming",
-                    "node_id": row[0],
-                    "relationship": row[1],
-                    "weight": row[2],
-                    "node_type": row[3],
-                    "properties": json.loads(row[4]) if row[4] else {}
-                })
-        
+                neighbors.append(
+                    {
+                        "direction": "incoming",
+                        "node_id": row[0],
+                        "relationship": row[1],
+                        "weight": row[2],
+                        "node_type": row[3],
+                        "properties": json.loads(row[4]) if row[4] else {},
+                    }
+                )
+
         if not neighbors:
             return f"No neighbors found for node '{node_id}'"
-        
+
         # Format output
         output = [f"Neighbors of '{node_id}' ({node_info[0]}):\n"]
         for n in neighbors:
             arrow = "<--" if n["direction"] == "incoming" else "-->"
-            output.append(f"  {node_id} {arrow}[{n['relationship']}]--> {n['node_id']} ({n['node_type']})")
+            output.append(
+                f"  {node_id} {arrow}[{n['relationship']}]--> {n['node_id']} ({n['node_type']})"
+            )
             if n["properties"]:
-                output.append(f"    Properties: {json.dumps(n['properties'], indent=6)[:100]}")
-        
+                output.append(
+                    f"    Properties: {json.dumps(n['properties'], indent=6)[:100]}"
+                )
+
         output.append(f"\nTotal neighbors: {len(neighbors)}")
         return "\n".join(output)
 
-    def _query_path(self, conn: sqlite3.Connection, start: str, end: str, 
-                   relationship: Optional[str]) -> str:
+    def _query_path(
+        self,
+        conn: sqlite3.Connection,
+        start: str,
+        end: str,
+        relationship: Optional[str],
+    ) -> str:
         """Find shortest path between two nodes using BFS."""
         cursor = conn.cursor()
-        
+
         # Check if nodes exist
         cursor.execute("SELECT id FROM nodes WHERE id IN (?, ?)", (start, end))
         existing = [row[0] for row in cursor.fetchall()]
@@ -312,224 +359,253 @@ Examples:
             return f"Error: Start node '{start}' not found"
         if end not in existing:
             return f"Error: End node '{end}' not found"
-        
+
         # BFS to find shortest path
         queue = deque([(start, [start])])
         visited = {start}
-        
+
         while queue:
             current, path = queue.popleft()
-            
+
             if current == end:
                 # Found path, get edge details
                 output = [f"Shortest path from '{start}' to '{end}':\n"]
-                
+
                 for i in range(len(path) - 1):
                     src, tgt = path[i], path[i + 1]
-                    
+
                     # Get edge details
                     query = "SELECT relationship, weight FROM edges WHERE source = ? AND target = ?"
                     cursor.execute(query, (src, tgt))
                     edge = cursor.fetchone()
-                    
+
                     if edge:
                         output.append(f"  {src} --[{edge[0]}]--> {tgt}")
                     else:
                         output.append(f"  {src} --> {tgt}")
-                
+
                 output.append(f"\nPath length: {len(path) - 1} edge(s)")
                 return "\n".join(output)
-            
+
             # Get neighbors
             query = "SELECT target FROM edges WHERE source = ?"
             params = [current]
             if relationship:
                 query += " AND relationship = ?"
                 params.append(relationship)
-                
+
             cursor.execute(query, params)
-            
+
             for (neighbor,) in cursor.fetchall():
                 if neighbor not in visited:
                     visited.add(neighbor)
                     queue.append((neighbor, path + [neighbor]))
-        
-        return f"No path found from '{start}' to '{end}'" + (f" with relationship '{relationship}'" if relationship else "")
 
-    def _query_subgraph(self, conn: sqlite3.Connection, node_id: str, depth: int,
-                       relationship: Optional[str], node_type: Optional[str], 
-                       direction: str) -> str:
+        return f"No path found from '{start}' to '{end}'" + (
+            f" with relationship '{relationship}'" if relationship else ""
+        )
+
+    def _query_subgraph(
+        self,
+        conn: sqlite3.Connection,
+        node_id: str,
+        depth: int,
+        relationship: Optional[str],
+        node_type: Optional[str],
+        direction: str,
+    ) -> str:
         """Get subgraph around a node up to specified depth."""
         cursor = conn.cursor()
-        
+
         # Check if node exists
         cursor.execute("SELECT type FROM nodes WHERE id = ?", (node_id,))
         if not cursor.fetchone():
             return f"Error: Node '{node_id}' not found"
-        
+
         # BFS to collect nodes and edges
         nodes = {node_id: 0}  # node_id -> depth
         edges = set()  # (source, target, relationship)
         queue = deque([(node_id, 0)])
-        
+
         while queue:
             current, current_depth = queue.popleft()
-            
+
             if current_depth >= depth:
                 continue
-            
+
             # Get edges based on direction
             if direction in ["both", "outgoing"]:
                 query = """SELECT e.target, e.relationship, n.type 
                           FROM edges e JOIN nodes n ON e.target = n.id
                           WHERE e.source = ?"""
                 params = [current]
-                
+
                 if relationship:
                     query += " AND e.relationship = ?"
                     params.append(relationship)
                 if node_type:
                     query += " AND n.type = ?"
                     params.append(node_type)
-                    
+
                 cursor.execute(query, params)
-                
+
                 for target, rel, _ in cursor.fetchall():
                     edges.add((current, target, rel))
                     if target not in nodes or nodes[target] > current_depth + 1:
                         nodes[target] = current_depth + 1
                         queue.append((target, current_depth + 1))
-            
+
             if direction in ["both", "incoming"]:
                 query = """SELECT e.source, e.relationship, n.type 
                           FROM edges e JOIN nodes n ON e.source = n.id
                           WHERE e.target = ?"""
                 params = [current]
-                
+
                 if relationship:
                     query += " AND e.relationship = ?"
                     params.append(relationship)
                 if node_type:
                     query += " AND n.type = ?"
                     params.append(node_type)
-                    
+
                 cursor.execute(query, params)
-                
+
                 for source, rel, _ in cursor.fetchall():
                     edges.add((source, current, rel))
                     if source not in nodes or nodes[source] > current_depth + 1:
                         nodes[source] = current_depth + 1
                         queue.append((source, current_depth + 1))
-        
+
         # Format output
         output = [f"Subgraph around '{node_id}' (depth={depth}):\n"]
         output.append(f"Nodes ({len(nodes)}):")
-        
+
         # Get node details
         for node, d in sorted(nodes.items(), key=lambda x: (x[1], x[0])):
             cursor.execute("SELECT type FROM nodes WHERE id = ?", (node,))
             node_type = cursor.fetchone()[0]
             output.append(f"  [{d}] {node} ({node_type})")
-        
+
         output.append(f"\nEdges ({len(edges)}):")
         for src, tgt, rel in sorted(edges):
             output.append(f"  {src} --[{rel}]--> {tgt}")
-        
+
         return "\n".join(output)
 
-    def _query_connected(self, conn: sqlite3.Connection, node_id: str,
-                        relationship: Optional[str], node_type: Optional[str], 
-                        direction: str) -> str:
+    def _query_connected(
+        self,
+        conn: sqlite3.Connection,
+        node_id: str,
+        relationship: Optional[str],
+        node_type: Optional[str],
+        direction: str,
+    ) -> str:
         """Find all nodes connected to a node (transitive closure)."""
         cursor = conn.cursor()
-        
+
         # Check if node exists
         cursor.execute("SELECT type FROM nodes WHERE id = ?", (node_id,))
         if not cursor.fetchone():
             return f"Error: Node '{node_id}' not found"
-        
+
         # BFS to find all connected nodes
         visited = {node_id}
         queue = deque([node_id])
         connections = []  # (node_id, node_type, distance)
         distance = {node_id: 0}
-        
+
         while queue:
             current = queue.popleft()
             current_dist = distance[current]
-            
+
             # Get edges based on direction
             neighbors = []
-            
+
             if direction in ["both", "outgoing"]:
                 query = """SELECT e.target, n.type FROM edges e 
                           JOIN nodes n ON e.target = n.id
                           WHERE e.source = ?"""
                 params = [current]
-                
+
                 if relationship:
                     query += " AND e.relationship = ?"
                     params.append(relationship)
                 if node_type:
                     query += " AND n.type = ?"
                     params.append(node_type)
-                    
+
                 cursor.execute(query, params)
                 neighbors.extend(cursor.fetchall())
-            
+
             if direction in ["both", "incoming"]:
                 query = """SELECT e.source, n.type FROM edges e 
                           JOIN nodes n ON e.source = n.id
                           WHERE e.target = ?"""
                 params = [current]
-                
+
                 if relationship:
                     query += " AND e.relationship = ?"
                     params.append(relationship)
                 if node_type:
                     query += " AND n.type = ?"
                     params.append(node_type)
-                    
+
                 cursor.execute(query, params)
                 neighbors.extend(cursor.fetchall())
-            
+
             for neighbor, n_type in neighbors:
                 if neighbor not in visited:
                     visited.add(neighbor)
                     queue.append(neighbor)
                     distance[neighbor] = current_dist + 1
                     connections.append((neighbor, n_type, current_dist + 1))
-        
+
         if not connections:
             return f"No connected nodes found for '{node_id}'"
-        
+
         # Format output
         output = [f"Nodes connected to '{node_id}' ({direction}):"]
         output.append(f"\nTotal connected: {len(connections)}\n")
-        
+
         # Group by distance
         by_distance = {}
         for node, n_type, dist in connections:
             if dist not in by_distance:
                 by_distance[dist] = []
             by_distance[dist].append((node, n_type))
-        
+
         for dist in sorted(by_distance.keys()):
             output.append(f"Distance {dist}:")
             for node, n_type in sorted(by_distance[dist]):
                 output.append(f"  {node} ({n_type})")
-        
+
         return "\n".join(output)
 
-    def _query_ancestors(self, conn: sqlite3.Connection, node_id: str, depth: int,
-                        relationship: Optional[str], node_type: Optional[str]) -> str:
+    def _query_ancestors(
+        self,
+        conn: sqlite3.Connection,
+        node_id: str,
+        depth: int,
+        relationship: Optional[str],
+        node_type: Optional[str],
+    ) -> str:
         """Find nodes that point TO this node (incoming edges only)."""
-        return self._query_subgraph(conn, node_id, depth, relationship, node_type, "incoming")
+        return self._query_subgraph(
+            conn, node_id, depth, relationship, node_type, "incoming"
+        )
 
-    def _query_descendants(self, conn: sqlite3.Connection, node_id: str, depth: int,
-                          relationship: Optional[str], node_type: Optional[str]) -> str:
+    def _query_descendants(
+        self,
+        conn: sqlite3.Connection,
+        node_id: str,
+        depth: int,
+        relationship: Optional[str],
+        node_type: Optional[str],
+    ) -> str:
         """Find nodes that this node points TO (outgoing edges only)."""
-        return self._query_subgraph(conn, node_id, depth, relationship, node_type, "outgoing")
+        return self._query_subgraph(
+            conn, node_id, depth, relationship, node_type, "outgoing"
+        )
 
     def register(self, mcp_server) -> None:
         """Register this tool with the MCP server."""