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

adbpg_mcp_server.py

@@ -0,0 +1,1143 @@
+import asyncio
+import logging
+import os
+import sys
+import json
+import psycopg
+import re
+import ast
+from psycopg import OperationalError as Error
+from psycopg import Connection
+from mcp.server import Server
+from mcp.types import Resource, Tool, TextContent, ResourceTemplate
+from pydantic import AnyUrl
+from dotenv import load_dotenv
+from mcp.server.stdio import stdio_server
+
+# 配置日志，输出到标准错误
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    stream=sys.stderr
+)
+logger = logging.getLogger("adbpg-mcp-server")
+
+# 检查环境变量Flag
+GRAPHRAG_ENV_IS_READY = True
+LLMEMORY_ENV_IS_READY = True
+# 加载环境变量
+try:
+    load_dotenv()
+    logger.info("Environment variables loaded")
+    
+    # 检查必要的环境变量
+    required_vars = ["ADBPG_HOST", "ADBPG_PORT", "ADBPG_USER", "ADBPG_PASSWORD", "ADBPG_DATABASE"]
+    missing_vars = [var for var in required_vars if not os.getenv(var)]
+    if missing_vars:
+        error_msg = f"Missing required environment variables: {', '.join(missing_vars)}"
+        logger.error(error_msg)
+        raise ValueError(error_msg)
+    logger.info("All ADBPG required environment variables are set")
+
+    # 检查graphrag/llmemory 环境变量
+    required_graphrag_vars = [
+        "GRAPHRAG_LLM_MODEL", 
+        "GRAPHRAG_API_KEY", 
+        "GRAPHRAG_BASE_URL", 
+        "GRAPHRAG_EMBEDDING_MODEL",
+        "GRAPHRAG_EMBEDDING_BASE_URL",
+        "GRAPHRAG_EMBEDDING_API_KEY"
+        ]
+    missing_graphrag_vars = [var for var in required_graphrag_vars if not os.getenv(var)]
+    if missing_graphrag_vars:
+        GRAPHRAG_ENV_IS_READY = False
+        error_msg = f"Missing required graphrag environment variables:{', '.join(missing_graphrag_vars)}"
+        logger.error(error_msg)
+    else:
+        logger.info("All graphRAG required environment variables are set")
+
+    required_llmemory_vars = ["LLMEMORY_LLM_MODEL", "LLMEMORY_API_KEY", "LLMEMORY_BASE_URL", "LLMEMORY_EMBEDDING_MODEL"]
+    missing_llmemory_vars = [var for var in required_llmemory_vars if not os.getenv(var)]
+    if missing_llmemory_vars:
+        LLMEMORY_ENV_IS_READY = False
+        error_msg = f"Missing required llm memory environment variables:{', '.join(missing_llmemory_vars)}"
+        logger.error(error_msg)
+    else:
+        logger.info("All llm memory required environment variables are set")
+    
+
+except Exception as e:
+    logger.error(f"Error loading environment variables: {e}")
+    sys.exit(1)
+
+SERVER_VERSION = "0.1.0"
+
+
+# 获得 graphrag 初始化配置
+def get_graphrag_config():
+    graphrag_config = {
+        "llm_model": os.getenv("GRAPHRAG_LLM_MODEL"),
+        "llm_api_key": os.getenv("GRAPHRAG_API_KEY"),
+        "llm_url": os.getenv("GRAPHRAG_BASE_URL"),
+        "embedding_model": os.getenv("GRAPHRAG_EMBEDDING_MODEL"),
+        "embedding_api_key": os.getenv("GRAPHRAG_EMBEDDING_API_KEY"),
+        "embedding_url": os.getenv("GRAPHRAG_EMBEDDING_BASE_URL"),
+        "language": os.getenv("GRAPHRAG_LANGUAGE", "English"),
+        "entity_types": os.getenv("GRAPHRAG_ENTITY_TYPES"),
+        "relationship_types": os.getenv("GRAPHRAG_RELATIONSHIP_TYPES"),
+        "postgres_password": os.getenv("ADBPG_PASSWORD")
+    }
+    return graphrag_config
+    
+# 获得llmemory 初始化配置
+def get_llmemory_config():
+    llm_memory_config = {
+        "llm": {
+            "provider": "openai",
+            "config": {
+                "model": os.getenv("LLMEMORY_LLM_MODEL"),
+                "openai_base_url": os.getenv("LLMEMORY_BASE_URL"),
+                "api_key": os.getenv("LLMEMORY_API_KEY")
+            }
+        },
+        "embedder": {
+            "provider": "openai",
+            "config": {
+                "model": os.getenv("LLMEMORY_EMBEDDING_MODEL"),
+                "embedding_dims": os.getenv("LLMEMORY_EMBEDDING_DIMS", 1024),
+                "api_key": os.getenv("LLMEMORY_API_KEY"),
+                "openai_base_url": os.getenv("LLMEMORY_BASE_URL")
+            }
+        },
+        "vector_store": {
+            "provider": "adbpg",
+            "config": {
+                "user": os.getenv("ADBPG_USER"),
+                "password": os.getenv("ADBPG_PASSWORD"),
+                "dbname": os.getenv("ADBPG_DATABASE"),
+                "hnsw": "True",
+                "embedding_model_dims": os.getenv("LLMEMORY_EMBEDDING_DIMS", 1024)
+            }
+        }
+    }
+    return llm_memory_config
+
+def get_db_config():
+    """从环境变量获取数据库配置信息"""
+    try:
+        config = {
+            "host": os.getenv("ADBPG_HOST", "localhost"),
+            "port": os.getenv("ADBPG_PORT"),
+            "user": os.getenv("ADBPG_USER"),
+            "password": os.getenv("ADBPG_PASSWORD"),
+            "dbname": os.getenv("ADBPG_DATABASE"),
+            "application_name": f"adbpg-mcp-server-{SERVER_VERSION}"
+        }
+        
+        # 记录配置信息（不包含密码）
+        logger.info(f"Database config: host={config['host']}, port={config['port']}, user={config['user']}, dbname={config['dbname']}")
+        return config
+    except Exception as e:
+        logger.error(f"Error getting database config: {str(e)}")
+        raise
+
+# 全局graphrag长连接 并 初始化
+GRAPHRAG_CONN: Connection | None = None
+def get_graphrag_tool_connection() -> Connection:
+    global GRAPHRAG_CONN
+    global GRAPHRAG_ENV_IS_READY
+    config = get_db_config()
+    # 如果未连接，或者连接失效 重新连接
+    if GRAPHRAG_CONN is None or GRAPHRAG_CONN.closed:
+        GRAPHRAG_CONN = psycopg.connect(**config)
+        GRAPHRAG_CONN.autocommit = True
+        try:
+            graphrag_conn = GRAPHRAG_CONN
+            with graphrag_conn.cursor() as cursor:
+                cursor.execute("SELECT adbpg_graphrag.initialize(%s::json);", (graphrag_config_str,))    
+                logger.info(f"[GraphRAG] Use the connection {id(graphrag_conn)} when executing the graphrag init")
+            logger.info("Successfully initialize the graphrag server\n")
+        except Exception as e:
+            GRAPHRAG_ENV_IS_READY = False
+            logger.error(f"Failed to initialize the graphrag server: {e}")
+        # 重新执行初始化
+    else:
+        # 发送一个轻量级查询 检测 连接是否健康
+        try:
+            with GRAPHRAG_CONN.cursor() as cur:
+                cur.execute("SELECT 1;")
+                _ = cur.fetchone()
+        except Exception:
+            # 重连
+            GRAPHRAG_CONN.close()
+            GRAPHRAG_CONN = psycopg.connect(**config)
+            GRAPHRAG_CONN.autocommit = True
+            # 重新执行初始化
+            try:
+                graphrag_conn = GRAPHRAG_CONN
+                with graphrag_conn.cursor() as cursor:
+                    cursor.execute("SELECT adbpg_graphrag.initialize(%s::json);", (graphrag_config_str,))
+                    logger.info(f"[GraphRAG] Use the connection {id(graphrag_conn)} when executing the graphrag init")
+                logger.info("Successfully initialize the graphrag server\n")
+            except Exception as e:
+                GRAPHRAG_ENV_IS_READY = False
+                logger.error(f"Failed to initialize the graphrag server: {e}")
+    
+    return GRAPHRAG_CONN
+
+LLM_MEMORY_CONN: Connection | None = None
+def get_llm_memory_tool_connection() -> Connection:
+    global LLMEMORY_ENV_IS_READY
+    global LLM_MEMORY_CONN
+    config = get_db_config()
+    # 如果未连接，或者连接失效 重新连接
+    if LLM_MEMORY_CONN is None or LLM_MEMORY_CONN.closed:
+        LLM_MEMORY_CONN = psycopg.connect(**config)
+        LLM_MEMORY_CONN.autocommit = True
+        try:
+            llm_memory_conn = LLM_MEMORY_CONN
+            with llm_memory_conn.cursor() as cursor:
+                cursor.execute("SELECT adbpg_llm_memory.config(%s::json)", (llm_memory_config_str,))
+                logger.info(f"[LLM Memory] Use the connection {id(llm_memory_conn)} when executing the llm_memory init")
+            logger.info("Successfully initialize the llm server\n")
+        except Exception as e:
+            LLMEMORY_ENV_IS_READY = False
+            logger.error(f"Failed to initialize the llm_memory server: {e}")
+    else:
+        # 发送一个轻量级查询 检测 连接是否健康
+        try:
+            with LLM_MEMORY_CONN.cursor() as cur:
+                cur.execute("SELECT 1;")
+                _ = cur.fetchone()
+        except Exception:
+            # 重连
+            LLM_MEMORY_CONN.close()
+            LLM_MEMORY_CONN = psycopg.connect(**config)
+            LLM_MEMORY_CONN.autocommit = True
+            try:
+                llm_memory_conn = LLM_MEMORY_CONN
+                with llm_memory_conn.cursor() as cursor:
+                    cursor.execute("SELECT adbpg_llm_memory.config(%s::json)", (llm_memory_config_str,))
+                    logger.info(f"[LLM Memory] Use the connection {id(llm_memory_conn)} when executing the llm_memory init")
+                logger.info("Successfully initialize the llm server\n")
+            except Exception as e:
+                LLMEMORY_ENV_IS_READY = False
+                logger.error(f"Failed to initialize the llm_memory server: {e}")
+
+    return LLM_MEMORY_CONN
+
+#### 初始化 
+if GRAPHRAG_ENV_IS_READY == True:
+    # 初始化 graphrag
+    logger.info("Starting graphRAG server...")
+    graphrag_config = get_graphrag_config()
+    graphrag_config_str = json.dumps(graphrag_config)
+    # 建立长连接 并 初始化
+    get_graphrag_tool_connection()
+if LLMEMORY_ENV_IS_READY == True:
+    # 初始化 llmemory
+    logger.info("Starting llmemory server...")
+    llm_memory_config = get_llmemory_config()
+    llm_memory_config_str = json.dumps(llm_memory_config)
+    # 建立长连接 并 初始化
+    get_llm_memory_tool_connection()
+
+# 初始化服务器
+try:
+    app = Server("adbpg-mcp-server")
+    logger.info("MCP server initialized")
+except Exception as e:
+    logger.error(f"Error initializing MCP server: {e}")
+    sys.exit(1)
+
+@app.list_resources()
+async def list_resources() -> list[Resource]:
+    """列出可用的基本资源"""
+    try:
+        return [
+            Resource(
+                uri="adbpg:///schemas",
+                name="All Schemas",
+                description="AnalyticDB PostgreSQL schemas. List all schemas in the database",
+                mimeType="text/plain"
+            )
+        ]
+    except Exception as e:
+        logger.error(f"Error listing resources: {str(e)}")
+        raise
+
+@app.list_resource_templates()
+async def list_resource_templates() -> list[ResourceTemplate]:
+    """
+    定义动态资源模板
+    
+    返回:
+        list[ResourceTemplate]: 资源模板列表
+        包含以下模板：
+        - 列出schema中的表
+        - 获取表DDL
+        - 获取表统计信息
+    """
+    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",  # 表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"
+        )
+    ]
+
+@app.read_resource()
+async def read_resource(uri: AnyUrl) -> str:
+    """
+    读取资源内容
+    
+    参数:
+        uri (AnyUrl): 资源URI
+        
+    返回:
+        str: 资源内容
+        
+    支持的URI格式：
+    - adbpg:///schemas: 列出所有schema
+    - adbpg:///{schema}/tables: 列出指定schema中的表
+    - adbpg:///{schema}/{table}/ddl: 获取表的DDL
+    - adbpg:///{schema}/{table}/statistics: 获取表的统计信息
+    """
+    config = get_db_config()
+    uri_str = str(uri)
+    
+    if not uri_str.startswith("adbpg:///"):
+        raise ValueError(f"Invalid URI scheme: {uri_str}")
+    
+    try:
+        with psycopg.connect(**config) as conn:  # 建立数据库连接
+            conn.autocommit = True  # 设置自动提交
+            with conn.cursor() as cursor:  # 创建游标
+                path_parts = uri_str[9:].split('/')  # 解析URI路径
+                
+                if path_parts[0] == "schemas":
+                    # 列出所有schema
+                    query = """
+                        SELECT schema_name 
+                        FROM information_schema.schemata 
+                        WHERE schema_name NOT IN ('pg_catalog', 'information_schema')
+                        ORDER BY schema_name;
+                    """
+                    cursor.execute(query)
+                    schemas = cursor.fetchall()
+                    return "\n".join([schema[0] for schema in schemas])
+                    
+                elif len(path_parts) == 2 and path_parts[1] == "tables":
+                    # 列出指定schema中的表
+                    schema = path_parts[0]
+                    query = f"""
+                        SELECT table_name, table_type
+                        FROM information_schema.tables
+                        WHERE table_schema = %s
+                        ORDER BY table_name;
+                    """
+                    cursor.execute(query, (schema,))
+                    tables = cursor.fetchall()
+                    return "\n".join([f"{table[0]} ({table[1]})" for table in tables])
+                    
+                elif len(path_parts) == 3 and path_parts[2] == "ddl":
+                    # 获取表的DDL
+                    schema = path_parts[0]
+                    table = 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 = path_parts[0]
+                    table = path_parts[1]
+                    query = """
+                        SELECT 
+                            schemaname,
+                            tablename,
+                            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 = []
+                    for row in rows:
+                        result.append(f"Column: {row[2]}")
+                        result.append(f"  Null fraction: {row[3]}")
+                        result.append(f"  Average width: {row[4]}")
+                        result.append(f"  Distinct values: {row[5]}")
+                        if row[6]:
+                            result.append(f"  Most common values: {row[6]}")
+                            result.append(f"  Most common frequencies: {row[7]}")
+                        result.append("")
+                    return "\n".join(result)
+                
+                raise ValueError(f"Invalid resource URI format: {uri_str}")
+      
+    except Error as e:
+        raise RuntimeError(f"Database error: {str(e)}")
+
+@app.list_tools()
+async def list_tools() -> list[Tool]:
+    """
+    列出可用的工具
+    
+    返回:
+        list[Tool]: 工具列表
+        包含以下工具：
+        - execute_select_sql: 执行SELECT查询
+        - execute_dml_sql: 执行DML操作
+        - execute_ddl_sql: 执行DDL操作
+        - analyze_table: 分析表统计信息
+        - explain_query: 获取查询执行计划
+
+        - adbpg_graphrag_upload: 执行 graphRAG upload 操作，上传文本
+        - adbpg_graphrag_query: 执行 graphRAG query 操作 
+        - adbpg_graphrag_upload_decision_tree: 上传一个决策树
+        - adbpg_graphrag_append_decision_tree: 在某个节点上新增子树
+        - adbpg_graphrag_delete_decision_tree: 根据节点id删除起下层子树
+
+        - adbpg_llm_memory_add: 执行新增记忆操作
+        - adbpg_llm_memory_get_all: 获取所有记忆
+        - adbpg_llm_memory_search: 根据查询检索记忆
+        - adbpg_llm_memory_delete_all: 删除所有记忆
+    """
+    return [
+        Tool(
+            name="execute_select_sql",
+            description="Execute SELECT SQL to query data from ADBPG database.",
+            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",
+                        "description": "Schema name"
+                    },
+                    "table": {
+                        "type": "string",
+                        "description": "Table name"
+                    }
+                },
+                "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"]
+            }
+        ),
+
+        #### graphrag & llm_memory tool list
+        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": []
+            }
+        ),
+        Tool(
+            name = "adbpg_llm_memory_add",
+            description = "Execute llm_memory add operation",
+            # 参数：messages json, user_id text, run_id text, agent_id text, metadata json
+            # 增加新的记忆
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "messages": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "role": {"type": "string"},
+                                "content": {"type": "string"}
+                            },
+                            "required": ["role", "content"]
+                        },
+                        "description": "List of messages objects (e.g., conversation history)"
+                    },
+                    "user_id": {
+                        "type": "string",
+                        "description": "the user_id"
+                    },
+                    "run_id": {
+                        "type": "string",
+                        "description": "the run_id"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "the agent_id"
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "description": "the metatdata json"
+                    },
+                    "memory_type": {
+                        "type": "string",
+                        "description": "the memory_type text"
+                    },
+                    "prompt": {
+                        "type": "string",
+                        "description": "the prompt"
+                    }
+                },
+                "required": ["messages"]
+            }
+        ),
+        Tool(
+            name = "adbpg_llm_memory_get_all",
+            description = "Execute llm_memory get_all operation",
+            # 参数：user_id text, run_id text, agent_id text
+            # 获取某个用户或者某个agent的所有记忆
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "user_id": {
+                        "type": "string",
+                        "description": "The user_id"
+                    },
+                    "run_id": {
+                        "type": "string",
+                        "description": "The run_id"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "The agent_id"
+                    }
+                },
+                "required": []
+            }
+        ),
+        Tool(
+            name = "adbpg_llm_memory_search",
+            description = "Execute llm_memory search operation",
+            # 参数：query text, user_id text, run_id text, agent_id text, filter json
+            # 获取与给定 query 相关的记忆
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "llm_memory relevant query"
+                    },
+                    "user_id": {
+                        "type": "string",
+                        "description": "The search of user_id"
+                    },
+                    "run_id": {
+                        "type": "string",
+                        "description": "The search of run_id"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "The search of agent_id"
+                    },
+                    "filter": {
+                        "type": "object",
+                        "description": "The search of filter"
+                    }
+                },
+                "required": ["query"]
+            }
+        )
+        ,
+        Tool(
+            name = "adbpg_llm_memory_delete_all",
+            description = "Execute llm_memory delete_all operation",
+            # 参数：user_id text, run_id text, agent_id text
+            # 删除某个用户或者agent的所有记忆
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "user_id": {
+                        "type": "string",
+                        "description": "The user_id"
+                    },
+                    "run_id": {
+                        "type": "string",
+                        "description": "The run_id"
+                    },
+                    "agent_id": {
+                        "type": "string",
+                        "description": "The agent_id"
+                    }
+                },
+                "required": []
+            }
+        )
+        
+    ]
+
+def get_graphrag_tool_result(wrapped_sql, params) -> list[TextContent]:
+    try:
+        conn = get_graphrag_tool_connection()
+        with conn.cursor() as cursor:
+            cursor.execute(wrapped_sql, params)
+            if cursor.description:
+                json_result = cursor.fetchone()[0]
+                return [TextContent(type="text", text=json_result)]
+            else:
+                return [TextContent(type="text", text="graphrag command executed successfully")]
+    except Exception as e:
+        return [TextContent(type="text", text=f"Error executing graphrag command: {str(e)}")]
+
+def get_llm_memory_tool_result(wrapped_sql, params) -> list[TextContent]:
+    try:
+        conn = get_llm_memory_tool_connection()
+        with conn.cursor() as cursor:
+
+            cursor.execute(wrapped_sql, params)
+            
+            if cursor.description:
+                json_result = cursor.fetchone()[0]
+                return [TextContent(type="text", text=json_result)]
+            else:
+                return [TextContent(type="text", text="llm_memory command executed successfully")]
+    except Exception as e:
+        return [TextContent(type="text", text=f"Error executing llm_memory command: {str(e)}")]
+            
+
+@app.call_tool()
+async def call_tool(name: str, arguments: dict) -> list[TextContent]:
+    """
+    执行工具操作
+    
+    参数:
+        name (str): 工具名称
+        arguments (dict): 工具参数
+        
+    返回:
+        list[TextContent]: 执行结果
+        
+    支持的工具：
+    - execute_select_sql: 执行SELECT查询
+    - execute_dml_sql: 执行DML操作
+    - execute_ddl_sql: 执行DDL操作
+    - analyze_table: 分析表统计信息
+    - explain_query: 获取查询执行计划
+    
+    - adbpg_graphrag_upload: 执行 graphRAG upload 操作，上传文本
+    - adbpg_graphrag_query: 执行 graphRAG query 操作 
+
+    - adbpg_llm_memory_add: 执行新增记忆操作
+    - adbpg_llm_memory_get_all: 获取所有记忆
+    - adbpg_llm_memory_search: 根据查询检索记忆
+    - adbpg_llm_memory_delete_all: 删除所有记忆
+    """
+    config = get_db_config()
+    global GRAPHRAG_ENV_IS_READY
+    # 根据工具名称处理不同的操作
+    if name == "execute_select_sql":
+        query = arguments.get("query")
+        if not query:
+            raise ValueError("Query is required")
+        if not query.strip().upper().startswith("SELECT"):
+            raise ValueError("Query must be a SELECT statement")
+        query = query.rstrip().rstrip(';')
+        query = f"""
+            SELECT json_agg(row_to_json(t))
+            FROM ({query}) AS t
+        """
+    elif name == "execute_dml_sql":
+        query = arguments.get("query")
+        if not query:
+            raise ValueError("Query is required")
+        if not any(query.strip().upper().startswith(keyword) for keyword 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:
+            raise ValueError("Query is required")
+        if not any(query.strip().upper().startswith(keyword) for keyword in ["CREATE", "ALTER", "DROP", "TRUNCATE"]):
+            raise ValueError("Query must be a DDL statement (CREATE, ALTER, DROP)")
+    elif name == "analyze_table":
+        schema = arguments.get("schema")
+        table = 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 = arguments.get("query")
+        if not query:
+            raise ValueError("Query is required")
+        query = f"EXPLAIN {query}"
+
+    # adbpg_graphrag tool
+    elif name == "adbpg_graphrag_upload":
+        # GraphRAG 服务初始化失败，工具不可用
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.") 
+        filename = arguments.get("filename")
+        context = arguments.get("context")
+        if not filename:
+            raise ValueError("Filename is required")
+        if not context:
+            raise ValueError("Context if required")
+        # 命令拼接
+        wrapped_sql = f"""
+                SELECT adbpg_graphrag.upload(%s::text, %s::text)
+            """
+        params = [filename, context]
+        return get_graphrag_tool_result(wrapped_sql, params)
+
+    elif name == "adbpg_graphrag_query":
+        # GraphRAG 服务初始化失败，工具不可用
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.") 
+        query_str = arguments.get("query_str")
+        query_mode = arguments.get("query_mode")
+        start_search_node_id = arguments.get("start_search_node_id")
+
+        if not query_str:
+            raise ValueError("Query is required")
+        if not query_mode:
+            # default mode
+            query_mode = "mix"
+        if not start_search_node_id:
+            start_search_node_id = None
+
+        # 命令拼接
+        wrapped_sql = f"""
+                SELECT adbpg_graphrag.query(%s::text, %s::text, %s::text)
+            """
+        params = [query_str, query_mode, start_search_node_id]
+        return get_graphrag_tool_result(wrapped_sql, params)
+    
+    elif name == "adbpg_graphrag_reset_tree_query":
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.")
+        wrapped_sql = f"""
+            SELECT adbpg_graphrag.reset_tree_query()
+        """
+        params = []
+        return get_graphrag_tool_result(wrapped_sql, params)
+    
+    elif name == "adbpg_graphrag_upload_decision_tree":
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.")
+        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
+        wrapped_sql = f"""
+            SELECT adbpg_graphrag.upload_decision_tree(%s::text, %s::text)
+        """
+        params = [context, root_node]
+        return get_graphrag_tool_result(wrapped_sql, params)
+
+    elif name == "adbpg_graphrag_append_decision_tree":
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.") 
+        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")
+        wrapped_sql = f"""
+            SELECT adbpg_graphrag.append_decision_tree(%s::text, %s::text)
+        """
+        params = [context, root_node]
+        return get_graphrag_tool_result(wrapped_sql, params)
+
+    elif name == "adbpg_graphrag_delete_decision_tree":
+        if GRAPHRAG_ENV_IS_READY == False:
+            raise ValueError("GraphRAG Server initialization failed. This tool cannot be used.") 
+        root_node = arguments.get("root_node_entity")
+        if not root_node:
+            raise ValueError("Root node entity is required")
+        wrapped_sql = f"""
+            SELECT adbpg_graphrag.delete_decision_tree(%s::text, %s::text)
+        """
+        params = [root_node]
+        return get_graphrag_tool_result(wrapped_sql, params)
+
+    # adbpg_llm_memory tool
+    elif name == "adbpg_llm_memory_add":
+        # LLMEMORY 服务初始化失败，工具不可用
+        if LLMEMORY_ENV_IS_READY == False:
+            raise ValueError("LLMEMORY Server initialization failed. This tool cannot be used.") 
+
+        messages = arguments.get("messages")
+        if not messages:
+            raise ValueError("messages is required")
+        messages_str = json.dumps(messages, ensure_ascii = False)
+
+        user_id = arguments.get("user_id")
+        if not user_id:
+            user_id = None
+        run_id = arguments.get("run_id")
+        if not run_id:
+            run_id = None
+        agent_id = arguments.get("agent_id")
+        if not agent_id:
+            agent_id = None
+        if user_id == None and run_id == None and agent_id == None:
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        
+        metadata = arguments.get("metadata")
+        metadata_str = None
+        if metadata:
+            metadata_str = json.dumps(metadata, ensure_ascii = False)
+        
+        memory_type = arguments.get("memory_type")
+        memory_prompt = arguments.get("prompt")
+        if not memory_type:
+            memory_type = None
+        if not memory_prompt:
+            memory_prompt = None
+
+
+        wrapped_sql = """
+        SELECT adbpg_llm_memory.add(
+            %s::json,
+            %s::text,
+            %s::text,
+            %s::text,
+            %s::json,
+            %s::text,
+            %s::text
+        )
+        """
+        params = [messages_str, user_id, run_id, agent_id, metadata_str, memory_type, memory_prompt]
+        return get_llm_memory_tool_result(wrapped_sql, params)
+        
+    elif name == "adbpg_llm_memory_get_all":
+        # LLMEMORY 服务初始化失败，工具不可用
+        if LLMEMORY_ENV_IS_READY == False:
+            raise ValueError("LLMEMORY Server initialization failed. This tool cannot be used.") 
+        
+        user_id = arguments.get("user_id")
+        if not user_id:
+            user_id = None
+        run_id = arguments.get("run_id")
+        if not run_id:
+            run_id = None
+        agent_id = arguments.get("agent_id")
+        if not agent_id:
+            agent_id = None
+        if user_id == None and run_id == None and agent_id == None:
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        wrapped_sql = f"""
+                SELECT adbpg_llm_memory.get_all(
+                    %s::text,
+                    %s::text,
+                    %s::text
+                )
+            """
+        params = [user_id, run_id, agent_id]
+        return get_llm_memory_tool_result(wrapped_sql, params)
+        
+
+    elif name == "adbpg_llm_memory_search":
+        # LLMEMORY 服务初始化失败，工具不可用
+        if LLMEMORY_ENV_IS_READY == False:
+            raise ValueError("LLMEMORY Server initialization failed. This tool cannot be used.") 
+        query = arguments.get("query")
+        if not query:
+            raise ValueError("Query is required")
+        
+        user_id = arguments.get("user_id")
+        if not user_id:
+            user_id = None
+        run_id = arguments.get("run_id")
+        if not run_id:
+            run_id = None
+        agent_id = arguments.get("agent_id")
+        if not agent_id:
+            agent_id = None
+        if user_id == None and run_id == None and agent_id == None:
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        
+        filter_json = arguments.get("filter")
+        filter_json_str = None
+        # json格式载入
+        if filter_json:
+            filter_json_str = json.dumps(filter_json, ensure_ascii = False)
+        # 拼接命令
+        wrapped_sql = f"""
+                SELECT adbpg_llm_memory.search(
+                    %s::text,
+                    %s::text,
+                    %s::text,
+                    %s::text,
+                    %s::json
+                )
+            """
+        params = [query, user_id, run_id, agent_id, filter_json_str]
+        return get_llm_memory_tool_result(wrapped_sql, params)
+
+    elif name == "adbpg_llm_memory_delete_all":
+        # LLMEMORY 服务初始化失败，工具不可用
+        if LLMEMORY_ENV_IS_READY == False:
+            raise ValueError("LLMEMORY Server initialization failed. This tool cannot be used.") 
+        
+        user_id = arguments.get("user_id")
+        if not user_id:
+            user_id = None
+        run_id = arguments.get("run_id")
+        if not run_id:
+            run_id = None
+        agent_id = arguments.get("agent_id")
+        if not agent_id:
+            agent_id = None
+        if user_id == None and run_id == None and agent_id == None:
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+
+        wrapped_sql = f"""
+                SELECT adbpg_llm_memory.delete_all(
+                    %s::text,
+                    %s::text,
+                    %s::text
+                )
+            """
+        params = [user_id, run_id, agent_id]
+        return get_llm_memory_tool_result(wrapped_sql, params)
+
+    else:
+        raise ValueError(f"Unknown tool: {name}")
+    
+    try:
+        with psycopg.connect(**config) as conn:
+            conn.autocommit = True
+            with conn.cursor() as cursor:
+                
+                cursor.execute(query)
+                
+                if name == "analyze_table":
+                    return [TextContent(type="text", text=f"Successfully analyzed table {schema}.{table}")]
+
+                if cursor.description:
+                    # 将返回结果存储为json格式
+                    json_result = cursor.fetchone()[0]
+                    json_str = json.dumps(json_result, ensure_ascii = False, indent = 2)
+                    result = [TextContent(type="text", text=json_str)]
+                    try:
+                        json.loads(result[0].text)
+                    except json.JSONDecodeError as e:
+                        raise Exception(f"JSON decode error: {e}\nRaw text: {result[0].text}") from e
+                    return result
+
+                else:
+                    return [TextContent(type="text", text="Query executed successfully")]
+    except Exception as e:
+        return [TextContent(type="text", text=f"Error executing query: {str(e)}")]
+
+async def main():
+    """服务器主入口点"""
+    try:
+        config = get_db_config()
+        logger.info("Starting ADBPG MCP server...")
+        
+        # 测试数据库连接
+        try:
+            with psycopg.connect(**config) as conn:
+                logger.info("Successfully connected to database")
+        except Exception as e:
+            logger.error(f"Failed to connect to database: {e}")
+            sys.exit(1)     
+        # 使用 stdio 传输
+        async with stdio_server() as (read_stream, write_stream):
+            try:
+                logger.info("Running MCP server with stdio transport...")
+                await app.run(
+                    read_stream=read_stream,
+                    write_stream=write_stream,
+                    initialization_options=app.create_initialization_options()
+                )
+            except Exception as e:
+                logger.error(f"Error running server: {str(e)}")
+                raise
+    except Exception as e:
+        logger.error(f"Server initialization error: {str(e)}")
+        raise
+
+def run():
+    """同步运行入口点"""
+    try:
+        asyncio.run(main())
+    except Exception as e:
+        logger.error(f"Fatal error: {e}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    run()