agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/mcp_convert/mcp_convert.py

@@ -0,0 +1,375 @@
+import os
+import inspect
+import logging
+import json
+from typing import Dict, Any, Optional, List, Callable, Awaitable
+from functools import wraps, partial
+from datetime import datetime
+from langfuse import propagate_attributes
+from langfuse.langchain import CallbackHandler
+from agent_api_server.shared.decode_token import decode_jwt
+from langgraph.config import get_stream_writer
+from langgraph.graph.state import CompiledStateGraph
+from fastmcp import FastMCP, Context
+from mcp import types
+from mcp.types import RequestParams, CallToolResult, TextContent
+from agent_api_server.shared.message import handle_stream_event
+from agent_api_server.shared.util_func import load_graph_config, load_graph, get_env
+
+logger = logging.getLogger(__name__)
+
+
+async def handle_messages(message):
+    writer = get_stream_writer()
+    if isinstance(message, types.ServerNotification):
+        notification = message.root
+        if isinstance(notification, types.ProgressNotification):
+            params = notification.params
+            logger.info(f"receive progress message: ({params.message})")
+            writer(params.message)
+
+async def create_mcp_tool_from_agent() -> FastMCP:
+    """Create and configure MCP tools from LangGraph agent configurations.
+
+    Returns:
+        FastMCP: Configured MCP instance with all graph tools registered.
+
+    Raises:
+        RuntimeError: If tool creation fails due to configuration or loading errors.
+    """
+    mcp = FastMCP(name="langgraph_mcp_tool")
+
+    try:
+        graph_cfg = await load_graph_config()
+
+        # Validate configurations
+        graphs = graph_cfg.get("graphs", {})
+        if not isinstance(graphs, dict):
+            raise RuntimeError(f"Invalid graphs config. Expected dict, got {type(graphs)}")
+
+        agent_descriptions = graph_cfg.get("agent_description", {})
+        if not isinstance(agent_descriptions, dict):
+            logger.warning("agent_description should be dict, got %s", type(agent_descriptions))
+            agent_descriptions = {}
+
+        for graph_name, graph_path in graphs.items():
+            _, graph_instance, _ = await load_graph(graph_name, graph_cfg, False)
+
+            if not isinstance(graph_instance, CompiledStateGraph):
+                raise RuntimeError(f"Graph {graph_name} is not CompiledStateGraph")
+
+            tool_impl = partial(
+                _execute_graph_tool,
+                mcp=mcp,
+                graph_name=graph_name,
+                graph_instance=graph_instance
+            )
+            tool_impl.__name__ = graph_name
+
+            tool = create_tool_from_schema(
+                schema=graph_instance.get_input_jsonschema(),
+                func_name=graph_name,
+                func_doc=agent_descriptions.get(graph_name, ""),
+                implementation=tool_impl
+            )
+
+            mcp.tool()(_add_metrics_wrapper(tool, graph_name))
+            logger.info("Registered tool: %s", graph_name)
+
+        return mcp
+    except Exception as e:
+        logger.error("MCP tool creation failed", exc_info=True)
+        raise RuntimeError(f"MCP tool creation failed: {str(e)}")
+
+
+async def _execute_graph_tool(
+        mcp: FastMCP,
+        graph_name: str,
+        graph_instance: CompiledStateGraph,
+        **kwargs: Any
+) -> Dict[str, Any]:
+    """Execute a LangGraph agent and stream results.
+
+    Args:
+        mcp: FastMCP instance
+        graph_name: Name of the graph to execute
+        graph_instance: Compiled graph instance
+        **kwargs: Input parameters
+
+    Returns:
+        Execution result dict
+    """
+    ctx = Context(fastmcp=mcp)
+    from fastmcp.server.dependencies import get_http_request
+    request = get_http_request()
+
+    use_sys_llm = request.headers.get("UseSysLLM", "")
+    ts_tenant = request.headers.get("TSTenant", "")
+    ei_token = request.headers.get("Authorization", "")
+    thread_id = request.headers.get('thread_id', '')
+    start_time = datetime.now()
+
+    logger.info(f"Executing graph '{graph_name}' for tenant '{ts_tenant}', ei_token is '{ei_token}', use_sys_llm is {use_sys_llm}")
+
+    # Validate input
+    schema = graph_instance.get_input_jsonschema()
+    input_dict = {}
+    validation_errors = []
+
+    for field, prop in schema.get("properties", {}).items():
+        if field in kwargs:
+            try:
+                input_dict[field] = _validate_field(field, kwargs[field], prop)
+            except ValueError as e:
+                validation_errors.append(str(e))
+        elif field in schema.get("required", []):
+            validation_errors.append(f"Missing required field: {field}")
+
+    if validation_errors:
+        error_msg = "Validation errors:\n" + "\n".join(validation_errors)
+        logger.error(error_msg)
+        return CallToolResult(
+            content=[TextContent(type="text", text=error_msg)],
+            isError=True
+        )
+
+    if thread_id is not None and thread_id != "":
+        logger.info(f"execute graph {graph_name} with thread id {thread_id} which get from context headers")
+    else:
+        thread_id = ctx.session_id
+        logger.info(f"execute graph {graph_name} with session id {ctx.session_id}")
+
+    user_id = None
+    if ei_token:
+        try:
+            header, payload, user_info = decode_jwt(ei_token)
+            if user_info and 'id' in user_info:
+                user_id = user_info['id']
+                logger.info(f"Extracted user_id from token: {user_id}")
+            else:
+                logger.warning("No user_id found in JWT token")
+        except Exception as e:
+            logger.warning(f"Failed to decode JWT token: {str(e)}")
+
+    # Build configurable parameters
+    configurable_params = {
+        "use_sys_llm": use_sys_llm,
+        "thread_id": thread_id,
+        "TSTenant": ts_tenant,
+        "EIToken": ei_token,
+        "graph_name": graph_name
+    }
+
+    if ts_tenant is not None and ts_tenant != "":
+        configurable_params = {
+            **dict(get_env(ts_tenant=ts_tenant)),
+            **configurable_params
+        }
+
+    logger.info(f"finally configurable parameter is {configurable_params}")
+
+    # Configure Langfuse callbacks if available
+    langfuse_keys = [
+        os.getenv("LANGFUSE_SECRET_KEY"),
+        os.getenv("LANGFUSE_PUBLIC_KEY"),
+        os.getenv("LANGFUSE_BASE_URL")
+    ]
+
+    callbacks_config = {}
+    if all(langfuse_keys):
+        langfuse_handler = CallbackHandler()
+        callbacks_config = {"callbacks": [langfuse_handler], "run_name": f"{graph_name}_MCP_Call"}
+
+    chunks = []
+    config = {"configurable": configurable_params, **callbacks_config}
+    try:
+        async def process_stream():
+            if user_id:
+                logger.info(f"execute graph {graph_name} with user_id: {user_id}")
+                with propagate_attributes(session_id=thread_id, user_id=user_id, trace_name=f"{graph_name}_MCP_Call"):
+                    async for stream_event in graph_instance.astream(
+                            input_dict,
+                            config=config,
+                            stream_mode=["updates"],
+                            subgraphs=True
+                    ):
+                        async for chunk in handle_stream_event(stream_event):
+                            yield chunk
+
+            else:
+                logger.info(f"execute graph {graph_name} without user_id")
+                with propagate_attributes(session_id=thread_id, trace_name=f"{graph_name}_MCP_Call"):
+                    async for stream_event in graph_instance.astream(
+                            input_dict,
+                            config=config,
+                            stream_mode=["updates"],
+                            subgraphs=True
+                    ):
+                        async for chunk in handle_stream_event(stream_event):
+                            yield chunk
+
+
+        async for chunk in process_stream():
+            ctx.request_context.meta = RequestParams.Meta(progressToken=ctx.request_id)
+            await ctx.report_progress(message=chunk, progress=len(chunks))
+            chunks.append(chunk)
+
+        if not chunks:
+            raise ValueError("No response from graph execution")
+
+        last_chunk = chunks[-1].strip()
+        if not last_chunk.startswith("data: "):
+            raise ValueError("Invalid response format")
+
+        data = json.loads(last_chunk[6:])
+        if not data.get("update_content", {}).get("content"):
+            raise ValueError("Missing content in response")
+
+        logger.info(
+            "Graph '%s' executed in %.2fs (%d chunks)",
+            graph_name,
+            (datetime.now() - start_time).total_seconds(),
+            len(chunks)
+        )
+
+        return CallToolResult(
+            content=[TextContent(type="text", text=data["update_content"]["content"])],
+            isError=False
+        )
+
+    except Exception as err:
+        logger.error("Graph execution failed: %s", str(err), exc_info=True)
+        return CallToolResult(
+            content=[TextContent(type="text", text=f"Execution failed: {str(err)}")],
+            isError=True
+        )
+
+
+def create_tool_from_schema(
+        schema: Dict[str, Any],
+        func_name: str,
+        func_doc: str = "",
+        implementation: Optional[Callable[..., Awaitable[Dict[str, Any]]]] = None
+) -> Callable[..., Awaitable[Dict[str, Any]]]:
+    properties = schema.get("properties", {})
+    required = schema.get("required", [])
+    type_map = {
+        "string": str,
+        "number": float,
+        "integer": int,
+        "boolean": bool,
+        "array": List,
+        "object": Dict
+    }
+
+    params = []
+    type_hints = {}
+    for param_name, param_info in properties.items():
+        param_type = type_map.get(param_info.get("type"), Any)
+
+        if param_name not in required:
+            if 'anyOf' in param_info:
+                first_non_null_type = None
+                for item in param_info['anyOf']:
+                    item_type = item.get('type')
+                    if item_type and item_type != 'null':
+                        first_non_null_type = item_type
+                        break
+
+                if first_non_null_type:
+                    non_null_param_type = type_map.get(first_non_null_type, Any)
+                    param_type = non_null_param_type
+
+        type_hints[param_name] = param_type
+        params.append(
+            inspect.Parameter(
+                param_name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                default=None if param_name not in required else inspect.Parameter.empty,
+                annotation=param_type,
+            )
+        )
+
+    if implementation is None:
+        async def default_impl(**kwargs: Any) -> Dict[str, Any]:
+            missing = [f for f in required if f not in kwargs]
+            if missing:
+                raise ValueError(f"Missing fields: {missing}")
+            return {"content": [f"Default response for {func_name}"]}
+
+        implementation = default_impl
+
+    implementation.__name__ = func_name
+    implementation.__qualname__ = func_name
+    implementation.__doc__ = func_doc or f"Tool from schema: {schema.get('title', '')}"
+    implementation.__signature__ = inspect.Signature(
+        parameters=params,
+        return_annotation=Dict[str, Any],
+    )
+    implementation.__annotations__ = type_hints
+
+    return implementation
+
+
+def _add_metrics_wrapper(
+        func: Callable[..., Awaitable[Dict[str, Any]]],
+        tool_name: str
+) -> Callable[..., Awaitable[Dict[str, Any]]]:
+
+    @wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Dict[str, Any]:
+        start_time = datetime.now()
+        logger.info("Tool '%s' started", tool_name)
+
+        try:
+            result = await func(*args, **kwargs)
+            logger.info(
+                "Tool '%s' completed in %.2fs",
+                tool_name,
+                (datetime.now() - start_time).total_seconds()
+            )
+            return result
+        except Exception as e:
+            logger.error(
+                "Tool '%s' failed after %.2fs: %s",
+                tool_name,
+                (datetime.now() - start_time).total_seconds(),
+                str(e),
+                exc_info=True
+            )
+            raise
+
+    return wrapper
+
+
+def _validate_field(
+        field_name: str,
+        value: Any,
+        schema: Dict[str, Any]
+) -> Any:
+    field_type = schema.get("type")
+
+    try:
+        if field_type == "string":
+            if not isinstance(value, str):
+                value = str(value)
+            if "enum" in schema and value not in schema["enum"]:
+                raise ValueError(f"Value not in {schema['enum']}")
+        elif field_type == "number":
+            value = float(value)
+        elif field_type == "integer":
+            value = int(value)
+        elif field_type == "boolean":
+            if isinstance(value, str):
+                value = value.lower() in ("true", "1", "yes")
+            value = bool(value)
+        elif field_type == "array":
+            if not isinstance(value, list):
+                raise ValueError("Expected list")
+        elif field_type == "object":
+            if not isinstance(value, dict):
+                raise ValueError("Expected dict")
+    except (ValueError, TypeError) as e:
+        raise ValueError(f"Invalid value for '{field_name}': {str(e)}")
+
+    return value

agent_api_server/memeory/postgres.py

@@ -0,0 +1,233 @@
+import os
+import asyncio
+import logging
+import weakref
+from typing import Optional, cast, Dict, Any
+from psycopg_pool import AsyncConnectionPool
+from psycopg.rows import dict_row
+from psycopg import OperationalError, AsyncConnection
+from agent_api_server.configs import global_config
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type,
+)
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.store.postgres import AsyncPostgresStore
+
+logger = logging.getLogger(__name__)
+
+class AsyncPostgresCheckpointer:
+    _instances = weakref.WeakSet()
+    _shutdown_lock = asyncio.Lock()
+    _worker_instances = weakref.WeakValueDictionary()
+
+    def __init__(self, max_retries: int = 3, retry_delay: float = 1.0):
+        self.worker_pid = os.getpid()
+        logger.info(
+            f"Initializing checkpointer for worker {self.worker_pid} "
+            f"(max_retries={max_retries}, retry_delay={retry_delay:.1f}s)"
+        )
+
+        self.conn_str = self._get_conn_str()
+        self.pool: Optional[AsyncConnectionPool] = None
+        self.saver: Optional[AsyncPostgresSaver] = None
+        self.max_retries = max(max_retries, 1)
+        self.retry_delay = max(retry_delay, 0.1)
+        self._lock = asyncio.Lock()
+        self._is_initialized = False
+
+        # Connection pool configuration with validation
+        self.pool_min_size = max(1, int(global_config.POSTGRES_POOL_MIN_SIZE))
+        self.pool_max_size = max(self.pool_min_size, int(global_config.POSTGRES_POOL_MAX_SIZE))
+        self.pool_timeout = max(1.0, float(global_config.POSTGRES_POOL_TIMEOUT))
+        self.pool_recycle = max(60, int(global_config.POSTGRES_POOL_RECYCLE))
+
+        logger.debug(
+            f"Worker {self.worker_pid} pool config: "
+            f"min_size={self.pool_min_size}, max_size={self.pool_max_size}, "
+            f"timeout={self.pool_timeout:.1f}, recycle={self.pool_recycle}"
+        )
+
+        AsyncPostgresCheckpointer._instances.add(self)
+        AsyncPostgresCheckpointer._worker_instances[self.worker_pid] = self
+        logger.debug(f"New checkpointer instance registered for worker {self.worker_pid}")
+
+    @classmethod
+    def get_worker_instance(cls) -> 'AsyncPostgresCheckpointer':
+        """Get or create instance for current worker process"""
+        worker_pid = os.getpid()
+        if worker_pid not in cls._worker_instances:
+            instance = cls()
+            cls._worker_instances[worker_pid] = instance
+            logger.debug(f"Created new checkpointer instance for worker {worker_pid}")
+        return cls._worker_instances[worker_pid]
+
+    @classmethod
+    async def close_worker_instance(cls):
+        """Close instance for current worker process"""
+        worker_pid = os.getpid()
+        if worker_pid in cls._worker_instances:
+            instance = cls._worker_instances.pop(worker_pid)
+            await instance.close()
+
+    @staticmethod
+    def _get_conn_str() -> str:
+        """获取安全的连接字符串（隐藏密码）"""
+        conn_str = global_config.POSTGRES_URL
+        safe_conn_str = conn_str.split('@')[0] + '@[REDACTED]' if '@' in conn_str else conn_str
+        logger.debug("Using connection string: %s", safe_conn_str)
+        return conn_str
+
+    @classmethod
+    async def close_all(cls):
+        logger.info("Initiating shutdown of all checkpointer instances")
+
+        async with cls._shutdown_lock:
+            instances = list(cls._instances)
+            if not instances:
+                logger.debug("No active instances to close")
+                return
+
+            logger.debug("Closing %d active instance(s)", len(instances))
+            for instance in instances:
+                try:
+                    await instance.close()
+                    logger.debug("Instance closed successfully")
+                except Exception as e:
+                    logger.warning(
+                        "Error closing instance: %s",
+                        str(e),
+                        exc_info=logger.isEnabledFor(logging.DEBUG)
+                    )
+
+    async def initialize(self) -> None:
+        """Initialize connection pool and database schema"""
+        async with self._lock:
+            if self._is_initialized:
+                logger.debug("Already initialized, skipping")
+                return
+
+            logger.info("Starting connection pool initialization")
+
+            try:
+                self.pool = AsyncConnectionPool(
+                    self.conn_str,
+                    min_size=self.pool_min_size,
+                    max_size=self.pool_max_size,
+                    timeout=self.pool_timeout,
+                    open=False,
+                    kwargs={"row_factory": dict_row},
+                )
+                logger.debug("Connection pool created")
+
+                await self.pool.open()
+                logger.debug("Connection pool opened")
+
+                pool_with_dict_conn = cast(AsyncConnectionPool[AsyncConnection[Dict[str, Any]]], self.pool)
+                self.saver = AsyncPostgresSaver(pool_with_dict_conn)
+                self._is_initialized = True
+                logger.info("Checkpointer initialized successfully")
+
+            except Exception as e:
+                logger.error("Initialization failed", exc_info=True)
+                await self._safe_close()
+                raise RuntimeError("Checkpointer initialization failed") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type(OperationalError),
+        before_sleep=lambda _: logger.warning("Retrying schema initialization..."),
+    )
+    async def initialize_schema_with_retry(self):
+        logger.debug("Starting schema initialization")
+
+        async with self.pool.connection() as conn:
+            try:
+                await conn.set_autocommit(True)
+                logger.debug("Autocommit enabled for schema setup")
+
+                await AsyncPostgresStore(conn).setup()
+                await AsyncPostgresSaver(conn).setup()
+                logger.debug("Schema setup completed")
+
+            finally:
+                await conn.set_autocommit(False)
+                logger.debug("Autocommit restored to default")
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type(OperationalError),
+    )
+    async def checkpointer(self) -> AsyncPostgresSaver | None:
+        if not self._is_initialized or self.saver is None:
+            logger.debug("Checkpointer not ready, initializing...")
+            await self.initialize()
+            return self.saver
+
+        try:
+            async with self.pool.connection() as conn:
+                await conn.execute("SELECT 1")  # 简单心跳检测
+        except OperationalError:
+            logger.warning("Connection pool invalid, reinitializing...")
+            await self._safe_close()
+            await self.initialize()
+
+        return self.saver
+
+    async def _safe_close(self) -> None:
+        logger.debug("Starting safe cleanup")
+
+        if self.pool:
+            try:
+                if not self.pool.closed:
+                    await self.pool.close()
+            except Exception as e:
+                logger.warning(f"Error closing pool with an error {str(e)}", exc_info=True)
+
+        self.pool = None
+        self.saver = None
+        self._is_initialized = False
+        logger.info("Resources cleaned up")
+
+    async def close(self) -> None:
+        """关闭所有连接并清理资源"""
+        async with self._lock:
+            if not self._is_initialized:
+                logger.debug("Already closed, skipping")
+                return
+
+            logger.info("Starting graceful shutdown")
+
+            try:
+                if self.pool and not self.pool.closed:
+                    logger.debug("Closing connection pool gracefully")
+                    await self.pool.close()
+                    logger.info("Connection pool closed gracefully")
+
+            except Exception as e:
+                logger.error(
+                    "Error during shutdown: %s",
+                    str(e),
+                    exc_info=logger.isEnabledFor(logging.DEBUG)
+                )
+            finally:
+                await self._safe_close()
+                AsyncPostgresCheckpointer._instances.discard(self)
+                logger.debug("Instance unregistered")
+
+    async def __aenter__(self) -> AsyncPostgresSaver:
+        logger.debug("Entering context manager")
+        await self.initialize()
+        return await self.checkpointer()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        logger.debug("Exiting context manager")
+        await self.close()
+
+    @property
+    def is_initialized(self):
+        return self._is_initialized

agent_api_server/register/register.py

@@ -0,0 +1,65 @@
+import logging
+from typing import List, Dict
+from agent_api_server.shared.common import process_model_from_config_dict
+from agent_api_server.shared.util_func import load_graph_config, load_graph
+from model_manage_client import ModelManageClient
+
+logger = logging.getLogger(__name__)
+
+class AgentRegistry:
+    def __init__(self, base_url:str, client_token:str):
+        self.client = ModelManageClient(
+            base_url=base_url,
+            client_token=client_token
+        )
+
+    @staticmethod
+    async def get_llm_model_from_agent_cfg(agent: Dict) -> Dict:
+        graph_cfg = await load_graph_config()
+        _, graph_instance, _ = await load_graph(agent['agent_name'], graph_cfg, False)
+
+        cfg = graph_instance.get_context_jsonschema()
+
+        logger.info(f"Get agent config {cfg} from {agent['agent_name']}")
+
+        return process_model_from_config_dict(cfg)
+
+    async def register_all(self, agents: List[Dict]):
+        for agent in agents:
+            try:
+                model_info = await self.get_llm_model_from_agent_cfg(agent)
+
+                extra_params = {
+                    "agent_description": agent["agent_description"],
+                    "agent_icon_url": agent["agent_icon_url"],
+                    "agent_api_version": agent["agent_api_version"],
+                    "agent_features": agent["agent_features"],
+                    "agent_labels": agent["agent_labels"],
+                    "support_models": model_info,
+                    "has_site": agent["has_site"],
+                    "is_system_agent": agent["is_system_agent"],
+                    "multilangs": agent.get("multilangs", {})
+                }
+
+                agent_find = self.client.get_agent(agent_name=agent["agent_name"])
+                if agent_find:
+                    logger.info(f"agent: {agent['agent_name']} already register successfully, update agent with extra_params {extra_params}")
+
+                    self.client.update_agent(
+                        agent_name=agent["agent_name"],
+                        agent_url=agent["agent_url"],
+                        **extra_params
+                    )
+                    continue
+
+                self.client.register_agent(
+                    agent_name=agent["agent_name"],
+                    agent_id=agent["agent_name"],
+                    agent_url=agent["agent_url"],
+                    **extra_params
+                )
+
+                logger.info(
+                    f"agent: {agent['agent_name']} register successfully, register agent with extra_params {extra_params}")
+            except Exception as e:
+                raise e