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

hanzo_mcp/tools/database/graph.py

@@ -1,17 +1,24 @@
 """Unified graph database tool."""
 
-from typing import Annotated, TypedDict, Unpack, final, override, Optional, List, Dict, Any
 import json
-from pathlib import Path
+from typing import (
+    Any,
+    Dict,
+    Unpack,
+    Optional,
+    Annotated,
+    TypedDict,
+    final,
+    override,
+)
 
-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.permissions import PermissionManager
 from hanzo_mcp.tools.database.database_manager import DatabaseManager
 
-
 # Parameter types
 Action = Annotated[
     str,
@@ -96,6 +103,7 @@ Limit = Annotated[
 
 class GraphParams(TypedDict, total=False):
     """Parameters for graph tool."""
+
     action: str
     node_id: Optional[str]
     node_type: Optional[str]
@@ -111,8 +119,10 @@ class GraphParams(TypedDict, total=False):
 @final
 class GraphTool(BaseTool):
     """Unified graph database tool."""
-    
-    def __init__(self, permission_manager: PermissionManager, db_manager: DatabaseManager):
+
+    def __init__(
+        self, permission_manager: PermissionManager, db_manager: DatabaseManager
+    ):
         """Initialize the graph tool."""
         super().__init__(permission_manager)
         self.db_manager = db_manager
@@ -153,7 +163,7 @@ graph --action search --pattern "John" --node-type User
 
         # Extract action
         action = params.get("action", "query")
-        
+
         # Route to appropriate handler
         if action == "query":
             return await self._handle_query(project_db, params, tool_ctx)
@@ -174,17 +184,18 @@ graph --action search --pattern "John" --node-type User
         node_type = params.get("node_type")
         depth = params.get("depth", 2)
         limit = params.get("limit", 50)
-        
+
         if not node_id and not node_type:
             return "Error: node_id or node_type required for query"
-        
+
         try:
             with project_db.get_graph_connection() as conn:
                 results = []
-                
+
                 if node_id:
                     # Query specific node and its relationships
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         WITH RECURSIVE
                         node_tree(id, type, properties, depth, path) AS (
                             SELECT id, type, properties, 0, id
@@ -206,52 +217,59 @@ graph --action search --pattern "John" --node-type User
                         SELECT DISTINCT * FROM node_tree
                         ORDER BY depth, id
                         LIMIT ?
-                    """, (node_id, depth, limit))
-                    
+                    """,
+                        (node_id, depth, limit),
+                    )
+
                     nodes = cursor.fetchall()
-                    
+
                     # Get edges
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         SELECT from_node, to_node, type, properties
                         FROM edges
                         WHERE from_node IN (SELECT id FROM node_tree)
                            OR to_node IN (SELECT id FROM node_tree)
-                    """)
-                    
+                    """
+                    )
+
                     edges = cursor.fetchall()
-                    
+
                 else:
                     # Query by type
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         SELECT id, type, properties
                         FROM nodes
                         WHERE type = ?
                         LIMIT ?
-                    """, (node_type, limit))
-                    
+                    """,
+                        (node_type, limit),
+                    )
+
                     nodes = cursor.fetchall()
                     edges = []
-                
+
                 # Format results
                 output = ["=== Graph Query Results ==="]
-                
+
                 if nodes:
                     output.append(f"\nNodes ({len(nodes)}):")
                     for node in nodes:
                         props = json.loads(node[2]) if node[2] else {}
                         output.append(f"  {node[0]} [{node[1]}] {props}")
-                
+
                 if edges:
                     output.append(f"\nEdges ({len(edges)}):")
                     for edge in edges:
                         props = json.loads(edge[3]) if edge[3] else {}
                         output.append(f"  {edge[0]} --[{edge[2]}]--> {edge[1]} {props}")
-                
+
                 if not nodes and not edges:
                     output.append("No results found")
-                
+
                 return "\n".join(output)
-                
+
         except Exception as e:
             await tool_ctx.error(f"Query failed: {str(e)}")
             return f"Error during query: {str(e)}"
@@ -261,50 +279,56 @@ graph --action search --pattern "John" --node-type User
         node_id = params.get("node_id")
         from_node = params.get("from_node")
         to_node = params.get("to_node")
-        
+
         if node_id:
             # Add node
             node_type = params.get("node_type")
             if not node_type:
                 return "Error: node_type required when adding node"
-            
+
             properties = params.get("properties", {})
-            
+
             try:
                 with project_db.get_graph_connection() as conn:
-                    conn.execute("""
+                    conn.execute(
+                        """
                         INSERT OR REPLACE INTO nodes (id, type, properties)
                         VALUES (?, ?, ?)
-                    """, (node_id, node_type, json.dumps(properties)))
+                    """,
+                        (node_id, node_type, json.dumps(properties)),
+                    )
                     conn.commit()
-                
+
                 await tool_ctx.info(f"Added node: {node_id}")
                 return f"Added node {node_id} [{node_type}]"
-                
+
             except Exception as e:
                 await tool_ctx.error(f"Failed to add node: {str(e)}")
                 return f"Error adding node: {str(e)}"
-        
+
         elif from_node and to_node:
             # Add edge
             edge_type = params.get("edge_type", "RELATED")
             properties = params.get("properties", {})
-            
+
             try:
                 with project_db.get_graph_connection() as conn:
-                    conn.execute("""
+                    conn.execute(
+                        """
                         INSERT OR REPLACE INTO edges (from_node, to_node, type, properties)
                         VALUES (?, ?, ?, ?)
-                    """, (from_node, to_node, edge_type, json.dumps(properties)))
+                    """,
+                        (from_node, to_node, edge_type, json.dumps(properties)),
+                    )
                     conn.commit()
-                
+
                 await tool_ctx.info(f"Added edge: {from_node} -> {to_node}")
                 return f"Added edge {from_node} --[{edge_type}]--> {to_node}"
-                
+
             except Exception as e:
                 await tool_ctx.error(f"Failed to add edge: {str(e)}")
                 return f"Error adding edge: {str(e)}"
-        
+
         else:
             return "Error: Either node_id (for node) or from_node + to_node (for edge) required"
 
@@ -313,68 +337,80 @@ graph --action search --pattern "John" --node-type User
         node_id = params.get("node_id")
         from_node = params.get("from_node")
         to_node = params.get("to_node")
-        
+
         if node_id:
             # Remove node and its edges
             try:
                 with project_db.get_graph_connection() as conn:
                     # Delete edges first
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         DELETE FROM edges
                         WHERE from_node = ? OR to_node = ?
-                    """, (node_id, node_id))
-                    
+                    """,
+                        (node_id, node_id),
+                    )
+
                     edges_deleted = cursor.rowcount
-                    
+
                     # Delete node
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         DELETE FROM nodes WHERE id = ?
-                    """, (node_id,))
-                    
+                    """,
+                        (node_id,),
+                    )
+
                     if cursor.rowcount == 0:
                         return f"Node {node_id} not found"
-                    
+
                     conn.commit()
-                
+
                 msg = f"Removed node {node_id}"
                 if edges_deleted > 0:
                     msg += f" and {edges_deleted} connected edges"
-                
+
                 await tool_ctx.info(msg)
                 return msg
-                
+
             except Exception as e:
                 await tool_ctx.error(f"Failed to remove node: {str(e)}")
                 return f"Error removing node: {str(e)}"
-        
+
         elif from_node and to_node:
             # Remove specific edge
             edge_type = params.get("edge_type")
-            
+
             try:
                 with project_db.get_graph_connection() as conn:
                     if edge_type:
-                        cursor = conn.execute("""
+                        cursor = conn.execute(
+                            """
                             DELETE FROM edges
                             WHERE from_node = ? AND to_node = ? AND type = ?
-                        """, (from_node, to_node, edge_type))
+                        """,
+                            (from_node, to_node, edge_type),
+                        )
                     else:
-                        cursor = conn.execute("""
+                        cursor = conn.execute(
+                            """
                             DELETE FROM edges
                             WHERE from_node = ? AND to_node = ?
-                        """, (from_node, to_node))
-                    
+                        """,
+                            (from_node, to_node),
+                        )
+
                     if cursor.rowcount == 0:
                         return f"Edge not found"
-                    
+
                     conn.commit()
-                
+
                 return f"Removed edge {from_node} --> {to_node}"
-                
+
             except Exception as e:
                 await tool_ctx.error(f"Failed to remove edge: {str(e)}")
                 return f"Error removing edge: {str(e)}"
-        
+
         else:
             return "Error: Either node_id or from_node + to_node required for remove"
 
@@ -383,43 +419,49 @@ graph --action search --pattern "John" --node-type User
         pattern = params.get("pattern")
         if not pattern:
             return "Error: pattern required for search"
-        
+
         node_type = params.get("node_type")
         limit = params.get("limit", 50)
-        
+
         try:
             with project_db.get_graph_connection() as conn:
                 # Search in properties
                 if node_type:
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         SELECT id, type, properties
                         FROM nodes
                         WHERE type = ? AND properties LIKE ?
                         LIMIT ?
-                    """, (node_type, f"%{pattern}%", limit))
+                    """,
+                        (node_type, f"%{pattern}%", limit),
+                    )
                 else:
-                    cursor = conn.execute("""
+                    cursor = conn.execute(
+                        """
                         SELECT id, type, properties
                         FROM nodes
                         WHERE properties LIKE ?
                         LIMIT ?
-                    """, (f"%{pattern}%", limit))
-                
+                    """,
+                        (f"%{pattern}%", limit),
+                    )
+
                 results = cursor.fetchall()
-                
+
                 if not results:
                     return f"No nodes found matching '{pattern}'"
-                
+
                 # Format results
                 output = [f"=== Graph Search Results for '{pattern}' ==="]
                 output.append(f"Found {len(results)} nodes\n")
-                
+
                 for node in results:
                     props = json.loads(node[2]) if node[2] else {}
                     output.append(f"{node[0]} [{node[1]}] {props}")
-                
+
                 return "\n".join(output)
-                
+
         except Exception as e:
             await tool_ctx.error(f"Search failed: {str(e)}")
             return f"Error during search: {str(e)}"
@@ -429,54 +471,58 @@ graph --action search --pattern "John" --node-type User
         try:
             with project_db.get_graph_connection() as conn:
                 # Node stats
-                cursor = conn.execute("""
+                cursor = conn.execute(
+                    """
                     SELECT type, COUNT(*) as count
                     FROM nodes
                     GROUP BY type
                     ORDER BY count DESC
-                """)
-                
+                """
+                )
+
                 node_stats = cursor.fetchall()
-                
+
                 # Edge stats
-                cursor = conn.execute("""
+                cursor = conn.execute(
+                    """
                     SELECT type, COUNT(*) as count
                     FROM edges
                     GROUP BY type
                     ORDER BY count DESC
-                """)
-                
+                """
+                )
+
                 edge_stats = cursor.fetchall()
-                
+
                 # Total counts
                 cursor = conn.execute("SELECT COUNT(*) FROM nodes")
                 total_nodes = cursor.fetchone()[0]
-                
+
                 cursor = conn.execute("SELECT COUNT(*) FROM edges")
                 total_edges = cursor.fetchone()[0]
-                
+
                 # Format output
                 output = [f"=== Graph Database Statistics ==="]
                 output.append(f"Project: {project_db.project_path}")
                 output.append(f"\nTotal nodes: {total_nodes}")
                 output.append(f"Total edges: {total_edges}")
-                
+
                 if node_stats:
                     output.append("\nNodes by type:")
                     for node_type, count in node_stats:
                         output.append(f"  {node_type}: {count}")
-                
+
                 if edge_stats:
                     output.append("\nEdges by type:")
                     for edge_type, count in edge_stats:
                         output.append(f"  {edge_type}: {count}")
-                
+
                 return "\n".join(output)
-                
+
         except Exception as e:
             await tool_ctx.error(f"Failed to get stats: {str(e)}")
             return f"Error getting stats: {str(e)}"
 
     def register(self, mcp_server) -> None:
         """Register this tool with the MCP server."""
-        pass
+        pass

hanzo_mcp/tools/database/graph_add.py

@@ -1,17 +1,16 @@
 """Graph add tool for adding nodes and edges to the graph database."""
 
 import json
-from typing import Annotated, Optional, TypedDict, Unpack, final, override
+from typing import Unpack, Optional, Annotated, TypedDict, final, override
 
-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
 
-
 NodeId = Annotated[
     Optional[str],
     Field(
@@ -94,7 +93,9 @@ class GraphAddParams(TypedDict, total=False):
 class GraphAddTool(BaseTool):
     """Tool for adding nodes and edges to graph database."""
 
-    def __init__(self, permission_manager: PermissionManager, db_manager: DatabaseManager):
+    def __init__(
+        self, permission_manager: PermissionManager, db_manager: DatabaseManager
+    ):
         """Initialize the graph add tool.
 
         Args:
@@ -181,8 +182,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,13 +208,16 @@ Examples:
                 properties_json = json.dumps(properties) if properties else None
 
                 # Insert or update node
-                graph_conn.execute("""
+                graph_conn.execute(
+                    """
                     INSERT OR REPLACE INTO nodes (id, type, properties)
                     VALUES (?, ?, ?)
-                """, (node_id, node_type, properties_json))
-                
+                """,
+                    (node_id, node_type, properties_json),
+                )
+
                 graph_conn.commit()
-                
+
                 return f"Successfully added node '{node_id}' of type '{node_type}'"
 
             else:
@@ -220,13 +225,17 @@ Examples:
                 if not relationship:
                     return "Error: relationship is required when adding an edge"
 
-                await tool_ctx.info(f"Adding edge: {source} --[{relationship}]--> {target}")
+                await tool_ctx.info(
+                    f"Adding edge: {source} --[{relationship}]--> {target}"
+                )
 
                 # Check if nodes exist
                 cursor = graph_conn.cursor()
-                cursor.execute("SELECT id FROM nodes WHERE id IN (?, ?)", (source, target))
+                cursor.execute(
+                    "SELECT id FROM nodes WHERE id IN (?, ?)", (source, target)
+                )
                 existing = [row[0] for row in cursor.fetchall()]
-                
+
                 if source not in existing:
                     return f"Error: Source node '{source}' does not exist"
                 if target not in existing:
@@ -236,16 +245,19 @@ Examples:
                 properties_json = json.dumps(properties) if properties else None
 
                 # Insert or update edge
-                graph_conn.execute("""
+                graph_conn.execute(
+                    """
                     INSERT OR REPLACE INTO edges (source, target, relationship, weight, properties)
                     VALUES (?, ?, ?, ?, ?)
-                """, (source, target, relationship, weight, properties_json))
-                
+                """,
+                    (source, target, relationship, weight, properties_json),
+                )
+
                 graph_conn.commit()
-                
+
                 # Save to disk
                 project_db._save_graph_to_disk()
-                
+
                 return f"Successfully added edge: {source} --[{relationship}]--> {target} (weight: {weight})"
 
         except Exception as e:
@@ -254,4 +266,4 @@ Examples:
 
     def register(self, mcp_server) -> None:
         """Register this tool with the MCP server."""
-        pass
+        pass