feat(auth): support auto auth for web search (#272)

* feat(auth): support auto auth for web search

* change version

Author: yaozheng-fang

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.16"
+version = "0.2.17"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"

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]

veadk/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-VERSION = "0.2.16"
+VERSION = "0.2.17"