Просмотр исходного кода

feat: docker环境操作工具(办公室4090机器)

tanjingyu 1 месяц назад
Родитель
Сommit
1271cf20c4
2 измененных файлов с 517 добавлено и 46 удалено
  1. 433 0
      agent/tools/builtin/sandbox.py
  2. 84 46
      agent/tools/builtin/search.py

+ 433 - 0
agent/tools/builtin/sandbox.py

@@ -0,0 +1,433 @@
+"""
+Sandbox Tools (Async)
+通过 HTTP 异步调用沙盒管理服务的客户端库。
+"""
+
+import json
+import httpx
+from typing import Optional, List, Dict, Any
+
+from agent import tool
+from agent.tools.models import ToolResult
+
+
+# 服务地址,可根据实际部署情况修改
+SANDBOX_SERVER_URL = "http://61.48.133.26:9999"
+
+# 默认超时时间(秒)
+DEFAULT_TIMEOUT = 300.0
+
+
+@tool(
+    display={
+        "zh": {
+            "name": "创建沙盒环境",
+            "params": {
+                "image": "Docker 镜像",
+                "mem_limit": "内存限制",
+                "nano_cpus": "CPU 限制",
+                "ports": "端口列表"
+            }
+        },
+        "en": {
+            "name": "Create Sandbox",
+            "params": {
+                "image": "Docker image",
+                "mem_limit": "Memory limit",
+                "nano_cpus": "CPU limit",
+                "ports": "Port list"
+            }
+        }
+    }
+)
+async def sandbox_create_environment(
+    image: str = "agent-sandbox:latest",
+    mem_limit: str = "512m",
+    nano_cpus: int = 500000000,
+    ports: Optional[List[int]] = None,
+    uid: str = ""
+) -> ToolResult:
+    """
+    创建一个隔离的 Docker 开发环境。
+
+    Args:
+        image: Docker 镜像名称,默认为 "agent-sandbox:latest"。
+               可以使用其他镜像如 "python:3.12-slim", "node:18-slim" 等。
+        mem_limit: 容器最大内存限制,默认为 "512m"。
+        nano_cpus: 容器最大 CPU 限制(纳秒),默认为 500000000(0.5 CPU)。
+        ports: 需要映射的端口列表,如 [8080, 3000]。
+        uid: 用户ID(自动注入)
+
+    Returns:
+        ToolResult 包含:
+            - sandbox_id (str): 沙盒唯一标识,后续操作需要用到
+            - message (str): 提示信息
+            - port_mapping (dict): 端口映射关系,如 {8080: 32001}
+            - access_urls (list): 访问 URL 列表
+    """
+    try:
+        url = f"{SANDBOX_SERVER_URL}/api/create_environment"
+        payload = {
+            "image": image,
+            "mem_limit": mem_limit,
+            "nano_cpus": nano_cpus
+        }
+        if ports:
+            payload["ports"] = ports
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        return ToolResult(
+            title="Sandbox Created",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Created sandbox: {data.get('sandbox_id', 'unknown')}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="Create Sandbox Failed",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="Create Sandbox Failed",
+            output="",
+            error=str(e)
+        )
+
+
+@tool(
+    display={
+        "zh": {
+            "name": "执行沙盒命令",
+            "params": {
+                "sandbox_id": "沙盒 ID",
+                "command": "Shell 命令",
+                "is_background": "后台执行",
+                "timeout": "超时时间"
+            }
+        },
+        "en": {
+            "name": "Run Shell in Sandbox",
+            "params": {
+                "sandbox_id": "Sandbox ID",
+                "command": "Shell command",
+                "is_background": "Run in background",
+                "timeout": "Timeout"
+            }
+        }
+    }
+)
+async def sandbox_run_shell(
+    sandbox_id: str,
+    command: str,
+    is_background: bool = False,
+    timeout: int = 120,
+    uid: str = ""
+) -> ToolResult:
+    """
+    在指定的沙盒中执行 Shell 命令。
+
+    Args:
+        sandbox_id: 沙盒 ID,由 create_environment 返回。
+        command: 要执行的 Shell 命令,如 "pip install flask" 或 "python app.py"。
+        is_background: 是否后台执行,默认为 False。
+            - False:前台执行,等待命令完成并返回输出
+            - True:后台执行,适合启动长期运行的服务
+        timeout: 前台命令的超时时间(秒),默认 120 秒。后台命令不受此限制。
+        uid: 用户ID(自动注入)
+
+    Returns:
+        ToolResult 包含:
+            前台执行:
+                - exit_code (int): 命令退出码
+                - stdout (str): 标准输出
+                - stderr (str): 标准错误
+            后台执行:
+                - status (str): 状态
+                - message (str): 提示信息
+                - log_file (str): 日志文件路径
+    """
+    try:
+        url = f"{SANDBOX_SERVER_URL}/api/run_shell"
+        payload = {
+            "sandbox_id": sandbox_id,
+            "command": command,
+            "is_background": is_background,
+            "timeout": timeout
+        }
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        return ToolResult(
+            title=f"Shell: {command[:50]}{'...' if len(command) > 50 else ''}",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Executed in sandbox {sandbox_id}: {command[:100]}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="Run Shell Failed",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="Run Shell Failed",
+            output="",
+            error=str(e)
+        )
+
+
+@tool(
+    display={
+        "zh": {
+            "name": "重建沙盒端口",
+            "params": {
+                "sandbox_id": "沙盒 ID",
+                "ports": "端口列表",
+                "mem_limit": "内存限制",
+                "nano_cpus": "CPU 限制"
+            }
+        },
+        "en": {
+            "name": "Rebuild Sandbox Ports",
+            "params": {
+                "sandbox_id": "Sandbox ID",
+                "ports": "Port list",
+                "mem_limit": "Memory limit",
+                "nano_cpus": "CPU limit"
+            }
+        }
+    }
+)
+async def sandbox_rebuild_with_ports(
+    sandbox_id: str,
+    ports: List[int],
+    mem_limit: str = "1g",
+    nano_cpus: int = 1000000000,
+    uid: str = ""
+) -> ToolResult:
+    """
+    重建沙盒并应用新的端口映射。
+
+    使用场景:先创建沙盒克隆项目,阅读 README 后才知道需要暴露哪些端口,
+    此时调用此函数重建沙盒,应用正确的端口映射。
+
+    注意:重建后会返回新的 sandbox_id,后续操作需要使用新 ID。
+    容器内的所有文件(克隆的代码、安装的依赖等)都会保留。
+
+    Args:
+        sandbox_id: 当前沙盒 ID。
+        ports: 需要映射的端口列表,如 [8080, 3306, 6379]。
+        mem_limit: 容器最大内存限制,默认为 "1g"。
+        nano_cpus: 容器最大 CPU 限制(纳秒),默认为 1000000000(1 CPU)。
+        uid: 用户ID(自动注入)
+
+    Returns:
+        ToolResult 包含:
+            - old_sandbox_id (str): 旧沙盒 ID(已销毁)
+            - new_sandbox_id (str): 新沙盒 ID(后续使用这个)
+            - port_mapping (dict): 端口映射关系
+            - access_urls (list): 访问 URL 列表
+    """
+    try:
+        url = f"{SANDBOX_SERVER_URL}/api/rebuild_with_ports"
+        payload = {
+            "sandbox_id": sandbox_id,
+            "ports": ports,
+            "mem_limit": mem_limit,
+            "nano_cpus": nano_cpus
+        }
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        return ToolResult(
+            title="Sandbox Rebuilt",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Rebuilt sandbox {sandbox_id} -> {data.get('new_sandbox_id', 'unknown')} with ports {ports}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="Rebuild Sandbox Failed",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="Rebuild Sandbox Failed",
+            output="",
+            error=str(e)
+        )
+
+
+@tool(
+    requires_confirmation=True,
+    display={
+        "zh": {
+            "name": "销毁沙盒环境",
+            "params": {
+                "sandbox_id": "沙盒 ID"
+            }
+        },
+        "en": {
+            "name": "Destroy Sandbox",
+            "params": {
+                "sandbox_id": "Sandbox ID"
+            }
+        }
+    }
+)
+async def sandbox_destroy_environment(
+    sandbox_id: str,
+    uid: str = ""
+) -> ToolResult:
+    """
+    销毁沙盒环境,释放资源。
+
+    Args:
+        sandbox_id: 沙盒 ID。
+        uid: 用户ID(自动注入)
+
+    Returns:
+        ToolResult 包含:
+            - status (str): 操作状态,如 "success"
+            - message (str): 提示信息
+            - removed_tools (list): 被移除的工具列表(如有关联的已注册工具)
+    """
+    try:
+        url = f"{SANDBOX_SERVER_URL}/api/destroy_environment"
+        payload = {
+            "sandbox_id": sandbox_id
+        }
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        return ToolResult(
+            title="Sandbox Destroyed",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Destroyed sandbox: {sandbox_id}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="Destroy Sandbox Failed",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="Destroy Sandbox Failed",
+            output="",
+            error=str(e)
+        )
+
+
+@tool(
+    display={
+        "zh": {
+            "name": "注册沙盒工具",
+            "params": {
+                "tool_name": "工具名称",
+                "description": "工具描述",
+                "input_schema": "参数定义",
+                "sandbox_id": "沙盒 ID",
+                "internal_port": "内部端口",
+                "endpoint_path": "API 路径",
+                "http_method": "HTTP 方法"
+            }
+        },
+        "en": {
+            "name": "Register Sandbox Tool",
+            "params": {
+                "tool_name": "Tool name",
+                "description": "Tool description",
+                "input_schema": "Input schema",
+                "sandbox_id": "Sandbox ID",
+                "internal_port": "Internal port",
+                "endpoint_path": "API path",
+                "http_method": "HTTP method"
+            }
+        }
+    }
+)
+async def sandbox_register_tool(
+    tool_name: str,
+    description: str,
+    input_schema: Dict[str, Any],
+    sandbox_id: str,
+    internal_port: int,
+    endpoint_path: str = "/",
+    http_method: str = "POST",
+    metadata: Optional[Dict[str, Any]] = None,
+    uid: str = ""
+) -> ToolResult:
+    """
+    将部署好的服务注册为工具。
+
+    注册后,该工具会出现在统一 MCP Server 的工具列表中,可被上游服务调用。
+
+    Args:
+        tool_name: 工具唯一标识(字母开头,只能包含字母、数字、下划线),
+                   如 "rendercv_api"。
+        description: 工具描述,描述该工具的功能。
+        input_schema: JSON Schema 格式的参数定义,定义工具接收的参数。
+        sandbox_id: 服务所在的沙盒 ID。
+        internal_port: 服务在容器内的端口。
+        endpoint_path: API 路径,默认 "/"。
+        http_method: HTTP 方法,默认 "POST"。
+        metadata: 额外元数据(可选)。
+        uid: 用户ID(自动注入)
+
+    Returns:
+        ToolResult 包含:
+            - status (str): 操作状态,"success" 或 "error"
+            - message (str): 提示信息
+            - tool_info (dict): 工具信息(成功时)
+    """
+    try:
+        url = f"{SANDBOX_SERVER_URL}/api/register_tool"
+        payload = {
+            "tool_name": tool_name,
+            "description": description,
+            "input_schema": input_schema,
+            "sandbox_id": sandbox_id,
+            "internal_port": internal_port,
+            "endpoint_path": endpoint_path,
+            "http_method": http_method
+        }
+        if metadata:
+            payload["metadata"] = metadata
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            data = response.json()
+
+        return ToolResult(
+            title=f"Tool Registered: {tool_name}",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Registered tool '{tool_name}' on sandbox {sandbox_id}:{internal_port}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="Register Tool Failed",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="Register Tool Failed",
+            output="",
+            error=str(e)
+        )

+ 84 - 46
agent/tools/builtin/search.py

@@ -8,12 +8,14 @@
 2. search_suggestions - 建议词搜索
 """
 
+import json
 from enum import Enum
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 import httpx
 
 from agent import tool
+from agent.tools.models import ToolResult
 
 
 # API 基础配置
@@ -77,7 +79,7 @@ async def search_posts(
     cursor: str = "0",
     max_count: int = 5,
     uid: str = "",
-) -> Dict[str, Any]:
+) -> ToolResult:
     """
     帖子搜索
 
@@ -100,7 +102,7 @@ async def search_posts(
         uid: 用户ID(自动注入)
 
     Returns:
-        API 返回的原始响应,结构如下
+        ToolResult 包含搜索结果
         {
             "code": 0,                    # 状态码,0 表示成功
             "message": "success",         # 状态消息
@@ -119,31 +121,48 @@ async def search_posts(
                 }
             ]
         }
-
-    Raises:
-        httpx.HTTPStatusError: HTTP 请求返回非 2xx 状态码
-        httpx.TimeoutException: 请求超时
-        httpx.RequestError: 其他请求错误
     """
-    # 处理 channel 参数,支持枚举和字符串
-    channel_value = channel.value if isinstance(channel, PostSearchChannel) else channel
-
-    url = f"{BASE_URL}/data"
-    payload = {
-        "type": channel_value,
-        "keyword": keyword,
-        "cursor": cursor,
-        "max_count": max_count,
-    }
+    try:
+        # 处理 channel 参数,支持枚举和字符串
+        channel_value = channel.value if isinstance(channel, PostSearchChannel) else channel
+
+        url = f"{BASE_URL}/data"
+        payload = {
+            "type": channel_value,
+            "keyword": keyword,
+            "cursor": cursor,
+            "max_count": max_count,
+        }
 
-    async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
-        response = await client.post(
-            url,
-            json=payload,
-            headers={"Content-Type": "application/json"},
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(
+                url,
+                json=payload,
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        # 计算结果数量
+        result_count = len(data.get("data", []))
+
+        return ToolResult(
+            title=f"搜索结果: {keyword} ({channel_value})",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Searched '{keyword}' on {channel_value}, found {result_count} posts"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="搜索失败",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="搜索失败",
+            output="",
+            error=str(e)
         )
-        response.raise_for_status()
-        return response.json()
 
 
 @tool(
@@ -168,7 +187,7 @@ async def search_suggestions(
     keyword: str,
     channel: str = "xhs",
     uid: str = "",
-) -> Dict[str, Any]:
+) -> ToolResult:
     """
     建议词搜索
 
@@ -187,7 +206,7 @@ async def search_suggestions(
         uid: 用户ID(自动注入)
 
     Returns:
-        API 返回的原始响应,结构如下
+        ToolResult 包含建议词数据
         {
             "code": 0,                    # 状态码,0 表示成功
             "message": "success",         # 状态消息
@@ -202,26 +221,45 @@ async def search_suggestions(
                 }
             ]
         }
-
-    Raises:
-        httpx.HTTPStatusError: HTTP 请求返回非 2xx 状态码
-        httpx.TimeoutException: 请求超时
-        httpx.RequestError: 其他请求错误
     """
-    # 处理 channel 参数,支持枚举和字符串
-    channel_value = channel.value if isinstance(channel, SuggestSearchChannel) else channel
+    try:
+        # 处理 channel 参数,支持枚举和字符串
+        channel_value = channel.value if isinstance(channel, SuggestSearchChannel) else channel
 
-    url = f"{BASE_URL}/suggest"
-    payload = {
-        "type": channel_value,
-        "keyword": keyword,
-    }
+        url = f"{BASE_URL}/suggest"
+        payload = {
+            "type": channel_value,
+            "keyword": keyword,
+        }
+
+        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
+            response = await client.post(
+                url,
+                json=payload,
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+            data = response.json()
 
-    async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
-        response = await client.post(
-            url,
-            json=payload,
-            headers={"Content-Type": "application/json"},
+        # 计算建议词数量
+        suggestion_count = 0
+        for item in data.get("data", []):
+            suggestion_count += len(item.get("list", []))
+
+        return ToolResult(
+            title=f"建议词: {keyword} ({channel_value})",
+            output=json.dumps(data, ensure_ascii=False, indent=2),
+            long_term_memory=f"Got {suggestion_count} suggestions for '{keyword}' on {channel_value}"
+        )
+    except httpx.HTTPStatusError as e:
+        return ToolResult(
+            title="获取建议词失败",
+            output="",
+            error=f"HTTP error {e.response.status_code}: {e.response.text}"
+        )
+    except Exception as e:
+        return ToolResult(
+            title="获取建议词失败",
+            output="",
+            error=str(e)
         )
-        response.raise_for_status()
-        return response.json()