feat: add opt-in VPC endpoints for knowledge base retrieval

Allow Bailian, ADB, and OTS knowledge base calls to use VPC endpoints via
AGENTRUN_KB_USE_VPC while keeping public network as the default behavior.

Change-Id: I8b51c2cf7f9089bc4e0a9b1609c119c126b7273a
Co-developed-by: Cursor <noreply@cursor.com>

Author: 灵轮

agentrun/knowledgebase/api/__data_async_template.py

@@ -661,8 +661,9 @@ def _build_agent_storage_client(
             raise ValueError("provider_settings is required for OTS retrieval")
 
         cfg = Config.with_configs(self.config, config)
-        region_id = cfg.get_region_id()
-        ots_endpoint = f"http://ots-{region_id}.aliyuncs.com"
+        ots_endpoint = cfg.get_ots_endpoint(
+            self.provider_settings.ots_instance_name
+        )
 
         return AgentStorageClient(
             access_key_id=cfg.get_access_key_id(),

agentrun/knowledgebase/api/data.py

@@ -941,8 +941,9 @@ def _build_agent_storage_client(
             raise ValueError("provider_settings is required for OTS retrieval")
 
         cfg = Config.with_configs(self.config, config)
-        region_id = cfg.get_region_id()
-        ots_endpoint = f"http://ots-{region_id}.aliyuncs.com"
+        ots_endpoint = cfg.get_ots_endpoint(
+            self.provider_settings.ots_instance_name
+        )
 
         return AgentStorageClient(
             access_key_id=cfg.get_access_key_id(),

agentrun/utils/config.py

@@ -12,6 +12,16 @@
 load_dotenv()
 
 
+_GPDB_PUBLIC_SHARED_REGIONS = (
+    "cn-beijing",
+    "cn-hangzhou",
+    "cn-shanghai",
+    "cn-shenzhen",
+    "cn-hongkong",
+    "ap-southeast-1",
+)
+
+
 def get_env_with_default(default: str, *key: str) -> str:
     """从环境变量获取值,支持多个候选键 / Get value from environment variables with multiple fallback keys
 
@@ -29,6 +39,15 @@ def get_env_with_default(default: str, *key: str) -> str:
     return default
 
 
+def get_env_flag(*key: str) -> bool:
+    """从环境变量读取布尔开关 / Read boolean flag from environment variables"""
+    for k in key:
+        v = os.getenv(k)
+        if v is not None:
+            return v.strip().lower() in ("1", "true", "yes", "on")
+    return False
+
+
 class Config:
     """AgentRun SDK 全局配置类 / AgentRun SDK Global Configuration Class
 
@@ -62,6 +81,7 @@ class Config:
         "_data_endpoint",
         "_devs_endpoint",
         "_bailian_endpoint",
+        "_use_vpc_endpoint",
         "_headers",
         "__weakref__",
     )
@@ -80,6 +100,7 @@ def __init__(
         data_endpoint: Optional[str] = None,
         devs_endpoint: Optional[str] = None,
         bailian_endpoint: Optional[str] = None,
+        use_vpc_endpoint: Optional[bool] = None,
         headers: Optional[Dict[str, str]] = None,
     ) -> None:
         """初始化配置 / Initialize configuration
@@ -103,7 +124,9 @@ def __init__(
             read_timeout: 读取超时时间(秒),默认 100000 / Read timeout in seconds, defaults to 100000
             control_endpoint: 自定义控制链路端点,可选 / Custom control endpoint, optional
             data_endpoint: 自定义数据链路端点,可选 / Custom data endpoint, optional
-            devs_endpoint: 自定义 DevS 端点,可选 / Custom DevS endpoint, optional
+            devs_endpoint: 自定义 Devs 端点,可选 / Custom DevS endpoint, optional
+            use_vpc_endpoint: 知识库检索是否使用 VPC 内网 endpoint,默认 false
+                未提供时从环境变量读取: AGENTRUN_KB_USE_VPC
             headers: 自定义请求头,可选 / Custom request headers, optional
         """
 
@@ -139,6 +162,8 @@ def __init__(
             devs_endpoint = get_env_with_default("", "DEVS_ENDPOINT")
         if bailian_endpoint is None:
             bailian_endpoint = get_env_with_default("", "BAILIAN_ENDPOINT")
+        if use_vpc_endpoint is None:
+            use_vpc_endpoint = get_env_flag("AGENTRUN_KB_USE_VPC")
 
         self._access_key_id = access_key_id
         self._access_key_secret = access_key_secret
@@ -152,6 +177,7 @@ def __init__(
         self._data_endpoint = data_endpoint
         self._devs_endpoint = devs_endpoint
         self._bailian_endpoint = bailian_endpoint
+        self._use_vpc_endpoint = use_vpc_endpoint
         self._headers = headers or {}
 
     @classmethod
@@ -263,8 +289,33 @@ def get_bailian_endpoint(self) -> str:
         if self._bailian_endpoint:
             return self._bailian_endpoint
 
+        if self._use_vpc_endpoint:
+            return f"bailian-vpc.{self.get_region_id()}.aliyuncs.com"
+
         return "https://bailian.cn-beijing.aliyuncs.com"
 
+    def get_gpdb_endpoint(self) -> str:
+        """获取 GPDB (ADB) OpenAPI 端点 / Get GPDB (ADB) OpenAPI endpoint"""
+        region_id = self.get_region_id()
+        if self._use_vpc_endpoint:
+            return f"gpdb-vpc.{region_id}.aliyuncs.com"
+        if region_id in _GPDB_PUBLIC_SHARED_REGIONS:
+            return "gpdb.aliyuncs.com"
+        return f"gpdb.{region_id}.aliyuncs.com"
+
+    def get_ots_endpoint(self, instance_name: str) -> str:
+        """获取 OTS endpoint / Get OTS endpoint"""
+        region_id = self.get_region_id()
+        if self._use_vpc_endpoint:
+            return (
+                f"https://{instance_name}.{region_id}.vpc.tablestore.aliyuncs.com"
+            )
+        return f"http://ots-{region_id}.aliyuncs.com"
+
+    def get_use_vpc_endpoint(self) -> bool:
+        """知识库检索是否使用 VPC 内网 endpoint"""
+        return self._use_vpc_endpoint
+
     def get_headers(self) -> Dict[str, str]:
         """获取自定义请求头"""
         return self._headers or {}

agentrun/utils/control_api.py

@@ -126,19 +126,7 @@ def _get_gpdb_client(self, config: Optional[Config] = None) -> "GPDBClient":
         from alibabacloud_gpdb20160503.client import Client as GPDBClient
 
         cfg = Config.with_configs(self.config, config)
-        # GPDB 使用区域级别的 endpoint / GPDB uses region-level endpoint
-        region_id = cfg.get_region_id()
-        if region_id in (
-            "cn-beijing",
-            "cn-hangzhou",
-            "cn-shanghai",
-            "cn-shenzhen",
-            "cn-hongkong",
-            "ap-southeast-1",
-        ):
-            endpoint = "gpdb.aliyuncs.com"
-        else:
-            endpoint = f"gpdb.{region_id}.aliyuncs.com"
+        endpoint = cfg.get_gpdb_endpoint()
 
         return GPDBClient(
             open_api_util_models.Config(

examples/knowledgebase.py

@@ -16,6 +16,7 @@
 - AGENTRUN_ACCESS_KEY_ID: 阿里云 AccessKey ID
 - AGENTRUN_ACCESS_KEY_SECRET: 阿里云 AccessKey Secret
 - AGENTRUN_REGION: 区域（默认 cn-hangzhou）
+- AGENTRUN_KB_USE_VPC: 设为 true 时，百炼/ADB/OTS 知识库检索走 VPC 内网 endpoint（默认 false，RagFlow 不受影响）
 
 百炼知识库额外配置 / Additional config for Bailian:
 - BAILIAN_WORKSPACE_ID: 百炼工作空间 ID

tests/unittests/knowledgebase/test_ots_knowledgebase.py

@@ -597,7 +597,9 @@ class TestOTSDataAPIBuildClient:
     def test_build_client(self, mock_client_class):
         """测试构建客户端"""
         mock_config = MagicMock(spec=Config)
-        mock_config.get_region_id.return_value = "cn-hangzhou"
+        mock_config.get_ots_endpoint.return_value = (
+            "http://ots-cn-hangzhou.aliyuncs.com"
+        )
         mock_config.get_access_key_id.return_value = "test-ak"
         mock_config.get_access_key_secret.return_value = "test-sk"
         mock_config.get_security_token.return_value = "test-sts"
@@ -619,6 +621,36 @@ def test_build_client(self, mock_client_class):
             ots_instance_name="test-instance",
         )
 
+    @patch("agentrun.knowledgebase.api.data.AgentStorageClient")
+    def test_build_client_vpc_mode(self, mock_client_class):
+        """测试 VPC 模式构建 OTS 客户端"""
+        mock_config = MagicMock(spec=Config)
+        mock_config.get_ots_endpoint.return_value = (
+            "https://test-instance.cn-hangzhou.vpc.tablestore.aliyuncs.com"
+        )
+        mock_config.get_access_key_id.return_value = "test-ak"
+        mock_config.get_access_key_secret.return_value = "test-sk"
+        mock_config.get_security_token.return_value = "test-sts"
+
+        with patch.object(Config, "with_configs", return_value=mock_config):
+            api = OTSDataAPI(
+                "test-kb",
+                provider_settings=OTSProviderSettings(
+                    ots_instance_name="test-instance"
+                ),
+            )
+            api._build_agent_storage_client()
+
+        mock_client_class.assert_called_once_with(
+            access_key_id="test-ak",
+            access_key_secret="test-sk",
+            sts_token="test-sts",
+            ots_endpoint=(
+                "https://test-instance.cn-hangzhou.vpc.tablestore.aliyuncs.com"
+            ),
+            ots_instance_name="test-instance",
+        )
+
     def test_build_client_without_provider_settings(self):
         """测试无 provider_settings 时构建客户端"""
         api = OTSDataAPI("test-kb")

tests/unittests/utils/test_config.py

@@ -19,3 +19,34 @@ def test_init_without_parameters(self):
             assert config._access_key_id == "mock-access-key-id"
             assert config._access_key_secret == "mock-access-key-secret"
             assert config._account_id == "mock-account-id"
+            assert config._use_vpc_endpoint is False
+
+    def test_use_vpc_endpoint_from_env(self):
+        with patch.dict(os.environ, {"AGENTRUN_KB_USE_VPC": "true"}, clear=False):
+            config = Config()
+            assert config.get_use_vpc_endpoint() is True
+
+    def test_kb_endpoints_default_public(self):
+        config = Config(region_id="cn-hangzhou")
+        assert config.get_bailian_endpoint() == "https://bailian.cn-beijing.aliyuncs.com"
+        assert config.get_gpdb_endpoint() == "gpdb.aliyuncs.com"
+        assert config.get_ots_endpoint("my-instance") == (
+            "http://ots-cn-hangzhou.aliyuncs.com"
+        )
+
+    def test_kb_endpoints_vpc_mode(self):
+        config = Config(region_id="cn-hangzhou", use_vpc_endpoint=True)
+        assert config.get_bailian_endpoint() == (
+            "bailian-vpc.cn-hangzhou.aliyuncs.com"
+        )
+        assert config.get_gpdb_endpoint() == "gpdb-vpc.cn-hangzhou.aliyuncs.com"
+        assert config.get_ots_endpoint("my-instance") == (
+            "https://my-instance.cn-hangzhou.vpc.tablestore.aliyuncs.com"
+        )
+
+    def test_bailian_endpoint_override_takes_precedence(self):
+        config = Config(
+            use_vpc_endpoint=True,
+            bailian_endpoint="custom.bailian.example.com",
+        )
+        assert config.get_bailian_endpoint() == "custom.bailian.example.com"

tests/unittests/utils/test_control_api.py

@@ -342,6 +342,26 @@ def test_get_bailian_client_strips_http_prefix(self, mock_client_class):
         config_arg = call_args[0][0]
         assert config_arg.endpoint == "bailian.custom.com"
 
+    @patch("alibabacloud_bailian20231229.client.Client")
+    def test_get_bailian_client_vpc_mode(self, mock_client_class):
+        """测试 VPC 模式使用 bailian-vpc endpoint"""
+        config = Config(
+            access_key_id="ak",
+            access_key_secret="sk",
+            region_id="cn-beijing",
+            use_vpc_endpoint=True,
+        )
+        api = ControlAPI(config=config)
+
+        mock_client = MagicMock()
+        mock_client_class.return_value = mock_client
+
+        api._get_bailian_client()
+
+        call_args = mock_client_class.call_args
+        config_arg = call_args[0][0]
+        assert config_arg.endpoint == "bailian-vpc.cn-beijing.aliyuncs.com"
+
 
 class TestControlAPIGetGPDBClient:
     """测试 ControlAPI._get_gpdb_client"""
@@ -413,3 +433,23 @@ def test_get_gpdb_client_all_known_regions(self, mock_client_class):
             assert (
                 config_arg.endpoint == "gpdb.aliyuncs.com"
             ), f"Region {region} should use gpdb.aliyuncs.com"
+
+    @patch("alibabacloud_gpdb20160503.client.Client")
+    def test_get_gpdb_client_vpc_mode(self, mock_client_class):
+        """测试 VPC 模式使用 gpdb-vpc endpoint"""
+        config = Config(
+            access_key_id="ak",
+            access_key_secret="sk",
+            region_id="cn-hangzhou",
+            use_vpc_endpoint=True,
+        )
+        api = ControlAPI(config=config)
+
+        mock_client = MagicMock()
+        mock_client_class.return_value = mock_client
+
+        api._get_gpdb_client()
+
+        call_args = mock_client_class.call_args
+        config_arg = call_args[0][0]
+        assert config_arg.endpoint == "gpdb-vpc.cn-hangzhou.aliyuncs.com"