Skip to content

Commit a784036

Browse files
author
hzzz2020
committed
fix: add OTS instance endpoint fallback
1 parent 40083c9 commit a784036

5 files changed

Lines changed: 173 additions & 17 deletions

File tree

‎agentrun/knowledgebase/api/__data_async_template.py‎

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -647,12 +647,15 @@ def __init__(
647647
self.retrieve_settings = retrieve_settings
648648

649649
def _build_agent_storage_client(
650-
self, config: Optional[Config] = None
650+
self,
651+
config: Optional[Config] = None,
652+
ots_endpoint: Optional[str] = None,
651653
) -> AgentStorageClient:
652654
"""构建 AgentStorageClient / Build AgentStorageClient
653655
654656
Args:
655657
config: 配置 / Configuration
658+
ots_endpoint: OTS 访问域名 / OTS endpoint
656659
657660
Returns:
658661
AgentStorageClient: OTS 存储客户端
@@ -661,18 +664,24 @@ def _build_agent_storage_client(
661664
raise ValueError("provider_settings is required for OTS retrieval")
662665

663666
cfg = Config.with_configs(self.config, config)
664-
ots_endpoint = cfg.get_ots_endpoint(
665-
self.provider_settings.ots_instance_name
666-
)
667+
ots_instance_name = self.provider_settings.ots_instance_name
668+
if ots_endpoint is None:
669+
ots_endpoint = cfg.get_ots_endpoint(ots_instance_name)
667670

668671
return AgentStorageClient(
669672
access_key_id=cfg.get_access_key_id(),
670673
access_key_secret=cfg.get_access_key_secret(),
671674
sts_token=cfg.get_security_token(),
672675
ots_endpoint=ots_endpoint,
673-
ots_instance_name=self.provider_settings.ots_instance_name,
676+
ots_instance_name=ots_instance_name,
674677
)
675678

679+
def _build_frontend_ots_endpoint(
680+
self, config: Optional[Config] = None
681+
) -> str:
682+
cfg = Config.with_configs(self.config, config)
683+
return f"http://ots-{cfg.get_region_id()}.aliyuncs.com"
684+
676685
def _build_retrieval_configuration(
677686
self, filter: Optional[Dict[str, Any]] = None
678687
) -> Optional[Dict[str, Any]]:
@@ -838,7 +847,20 @@ async def retrieve_async(
838847
request["retrievalConfiguration"] = retrieval_config
839848

840849
logger.debug(f"OTS retrieve request: {request}")
841-
response = client.retrieve(request)
850+
try:
851+
response = client.retrieve(request)
852+
except Exception as instance_error:
853+
frontend_endpoint = self._build_frontend_ots_endpoint(config)
854+
logger.warning(
855+
"Failed to retrieve from OTS knowledge base "
856+
f"'{self.knowledge_base_name}' with instance endpoint, "
857+
f"fallback to frontend endpoint {frontend_endpoint}: "
858+
f"{instance_error}"
859+
)
860+
client = self._build_agent_storage_client(
861+
config, ots_endpoint=frontend_endpoint
862+
)
863+
response = client.retrieve(request)
842864
logger.debug(f"OTS retrieve response: {response}")
843865

844866
return self._parse_retrieve_response(response, query)

‎agentrun/knowledgebase/api/data.py‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -927,12 +927,15 @@ def __init__(
927927
self.retrieve_settings = retrieve_settings
928928

929929
def _build_agent_storage_client(
930-
self, config: Optional[Config] = None
930+
self,
931+
config: Optional[Config] = None,
932+
ots_endpoint: Optional[str] = None,
931933
) -> AgentStorageClient:
932934
"""构建 AgentStorageClient / Build AgentStorageClient
933935
934936
Args:
935937
config: 配置 / Configuration
938+
ots_endpoint: OTS 访问域名 / OTS endpoint
936939
937940
Returns:
938941
AgentStorageClient: OTS 存储客户端
@@ -941,18 +944,24 @@ def _build_agent_storage_client(
941944
raise ValueError("provider_settings is required for OTS retrieval")
942945

943946
cfg = Config.with_configs(self.config, config)
944-
ots_endpoint = cfg.get_ots_endpoint(
945-
self.provider_settings.ots_instance_name
946-
)
947+
ots_instance_name = self.provider_settings.ots_instance_name
948+
if ots_endpoint is None:
949+
ots_endpoint = cfg.get_ots_endpoint(ots_instance_name)
947950

948951
return AgentStorageClient(
949952
access_key_id=cfg.get_access_key_id(),
950953
access_key_secret=cfg.get_access_key_secret(),
951954
sts_token=cfg.get_security_token(),
952955
ots_endpoint=ots_endpoint,
953-
ots_instance_name=self.provider_settings.ots_instance_name,
956+
ots_instance_name=ots_instance_name,
954957
)
955958

959+
def _build_frontend_ots_endpoint(
960+
self, config: Optional[Config] = None
961+
) -> str:
962+
cfg = Config.with_configs(self.config, config)
963+
return f"http://ots-{cfg.get_region_id()}.aliyuncs.com"
964+
956965
def _build_retrieval_configuration(
957966
self, filter: Optional[Dict[str, Any]] = None
958967
) -> Optional[Dict[str, Any]]:
@@ -1118,7 +1127,20 @@ async def retrieve_async(
11181127
request["retrievalConfiguration"] = retrieval_config
11191128

11201129
logger.debug(f"OTS retrieve request: {request}")
1121-
response = client.retrieve(request)
1130+
try:
1131+
response = client.retrieve(request)
1132+
except Exception as instance_error:
1133+
frontend_endpoint = self._build_frontend_ots_endpoint(config)
1134+
logger.warning(
1135+
"Failed to retrieve from OTS knowledge base "
1136+
f"'{self.knowledge_base_name}' with instance endpoint, "
1137+
f"fallback to frontend endpoint {frontend_endpoint}: "
1138+
f"{instance_error}"
1139+
)
1140+
client = self._build_agent_storage_client(
1141+
config, ots_endpoint=frontend_endpoint
1142+
)
1143+
response = client.retrieve(request)
11221144
logger.debug(f"OTS retrieve response: {response}")
11231145

11241146
return self._parse_retrieve_response(response, query)
@@ -1174,7 +1196,20 @@ def retrieve(
11741196
request["retrievalConfiguration"] = retrieval_config
11751197

11761198
logger.debug(f"OTS retrieve request: {request}")
1177-
response = client.retrieve(request)
1199+
try:
1200+
response = client.retrieve(request)
1201+
except Exception as instance_error:
1202+
frontend_endpoint = self._build_frontend_ots_endpoint(config)
1203+
logger.warning(
1204+
"Failed to retrieve from OTS knowledge base "
1205+
f"'{self.knowledge_base_name}' with instance endpoint, "
1206+
f"fallback to frontend endpoint {frontend_endpoint}: "
1207+
f"{instance_error}"
1208+
)
1209+
client = self._build_agent_storage_client(
1210+
config, ots_endpoint=frontend_endpoint
1211+
)
1212+
response = client.retrieve(request)
11781213
logger.debug(f"OTS retrieve response: {response}")
11791214

11801215
return self._parse_retrieve_response(response, query)

‎agentrun/utils/config.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ def get_ots_endpoint(self, instance_name: str) -> str:
343343
return (
344344
f"https://{instance_name}.{region_id}.vpc.tablestore.aliyuncs.com"
345345
)
346-
return f"http://ots-{region_id}.aliyuncs.com"
346+
return f"https://{instance_name}.{region_id}.ots.aliyuncs.com"
347347

348348
def get_use_vpc_endpoint(self) -> bool:
349349
"""知识库检索是否使用 VPC 内网 endpoint"""

‎tests/unittests/knowledgebase/test_ots_knowledgebase.py‎

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,55 @@ def test_retrieve_error_handling(self, mock_build_client):
493493
assert result["query"] == "test query"
494494
assert result["knowledge_base_name"] == "test-kb"
495495

496+
@patch("agentrun.knowledgebase.api.data.AgentStorageClient")
497+
def test_retrieve_falls_back_to_frontend_endpoint(
498+
self, mock_client_class
499+
):
500+
"""测试实例域名检索失败后回退到大前端域名"""
501+
mock_config = MagicMock(spec=Config)
502+
mock_config.get_region_id.return_value = "cn-hangzhou"
503+
mock_config.get_ots_endpoint.return_value = (
504+
"https://test-instance.cn-hangzhou.ots.aliyuncs.com"
505+
)
506+
mock_config.get_access_key_id.return_value = "test-ak"
507+
mock_config.get_access_key_secret.return_value = "test-sk"
508+
mock_config.get_security_token.return_value = "test-sts"
509+
510+
instance_client = MagicMock()
511+
instance_client.retrieve.side_effect = RuntimeError(
512+
"instance endpoint unavailable"
513+
)
514+
fallback_client = MagicMock()
515+
fallback_client.retrieve.return_value = {
516+
"code": "SUCCESS",
517+
"data": {
518+
"retrievalResults": [
519+
{"content": "fallback result", "score": 0.8}
520+
]
521+
},
522+
}
523+
mock_client_class.side_effect = [instance_client, fallback_client]
524+
525+
with patch.object(Config, "with_configs", return_value=mock_config):
526+
api = OTSDataAPI(
527+
"test-kb",
528+
provider_settings=OTSProviderSettings(
529+
ots_instance_name="test-instance"
530+
),
531+
)
532+
result = api.retrieve("test query")
533+
534+
assert "error" not in result
535+
assert result["data"][0]["content"] == "fallback result"
536+
assert mock_client_class.call_args_list[0].kwargs["ots_endpoint"] == (
537+
"https://test-instance.cn-hangzhou.ots.aliyuncs.com"
538+
)
539+
assert mock_client_class.call_args_list[1].kwargs["ots_endpoint"] == (
540+
"http://ots-cn-hangzhou.aliyuncs.com"
541+
)
542+
instance_client.retrieve.assert_called_once()
543+
fallback_client.retrieve.assert_called_once()
544+
496545
@patch(
497546
"agentrun.knowledgebase.api.data.OTSDataAPI._build_agent_storage_client"
498547
)
@@ -512,6 +561,56 @@ async def test_retrieve_async_error_handling(self, mock_build_client):
512561
assert result["error"] is True
513562
assert "Failed to retrieve" in result["data"]
514563

564+
@patch("agentrun.knowledgebase.api.data.AgentStorageClient")
565+
@pytest.mark.asyncio
566+
async def test_retrieve_async_falls_back_to_frontend_endpoint(
567+
self, mock_client_class
568+
):
569+
"""测试异步实例域名检索失败后回退到大前端域名"""
570+
mock_config = MagicMock(spec=Config)
571+
mock_config.get_region_id.return_value = "cn-hangzhou"
572+
mock_config.get_ots_endpoint.return_value = (
573+
"https://test-instance.cn-hangzhou.ots.aliyuncs.com"
574+
)
575+
mock_config.get_access_key_id.return_value = "test-ak"
576+
mock_config.get_access_key_secret.return_value = "test-sk"
577+
mock_config.get_security_token.return_value = "test-sts"
578+
579+
instance_client = MagicMock()
580+
instance_client.retrieve.side_effect = RuntimeError(
581+
"instance endpoint unavailable"
582+
)
583+
fallback_client = MagicMock()
584+
fallback_client.retrieve.return_value = {
585+
"code": "SUCCESS",
586+
"data": {
587+
"retrievalResults": [
588+
{"content": "async fallback result", "score": 0.8}
589+
]
590+
},
591+
}
592+
mock_client_class.side_effect = [instance_client, fallback_client]
593+
594+
with patch.object(Config, "with_configs", return_value=mock_config):
595+
api = OTSDataAPI(
596+
"test-kb",
597+
provider_settings=OTSProviderSettings(
598+
ots_instance_name="test-instance"
599+
),
600+
)
601+
result = await api.retrieve_async("test query")
602+
603+
assert "error" not in result
604+
assert result["data"][0]["content"] == "async fallback result"
605+
assert mock_client_class.call_args_list[0].kwargs["ots_endpoint"] == (
606+
"https://test-instance.cn-hangzhou.ots.aliyuncs.com"
607+
)
608+
assert mock_client_class.call_args_list[1].kwargs["ots_endpoint"] == (
609+
"http://ots-cn-hangzhou.aliyuncs.com"
610+
)
611+
instance_client.retrieve.assert_called_once()
612+
fallback_client.retrieve.assert_called_once()
613+
515614
def test_retrieve_without_provider_settings(self):
516615
"""测试无 provider_settings 时检索"""
517616
api = OTSDataAPI("test-kb")
@@ -598,7 +697,7 @@ def test_build_client(self, mock_client_class):
598697
"""测试构建客户端"""
599698
mock_config = MagicMock(spec=Config)
600699
mock_config.get_ots_endpoint.return_value = (
601-
"http://ots-cn-hangzhou.aliyuncs.com"
700+
"https://test-instance.cn-hangzhou.ots.aliyuncs.com"
602701
)
603702
mock_config.get_access_key_id.return_value = "test-ak"
604703
mock_config.get_access_key_secret.return_value = "test-sk"
@@ -617,7 +716,7 @@ def test_build_client(self, mock_client_class):
617716
access_key_id="test-ak",
618717
access_key_secret="test-sk",
619718
sts_token="test-sts",
620-
ots_endpoint="http://ots-cn-hangzhou.aliyuncs.com",
719+
ots_endpoint="https://test-instance.cn-hangzhou.ots.aliyuncs.com",
621720
ots_instance_name="test-instance",
622721
)
623722

‎tests/unittests/utils/test_config.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def test_kb_endpoints_default_public(self):
3535
assert config.get_bailian_endpoint() == "https://bailian.cn-beijing.aliyuncs.com"
3636
assert config.get_gpdb_endpoint() == "gpdb.aliyuncs.com"
3737
assert config.get_ots_endpoint("my-instance") == (
38-
"http://ots-cn-hangzhou.aliyuncs.com"
38+
"https://my-instance.cn-hangzhou.ots.aliyuncs.com"
3939
)
4040

4141
def test_kb_endpoints_vpc_mode(self):

0 commit comments

Comments
 (0)