fix: support byteplus in viking, web_search and mem0

Author: BhAem

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -705,6 +705,7 @@ def _do_request(
             session_token=sts_token or self.session_token,
             method=method,
             data=body,
+            region=self.region,
         )
         response = requests.request(
             method=method,

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -150,7 +150,14 @@ def _get_ak_sk_sts(self) -> tuple[str, str, str]:
 
     def _get_client(self) -> VikingDBMemoryClient:
         ak, sk, sts_token = self._get_ak_sk_sts()
+        if self.cloud_provider.lower() == "byteplus":
+            host = f"api-knowledgebase.mlp.{self.region}.bytepluses.com"
+
+        logger.info(f"Cloud provider: {self.cloud_provider.lower()}")
+        logger.info(f"VikingDBLTMBackend: region={self.region}, host={host}")
+
         return VikingDBMemoryClient(
+            host=host,
             ak=ak,
             sk=sk,
             sts_token=sts_token,
@@ -159,23 +166,23 @@ def _get_client(self) -> VikingDBMemoryClient:
 
     def _get_sdk_client(self) -> VikingMem:
         ak, sk, sts_token = self._get_ak_sk_sts()
-        client = VikingDBMemoryClient(
-            ak=ak,
-            sk=sk,
-            sts_token=sts_token,
-            region=self.region,
-        )
-
-        host = client.get_host()
         if self.cloud_provider.lower() == "byteplus":
             host = f"api-knowledgebase.mlp.{self.region}.bytepluses.com"
 
         logger.info(f"Cloud provider: {self.cloud_provider.lower()}")
         logger.info(f"VikingDBLTMBackend: region={self.region}, host={host}")
 
-        return VikingMem(
+        client = VikingDBMemoryClient(
             host=host,
             region=self.region,
+            ak=ak,
+            sk=sk,
+            sts_token=sts_token,
+        )
+
+        return VikingMem(
+            host=client.get_host(),
+            region=self.region,
             auth=IAM(
                 ak=ak,
                 sk=sk,

veadk/tools/builtin_tools/web_search.py

@@ -17,6 +17,7 @@
 """
 
 import os
+import requests
 
 from google.adk.tools import ToolContext
 
@@ -71,33 +72,40 @@ def web_search(query: str, tool_context: ToolContext | None = None) -> list[str]
             "Query": query,
             "Count": 5,
         }
-        host = "torchlight.byteintlapi.com/search_api/web_search"
         api_key = os.getenv("BYTEPLUS_WEB_SEARCH_API_KEY")
+        if not api_key:
+            logger.error("BYTEPLUS_WEB_SEARCH_API_KEY is not set.")
+            return ["Web search failed: API key is not set."]
         header = {
-            "X-Security-Token": session_token,
             "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
         }
-    else:
-        request_body = {
-            "Query": query,
-            "SearchType": "web",
-            "Count": 5,
-            "NeedSummary": True,
-        }
-        host = "mercury.volcengineapi.com"
-        header = {"X-Security-Token": session_token}
 
-    response = ve_request(
-        request_body=request_body,
-        action="WebSearch",
-        ak=ak,
-        sk=sk,
-        service="volc_torchlight_api",
-        version="2025-01-01",
-        region="cn-beijing",
-        host=host,
-        header=header,
-    )
+        response = requests.post(
+            url="https://torchlight.byteintlapi.com/search_api/web_search",
+            headers=header,
+            json=request_body,
+            timeout=60,
+        )
+        response.raise_for_status()
+        response = response.json()
+    else:
+        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["Result"]["WebResults"]