feat: support byteplus in viking, web_search and mem0 (#482)

* feat: support byteplus in viking, web_search and mem0

* fix: support byteplus in viking, web_search and mem0

Author: BhAem

veadk/configs/database_configs.py

@@ -118,6 +118,13 @@ class TOSConfig(BaseSettings):
 
     region: str = "cn-beijing"
 
+    def model_post_init(self, __context) -> None:
+        cloud_provider = os.getenv("CLOUD_PROVIDER", "volces").lower()
+
+        if cloud_provider == "byteplus":
+            self.endpoint = "tos-ap-southeast-1.bytepluses.com"
+            self.region = "ap-southeast-1"
+
     @cached_property
     def bucket(self) -> str:
         _bucket = os.getenv("DATABASE_TOS_BUCKET") or DEFAULT_TOS_BUCKET_NAME

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -169,19 +169,37 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         default_factory=lambda: os.getenv("DATABASE_VIKING_VERSION", "2")
     )
 
-    region: str = Field(
-        default_factory=lambda: os.getenv("DATABASE_VIKING_REGION", "cn-beijing")
+    cloud_provider: str = Field(
+        default_factory=lambda: os.getenv("CLOUD_PROVIDER", "volces")
     )
 
-    base_url: str = "https://api-knowledgebase.mlp.cn-beijing.volces.com"
-    host: str = "api-knowledgebase.mlp.cn-beijing.volces.com"
+    region: str = Field(default="")
+    base_url: str = Field(default="")
+    host: str = Field(default="")
     schema: str = "https"
 
     tos_config: TOSConfig | NormalTOSConfig = Field(default_factory=TOSConfig)
 
     _viking_sdk_client = None
 
     def model_post_init(self, __context: Any) -> None:
+        if not self.region:
+            if self.cloud_provider.lower() == "byteplus":
+                self.region = os.getenv("DATABASE_VIKING_REGION", "cn-hongkong")
+                self.base_url = (
+                    f"https://api-knowledgebase.mlp.{self.region}.bytepluses.com"
+                )
+                self.host = f"api-knowledgebase.mlp.{self.region}.bytepluses.com"
+            else:
+                self.region = os.getenv("DATABASE_VIKING_REGION", "cn-beijing")
+                self.base_url = (
+                    f"https://api-knowledgebase.mlp.{self.region}.volces.com"
+                )
+                self.host = f"api-knowledgebase.mlp.{self.region}.volces.com"
+
+        logger.info(f"Cloud provider: {self.cloud_provider.lower()}")
+        logger.info(f"VikingDBKnowledgeBackend: region={self.region}, host={self.host}")
+
         self.precheck_index_naming()
 
         # check whether collection exist, if not, create it
@@ -687,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

@@ -48,9 +48,11 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
 
     session_token: str = ""
 
-    region: str = Field(
-        default_factory=lambda: os.getenv("DATABASE_VIKINGMEM_REGION") or "cn-beijing"
+    cloud_provider: str = Field(
+        default_factory=lambda: os.getenv("CLOUD_PROVIDER", "volces")
     )
+
+    region: str = Field(default="")
     """VikingDB memory region"""
 
     volcengine_project: str = Field(
@@ -61,6 +63,12 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
     memory_type: list[str] = Field(default_factory=list)
 
     def model_post_init(self, __context: Any) -> None:
+        if not self.region:
+            if self.cloud_provider.lower() == "byteplus":
+                self.region = os.getenv("DATABASE_VIKING_REGION", "cn-hongkong")
+            else:
+                self.region = os.getenv("DATABASE_VIKING_REGION", "cn-beijing")
+
         # We get memory type from:
         # 1. user input
         # 2. environment variable
@@ -142,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,
@@ -151,12 +166,20 @@ def _get_client(self) -> VikingDBMemoryClient:
 
     def _get_sdk_client(self) -> VikingMem:
         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}")
+
         client = VikingDBMemoryClient(
+            host=host,
+            region=self.region,
             ak=ak,
             sk=sk,
             sts_token=sts_token,
-            region=self.region,
         )
+
         return VikingMem(
             host=client.get_host(),
             region=self.region,

veadk/tools/builtin_tools/web_search.py

@@ -17,6 +17,7 @@
 """
 
 import os
+import requests
 
 from google.adk.tools import ToolContext
 
@@ -63,22 +64,48 @@ def web_search(query: str, tool_context: ToolContext | None = None) -> list[str]
     else:
         logger.debug("Successfully get AK/SK from tool context.")
 
-    response = ve_request(
-        request_body={
+    provider = (os.getenv("CLOUD_PROVIDER") or "").lower()
+    logger.info(f"Cloud provider: {provider}")
+
+    if provider == "byteplus":
+        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},
-    )
+        }
+        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 = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+
+        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"]