veadk-python 0.2.16__py3-none-any.whl → 0.2.17__py3-none-any.whl

veadk/tools/builtin_tools/llm_shield.py

@@ -0,0 +1,381 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import requests
+from typing import Optional, List, Dict, Any, Union
+from volcenginesdkllmshield.models.llm_shield_sign import request_sign
+
+from google.adk.plugins import BasePlugin
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.tools.tool_context import ToolContext
+from google.adk.models import LlmRequest, LlmResponse
+from google.genai import types
+from google.adk.tools.base_tool import BaseTool
+
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+
+logger = get_logger(__name__)
+
+
+class LLMShieldPlugin(BasePlugin):
+    """
+    LLM Shield Plugin for content moderation and safety.
+
+    This plugin integrates with Volcano Engine's LLM Shield service to provide
+    real-time content moderation for user inputs, model outputs, and tool interactions.
+    It helps detect and block potentially harmful content including sensitive information,
+    prompt injection attacks, and policy violations.
+
+    Examples:
+        Basic usage with default settings:
+        ```python
+        from veadk.tools.builtin_tools.llm_shield import content_safety
+        agent = Agent(
+            before_model_callback=content_safety.before_model_callback,
+            after_model_callback=content_safety.after_model_callback,
+            before_tool_callback=content_safety.before_tool_callback,
+            after_tool_callback=content_safety.after_tool_callback,
+        )
+        ```
+    """
+
+    def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None:
+        """
+        Initialize the LLM Shield Plugin.
+
+        Args:
+            region (str, optional): The service region. Defaults to "cn-beijing".
+            timeout (int, optional): Request timeout in seconds. Defaults to 50.
+        """
+        self.name = "LLMShieldPlugin"
+        super().__init__(name=self.name)
+
+        self.appid = getenv("TOOL_LLM_SHIELD_APP_ID")
+        self.region = region
+        self.timeout = timeout
+
+        self.category_map = {
+            101: "Model Misuse",
+            103: "Sensitive Information",
+            104: "Prompt Injection",
+            106: "General Topic Control",
+            107: "Computational Resource Consumption",
+        }
+
+    def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
+        """
+        Make a request to the LLM Shield service for content moderation.
+
+        This method sends a message to the LLM Shield API for security analysis.
+        If the content is deemed risky, it returns a blocking message explaining
+        the violation. Otherwise, it returns None to allow the content through.
+
+        Args:
+            message (str): The content to be moderated
+            role (str): The role of the message sender ("user" or "assistant")
+
+        Returns:
+            Optional[str]: A blocking message if content violates policies,
+                         None if content is safe or on error
+        """
+        if not self.appid:
+            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,
+                "Content": message,
+                "ContentType": 1,
+            },
+            "Scene": self.appid,
+        }
+
+        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"
+        path = "/v2/moderate"
+        action = "Moderate"
+        version = "2025-08-31"
+
+        signed_header = request_sign(
+            header, ak, sk, self.region, 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(
+                url + path,
+                headers=signed_header,
+                data=body_json,
+                params={"Action": action, "Version": version},
+                timeout=self.timeout,
+            )
+
+            if response.status_code != 200:
+                logger.error(
+                    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")
+            return None
+        except requests.exceptions.RequestException as e:
+            logger.error(f"LLM Shield network request failed: {e}")
+            return None
+        except json.JSONDecodeError as e:
+            logger.error(f"LLM Shield response JSON decode failed: {e}")
+            return None
+        except Exception as e:
+            logger.error(f"LLM Shield request failed: {e}")
+            return None
+
+        # Process risk detection results
+        result = response.get("Result", None)
+        if result:
+            decision = result.get("Decision", None)
+            decision_type = decision.get("DecisionType", None)
+            risk_info = result.get("RiskInfo", None)
+            if decision_type is not None and int(decision_type) == 2 and risk_info:
+                risks = risk_info.get("Risks", [])
+                if risks:
+                    # Extract risk categories for user-friendly error message
+                    risk_reasons = set()
+                    for risk in risks:
+                        category = risk.get("Category", None)
+                        if category:
+                            category_name = self.category_map.get(
+                                int(category), f"Category {category}"
+                            )
+                            risk_reasons.add(category_name)
+
+                    # Generate blocking response
+                    reason_text = (
+                        ", ".join(risk_reasons)
+                        if risk_reasons
+                        else "security policy violation"
+                    )
+                    response_text = (
+                        f"Your request has been blocked due to: {reason_text}. "
+                        f"Please modify your input and try again."
+                    )
+
+                    return response_text
+
+        return None
+
+    def before_agent_callback(
+        self, callback_context: CallbackContext, **kwargs
+    ) -> None:
+        # TODO: Implement agent-level input validation and context analysis
+        return None
+
+    def after_agent_callback(self, callback_context: CallbackContext, **kwargs) -> None:
+        # TODO: Implement post-agent analysis and context analysis
+        return None
+
+    def before_model_callback(
+        self, callback_context: CallbackContext, llm_request: LlmRequest, **kwargs
+    ) -> Optional[LlmResponse]:
+        """
+        Moderate user input before sending to the language model.
+
+        Extracts the last user message from the LLM request and checks it
+        against the LLM Shield service. If the content violates safety policies,
+        returns a blocking response instead of allowing the request to proceed.
+
+        Args:
+            callback_context (CallbackContext): The callback execution context
+            llm_request (LlmRequest): The incoming LLM request to moderate
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Optional[LlmResponse]: A blocking response if content is unsafe,
+                                 None if content is safe to proceed
+        """
+        # Extract the last user message for moderation
+        last_user_message = None
+        contents = getattr(llm_request, "contents", [])
+
+        if contents:
+            last_content = contents[-1]
+            last_role = getattr(last_content, "role", "")
+            last_parts = getattr(last_content, "parts", [])
+
+            if last_role == "user" and last_parts:
+                last_user_message = getattr(last_parts[0], "text", "")
+
+        # Skip moderation if message is empty
+        if not last_user_message:
+            return None
+
+        response = self._request_llm_shield(message=last_user_message, role="user")
+        if response:
+            logger.debug("LLM Shield triggered in before_model_callback.")
+            return LlmResponse(
+                content=types.Content(
+                    role="model",
+                    parts=[types.Part(text=response)],
+                )
+            )
+        return None
+
+    def after_model_callback(
+        self, callback_context: CallbackContext, llm_response: LlmResponse, **kwargs
+    ) -> Optional[LlmResponse]:
+        """
+        Moderate model output before returning to the user.
+
+        Extracts the model's response and checks it against the LLM Shield service.
+        If the model's output violates safety policies, returns a blocking response
+        instead of the original model output.
+
+        Args:
+            callback_context (CallbackContext): The callback execution context
+            llm_response (LlmResponse): The model's response to moderate
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Optional[LlmResponse]: A blocking response if content is unsafe,
+                                 None if content is safe to return
+        """
+        # Extract the model's response for moderation
+        last_model_message = None
+        content = getattr(llm_response, "content", [])
+
+        if content:
+            last_role = getattr(content, "role", "")
+            last_parts = getattr(content, "parts", [])
+
+            if last_role == "model" and last_parts:
+                last_model_message = getattr(last_parts[0], "text", "")
+
+        # Skip moderation if message is empty
+        if not last_model_message:
+            return None
+
+        response = self._request_llm_shield(
+            message=last_model_message, role="assistant"
+        )
+        if response:
+            logger.debug("LLM Shield triggered in after_model_callback.")
+            return LlmResponse(
+                content=types.Content(
+                    role="model",
+                    parts=[types.Part(text=response)],
+                )
+            )
+        return None
+
+    def before_tool_callback(
+        self, tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext, **kwargs
+    ) -> Optional[Dict]:
+        """
+        Moderate tool arguments before tool execution.
+
+        Combines all tool arguments into a message and checks it against
+        the LLM Shield service. If the arguments contain unsafe content,
+        returns a blocking result instead of allowing tool execution.
+
+        Args:
+            tool (BaseTool): The tool to be executed
+            args (Dict[str, Any]): The arguments passed to the tool
+            tool_context (ToolContext): The tool execution context
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Optional[Dict]: A blocking result if arguments are unsafe,
+                          None if arguments are safe to proceed
+        """
+        args_list = []
+
+        for key, value in args.items():
+            args_list.append(f"{key}: {value}")
+
+        message = "\n".join(args_list)
+        response = self._request_llm_shield(message=message, role="user")
+        if response:
+            logger.debug("LLM Shield triggered in before_tool_callback.")
+            return {"result": response}
+        return None
+
+    def after_tool_callback(
+        self,
+        tool: BaseTool,
+        args: Dict[str, Any],
+        tool_context: CallbackContext,
+        tool_response: Union[str, Dict[str, Any], List[Any]],
+        **kwargs,
+    ) -> Optional[Dict]:
+        """
+        Moderate tool output after tool execution.
+
+        Processes the tool's response (string, dict, or list) into a message
+        and checks it against the LLM Shield service. If the tool's output
+        violates safety policies, returns a blocking result.
+
+        Args:
+            tool (BaseTool): The tool that was executed
+            args (Dict[str, Any]): The arguments that were passed to the tool
+            tool_context (CallbackContext): The tool execution context
+            tool_response (Union[str, Dict[str, Any], List[Any]]): The tool's response
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Optional[Dict]: A blocking result if tool output is unsafe,
+                          None if output is safe to return
+        """
+        message = ""
+        if isinstance(tool_response, str):
+            message = tool_response
+        elif isinstance(tool_response, dict):
+            for key, value in tool_response.items():
+                message += f"{value}\n"
+        elif isinstance(tool_response, list):
+            for item in tool_response:
+                message += f"{item}\n"
+
+        response = self._request_llm_shield(message=message, role="assistant")
+        if response:
+            logger.debug("LLM Shield triggered in after_tool_callback.")
+            return {"result": response}
+        return None
+
+
+content_safety = LLMShieldPlugin()

veadk/tools/builtin_tools/mcp_router.py

@@ -15,8 +15,15 @@
 from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
 from veadk.config import getenv
-from veadk.utils.mcp_utils import get_mcp_params
+from google.adk.tools.mcp_tool.mcp_session_manager import (
+    StreamableHTTPConnectionParams,
+)
 
 url = getenv("TOOL_MCP_ROUTER_URL")
+api_key = getenv("TOOL_MCP_ROUTER_API_KEY")
 
-mcp_router = MCPToolset(connection_params=get_mcp_params(url=url))
+mcp_router = MCPToolset(
+    connection_params=StreamableHTTPConnectionParams(
+        url=url, headers={"Authorization": f"Bearer {api_key}"}
+    ),
+)

veadk/tools/builtin_tools/run_code.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 
 import json
+import os
 
 from google.adk.tools import ToolContext
 
 from veadk.config import getenv
 from veadk.utils.logger import get_logger
 from veadk.utils.volcengine_sign import ve_request
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 
 logger = get_logger(__name__)
 
@@ -51,8 +53,26 @@ def run_code(code: str, language: str, tool_context: ToolContext) -> str:
         f"Running code in language: {language}, session_id={session_id}, code={code}, tool_id={tool_id}, host={host}, service={service}, region={region}"
     )
 
-    access_key = getenv("VOLCENGINE_ACCESS_KEY")
-    secret_key = getenv("VOLCENGINE_SECRET_KEY")
+    ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
+    sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
+    header = {}
+
+    if not (ak and sk):
+        logger.debug("Get AK/SK from tool context failed.")
+        ak = os.getenv("VOLCENGINE_ACCESS_KEY")
+        sk = os.getenv("VOLCENGINE_SECRET_KEY")
+        if not (ak and sk):
+            logger.debug(
+                "Get AK/SK from environment variables failed. Try to use credential from Iam."
+            )
+            credential = get_credential_from_vefaas_iam()
+            ak = credential.access_key_id
+            sk = credential.secret_access_key
+            header = {"X-Security-Token": credential.session_token}
+        else:
+            logger.debug("Successfully get AK/SK from environment variables.")
+    else:
+        logger.debug("Successfully get AK/SK from tool context.")
 
     res = ve_request(
         request_body={
@@ -68,14 +88,14 @@ def run_code(code: str, language: str, tool_context: ToolContext) -> str:
             ),
         },
         action="InvokeTool",
-        ak=access_key,
-        sk=secret_key,
+        ak=ak,
+        sk=sk,
         service=service,
         version="2025-10-30",
         region=region,
         host=host,
+        header=header,
     )
-
     logger.debug(f"Invoke run code response: {res}")
 
     try:

veadk/tools/builtin_tools/web_search.py

@@ -16,142 +16,17 @@
 The document of this tool see: https://www.volcengine.com/docs/85508/1650263
 """
 
-import datetime
-import hashlib
-import hmac
-import json
-from urllib.parse import quote
+import os
 
-import requests
 from google.adk.tools import ToolContext
 
-from veadk.config import getenv
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.utils.logger import get_logger
+from veadk.utils.volcengine_sign import ve_request
 
 logger = get_logger(__name__)
 
 
-Service = "volc_torchlight_api"
-Version = "2025-01-01"
-Region = "cn-beijing"
-Host = "mercury.volcengineapi.com"
-ContentType = "application/json"
-
-
-def norm_query(params):
-    query = ""
-    for key in sorted(params.keys()):
-        if isinstance(params[key], list):
-            for k in params[key]:
-                query = (
-                    query + quote(key, safe="-_.~") + "=" + quote(k, safe="-_.~") + "&"
-                )
-        else:
-            query = (
-                query
-                + quote(key, safe="-_.~")
-                + "="
-                + quote(params[key], safe="-_.~")
-                + "&"
-            )
-    query = query[:-1]
-    return query.replace("+", "%20")
-
-
-def hmac_sha256(key: bytes, content: str):
-    return hmac.new(key, content.encode("utf-8"), hashlib.sha256).digest()
-
-
-def hash_sha256(content: str):
-    return hashlib.sha256(content.encode("utf-8")).hexdigest()
-
-
-def request(method, date, query, header, ak, sk, action, body):
-    credential = {
-        "access_key_id": ak,
-        "secret_access_key": sk,
-        "service": Service,
-        "region": Region,
-    }
-    request_param = {
-        "body": body,
-        "host": Host,
-        "path": "/",
-        "method": method,
-        "content_type": ContentType,
-        "date": date,
-        "query": {"Action": action, "Version": Version, **query},
-    }
-    if body is None:
-        request_param["body"] = ""
-    # 第四步：接下来开始计算签名。在计算签名前，先准备好用于接收签算结果的 signResult 变量，并设置一些参数。
-    # 初始化签名结果的结构体
-    x_date = request_param["date"].strftime("%Y%m%dT%H%M%SZ")
-    short_x_date = x_date[:8]
-    x_content_sha256 = hash_sha256(request_param["body"])
-    sign_result = {
-        "Host": request_param["host"],
-        "X-Content-Sha256": x_content_sha256,
-        "X-Date": x_date,
-        "Content-Type": request_param["content_type"],
-    }
-    signed_headers_str = ";".join(
-        ["content-type", "host", "x-content-sha256", "x-date"]
-    )
-    # signed_headers_str = signed_headers_str + ";x-security-token"
-    canonical_request_str = "\n".join(
-        [
-            request_param["method"].upper(),
-            request_param["path"],
-            norm_query(request_param["query"]),
-            "\n".join(
-                [
-                    "content-type:" + request_param["content_type"],
-                    "host:" + request_param["host"],
-                    "x-content-sha256:" + x_content_sha256,
-                    "x-date:" + x_date,
-                ]
-            ),
-            "",
-            signed_headers_str,
-            x_content_sha256,
-        ]
-    )
-
-    hashed_canonical_request = hash_sha256(canonical_request_str)
-
-    credential_scope = "/".join(
-        [short_x_date, credential["region"], credential["service"], "request"]
-    )
-    string_to_sign = "\n".join(
-        ["HMAC-SHA256", x_date, credential_scope, hashed_canonical_request]
-    )
-
-    k_date = hmac_sha256(credential["secret_access_key"].encode("utf-8"), short_x_date)
-    k_region = hmac_sha256(k_date, credential["region"])
-    k_service = hmac_sha256(k_region, credential["service"])
-    k_signing = hmac_sha256(k_service, "request")
-    signature = hmac_sha256(k_signing, string_to_sign).hex()
-
-    sign_result["Authorization"] = (
-        "HMAC-SHA256 Credential={}, SignedHeaders={}, Signature={}".format(
-            credential["access_key_id"] + "/" + credential_scope,
-            signed_headers_str,
-            signature,
-        )
-    )
-    header = {**header, **sign_result}
-    # header = {**header, **{"X-Security-Token": SessionToken}}
-    r = requests.request(
-        method=method,
-        url="https://{}{}".format(request_param["host"], request_param["path"]),
-        headers=header,
-        params=request_param["query"],
-        data=request_param["body"],
-    )
-    return r.json()
-
-
 def web_search(query: str, tool_context: ToolContext) -> list[str]:
     """Search a query in websites.
 
@@ -161,39 +36,48 @@ def web_search(query: str, tool_context: ToolContext) -> list[str]:
     Returns:
         A list of result documents.
     """
-    req = {
-        "Query": query,
-        "SearchType": "web",
-        "Count": 5,
-        "NeedSummary": True,
-    }
-
     ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
-    if not ak:
-        ak = getenv("VOLCENGINE_ACCESS_KEY")
-
     sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
-    if not sk:
-        sk = getenv("VOLCENGINE_SECRET_KEY")
-
-    now = datetime.datetime.utcnow()
-    response_body = request(
-        "POST",
-        now,
-        {},
-        {},
-        ak,
-        sk,
-        "WebSearch",
-        json.dumps(req),
+    session_token = ""
+
+    if not (ak and sk):
+        logger.debug("Get AK/SK from tool context failed.")
+        ak = os.getenv("VOLCENGINE_ACCESS_KEY")
+        sk = os.getenv("VOLCENGINE_SECRET_KEY")
+        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.")
+    else:
+        logger.debug("Successfully get AK/SK from tool context.")
+
+    response = ve_request(
+        request_body={
+            "Query": query,
+            "SearchType": "web",
+            "Count": 5,
+            "NeedSummary": True,
+        },
+        action="WebSearch",
+        ak=ak,
+        sk=sk,
+        service="volc_torchlight_api",
+        version="2025-01-01",
+        region="cn-beijing",
+        host="mercury.volcengineapi.com",
+        header={"X-Security-Token": session_token},
     )
 
     try:
-        results: list = response_body["Result"]["WebResults"]
+        results: list = response["Result"]["WebResults"]
         final_results = []
         for result in results:
             final_results.append(result["Summary"].strip())
         return final_results
     except Exception as e:
-        logger.error(f"Web search failed {e}, response body: {response_body}")
-        return [response_body]
+        logger.error(f"Web search failed {e}, response body: {response}")
+        return [response]