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

Author: floritange

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
@@ -51,10 +51,25 @@ def __init__(
         ak="",
         sk="",
         sts_token="",
-        scheme="http",
+        scheme="https",
         connection_timeout=30,
         socket_timeout=30,
     ):
+        env_host = getenv(
+            "DATABASE_VIKINGMEM_BASE_URL", default_value=None, allow_false_values=True
+        )
+        if env_host:
+            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:
+                raise ValueError(
+                    "DATABASE_VIKINGMEM_BASE_URL must start with http:// or https://"
+                )
+
         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
@@ -557,7 +557,21 @@ def _do_request(
         path: str,
         method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
     ) -> dict:
-        VIKINGDB_KNOWLEDGEBASE_BASE_URL = "api-knowledgebase.mlp.cn-beijing.volces.com"
+        VIKINGDB_KNOWLEDGEBASE_BASE_URL = (
+            "https://api-knowledgebase.mlp.cn-beijing.volces.com"
+        )
+        full_path = f"{VIKINGDB_KNOWLEDGEBASE_BASE_URL}{path}"
+
+        env_host = getenv(
+            "DATABASE_VIKING_BASE_URL", default_value=None, allow_false_values=True
+        )
+        if env_host:
+            if env_host.startswith("http://") or env_host.startswith("https://"):
+                full_path = f"{env_host}{path}"
+            else:
+                raise ValueError(
+                    "DATABASE_VIKING_BASE_URL must start with http:// or https://"
+                )
 
         volcengine_access_key = self.volcengine_access_key
         volcengine_secret_key = self.volcengine_secret_key
@@ -579,7 +593,7 @@ def _do_request(
         )
         response = requests.request(
             method=method,
-            url=f"https://{VIKINGDB_KNOWLEDGEBASE_BASE_URL}{path}",
+            url=full_path,
             headers=request.headers,
             data=request.body,
         )