|
| 1 | +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import json |
| 16 | +import os |
| 17 | +from typing import List, Optional |
| 18 | + |
| 19 | +from langchain.tools import ToolRuntime, tool |
| 20 | + |
| 21 | +from veadk.auth.veauth.utils import get_credential_from_vefaas_iam |
| 22 | +from veadk.config import getenv |
| 23 | +from veadk.utils.logger import get_logger |
| 24 | +from veadk.utils.volcengine_sign import ve_request |
| 25 | + |
| 26 | +logger = get_logger(__name__) |
| 27 | + |
| 28 | + |
| 29 | +def _clean_ansi_codes(text: str) -> str: |
| 30 | + """Remove ANSI escape sequences (color codes, etc.)""" |
| 31 | + import re |
| 32 | + |
| 33 | + ansi_escape = re.compile(r"\x1b\[[0-9;]*m") |
| 34 | + return ansi_escape.sub("", text) |
| 35 | + |
| 36 | + |
| 37 | +def _format_execution_result(result_str: str) -> str: |
| 38 | + """Format the execution results, handle escape characters and JSON structures""" |
| 39 | + try: |
| 40 | + result_json = json.loads(result_str) |
| 41 | + |
| 42 | + if not result_json.get("success"): |
| 43 | + message = result_json.get("message", "Unknown error") |
| 44 | + outputs = result_json.get("data", {}).get("outputs", []) |
| 45 | + if outputs and isinstance(outputs[0], dict): |
| 46 | + error_msg = outputs[0].get("ename", "Unknown error") |
| 47 | + return f"Execution failed: {message}, {error_msg}" |
| 48 | + |
| 49 | + outputs = result_json.get("data", {}).get("outputs", []) |
| 50 | + if not outputs: |
| 51 | + return "No output generated" |
| 52 | + |
| 53 | + formatted_lines = [] |
| 54 | + for output in outputs: |
| 55 | + if output and isinstance(output, dict) and "text" in output: |
| 56 | + text = output["text"] |
| 57 | + text = _clean_ansi_codes(text) |
| 58 | + text = text.replace("\\n", "\n") |
| 59 | + formatted_lines.append(text) |
| 60 | + |
| 61 | + return "".join(formatted_lines).strip() |
| 62 | + |
| 63 | + except json.JSONDecodeError: |
| 64 | + return _clean_ansi_codes(result_str) |
| 65 | + except Exception as e: |
| 66 | + logger.warning(f"Error formatting result: {e}, returning raw result") |
| 67 | + return result_str |
| 68 | + |
| 69 | + |
| 70 | +@tool |
| 71 | +def execute_skills( |
| 72 | + workflow_prompt: str, |
| 73 | + runtime: ToolRuntime, |
| 74 | + skills: Optional[List[str]] = None, |
| 75 | + timeout: int = 900, |
| 76 | +) -> str: |
| 77 | + """execute skills in a code sandbox and return the output. |
| 78 | + For C++ code, don't execute it directly, compile and execute via Python; write sources and object files to /tmp. |
| 79 | +
|
| 80 | + Args: |
| 81 | + workflow_prompt (str): instruction of workflow |
| 82 | + skills (Optional[List[str]]): The skills will be invoked |
| 83 | + timeout (int, optional): The timeout in seconds for the code execution, less than or equal to 900. Defaults to 900. |
| 84 | +
|
| 85 | + Returns: |
| 86 | + str: The output of the code execution. |
| 87 | + """ |
| 88 | + |
| 89 | + tool_id = getenv("AGENTKIT_TOOL_ID") |
| 90 | + |
| 91 | + service = getenv( |
| 92 | + "AGENTKIT_TOOL_SERVICE_CODE", "agentkit" |
| 93 | + ) # temporary service for code run tool |
| 94 | + region = getenv("AGENTKIT_TOOL_REGION", "cn-beijing") |
| 95 | + host = getenv( |
| 96 | + "AGENTKIT_TOOL_HOST", service + "." + region + ".volces.com" |
| 97 | + ) # temporary host for code run tool |
| 98 | + logger.debug(f"tools endpoint: {host}") |
| 99 | + |
| 100 | + session_id = runtime.session_id # type: ignore |
| 101 | + agent_name = runtime.context.agent_name # type: ignore |
| 102 | + user_id = runtime.context.user_id # type: ignore |
| 103 | + tool_user_session_id = agent_name + "_" + user_id + "_" + session_id |
| 104 | + logger.debug(f"tool_user_session_id: {tool_user_session_id}") |
| 105 | + |
| 106 | + logger.debug( |
| 107 | + f"Execute skills in session_id={session_id}, tool_id={tool_id}, host={host}, service={service}, region={region}, timeout={timeout}" |
| 108 | + ) |
| 109 | + |
| 110 | + header = {} |
| 111 | + |
| 112 | + ak = os.getenv("VOLCENGINE_ACCESS_KEY") |
| 113 | + sk = os.getenv("VOLCENGINE_SECRET_KEY") |
| 114 | + if not (ak and sk): |
| 115 | + logger.debug( |
| 116 | + "Get AK/SK from environment variables failed. Try to use credential from Iam." |
| 117 | + ) |
| 118 | + credential = get_credential_from_vefaas_iam() |
| 119 | + ak = credential.access_key_id |
| 120 | + sk = credential.secret_access_key |
| 121 | + header = {"X-Security-Token": credential.session_token} |
| 122 | + else: |
| 123 | + logger.debug("Successfully get AK/SK from environment variables.") |
| 124 | + |
| 125 | + cmd = ["python", "agent.py", workflow_prompt] |
| 126 | + if skills: |
| 127 | + cmd.extend(["--skills"] + skills) |
| 128 | + |
| 129 | + # TODO: remove after agentkit supports custom environment variables setting |
| 130 | + res = ve_request( |
| 131 | + request_body={}, |
| 132 | + action="GetCallerIdentity", |
| 133 | + ak=ak, |
| 134 | + sk=sk, |
| 135 | + service="sts", |
| 136 | + version="2018-01-01", |
| 137 | + region=region, |
| 138 | + host="sts.volcengineapi.com", |
| 139 | + header=header, |
| 140 | + ) |
| 141 | + try: |
| 142 | + account_id = res["Result"]["AccountId"] |
| 143 | + except KeyError as e: |
| 144 | + logger.error(f"Error occurred while getting account id: {e}, response is {res}") |
| 145 | + return res |
| 146 | + |
| 147 | + env_vars = { |
| 148 | + "TOS_SKILLS_DIR": f"tos://agentkit-platform-{account_id}/skills/", |
| 149 | + "TOOL_USER_SESSION_ID": tool_user_session_id, |
| 150 | + } |
| 151 | + |
| 152 | + code = f""" |
| 153 | +import subprocess |
| 154 | +import os |
| 155 | +import time |
| 156 | +import select |
| 157 | +import sys |
| 158 | +
|
| 159 | +env = os.environ.copy() |
| 160 | +for key, value in {env_vars!r}.items(): |
| 161 | + if key not in env: |
| 162 | + env[key] = value |
| 163 | +
|
| 164 | +process = subprocess.Popen( |
| 165 | + {cmd!r}, |
| 166 | + cwd='/home/gem/veadk_skills', |
| 167 | + stdout=subprocess.PIPE, |
| 168 | + stderr=subprocess.PIPE, |
| 169 | + text=True, |
| 170 | + env=env, |
| 171 | + bufsize=1, |
| 172 | + universal_newlines=True |
| 173 | +) |
| 174 | +
|
| 175 | +start_time = time.time() |
| 176 | +timeout = {timeout - 10} |
| 177 | +
|
| 178 | +with open('/tmp/agent.log', 'w') as log_file: |
| 179 | + while True: |
| 180 | + if time.time() - start_time > timeout: |
| 181 | + process.kill() |
| 182 | + log_file.write('log_type=stderr request_id=x function_id=y revision_number=1 Process timeout\\n') |
| 183 | + break |
| 184 | + |
| 185 | + reads = [process.stdout.fileno(), process.stderr.fileno()] |
| 186 | + ret = select.select(reads, [], [], 1) |
| 187 | + |
| 188 | + for fd in ret[0]: |
| 189 | + if fd == process.stdout.fileno(): |
| 190 | + line = process.stdout.readline() |
| 191 | + if line: |
| 192 | + log_file.write(f'log_type=stdout request_id=x function_id=y revision_number=1 {{line}}') |
| 193 | + log_file.flush() |
| 194 | + if fd == process.stderr.fileno(): |
| 195 | + line = process.stderr.readline() |
| 196 | + if line: |
| 197 | + log_file.write(f'log_type=stderr request_id=x function_id=y revision_number=1 {{line}}') |
| 198 | + log_file.flush() |
| 199 | + |
| 200 | + if process.poll() is not None: |
| 201 | + break |
| 202 | + |
| 203 | + for line in process.stdout: |
| 204 | + log_file.write(f'log_type=stdout request_id=x function_id=y revision_number=1 {{line}}') |
| 205 | + for line in process.stderr: |
| 206 | + log_file.write(f'log_type=stderr request_id=x function_id=y revision_number=1 {{line}}') |
| 207 | +
|
| 208 | +with open('/tmp/agent.log', 'r') as log_file: |
| 209 | + output = log_file.read() |
| 210 | + print(output) |
| 211 | + """ |
| 212 | + |
| 213 | + res = ve_request( |
| 214 | + request_body={ |
| 215 | + "ToolId": tool_id, |
| 216 | + "UserSessionId": tool_user_session_id, |
| 217 | + "OperationType": "RunCode", |
| 218 | + "OperationPayload": json.dumps( |
| 219 | + { |
| 220 | + "code": code, |
| 221 | + "timeout": timeout, |
| 222 | + "kernel_name": "python3", |
| 223 | + } |
| 224 | + ), |
| 225 | + }, |
| 226 | + action="InvokeTool", |
| 227 | + ak=ak, |
| 228 | + sk=sk, |
| 229 | + service=service, |
| 230 | + version="2025-10-30", |
| 231 | + region=region, |
| 232 | + host=host, |
| 233 | + header=header, |
| 234 | + ) |
| 235 | + logger.debug(f"Invoke run code response: {res}") |
| 236 | + |
| 237 | + try: |
| 238 | + return _format_execution_result(res["Result"]["Result"]) |
| 239 | + except KeyError as e: |
| 240 | + logger.error(f"Error occurred while running code: {e}, response is {res}") |
| 241 | + return res |
0 commit comments