adbpg-mcp-server 1.0.7__py3-none-any.whl → 2.0.0__py3-none-any.whl

adbpg_mcp_server/__main__.py

@@ -0,0 +1,45 @@
+import argparse
+import sys
+
+from .server_http import run_http_server
+from .server_stdio import run_stdio_server
+
+def main():
+    """
+    通过参数选择使用 stdio 或 http 服务器
+    """
+    parser = argparse.ArgumentParser(
+        description="ADBPG Management & Control Plane MCP Server"
+    )
+    parser.add_argument(
+        "--transport",
+        choices = ["stdio", "http"],
+        default = "stdio",
+        help = "Transport protocol to use (stdio or streamable-http)"
+    )
+    http_group = parser.add_argument_group("HTTP Server Options")
+    http_group.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help = "Host to bind the HTTP server to (default: 127.0.0.1)."
+    )
+    http_group.add_argument(
+        "--port",
+        type = int,
+        default = 3000,
+        help = "Port to bind the HTTP server to (default: 3000)."
+    )
+    
+    args = parser.parse_args()
+    print(f"Starting ADBPG MCP Server with transport {args.transport}")
+
+    if args.transport == "http":
+        run_http_server(host=args.host, port=args.port)
+    elif args.transport == "stdio":
+        run_stdio_server()
+    else:
+        raise ValueError(f"Invalid transport: {args.transport}")
+
+if __name__ == "__main__":
+    main()
+

adbpg_mcp_server/adbpg.py

@@ -0,0 +1,113 @@
+# mcp_server/db.py
+import psycopg
+import logging
+import json
+from psycopg import Connection
+from .adbpg_config import AppConfig
+from typing import Optional, Callable
+from .adbpg_config import settings
+
+logger = logging.getLogger(__name__)
+
+class DatabaseManager:
+    def __init__(self, config: AppConfig):
+        self._config_info = config.get_db_connection_info()
+        self._conn = None
+        self._graphrag_conn = None
+        self._llm_memory_conn = None
+        self.db_master_port = self.get_master_port()
+
+    def _get_connection(self, conn_attr: str, initializer: Optional[Callable[[Connection],None]] = None) -> Connection:
+        """通用连接获取和健康检查逻辑"""
+
+        def _create_and_initialize() -> Connection:
+            new_conn = psycopg.connect(**self._config_info)
+            new_conn.autocommit = True
+
+            if initializer:
+                logger.info(f"Running initializer for {conn_attr} on new connection...")
+                try:
+                    initializer(new_conn)
+                except Exception as e:
+                    logger.error(f"Error initializing {conn_attr}: {e}")
+                    new_conn.close()
+                    raise e
+            setattr(self, conn_attr, new_conn)
+            return new_conn
+        
+        conn = getattr(self, conn_attr)
+        
+        if conn is None or conn.closed:
+            logger.info(f"Connecting to database for {conn_attr}...")
+            conn = _create_and_initialize()
+            logger.info(f"New database connection established and initialized for {conn_attr} (id: {id(conn)})")
+        else:
+            try:
+                with conn.cursor() as cur:
+                    cur.execute("SELECT 1;")
+            except psycopg.Error:
+                logger.warning(f"Connection for {conn_attr} is stale. Reconnecting...")
+                conn.close()
+                conn = psycopg.connect(**self._config_info)
+                conn.autocommit = True
+                setattr(self, conn_attr, conn)
+                logger.info(f"Reconnected for {conn_attr} (id: {id(conn)})")
+        return conn
+
+    def get_basic_connection(self) -> Connection:
+        """获取用于基本操作的连接"""
+        return self._get_connection('_conn')
+
+    def get_graphrag_connection(self) -> Connection:
+        """
+        获取 GraphRAG 专用的长连接
+        每次重新获得连接就初始化
+        """
+        def initializer(conn: Connection):
+            try:
+                config_json = json.dumps(settings.get_graphrag_init_config())
+                with conn.cursor() as cursor:
+                    cursor.execute("SELECT adbpg_graphrag.initialize(%s::json);", (config_json,))
+            except (psycopg.Error, json.JSONDecodeError, AttributeError) as e:
+                logger.error(f"ADBPG GraphRAG initialization failed: {e}")
+                raise RuntimeError(f"ADBPG GraphRAG initialization failed: {e}") from e
+        return self._get_connection('_graphrag_conn', initializer=initializer)
+
+    def get_llm_memory_connection(self) -> Connection:
+        """
+        获取 LLM Memory 专用的长连接
+        每次重新连接就初始化
+        """
+        def initializer(conn: Connection):
+            try:
+                config_json = json.dumps(settings.get_memory_init_config(self.db_master_port))
+                with conn.cursor() as cursor:
+                    cursor.execute("SELECT adbpg_llm_memory.config(%s::json);", (config_json,))
+            except (psycopg.Error, json.JSONDecodeError, AttributeError) as e:
+                logger.error(f"ADBPG LLM Memory initialization failed: {e}")
+                raise RuntimeError(f"ADBPG LLM Memory initialization failed: {e}") from e
+
+        return self._get_connection('_llm_memory_conn',initializer=initializer)
+
+    def get_master_port(self) -> int:
+        """获取 master 节点的端口，用于 llm_memory 配置"""
+        sql = "SELECT port FROM gp_segment_configuration WHERE content = -1 AND role = 'p';"
+        try:
+            with self.get_basic_connection().cursor() as cursor:
+                cursor.execute(sql)
+                port = cursor.fetchone()[0]
+                return port
+        except psycopg.Error as e:
+            logger.error(f"Database error while getting master port: {e}")
+            raise RuntimeError(f"Database error: {str(e)}") from e
+
+    def close_all(self):
+        """关闭所有连接"""
+        if self._conn and not self._conn.closed:
+            self._conn.close()
+        if self._graphrag_conn and not self._graphrag_conn.closed:
+            self._graphrag_conn.close()
+        if self._llm_memory_conn and not self._llm_memory_conn.closed:
+            self._llm_memory_conn.close()
+        logger.info("All database connections closed.")
+

adbpg_mcp_server/adbpg_basic_operation.py

@@ -0,0 +1,160 @@
+# mcp_server/basic_operation.py
+import psycopg
+import logging
+from pydantic import AnyUrl
+from mcp.types import Resource, ResourceTemplate, Tool
+from .adbpg import DatabaseManager
+from typing import Tuple
+
+
+logger = logging.getLogger(__name__)
+
+async def list_resources() -> list[Resource]:
+    """列出可用的基本资源"""
+    return [
+        Resource(
+            uri="adbpg:///schemas",
+            name="All Schemas",
+            description="AnalyticDB PostgreSQL schemas. List all schemas in the database",
+            mimeType="text/plain"
+        )
+    ]
+
+async def list_resource_templates() -> list[ResourceTemplate]:
+    """定义动态资源模板"""
+    return [
+        ResourceTemplate(
+            uriTemplate="adbpg:///{schema}/tables",
+            name="Schema Tables",
+            description="List all tables in a specific schema",
+            mimeType="text/plain"
+        ),
+        ResourceTemplate(
+            uriTemplate="adbpg:///{schema}/{table}/ddl",
+            name="Table DDL",
+            description="Get the DDL script of a table in a specific schema",
+            mimeType="text/plain"
+        ),
+        ResourceTemplate(
+            uriTemplate="adbpg:///{schema}/{table}/statistics",
+            name="Table Statistics",
+            description="Get statistics information of a table",
+            mimeType="text/plain"
+        )
+    ]
+
+async def read_resource(uri: AnyUrl, db: DatabaseManager) -> str:
+    """读取资源内容"""
+    uri_str = str(uri)
+    if not uri_str.startswith("adbpg:///"):
+        raise ValueError(f"Invalid URI scheme: {uri_str}")
+
+    try:
+        conn = db.get_basic_connection()
+        with conn.cursor() as cursor:
+            path_parts = uri_str[9:].split('/')
+            
+            if path_parts[0] == "schemas":
+                query = "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('pg_catalog', 'information_schema') ORDER BY schema_name;"
+                cursor.execute(query)
+                return "\n".join([schema[0] for schema in cursor.fetchall()])
+                
+            elif len(path_parts) == 2 and path_parts[1] == "tables":
+                schema = path_parts[0]
+                query = "SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = %s ORDER BY table_name;"
+                cursor.execute(query, (schema,))
+                return "\n".join([f"{table[0]} ({table[1]})" for table in cursor.fetchall()])
+                
+            elif len(path_parts) == 3 and path_parts[2] == "ddl":
+                schema, table = path_parts[0], path_parts[1]
+                query = f"SELECT pg_get_ddl('{schema}.{table}'::regclass);"
+                cursor.execute(query)
+                ddl = cursor.fetchone()
+                return ddl[0] if ddl else f"No DDL found for {schema}.{table}"
+                
+            elif len(path_parts) == 3 and path_parts[2] == "statistics":
+                schema, table = path_parts[0], path_parts[1]
+                query = "SELECT attname, null_frac, avg_width, n_distinct, most_common_vals, most_common_freqs FROM pg_stats WHERE schemaname = %s AND tablename = %s ORDER BY attname;"
+                cursor.execute(query, (schema, table))
+                rows = cursor.fetchall()
+                if not rows: return f"No statistics found for {schema}.{table}"
+                
+                result = [f"Statistics for {schema}.{table}:\n"]
+                for row in rows:
+                    result.append(f"Column: {row[0]}, Null fraction: {row[1]}, Avg width: {row[2]}, Distinct values: {row[3]}")
+                return "\n".join(result)
+            
+            raise ValueError(f"Invalid resource URI format: {uri_str}")
+  
+    except psycopg.Error as e:
+        logger.error(f"Database error in read_resource: {e}")
+        raise RuntimeError(f"Database error: {str(e)}")
+
+def get_basic_tools() -> list[Tool]:
+    """返回基础数据库操作工具列表"""
+    return [
+        Tool(
+            name="execute_select_sql",
+            description="Execute SELECT SQL to query data from ADBPG database. Returns data in JSON format.",
+            inputSchema={ "type": "object", "properties": {"query": {"type": "string", "description": "The (SELECT) SQL query to execute"}}, "required": ["query"]}
+        ),
+        Tool(
+            name="execute_dml_sql",
+            description="Execute (INSERT, UPDATE, DELETE) SQL to modify data in ADBPG database.",
+            inputSchema={ "type": "object", "properties": {"query": {"type": "string", "description": "The DML SQL query to execute"}}, "required": ["query"]}
+        ),
+        Tool(
+            name="execute_ddl_sql",
+            description="Execute (CREATE, ALTER, DROP) SQL statements to manage database objects.",
+            inputSchema={ "type": "object", "properties": {"query": {"type": "string", "description": "The DDL SQL query to execute"}}, "required": ["query"]}
+        ),
+        Tool(
+            name="analyze_table",
+            description="Execute ANALYZE command to collect table statistics.",
+            inputSchema={ "type": "object", "properties": {"schema": {"type": "string"}, "table": {"type": "string"}}, "required": ["schema", "table"]}
+        ),
+        Tool(
+            name="explain_query",
+            description="Get query execution plan.",
+            inputSchema={ "type": "object", "properties": {"query": {"type": "string", "description": "The SQL query to analyze"}}, "required": ["query"]}
+        ),
+    ]
+
+async def call_basic_tool(name: str, arguments: dict, db: DatabaseManager) -> Tuple[str, dict, bool]:
+    """
+    准备执行基础工具的SQL和参数。
+    返回 (query_string, params, needs_json_agg)
+    """
+    query, params, needs_json_agg = None, None, False
+    
+    if name == "execute_select_sql":
+        query_text = arguments.get("query")
+        if not query_text or not query_text.strip().upper().startswith("SELECT"):
+            raise ValueError("Query must be a SELECT statement")
+        query = f"SELECT json_agg(row_to_json(t)) FROM ({query_text.rstrip(';')}) AS t"
+        needs_json_agg = True
+    elif name == "execute_dml_sql":
+        query = arguments.get("query")
+        if not query or not any(query.strip().upper().startswith(k) for k in ["INSERT", "UPDATE", "DELETE"]):
+            raise ValueError("Query must be a DML statement (INSERT, UPDATE, DELETE)")
+    elif name == "execute_ddl_sql":
+        query = arguments.get("query")
+        if not query or not any(query.strip().upper().startswith(k) for k in ["CREATE", "ALTER", "DROP", "TRUNCATE"]):
+            raise ValueError("Query must be a DDL statement (CREATE, ALTER, DROP)")
+    elif name == "analyze_table":
+        schema, table = arguments.get("schema"), arguments.get("table")
+        if not all([schema, table]):
+            raise ValueError("Schema and table are required")
+        query = f"ANALYZE {schema}.{table}"
+    elif name == "explain_query":
+        query_text = arguments.get("query")
+        if not query_text:
+            raise ValueError("Query is required")
+        query = f"EXPLAIN (FORMAT JSON) {query_text}"
+        needs_json_agg = True # The output is already a single JSON value in a single row
+
+    if query is None:
+        raise ValueError(f"Unknown basic tool: {name}")
+
+    return (query, params, needs_json_agg)
+

adbpg_mcp_server/adbpg_config.py

@@ -0,0 +1,155 @@
+import os
+import sys
+import logging
+from dotenv import load_dotenv
+
+logger = logging.getLogger(__name__)
+
+SERVER_VERSION = "0.2.0"
+
+try:
+    load_dotenv()
+    logger.info("Environment variables loaded")
+except Exception as e:
+    logger.error(f"Error loading .env file: {str(e)}")
+    sys.exit(1)
+
+class AppConfig:
+    def __init__(self):
+        # Database Config
+        self.db_host = os.getenv("ADBPG_HOST")
+        self.db_port = os.getenv("ADBPG_PORT")
+        self.db_user = os.getenv("ADBPG_USER")
+        self.db_password = os.getenv("ADBPG_PASSWORD")
+        self.db_name = os.getenv("ADBPG_DATABASE")
+
+        # GraphRAG Config
+        self.graphrag_llm_model = os.getenv("GRAPHRAG_LLM_MODEL")
+        self.graphrag_api_key = os.getenv("GRAPHRAG_API_KEY")
+        self.graphrag_base_url = os.getenv("GRAPHRAG_BASE_URL")
+        self.graphrag_embedding_model = os.getenv("GRAPHRAG_EMBEDDING_MODEL")
+        self.graphrag_embedding_api_key = os.getenv("GRAPHRAG_EMBEDDING_API_KEY")
+        self.graphrag_embedding_base_url = os.getenv("GRAPHRAG_EMBEDDING_BASE_URL")
+        self.graphrag_language = os.getenv("GRAPHRAG_LANGUAGE", "English")
+        self.graphrag_entity_types = os.getenv("GRAPHRAG_ENTITY_TYPES", '["Organization", "Person", "Location", "Event", "Technology", "Equipment", "Product", "Document", "Category"]')
+        self.graphrag_relationship_types = os.getenv("GRAPHRAG_RELATIONSHIP_TYPES", '["Causes", "Used For", "Helps In", "Includes", "Originated From", "Seasonal", "Coreference Of", "Synonym Of", "Conjunction", "Affects"]')
+
+        # LLM Memory Config
+        self.memory_llm_model = os.getenv("LLMEMORY_LLM_MODEL")
+        self.memory_api_key = os.getenv("LLMEMORY_API_KEY")
+        self.memory_base_url = os.getenv("LLMEMORY_BASE_URL")
+        self.memory_embedding_model = os.getenv("LLMEMORY_EMBEDDING_MODEL")
+        self.memory_embedding_dims = os.getenv("LLMEMORY_EMBEDDING_DIMS", 1024)
+        self.memory_enable_graph = os.getenv("LLMEMORY_ENABLE_GRAPH", "False").lower() in ('true', '1', 't')
+
+        self.db_env_ready = self._check_db_env()
+        self.graphrag_env_ready = self._check_graphrag_env()
+        self.memory_env_ready = self._check_memory_env()
+
+    def _check_db_env(self):
+        required = ["db_host", "db_port", "db_user", "db_password", "db_name"]
+        missing = [var.replace('db_', 'ADBPG_').upper() for var in required if not getattr(self, var)]
+        if missing:
+            logger.error(f"Missing required ADBPG environment variables: {', '.join(missing)}")
+            return False
+        logger.info("All ADBPG required environment variables are set")
+        return True
+    
+    def _check_graphrag_env(self):
+        required = [
+            "graphrag_llm_model", 
+            "graphrag_api_key", 
+            "graphrag_base_url",
+            "graphrag_embedding_model", 
+            "graphrag_embedding_api_key", 
+            "graphrag_embedding_base_url"
+        ]
+        missing = [var.replace('graphrag_', 'GRAPHRAG_').upper() for var in required if not getattr(self, var)]
+        if missing:
+            logger.warning(f"Missing GraphRAG environment variables: {', '.join(missing)}. GraphRAG tools will be disabled.")
+            return False
+        logger.info("All GraphRAG required environment variables are set")
+        return True
+    
+    def _check_memory_env(self):
+        required = [
+            "memory_llm_model", "memory_api_key", "memory_base_url", "memory_embedding_model"
+        ]
+        missing = [var.replace('memory_', 'LLMEMORY_').upper() for var in required if not getattr(self, var)]
+        if missing:
+            logger.warning(f"Missing LLM Memory environment variables: {', '.join(missing)}. LLM Memory tools will be disabled.")
+            return False
+        logger.info("All LLM Memory required environment variables are set")
+        return True
+    
+    def get_db_connection_info(self):
+        if not self.db_env_ready:
+            raise ValueError("Database environment variables are not set.")
+        return {
+            "host": self.db_host,
+            "port": self.db_port,
+            "user": self.db_user,
+            "password": self.db_password,
+            "dbname": self.db_name,
+            "application_name": f"adbpg-mcp-server-{SERVER_VERSION}"
+        }
+
+    def get_graphrag_init_config(self):
+        if not self.graphrag_env_ready:
+            return None
+        return {
+            "llm_model": self.graphrag_llm_model,
+            "llm_api_key": self.graphrag_api_key,
+            "llm_url": self.graphrag_base_url,
+            "embedding_model": self.graphrag_embedding_model,
+            "embedding_api_key": self.graphrag_embedding_api_key,
+            "embedding_url": self.graphrag_embedding_base_url,
+            "language": self.graphrag_language,
+            "entity_types": self.graphrag_entity_types,
+            "relationship_types": self.graphrag_relationship_types,
+            "postgres_password": self.db_password
+        }
+    def get_memory_init_config(self, db_master_port: int):
+        if not self.memory_env_ready:
+            return None
+        config = {
+            "llm": {
+                "provider": "openai", 
+                "config": {
+                    "model": self.memory_llm_model,
+                    "openai_base_url": self.memory_base_url,
+                    "api_key": self.memory_api_key
+                }},
+            "embedder": {
+                "provider": "openai", 
+                "config": {
+                    "model": self.memory_embedding_model,
+                    "embedding_dims": self.memory_embedding_dims,
+                    "api_key": self.memory_api_key,
+                    "openai_base_url": self.memory_base_url
+                }},
+            "vector_store": {
+                "provider": "adbpg", 
+                "config": {
+                    "user": self.db_user,
+                    "password": self.db_password,
+                    "dbname": self.db_name,
+                    "hnsw": "True",
+                    "embedding_model_dims": self.memory_embedding_dims,
+                    "port": db_master_port
+                }}
+        }
+        if self.memory_enable_graph:
+            config["graph_store"] = {
+                "provider": "adbpg", 
+                "config": {
+                    "url": "http://localhost",
+                    "username": self.db_user,
+                    "password": self.db_password,
+                    "database": self.db_name,
+                    "port": db_master_port
+                }
+            }
+        return config
+
+settings = AppConfig()

adbpg_mcp_server/adbpg_graphrag.py

@@ -0,0 +1,197 @@
+import logging
+from mcp.types import Tool, TextContent
+from .adbpg import DatabaseManager
+
+logger = logging.getLogger(__name__)
+
+def get_graphrag_tools() -> list[Tool]:
+    """
+    返回 ADBPG GraphRAG 工具列表
+    """
+    return [
+        Tool(
+            name = "adbpg_graphrag_upload",
+            description = "Execute graphrag upload operation",
+            # 参数：filename text， context text
+            # filename 表示文件名称， context 表示文件内容
+            inputSchema = {
+                "type": "object",
+                "properties": {
+                    "filename": {
+                        "type": "string",
+                        "description": "The file name need to upload"
+                    },
+                    "context": {
+                        "type": "string",
+                        "description": "the context of your file"
+                    }
+                },
+                "required": ["filename", "context"]
+            }
+        ),
+        Tool(
+            name = "adbpg_graphrag_query",
+            description = "Execute graphrag query operation",
+            # 参数：query_str text, [query_mode text]
+            # query_str 是询问的问题，query_mode 选择查询模式
+            inputSchema = {
+                "type": "object",
+                "properties": {
+                    "query_str": {
+                        "type": "string",
+                        "description": "The query you want to ask"
+                    },
+                    "query_mode": {
+                        "type": "string",
+                        "description": "The query mode you need to choose [ bypass,naive, local, global, hybrid, mix[default], tree ]."
+                    },
+                    "start_search_node_id": {
+                        "type": "string",
+                        "description": "If using 'tree' query mode, set the start node ID of tree."
+                    }
+                },
+                "required": ["query_str"]
+            }
+        ),
+        Tool(
+            name = "adbpg_graphrag_upload_decision_tree",
+            description = " Upload a decision tree with the specified root_node. If the root_node does not exist, a new decision tree will be created. ",
+            # context text, root_node text
+            inputSchema = {
+                "type": "object",
+                "properties": {
+                    "root_node": {
+                        "type": "string",
+                        "description": "the root_noot (optional)"
+                    },
+                    "context": {
+                        "type": "string",
+                        "description": "the context of decision"
+                    }
+                },
+                "required": ["context"]
+            }
+        ),
+        Tool(
+            name = "adbpg_graphrag_append_decision_tree",
+            description = "Append a subtree to an existing decision tree at the node specified by root_node_id. ",
+            # para: context text, root_node_id text
+            inputSchema = {
+                "type": "object",
+                "properties": {
+                    "root_node_id": {
+                        "type": "string",
+                        "description": "the root_noot_id"
+                    },
+                    "context": {
+                        "type": "string",
+                        "description": "the context of decision"
+                    }
+                },
+                "required": ["context", "root_node_id"]
+            }
+        ),
+        Tool(
+            name = "adbpg_graphrag_delete_decision_tree",
+            description = " Delete a sub-decision tree under the node specified by root_node_entity. ",
+            # para: root_node_entity text
+            inputSchema = {
+                "type": "object",
+                "properties": {
+                    "root_node_entity": {
+                        "type": "string",
+                        "description": "the root_noot_entity"
+                        
+                    }
+                },
+                "required": ["root_node_entity"]
+            }
+        ),
+        Tool(
+            name = "adbpg_graphrag_reset_tree_query",
+            description = " Reset the decision tree in the tree query mode",
+            # para: 
+            inputSchema = {
+                "type": "object",
+                "required": []
+            }
+        ),
+    ]
+
+def _execute_graphrag_tool(sql: str, params: list, db: DatabaseManager) -> list[TextContent]:
+    """
+    执行 ADBPG GraphRAG 工具并返回结果
+    """
+    try:
+        conn = db.get_graphrag_connection()
+        with conn.cursor() as cursor:
+            cursor.execute(sql, params)
+            if cursor.description:
+                json_result = cursor.fetchone()[0]
+                return [TextContent(type="text", text=json_result)]
+            else:
+                return [TextContent(type="text", text="ADBPG GraphRAG Tool executed successfully")]
+    except Exception as e:
+        logger.error(f"Error executing ADBPG GraphRAG Tool: {e}")
+        return [TextContent(type="text", text="Error executing ADBPG GraphRAG Tool")]
+    
+async def call_graphrag_tool(name: str, arguments: dict, db: DatabaseManager) -> list[TextContent]:
+    """
+    调用 ADBPG GraphRAG 工具
+    """
+    if name == "adbpg_graphrag_upload":
+        filename, context = arguments.get("filename"), arguments.get("context")
+        if not all([filename, context]):
+            raise ValueError("Filename and context are required.")
+        sql = "SELECT adbpg_graphrag.upload(%s::text, %s::text)"
+        params = [filename, context]
+    
+    elif name == "adbpg_graphrag_query":
+        query_str = arguments.get("query_str")
+        if not query_str:
+            raise ValueError("Query string is required.")
+        query_mode = arguments.get("query_mode", "mix")
+        start_node = arguments.get("start_search_node_id")
+
+        if start_node:
+            sql = "SELECT adbpg_graphrag.query(%s::text, %s::text, %s::text)"
+            params = [query_str, query_mode, start_node]
+        else:
+            sql = "SELECT adbpg_graphrag.query(%s::text, %s::text)"
+            params = [query_str, query_mode]
+    
+    elif name == "adbpg_graphrag_reset_tree_query":
+        sql = "SELECT adbpg_graphrag.reset_tree_query()"
+        params = []
+    
+    elif name == "adbpg_graphrag_upload_decision_tree":
+        root_node = arguments.get("root_node")
+        context = arguments.get("context")
+        if not context:
+            raise ValueError("Decision Tree Context is required")
+        if not root_node:
+            root_node = None
+        sql = "SELECT adbpg_graphrag.upload_decision_tree(%s::text, %s::text)"
+        params = [context, root_node]
+    
+    elif name == "adbpg_graphrag_append_decision_tree":
+        root_node = arguments.get("root_node_id")
+        context = arguments.get("context")
+        if not context:
+            raise ValueError("Decision Tree Context is required")
+        if not root_node:
+            raise ValueError("Root node id is required")
+        sql = "SELECT adbpg_graphrag.append_decision_tree(%s::text, %s::text)"
+        params = [context, root_node]
+
+    elif name == "adbpg_graphrag_delete_decision_tree":
+        root_node = arguments.get("root_node_entity")
+        if not root_node:
+            raise ValueError("Root node entity is required")
+        sql = "SELECT adbpg_graphrag.delete_decision_tree(%s::text)"
+        params = [root_node]
+
+    else:
+        raise ValueError(f"Unknown graphrag tool: {name}")
+    
+    return _execute_graphrag_tool(sql, params, db)