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

adbpg_mcp_server/adbpg_memory.py

@@ -0,0 +1,201 @@
+import json
+import logging
+from mcp.types import Tool, TextContent
+from .adbpg import DatabaseManager
+
+logger = logging.getLogger(__name__)
+
+def get_memory_tools() -> list[Tool]:
+    return [
+        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 _execute_memory_tool(sql: str, params: list, db: DatabaseManager) -> list[TextContent]:
+    try:
+        conn = db.get_llm_memory_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="LLM memory tool executed successfully.")]
+    except Exception as e:
+        return [TextContent(type="text", text=f"Error executing LLM memory tool: {str(e)}")]
+    
+
+async def call_memory_tool(name: str, arguments: dict, db: DatabaseManager) -> list[TextContent]:
+    """调用 LLM Memory 工具"""
+    user_id = arguments.get("user_id")
+    run_id = arguments.get("run_id")
+    agent_id = arguments.get("agent_id")
+
+    if name == "adbpg_llm_memory_add":
+        messages = arguments.get("messages")
+        if not messages:
+            raise ValueError("messages is required")
+        if not any([user_id, run_id, agent_id]):
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        sql = "SELECT adbpg_llm_memory.add(%s::json, %s::text, %s::text, %s::text, %s::json, %s::text, %s::text)"
+        params = [
+            json.dumps(messages, ensure_ascii=False),
+            user_id, run_id, agent_id,
+            json.dumps(arguments.get("metadata"),ensure_ascii=False) if arguments.get("metadata") else None,
+            arguments.get("memory_type"),
+            arguments.get("prompt")
+        ]
+    
+    elif name == "adbpg_llm_memory_search":
+        query = arguments.get("query")
+        if not query:
+            raise ValueError("query is required")
+        if not any([user_id, run_id, agent_id]):
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        sql = "SELECT adbpg_llm_memory.search(%s::text, %s::text, %s::text, %s::text, %s::json)"
+        params = [
+            query, user_id, run_id, agent_id,
+            json.dumps(arguments.get("filter"), ensure_ascii=False) if arguments.get("filter") else None
+        ]
+    
+    elif name == "adbpg_llm_memory_get_all":
+        if not any([user_id, run_id, agent_id]):
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        sql = "SELECT adbpg_llm_memory.get_all(%s::text, %s::text, %s::text)"
+        params = [user_id, run_id, agent_id]
+    
+    elif name == "adbpg_llm_memory_delete_all":
+        if not any([user_id, run_id, agent_id]):
+            raise ValueError("At least one of user_id, run_id, or agent_id must be provided.")
+        sql = "SELECT adbpg_llm_memory.delete_all(%s::text, %s::text, %s::text)"
+        params = [user_id, run_id, agent_id]
+    else:
+        raise ValueError(f"Unknown memory tool: {name}")
+    
+    return _execute_memory_tool(sql, params, db)

adbpg_mcp_server/server_http.py

@@ -0,0 +1,206 @@
+import requests
+import logging
+import json
+import click
+import contextlib
+import uvicorn
+import sys
+import mcp.types as types
+from pydantic import AnyUrl
+from collections.abc import AsyncIterator
+from mcp.server.lowlevel import Server
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.types import Resource, Tool, TextContent, ResourceTemplate
+from starlette.applications import Starlette
+from starlette.routing import Mount
+from .adbpg import DatabaseManager
+from . import adbpg_basic_operation, adbpg_graphrag, adbpg_memory
+from .adbpg_config import settings
+
+
+db_manager: DatabaseManager | None = None
+graphrag_is_available = False
+llm_memory_is_available = False
+
+logging.basicConfig(
+        level=logging.DEBUG,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+logger = logging.getLogger("ADBPG MCP Server")
+
+def initialize_services():
+    """初始化数据库、GraphRAG 和 LLM Memory"""
+    global db_manager, graphrag_is_available, llm_memory_is_available
+
+    if not settings.db_env_ready:
+        logger.error("Cannot start server: Database environment is not configured.")
+        sys.exit(1)
+
+    db_manager = DatabaseManager(settings)
+
+    # 测试主连接
+    try:
+        db_manager.get_basic_connection()
+        logger.info("Successfully connected to database.")
+    except Exception as e:
+        logger.error(f"Failed to connect to database: {e}")
+        sys.exit(1)
+
+    # 初始化 GraphRAG
+    if settings.graphrag_env_ready:
+        try:
+            db_manager.get_graphrag_connection()
+            graphrag_is_available = True
+            logger.info("GraphRAG initialized successfully.")
+        except Exception as e:
+            logger.error(f"Failed to initialize GraphRAG: {e}")
+            graphrag_is_available = False
+            
+    # 初始化 LLM Memory
+    if settings.memory_env_ready:
+        try:
+            db_manager.get_llm_memory_connection()
+            llm_memory_is_available = True
+            logger.info("LLM Memory initialized successfully.")
+        except Exception as e:
+            logger.error(f"Failed to initialize LLM Memory: {e}")
+            llm_memory_is_available = False
+
+
+
+def run_http_server(host, port):
+
+    # 创建MCP服务端
+    app = Server("adbpg-mcp-server")
+    initialize_services()
+    @app.list_resources()
+    async def list_resources() -> list[Resource]:
+        """列出可用的基本资源"""
+        return await adbpg_basic_operation.list_resources()
+
+    @app.list_resource_templates()
+    async def list_resource_templates() -> list[ResourceTemplate]:
+        """
+        定义动态资源模板
+        
+        返回:
+            list[ResourceTemplate]: 资源模板列表
+            包含以下模板：
+            - 列出schema中的表
+            - 获取表DDL
+            - 获取表统计信息
+        """
+        return await adbpg_basic_operation.list_resource_templates()
+
+    @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: 获取表的统计信息
+        """
+        if not db_manager:
+            raise Exception("Database connection not established")
+        return await adbpg_basic_operation.read_resource(uri, db_manager)
+
+    # 工具调用
+    @app.call_tool()
+    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
+        """
+        执行工具操作
+
+        参数:
+            name (str): 工具名称
+            arguments (dict): 工具参数
+
+        返回:
+            list[TextContent]: 执行结果
+        """
+        if not db_manager:
+            raise Exception("Database manager not initialized")
+        
+        try:
+            # 分发到 Basic Operation
+            if name in [t.name for t in adbpg_basic_operation.get_basic_tools()]:
+                query, params, needs_json_agg = await adbpg_basic_operation.call_basic_tool(name, arguments, db_manager)
+                conn = db_manager.get_basic_connection()
+                with conn.cursor() as cursor:
+                    cursor.execute(query, params)
+                    if needs_json_agg:
+                        json_result = cursor.fetchone()[0]
+                        return [TextContent(type="text", text=json.dumps(json_result, ensure_ascii=False, indent=2))]
+                    else:
+                        return [TextContent(type="text", text="Tool executed successfully.")]
+
+            # 分发到 GraphRAG
+            elif name in [t.name for t in adbpg_graphrag.get_graphrag_tools()]:
+                if not graphrag_is_available:
+                    raise ValueError("GraphRAG tool is not available due to configuration or initialization errors.")
+                return await adbpg_graphrag.call_graphrag_tool(name, arguments, db_manager)
+
+            # 分发到 LLM Memory
+            elif name in [t.name for t in adbpg_memory.get_memory_tools()]:
+                if not llm_memory_is_available:
+                    raise ValueError("LLM Memory tool is not available due to configuration or initialization errors.")
+                return await adbpg_memory.call_memory_tool(name, arguments, db_manager)
+
+            else:
+                raise ValueError(f"Unknown tool: {name}")
+
+        except Exception as e:
+            logger.error(f"Error calling tool '{name}': {e}", exc_info=True)
+            return [TextContent(type="text", text=f"Error executing tool '{name}': {str(e)}")]
+
+    # 工具列表
+    @app.list_tools()
+    async def list_tools() -> list[Tool]:
+        """
+        列出可用的工具
+        """
+        tools = adbpg_basic_operation.get_basic_tools()
+
+        if graphrag_is_available:
+            tools.extend(adbpg_graphrag.get_graphrag_tools())
+        if llm_memory_is_available:
+            tools.extend(adbpg_memory.get_memory_tools())
+        return tools
+        
+    #----------管理请求会话--------------
+    session_manager = StreamableHTTPSessionManager(
+        app=app,
+        event_store=None, #无状态，不保存历史事件
+        stateless=True
+    )
+    async def handle_streamable_http(scope, receive, send):
+        await session_manager.handle_request(scope, receive, send)
+
+    @contextlib.asynccontextmanager
+    async def lifespan(app):
+        async with session_manager.run():
+            logger.info("ADBPG MCP Server Started! ")
+            try:
+                yield
+            finally:
+                logger.info("ADBPG MCP Server shutting down…")
+    
+    # 将MCP服务挂载到/mcp路径上，用户访问整个路径时，就会进入刚才创建的MCP HTTP会话管理器
+    starlette_app = Starlette(
+        debug=False,
+        routes=[Mount("/mcp", app=handle_streamable_http)],
+        lifespan=lifespan,
+    )
+
+    # 利用uvicorn启动ASGI服务器并监听指定端口
+    uvicorn.run(starlette_app, host=host, port=port)
+
+    return 0

adbpg_mcp_server/server_stdio.py

@@ -0,0 +1,216 @@
+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")
+
+SERVER_VERSION = "0.2.0"
+
+from .adbpg_config import settings
+from .adbpg import DatabaseManager
+from . import adbpg_basic_operation, adbpg_graphrag, adbpg_memory
+
+
+# 初始化服务器
+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)
+db_manager: DatabaseManager | None = None
+graphrag_is_available = False
+llm_memory_is_available = False
+
+@app.list_resources()
+async def list_resources() -> list[Resource]:
+    """列出可用的基本资源"""
+    return await adbpg_basic_operation.list_resources()
+
+@app.list_resource_templates()
+async def list_resource_templates() -> list[ResourceTemplate]:
+    """
+    定义动态资源模板
+    
+    返回:
+        list[ResourceTemplate]: 资源模板列表
+        包含以下模板：
+        - 列出schema中的表
+        - 获取表DDL
+        - 获取表统计信息
+    """
+    return await adbpg_basic_operation.list_resource_templates()
+
+
+@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: 获取表的统计信息
+    """
+    if not db_manager:
+        raise Exception("Database connection not established")
+    return await adbpg_basic_operation.read_resource(uri, db_manager)
+
+@app.list_tools()
+async def list_tools() -> list[Tool]:
+    """
+    列出可用的工具
+    """
+    tools = adbpg_basic_operation.get_basic_tools()
+    #tools.extend(adbpg_graphrag.get_graphrag_tools())
+    #tools.extend(adbpg_memory.get_memory_tools())
+
+    if graphrag_is_available:
+        tools.extend(adbpg_graphrag.get_graphrag_tools())
+    if llm_memory_is_available:
+        tools.extend(adbpg_memory.get_memory_tools())
+    return tools
+@app.call_tool()
+async def call_tool(name: str, arguments: dict) -> list[TextContent]:
+    """
+    执行工具操作
+
+    参数:
+        name (str): 工具名称
+        arguments (dict): 工具参数
+
+    返回:
+        list[TextContent]: 执行结果
+    """
+    if not db_manager:
+        raise Exception("Database manager not initialized")
+    
+    try:
+        # 分发到 Basic Operation
+        if name in [t.name for t in adbpg_basic_operation.get_basic_tools()]:
+            query, params, needs_json_agg = await adbpg_basic_operation.call_basic_tool(name, arguments, db_manager)
+            conn = db_manager.get_basic_connection()
+            with conn.cursor() as cursor:
+                cursor.execute(query, params)
+                if needs_json_agg:
+                    json_result = cursor.fetchone()[0]
+                    return [TextContent(type="text", text=json.dumps(json_result, ensure_ascii=False, indent=2))]
+                else:
+                    return [TextContent(type="text", text="Tool executed successfully.")]
+
+        # 分发到 GraphRAG
+        elif name in [t.name for t in adbpg_graphrag.get_graphrag_tools()]:
+            if not graphrag_is_available:
+                raise ValueError("GraphRAG tool is not available due to configuration or initialization errors.")
+            return await adbpg_graphrag.call_graphrag_tool(name, arguments, db_manager)
+
+        # 分发到 LLM Memory
+        elif name in [t.name for t in adbpg_memory.get_memory_tools()]:
+            if not llm_memory_is_available:
+                raise ValueError("LLM Memory tool is not available due to configuration or initialization errors.")
+            return await adbpg_memory.call_memory_tool(name, arguments, db_manager)
+
+        else:
+            raise ValueError(f"Unknown tool: {name}")
+
+    except Exception as e:
+        logger.error(f"Error calling tool '{name}': {e}", exc_info=True)
+        return [TextContent(type="text", text=f"Error executing tool '{name}': {str(e)}")]
+
+# --- 服务器生命周期 ---
+def initialize_services():
+    """初始化数据库、GraphRAG 和 LLM Memory"""
+    global db_manager, graphrag_is_available, llm_memory_is_available
+
+    if not settings.db_env_ready:
+        logger.error("Cannot start server: Database environment is not configured.")
+        sys.exit(1)
+
+    db_manager = DatabaseManager(settings)
+
+    # 测试主连接
+    try:
+        db_manager.get_basic_connection()
+        logger.info("Successfully connected to database.")
+    except Exception as e:
+        logger.error(f"Failed to connect to database: {e}")
+        sys.exit(1)
+
+    # 初始化 GraphRAG
+    if settings.graphrag_env_ready:
+        try:
+            db_manager.get_graphrag_connection()
+            graphrag_is_available = True
+            logger.info("GraphRAG initialized successfully.")
+        except Exception as e:
+            logger.error(f"Failed to initialize GraphRAG: {e}")
+            graphrag_is_available = False
+            
+    # 初始化 LLM Memory
+    if settings.memory_env_ready:
+        try:
+            db_manager.get_llm_memory_connection()
+            llm_memory_is_available = True
+            logger.info("LLM Memory initialized successfully.")
+        except Exception as e:
+            logger.error(f"Failed to initialize LLM Memory: {e}")
+            llm_memory_is_available = False
+
+async def main():
+    """服务器主入口点"""
+    try:
+        initialize_services()
+        logger.info("Starting ADBPG MCP server...")
+        
+        # 使用 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
+
+    finally:
+        if db_manager:
+            db_manager.close_all()
+
+def run_stdio_server():
+    """同步运行入口点"""
+    try:
+        asyncio.run(main())
+    except Exception as e:
+        logger.error(f"Fatal error: {e}")
+        sys.exit(1)