Skip to content

Commit 647e161

Browse files
committed
feat: add execute_skills tool and add download_directory in VeTOS
1 parent cc50209 commit 647e161

File tree

2 files changed

+208
-0
lines changed

2 files changed

+208
-0
lines changed

veadk/integrations/ve_tos/ve_tos.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -705,6 +705,80 @@ def download(self, bucket_name: str, object_key: str, save_path: str) -> bool:
705705
logger.error(f"Image download failed: {str(e)}")
706706
return False
707707

708+
def download_directory(
709+
self, bucket_name: str, prefix: str, local_dir: str = "/tmp"
710+
) -> bool:
711+
"""Download entire directory from TOS bucket to local directory
712+
713+
Args:
714+
bucket_name: TOS bucket name
715+
prefix: Directory prefix in TOS (e.g., "skills/pdf/")
716+
local_dir: Local directory path to save files
717+
718+
Returns:
719+
bool: True if download succeeds, False otherwise
720+
"""
721+
bucket_name = self._check_bucket_name(bucket_name)
722+
723+
if not self._client:
724+
logger.error("TOS client is not initialized")
725+
return False
726+
727+
try:
728+
# Ensure prefix ends with /
729+
if prefix and not prefix.endswith("/"):
730+
prefix += "/"
731+
732+
# Create local directory if not exists
733+
os.makedirs(local_dir, exist_ok=True)
734+
735+
# List all objects with the prefix
736+
is_truncated = True
737+
next_continuation_token = ""
738+
downloaded_count = 0
739+
740+
while is_truncated:
741+
out = self._client.list_objects_type2(
742+
bucket_name,
743+
prefix=prefix,
744+
continuation_token=next_continuation_token,
745+
)
746+
is_truncated = out.is_truncated
747+
next_continuation_token = out.next_continuation_token
748+
749+
# Download each object
750+
for content in out.contents:
751+
object_key = content.key
752+
753+
# Skip directory markers (objects ending with /)
754+
if object_key.endswith("/"):
755+
continue
756+
757+
# Calculate relative path and local file path
758+
relative_path = object_key[len(prefix) :]
759+
local_file_path = os.path.join(local_dir, relative_path)
760+
761+
# Create subdirectories if needed
762+
local_file_dir = os.path.dirname(local_file_path)
763+
if local_file_dir:
764+
os.makedirs(local_file_dir, exist_ok=True)
765+
766+
# Download the file
767+
if self.download(bucket_name, object_key, local_file_path):
768+
downloaded_count += 1
769+
logger.debug(f"Downloaded: {object_key} -> {local_file_path}")
770+
else:
771+
logger.warning(f"Failed to download: {object_key}")
772+
773+
logger.info(
774+
f"Downloaded {downloaded_count} files from {bucket_name}/{prefix} to {local_dir}"
775+
)
776+
return downloaded_count > 0
777+
778+
except Exception as e:
779+
logger.error(f"Failed to download directory: {str(e)}")
780+
return False
781+
708782
def close(self):
709783
if self._client:
710784
self._client.close()
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
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 Optional, List
18+
19+
from google.adk.tools import ToolContext
20+
21+
from veadk.config import getenv
22+
from veadk.utils.logger import get_logger
23+
from veadk.utils.volcengine_sign import ve_request
24+
from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
25+
26+
logger = get_logger(__name__)
27+
28+
29+
def execute_skills(
30+
workflow_prompt: str,
31+
skills: Optional[List[str]] = None,
32+
tool_context: ToolContext = None,
33+
timeout: int = 30,
34+
) -> str:
35+
"""execute skills in a code sandbox and return the output.
36+
For C++ code, don't execute it directly, compile and execute via Python; write sources and object files to /tmp.
37+
38+
Args:
39+
workflow_prompt (str): instruction of workflow
40+
skills (Optional[List[str]]): The skills will be invoked
41+
timeout (int, optional): The timeout in seconds for the code execution. Defaults to 30.
42+
43+
Returns:
44+
str: The output of the code execution.
45+
"""
46+
47+
tool_id = getenv("AGENTKIT_TOOL_ID")
48+
49+
service = getenv(
50+
"AGENTKIT_TOOL_SERVICE_CODE", "agentkit"
51+
) # temporary service for code run tool
52+
region = getenv("AGENTKIT_TOOL_REGION", "cn-beijing")
53+
host = getenv(
54+
"AGENTKIT_TOOL_HOST", service + "." + region + ".volces.com"
55+
) # temporary host for code run tool
56+
logger.debug(f"tools endpoint: {host}")
57+
58+
session_id = tool_context._invocation_context.session.id
59+
agent_name = tool_context._invocation_context.agent.name
60+
user_id = tool_context._invocation_context.user_id
61+
tool_user_session_id = agent_name + "_" + user_id + "_" + session_id
62+
logger.debug(f"tool_user_session_id: {tool_user_session_id}")
63+
64+
logger.debug(
65+
f"Execute skills in session_id={session_id}, tool_id={tool_id}, host={host}, service={service}, region={region}, timeout={timeout}"
66+
)
67+
68+
ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
69+
sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
70+
header = {}
71+
72+
if not (ak and sk):
73+
logger.debug("Get AK/SK from tool context failed.")
74+
ak = os.getenv("VOLCENGINE_ACCESS_KEY")
75+
sk = os.getenv("VOLCENGINE_SECRET_KEY")
76+
if not (ak and sk):
77+
logger.debug(
78+
"Get AK/SK from environment variables failed. Try to use credential from Iam."
79+
)
80+
credential = get_credential_from_vefaas_iam()
81+
ak = credential.access_key_id
82+
sk = credential.secret_access_key
83+
header = {"X-Security-Token": credential.session_token}
84+
else:
85+
logger.debug("Successfully get AK/SK from environment variables.")
86+
else:
87+
logger.debug("Successfully get AK/SK from tool context.")
88+
89+
cmd = ["python", "agent.py", workflow_prompt]
90+
if skills:
91+
cmd.extend(["--skills"] + skills)
92+
93+
code = f"""
94+
import subprocess
95+
result = subprocess.run(
96+
{cmd!r},
97+
cwd='/home/gem/veadk_skills',
98+
capture_output=True,
99+
text=True
100+
)
101+
print(result.stdout)
102+
if result.stderr:
103+
print(result.stderr)
104+
"""
105+
106+
res = ve_request(
107+
request_body={
108+
"ToolId": tool_id,
109+
"UserSessionId": tool_user_session_id,
110+
"OperationType": "RunCode",
111+
"OperationPayload": json.dumps(
112+
{
113+
"code": code,
114+
"timeout": timeout,
115+
"kernel_name": "python3",
116+
}
117+
),
118+
},
119+
action="InvokeTool",
120+
ak=ak,
121+
sk=sk,
122+
service=service,
123+
version="2025-10-30",
124+
region=region,
125+
host=host,
126+
header=header,
127+
)
128+
logger.debug(f"Invoke run code response: {res}")
129+
130+
try:
131+
return res["Result"]["Result"]
132+
except KeyError as e:
133+
logger.error(f"Error occurred while running code: {e}, response is {res}")
134+
return res

0 commit comments

Comments
 (0)