feat(viking_url): viking mem and db envs (DATABASE_VIKINGMEM_BASE_URL and DATABASE_VIKING_BASE_URL)

Author: floritange

veadk/consts.py

@@ -72,3 +72,6 @@
 DEFAULT_IMAGE_GENERATE_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
 
 VEFAAS_IAM_CRIDENTIAL_PATH = "/var/run/secrets/iam/credential"
+
+DEFAULT_VIKING_BASE_URL = "api-knowledgebase.mlp.cn-beijing.volces.com"
+DEFAULT_VIKINGMEM_BASE_URL = "api-knowledgebase.mlp.cn-beijing.volces.com"

veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py

@@ -14,7 +14,7 @@
 
 import json
 import threading
-
+from veadk.utils.misc import getenv
 from volcengine.ApiInfo import ApiInfo
 from volcengine.auth.SignerV4 import SignerV4
 from volcengine.base.Service import Service
@@ -55,6 +55,21 @@ def __init__(
         connection_timeout=30,
         socket_timeout=30,
     ):
+        env_host = getenv(
+            "DATABASE_VIKINGMEM_BASE_URL", "api-knowledgebase.mlp.cn-beijing.volces.com"
+        )
+        # check if the environment variable contains protocol
+        if env_host.startswith("http://"):
+            host = env_host.replace("http://", "")
+            scheme = "http"
+        elif env_host.startswith("https://"):
+            host = env_host.replace("https://", "")
+            scheme = "https"
+        else:
+            # default http
+            host = env_host
+            scheme = "http"
+
         self.service_info = VikingDBMemoryClient.get_service_info(
             host, region, scheme, connection_timeout, socket_timeout
         )

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -22,7 +22,7 @@
 import requests
 from pydantic import Field
 from typing_extensions import override
-
+from veadk.utils.misc import getenv
 import veadk.config  # noqa E401
 from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
@@ -58,6 +58,23 @@ def get_files_in_directory(directory: str):
     return file_paths
 
 
+def _normalize_base_url(base_url: str) -> str:
+    """Normalize base URL: if already has scheme, return as-is; otherwise add http://
+
+    Args:
+        base_url: Base URL string
+
+    Returns:
+        Normalized URL with scheme
+    """
+    if not base_url:
+        return base_url
+    if base_url.startswith(("http://", "https://")):
+        return base_url
+    # PrivateLink URLs use http only
+    return f"http://{base_url}"
+
+
 class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
     """Volcengine Viking DB knowledgebase backend.
 
@@ -557,7 +574,18 @@ def _do_request(
         path: str,
         method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
     ) -> dict:
-        VIKINGDB_KNOWLEDGEBASE_BASE_URL = "api-knowledgebase.mlp.cn-beijing.volces.com"
+        base_url = getenv(
+            "DATABASE_VIKING_BASE_URL", "api-knowledgebase.mlp.cn-beijing.volces.com"
+        )
+        if base_url.startswith(("http://", "https://")):
+            full_url = f"{base_url}{path}"
+        else:
+            scheme = (
+                "https"
+                if base_url == "api-knowledgebase.mlp.cn-beijing.volces.com"
+                else "http"
+            )
+            full_url = f"{scheme}://{base_url}{path}"
 
         volcengine_access_key = self.volcengine_access_key
         volcengine_secret_key = self.volcengine_secret_key
@@ -577,9 +605,10 @@ def _do_request(
             method=method,
             data=body,
         )
+
         response = requests.request(
             method=method,
-            url=f"https://{VIKINGDB_KNOWLEDGEBASE_BASE_URL}{path}",
+            url=full_url,
             headers=request.headers,
             data=request.body,
         )