From 6607fcdf3d837c910a170636a8746dd694f7ca0f Mon Sep 17 00:00:00 2001 From: "zhenguo.li" Date: Tue, 16 Dec 2025 15:09:42 +0800 Subject: [PATCH 1/4] feat: support custom url and header for llm shield service --- veadk/tools/builtin_tools/llm_shield.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/veadk/tools/builtin_tools/llm_shield.py b/veadk/tools/builtin_tools/llm_shield.py index f0583b63..d893596e 100644 --- a/veadk/tools/builtin_tools/llm_shield.py +++ b/veadk/tools/builtin_tools/llm_shield.py @@ -67,6 +67,8 @@ def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None: self.appid = getenv("TOOL_LLM_SHIELD_APP_ID") self.region = region self.timeout = timeout + self.url = getenv("TOOL_LLM_SHIELD_URL", f"https://{self.region}.sdk.access.llm-shield.omini-shield.com") + self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY") self.category_map = { 101: "Model Misuse", @@ -120,13 +122,15 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]: body_json = json.dumps(body).encode("utf-8") header = {"X-Security-Token": session_token} - url = f"https://{self.region}.sdk.access.llm-shield.omini-shield.com" + # Add x-api-key header if API key is provided + if self.api_key: + header["x-api-key"] = self.api_key path = "/v2/moderate" action = "Moderate" version = "2025-08-31" signed_header = request_sign( - header, ak, sk, self.region, url, path, action, body_json + header, ak, sk, self.region, self.url, path, action, body_json ) signed_header.update( @@ -139,7 +143,7 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]: try: response = requests.post( - url + path, + self.url + path, headers=signed_header, data=body_json, params={"Action": action, "Version": version}, From 1c0adec04b290a91339faff462768be1c1953717 Mon Sep 17 00:00:00 2001 From: "zhenguo.li" Date: Tue, 16 Dec 2025 16:18:24 +0800 Subject: [PATCH 2/4] fix: pre commit linter --- veadk/tools/builtin_tools/llm_shield.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/veadk/tools/builtin_tools/llm_shield.py b/veadk/tools/builtin_tools/llm_shield.py index d893596e..62bfa0cc 100644 --- a/veadk/tools/builtin_tools/llm_shield.py +++ b/veadk/tools/builtin_tools/llm_shield.py @@ -67,7 +67,10 @@ def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None: self.appid = getenv("TOOL_LLM_SHIELD_APP_ID") self.region = region self.timeout = timeout - self.url = getenv("TOOL_LLM_SHIELD_URL", f"https://{self.region}.sdk.access.llm-shield.omini-shield.com") + self.url = getenv( + "TOOL_LLM_SHIELD_URL", + f"https://{self.region}.sdk.access.llm-shield.omini-shield.com", + ) self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY") self.category_map = { From 6d367b112014acf26b185c2c9c0b82589331e174 Mon Sep 17 00:00:00 2001 From: "zhenguo.li" Date: Wed, 17 Dec 2025 20:13:20 +0800 Subject: [PATCH 3/4] feat: disabled ak/sk auth when use api key auth for llm shield service --- veadk/tools/builtin_tools/llm_shield.py | 62 ++++++++++++++----------- 1 file changed, 35 insertions(+), 27 deletions(-) diff --git a/veadk/tools/builtin_tools/llm_shield.py b/veadk/tools/builtin_tools/llm_shield.py index 62bfa0cc..1a5eca7c 100644 --- a/veadk/tools/builtin_tools/llm_shield.py +++ b/veadk/tools/builtin_tools/llm_shield.py @@ -71,7 +71,7 @@ def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None: "TOOL_LLM_SHIELD_URL", f"https://{self.region}.sdk.access.llm-shield.omini-shield.com", ) - self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY") + self.api_key = getenv("TOOL_LLM_SHIELD_API_KEY", allow_false_values=True) self.category_map = { 101: "Model Misuse", @@ -101,18 +101,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]: logger.error("LLM Shield app ID not configured") return None - ak = os.getenv("VOLCENGINE_ACCESS_KEY") - sk = os.getenv("VOLCENGINE_SECRET_KEY") - session_token = "" - if not (ak and sk): - logger.debug("Get AK/SK from environment variables failed.") - credential = get_credential_from_vefaas_iam() - ak = credential.access_key_id - sk = credential.secret_access_key - session_token = credential.session_token - else: - logger.debug("Successfully get AK/SK from environment variables.") - body = { "Message": { "Role": role, @@ -124,25 +112,46 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]: body_json = json.dumps(body).encode("utf-8") - header = {"X-Security-Token": session_token} - # Add x-api-key header if API key is provided - if self.api_key: - header["x-api-key"] = self.api_key path = "/v2/moderate" action = "Moderate" version = "2025-08-31" - signed_header = request_sign( - header, ak, sk, self.region, self.url, path, action, body_json - ) - - signed_header.update( - { + # Check if using API key authentication + logger.debug(f"API key value: {self.api_key}, type: {type(self.api_key)}") + if self.api_key and self.api_key != "": + logger.debug("Using API key authentication (no AK/SK signature)") + # Use API key authentication only - match curl command headers exactly + signed_header = { "Content-Type": "application/json", - "X-Top-Service": "llmshield", - "X-Top-Region": self.region, + "x-api-key": self.api_key, } - ) + else: + logger.debug("Using AK/SK signature authentication") + # Use AK/SK signature authentication + ak = os.getenv("VOLCENGINE_ACCESS_KEY") + sk = os.getenv("VOLCENGINE_SECRET_KEY") + session_token = "" + if not (ak and sk): + logger.debug("Get AK/SK from environment variables failed.") + credential = get_credential_from_vefaas_iam() + ak = credential.access_key_id + sk = credential.secret_access_key + session_token = credential.session_token + else: + logger.debug("Successfully get AK/SK from environment variables.") + + header = {"X-Security-Token": session_token} + signed_header = request_sign( + header, ak, sk, self.region, self.url, path, action, body_json + ) + + signed_header.update( + { + "Content-Type": "application/json", + "X-Top-Service": "llmshield", + "X-Top-Region": self.region, + } + ) try: response = requests.post( @@ -158,7 +167,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]: f"LLM Shield HTTP error: {response.status_code} - {response.text}" ) return None - response = response.json() except requests.exceptions.Timeout: logger.error("LLM Shield request timeout") From 56a9cf056397f3e3c4e632e8fd885a954aa2f759 Mon Sep 17 00:00:00 2001 From: "zhenguo.li" Date: Thu, 18 Dec 2025 21:05:05 +0800 Subject: [PATCH 4/4] feat: add web search tool-specific ak/sk --- veadk/tools/builtin_tools/web_search.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/veadk/tools/builtin_tools/web_search.py b/veadk/tools/builtin_tools/web_search.py index bd011d5a..56266097 100644 --- a/veadk/tools/builtin_tools/web_search.py +++ b/veadk/tools/builtin_tools/web_search.py @@ -38,7 +38,12 @@ def web_search(query: str, tool_context: ToolContext | None = None) -> list[str] """ ak = None sk = None - if tool_context: + # First try to get tool-specific AK/SK + ak = os.getenv("TOOL_WEB_SEARCH_ACCESS_KEY") + sk = os.getenv("TOOL_WEB_SEARCH_SECRET_KEY") + if ak and sk: + logger.debug("Successfully get tool-specific AK/SK.") + elif tool_context: ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY") sk = tool_context.state.get("VOLCENGINE_SECRET_KEY") session_token = ""