adbpg-mcp-server 1.0.6__tar.gz → 1.0.8__tar.gz

{adbpg_mcp_server-1.0.6 → adbpg_mcp_server-1.0.8}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: adbpg-mcp-server
-Version: 1.0.6
-Summary: MCP server for AnalyticDB PostgreSQL
+Version: 1.0.8
+Summary: MCP Server for AnalyticDB PostgreSQL
 License-File: LICENSE
 Requires-Python: >=3.10
 Requires-Dist: mcp>=1.4.0

{adbpg_mcp_server-1.0.6 → adbpg_mcp_server-1.0.8}/README.md

@@ -43,7 +43,8 @@ Add the following configuration to the MCP client configuration file:
       "LLMEMORY_API_KEY": "llm memory api_key",
       "LLMEMORY_BASE_URL": "llm memory base_url",
       "LLMEMORY_LLM_MODEL": "llm memory model name",
-      "LLMEMORY_EMBEDDING_MODEL": "llm memory embedding model name"
+      "LLMEMORY_EMBEDDING_MODEL": "llm memory embedding model name",
+      "LLMEMORY_ENABLE_GRAPH": "enable graph engine for llm memory (Default: false)"
     }
   }
 }
@@ -78,6 +79,7 @@ pip install adbpg_mcp_server
       "LLMEMORY_BASE_URL": "llm memory base_url",
       "LLMEMORY_LLM_MODEL": "llm memory model name",
       "LLMEMORY_EMBEDDING_MODEL": "llm memory embedding model name"
+      "LLMEMORY_ENABLE_GRAPH": "enable graph engine for llm memory (Default: false)"
     }
   }
 }
@@ -206,6 +208,7 @@ MCP Server requires the following environment variables to initialize graphRAG a
 - `LLMEMORY_BASE_URL`: Base URL for LLM or embedding service endpoint
 - `LLMEMORY_LLM_MODEL`: LLM model name or identifier
 - `LLMEMORY_EMBEDDING_MODEL`: Embedding model name or identifier
+- `LLMEMORY_ENABLE_GRAPH`: (Optional) Enable graph engine for llm memory (Default: false)
 
 ## Dependencies
 

{adbpg_mcp_server-1.0.6 → adbpg_mcp_server-1.0.8}/adbpg_mcp_server.py

@@ -70,7 +70,7 @@ except Exception as e:
     logger.error(f"Error loading environment variables: {e}")
     sys.exit(1)
 
-SERVER_VERSION = "0.1.0"
+SERVER_VERSION = "0.2.0"
 
 
 # 获得 graphrag 初始化配置
@@ -84,12 +84,29 @@ def get_graphrag_config():
         "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")
+        "relationship_types": os.getenv("GRAPHRAG_RELATIONSHIP_TYPES"),
+        "postgres_password": os.getenv("ADBPG_PASSWORD")
     }
     return graphrag_config
     
 # 获得llmemory 初始化配置
 def get_llmemory_config():
+    config = get_db_config()
+    port = 3000
+    sql = """
+        select port from gp_segment_configuration where content = -1 and role = 'p';
+    """
+    try:
+        with psycopg.connect(**config) as conn:
+            conn.autocommit = True
+            with conn.cursor() as cursor:
+                cursor.execute(sql)
+                port = cursor.fetchone()[0]
+    except Error as e:
+        raise RuntimeError(f"Database error: {str(e)}")
+    llmemory_enable_graph = os.getenv("LLMEMORY_ENABLE_GRAPH", "False") 
+    
+
     llm_memory_config = {
         "llm": {
             "provider": "openai",
@@ -115,10 +132,22 @@ def get_llmemory_config():
                 "password": os.getenv("ADBPG_PASSWORD"),
                 "dbname": os.getenv("ADBPG_DATABASE"),
                 "hnsw": "True",
-                "embedding_model_dims": os.getenv("LLMEMORY_EMBEDDING_DIMS", 1024)
+                "embedding_model_dims": os.getenv("LLMEMORY_EMBEDDING_DIMS", 1024),
+                "port": port
             }
         }
     }
+    if llmemory_enable_graph == "True" or llmemory_enable_graph == "true":
+        llm_memory_config["graph_store"] = {
+            "provider": "adbpg",
+            "config": {
+                "url": "http://localhost",
+                "username": os.getenv("ADBPG_USER"),
+                "password": os.getenv("ADBPG_PASSWORD"),
+                "database": os.getenv("ADBPG_DATABASE"),
+                "port": port
+            }
+        }
     return llm_memory_config
 
 def get_db_config():
@@ -186,8 +215,8 @@ def get_graphrag_tool_connection() -> Connection:
 
 LLM_MEMORY_CONN: Connection | None = None
 def get_llm_memory_tool_connection() -> Connection:
-    global LLM_MEMORY_CONN
     global LLMEMORY_ENV_IS_READY
+    global LLM_MEMORY_CONN
     config = get_db_config()
     # 如果未连接，或者连接失效 重新连接
     if LLM_MEMORY_CONN is None or LLM_MEMORY_CONN.closed:
@@ -507,7 +536,7 @@ async def list_tools() -> list[Tool]:
         #### graphrag & llm_memory tool list
         Tool(
             name = "adbpg_graphrag_upload",
-            description = "Upload a text file (with its name) and file content to graphrag to generate a knowledge graph.",
+            description = "Execute graphrag upload operation",
             # 参数：filename text， context text
             # filename 表示文件名称， context 表示文件内容
             inputSchema = {
@@ -515,11 +544,11 @@ async def list_tools() -> list[Tool]:
                 "properties": {
                     "filename": {
                         "type": "string",
-                        "description": "The name of the file to be uploaded"
+                        "description": "The file name need to upload"
                     },
                     "context": {
                         "type": "string",
-                        "description": "The textual content of the file."
+                        "description": "the context of your file"
                     }
                 },
                 "required": ["filename", "context"]
@@ -527,7 +556,7 @@ async def list_tools() -> list[Tool]:
         ),
         Tool(
             name = "adbpg_graphrag_query",
-            description = "Query the graphrag using the specified query string and mode.",
+            description = "Execute graphrag query operation",
             # 参数：query_str text, [query_mode text]
             # query_str 是询问的问题，query_mode 选择查询模式
             inputSchema = {
@@ -535,11 +564,15 @@ async def list_tools() -> list[Tool]:
                 "properties": {
                     "query_str": {
                         "type": "string",
-                        "description": "The query content."
+                        "description": "The query you want to ask"
                     },
                     "query_mode": {
                         "type": "string",
-                        "description": "The query mode, choose from [bypass, naive, local, global, hybrid, mix]. If null, defaults to mix."
+                        "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"]
@@ -593,24 +626,40 @@ async def list_tools() -> list[Tool]:
                     "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 = "Add LLM long memory with a specific user, run or agent.",
+            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": "object",
-                        "description": "llm_memory 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",
@@ -642,7 +691,7 @@ async def list_tools() -> list[Tool]:
         ),
         Tool(
             name = "adbpg_llm_memory_get_all",
-            description = "Retrieves all memory records associated with a specific user, run or agent.",
+            description = "Execute llm_memory get_all operation",
             # 参数：user_id text, run_id text, agent_id text
             # 获取某个用户或者某个agent的所有记忆
             inputSchema={
@@ -666,7 +715,7 @@ async def list_tools() -> list[Tool]:
         ),
         Tool(
             name = "adbpg_llm_memory_search",
-            description = "Retrieves memories relevant to the given query for a specific user, run, or agent.",
+            description = "Execute llm_memory search operation",
             # 参数：query text, user_id text, run_id text, agent_id text, filter json
             # 获取与给定 query 相关的记忆
             inputSchema={
@@ -699,7 +748,7 @@ async def list_tools() -> list[Tool]:
         ,
         Tool(
             name = "adbpg_llm_memory_delete_all",
-            description = "Delete all memory records associated with a specific user, run or agent.",
+            description = "Execute llm_memory delete_all operation",
             # 参数：user_id text, run_id text, agent_id text
             # 删除某个用户或者agent的所有记忆
             inputSchema={
@@ -741,6 +790,7 @@ 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:
@@ -788,6 +838,11 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
             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:
@@ -798,7 +853,7 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
         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"]):
+        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")
@@ -836,18 +891,36 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
             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"
-        # 命令拼接
+            
         wrapped_sql = f"""
-                SELECT adbpg_graphrag.query(%s::text, %s::text)
-            """
+            SELECT adbpg_graphrag.query(%s::text, %s::text)
+        """
         params = [query_str, query_mode]
-        return get_graphrag_tool_result(wrapped_sql, params)
 
+        if start_search_node_id:
+            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.")
@@ -1040,15 +1113,8 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
             conn.autocommit = True
             with conn.cursor() as cursor:
                 
-                #cursor.execute("SET statement_timeout = 300000")
-                #过滤子查询的分号
-                query = query.rstrip().rstrip(';')
-                wrapped_query = f"""
-                SELECT json_agg(row_to_json(t))
-                FROM ({query}) AS t
-                """
-                cursor.execute(wrapped_query)
-
+                cursor.execute(query)
+                
                 if name == "analyze_table":
                     return [TextContent(type="text", text=f"Successfully analyzed table {schema}.{table}")]
 
@@ -1057,7 +1123,6 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
                     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:

{adbpg_mcp_server-1.0.6 → adbpg_mcp_server-1.0.8}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "hatchling.build"
 
 [project]
 name = "adbpg-mcp-server"
-version = "1.0.6"
-description = "MCP server for AnalyticDB PostgreSQL"
+version = "1.0.8"
+description = "MCP Server for AnalyticDB PostgreSQL"
 requires-python = ">=3.10"
 dependencies = [
     "mcp>=1.4.0",