Skip to content

Commit 6d367b1

Browse files
author
zhenguo.li
committed
feat: disabled ak/sk auth when use api key auth for llm shield service
1 parent 1c0adec commit 6d367b1

File tree

1 file changed

+35
-27
lines changed

1 file changed

+35
-27
lines changed

veadk/tools/builtin_tools/llm_shield.py

Lines changed: 35 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None:
7171
"TOOL_LLM_SHIELD_URL",
7272
f"https://{self.region}.sdk.access.llm-shield.omini-shield.com",
7373
)
74-
self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY")
74+
self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY", allow_false_values=True)
7575

7676
self.category_map = {
7777
101: "Model Misuse",
@@ -101,18 +101,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
101101
logger.error("LLM Shield app ID not configured")
102102
return None
103103

104-
ak = os.getenv("VOLCENGINE_ACCESS_KEY")
105-
sk = os.getenv("VOLCENGINE_SECRET_KEY")
106-
session_token = ""
107-
if not (ak and sk):
108-
logger.debug("Get AK/SK from environment variables failed.")
109-
credential = get_credential_from_vefaas_iam()
110-
ak = credential.access_key_id
111-
sk = credential.secret_access_key
112-
session_token = credential.session_token
113-
else:
114-
logger.debug("Successfully get AK/SK from environment variables.")
115-
116104
body = {
117105
"Message": {
118106
"Role": role,
@@ -124,25 +112,46 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
124112

125113
body_json = json.dumps(body).encode("utf-8")
126114

127-
header = {"X-Security-Token": session_token}
128-
# Add x-api-key header if API key is provided
129-
if self.api_key:
130-
header["x-api-key"] = self.api_key
131115
path = "/v2/moderate"
132116
action = "Moderate"
133117
version = "2025-08-31"
134118

135-
signed_header = request_sign(
136-
header, ak, sk, self.region, self.url, path, action, body_json
137-
)
138-
139-
signed_header.update(
140-
{
119+
# Check if using API key authentication
120+
logger.debug(f"API key value: {self.api_key}, type: {type(self.api_key)}")
121+
if self.api_key and self.api_key != "":
122+
logger.debug("Using API key authentication (no AK/SK signature)")
123+
# Use API key authentication only - match curl command headers exactly
124+
signed_header = {
141125
"Content-Type": "application/json",
142-
"X-Top-Service": "llmshield",
143-
"X-Top-Region": self.region,
126+
"x-api-key": self.api_key,
144127
}
145-
)
128+
else:
129+
logger.debug("Using AK/SK signature authentication")
130+
# Use AK/SK signature authentication
131+
ak = os.getenv("VOLCENGINE_ACCESS_KEY")
132+
sk = os.getenv("VOLCENGINE_SECRET_KEY")
133+
session_token = ""
134+
if not (ak and sk):
135+
logger.debug("Get AK/SK from environment variables failed.")
136+
credential = get_credential_from_vefaas_iam()
137+
ak = credential.access_key_id
138+
sk = credential.secret_access_key
139+
session_token = credential.session_token
140+
else:
141+
logger.debug("Successfully get AK/SK from environment variables.")
142+
143+
header = {"X-Security-Token": session_token}
144+
signed_header = request_sign(
145+
header, ak, sk, self.region, self.url, path, action, body_json
146+
)
147+
148+
signed_header.update(
149+
{
150+
"Content-Type": "application/json",
151+
"X-Top-Service": "llmshield",
152+
"X-Top-Region": self.region,
153+
}
154+
)
146155

147156
try:
148157
response = requests.post(
@@ -158,7 +167,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
158167
f"LLM Shield HTTP error: {response.status_code} - {response.text}"
159168
)
160169
return None
161-
162170
response = response.json()
163171
except requests.exceptions.Timeout:
164172
logger.error("LLM Shield request timeout")

0 commit comments

Comments
 (0)