Skip to content

Commit e8984ec

Browse files
committed
refactor(utils): Introduce lazy imports for BailianClient/GPDBClient and memory_collection module
This change refactors the code to use lazy imports for `BailianClient`, `GPDBClient`, and the entire `memory_collection` module. By doing so, we defer the loading of these components until they're actually needed, which helps reduce startup time and avoids importing unnecessary dependencies like tablestore, mem0ai, numpy, fastapi, uvicorn, etc., when simply importing the `agentrun` package. The key changes include: - Moving `BailianClient` and `GPDBClient` imports inside their respective methods using local imports. - Adding conditional imports (`TYPE_CHECKING`) for `memory_collection` items. - Implementing a custom `__getattr__` function that lazily loads modules on demand rather than during initial import. Co-developed-by: Aone Copilot <noreply@alibaba-inc.com> Signed-off-by: Sodawyx <sodawyx@126.com>
1 parent 38b7225 commit e8984ec

File tree

4 files changed

+72
-28
lines changed

4 files changed

+72
-28
lines changed

agentrun/__init__.py

Lines changed: 49 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -57,22 +57,26 @@
5757
CredentialUpdateInput,
5858
RelatedResource,
5959
)
60-
# Memory Collection
61-
from agentrun.memory_collection import (
62-
EmbedderConfig,
63-
EmbedderConfigConfig,
64-
LLMConfig,
65-
LLMConfigConfig,
66-
MemoryCollection,
67-
MemoryCollectionClient,
68-
MemoryCollectionCreateInput,
69-
MemoryCollectionListInput,
70-
MemoryCollectionListOutput,
71-
MemoryCollectionUpdateInput,
72-
NetworkConfiguration,
73-
VectorStoreConfig,
74-
VectorStoreConfigConfig,
75-
)
60+
61+
# Memory Collection - 延迟导入以避免 tablestore/mem0ai 等重型依赖
62+
# Lazy import to avoid heavy dependencies (tablestore, mem0ai, numpy, etc.)
63+
# Type hints for IDE and type checkers
64+
if TYPE_CHECKING:
65+
from agentrun.memory_collection import (
66+
EmbedderConfig,
67+
EmbedderConfigConfig,
68+
LLMConfig,
69+
LLMConfigConfig,
70+
MemoryCollection,
71+
MemoryCollectionClient,
72+
MemoryCollectionCreateInput,
73+
MemoryCollectionListInput,
74+
MemoryCollectionListOutput,
75+
MemoryCollectionUpdateInput,
76+
NetworkConfiguration,
77+
VectorStoreConfig,
78+
VectorStoreConfigConfig,
79+
)
7680
# Model Service
7781
from agentrun.model import (
7882
BackendType,
@@ -304,6 +308,24 @@
304308
"Config",
305309
]
306310

311+
# Memory Collection 模块的所有导出(延迟加载)
312+
# Memory Collection module exports (lazy loaded)
313+
_MEMORY_COLLECTION_EXPORTS = {
314+
"MemoryCollection",
315+
"MemoryCollectionClient",
316+
"EmbedderConfig",
317+
"EmbedderConfigConfig",
318+
"LLMConfig",
319+
"LLMConfigConfig",
320+
"NetworkConfiguration",
321+
"VectorStoreConfig",
322+
"VectorStoreConfigConfig",
323+
"MemoryCollectionCreateInput",
324+
"MemoryCollectionUpdateInput",
325+
"MemoryCollectionListInput",
326+
"MemoryCollectionListOutput",
327+
}
328+
307329
# Server 模块的所有导出
308330
_SERVER_EXPORTS = {
309331
"AgentRunServer",
@@ -346,11 +368,19 @@
346368

347369

348370
def __getattr__(name: str):
349-
"""延迟加载 server 模块的导出,避免可选依赖导致导入失败
371+
"""延迟加载 server / memory_collection 模块的导出,避免重型依赖在
372+
import agentrun 时被立即加载。
350373
351-
当用户访问 server 相关的类时,才尝试导入 server 模块。
352-
如果 server 可选依赖未安装,会抛出清晰的错误提示。
374+
Lazy-load server / memory_collection module exports to avoid pulling in
375+
heavy dependencies (tablestore, mem0ai, fastapi, etc.) at import time.
353376
"""
377+
# Memory Collection 模块(延迟加载以避免 tablestore/mem0ai 依赖)
378+
if name in _MEMORY_COLLECTION_EXPORTS:
379+
from agentrun import memory_collection
380+
381+
return getattr(memory_collection, name)
382+
383+
# Server 模块(延迟加载以避免 fastapi/uvicorn 依赖)
354384
if name in _SERVER_EXPORTS:
355385
try:
356386
from agentrun import server

agentrun/utils/control_api.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,22 @@
44
This module defines the base class for control API.
55
"""
66

7-
from typing import Optional
7+
from typing import Optional, TYPE_CHECKING
88

99
from alibabacloud_agentrun20250910.client import Client as AgentRunClient
10-
from alibabacloud_bailian20231229.client import Client as BailianClient
1110
from alibabacloud_devs20230714.client import Client as DevsClient
12-
from alibabacloud_gpdb20160503.client import Client as GPDBClient
1311
from alibabacloud_tea_openapi import utils_models as open_api_util_models
1412

1513
from agentrun.utils.config import Config
1614

15+
# 延迟导入:BailianClient 和 GPDBClient 仅在 knowledgebase 模块使用,
16+
# 不在顶层导入以减少非 knowledgebase 场景的依赖加载。
17+
# Lazy import: BailianClient and GPDBClient are only used by the knowledgebase
18+
# module. Deferring import to reduce dependency loading for non-KB scenarios.
19+
if TYPE_CHECKING:
20+
from alibabacloud_bailian20231229.client import Client as BailianClient
21+
from alibabacloud_gpdb20160503.client import Client as GPDBClient
22+
1723

1824
class ControlAPI:
1925
"""控制链路客户端基类 / Control API Client Base Class
@@ -88,6 +94,7 @@ def _get_bailian_client(
8894
Returns:
8995
BailianClient: 百炼 API 客户端实例 / Bailian API client instance
9096
"""
97+
from alibabacloud_bailian20231229.client import Client as BailianClient
9198

9299
cfg = Config.with_configs(self.config, config)
93100
endpoint = cfg.get_bailian_endpoint()
@@ -116,6 +123,7 @@ def _get_gpdb_client(self, config: Optional[Config] = None) -> "GPDBClient":
116123
Returns:
117124
GPDBClient: GPDB API 客户端实例 / GPDB API client instance
118125
"""
126+
from alibabacloud_gpdb20160503.client import Client as GPDBClient
119127

120128
cfg = Config.with_configs(self.config, config)
121129
# GPDB 使用区域级别的 endpoint / GPDB uses region-level endpoint

pyproject.toml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,12 @@ dependencies = [
2222
]
2323

2424
[project.optional-dependencies]
25+
# CLI 最小依赖集 / Minimal dependencies for CLI binary packaging
26+
core = [
27+
"click>=8.0.0",
28+
"rich>=13.0.0",
29+
]
30+
2531
server = [
2632
"fastapi>=0.104.0",
2733
"uvicorn>=0.24.0",

tests/unittests/utils/test_control_api.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ def test_get_devs_client_with_read_timeout(self, mock_client_class):
282282
class TestControlAPIGetBailianClient:
283283
"""测试 ControlAPI._get_bailian_client"""
284284

285-
@patch("agentrun.utils.control_api.BailianClient")
285+
@patch("alibabacloud_bailian20231229.client.Client")
286286
def test_get_bailian_client_basic(self, mock_client_class):
287287
"""测试获取基本百炼客户端"""
288288
config = Config(
@@ -304,7 +304,7 @@ def test_get_bailian_client_basic(self, mock_client_class):
304304
assert config_arg.access_key_secret == "sk"
305305
assert config_arg.region_id == "cn-hangzhou"
306306

307-
@patch("agentrun.utils.control_api.BailianClient")
307+
@patch("alibabacloud_bailian20231229.client.Client")
308308
def test_get_bailian_client_strips_https_prefix(self, mock_client_class):
309309
"""测试获取百炼客户端时去除 https:// 前缀"""
310310
config = Config(
@@ -323,7 +323,7 @@ def test_get_bailian_client_strips_https_prefix(self, mock_client_class):
323323
config_arg = call_args[0][0]
324324
assert config_arg.endpoint == "bailian.cn-hangzhou.aliyuncs.com"
325325

326-
@patch("agentrun.utils.control_api.BailianClient")
326+
@patch("alibabacloud_bailian20231229.client.Client")
327327
def test_get_bailian_client_strips_http_prefix(self, mock_client_class):
328328
"""测试获取百炼客户端时去除 http:// 前缀"""
329329
config = Config(
@@ -346,7 +346,7 @@ def test_get_bailian_client_strips_http_prefix(self, mock_client_class):
346346
class TestControlAPIGetGPDBClient:
347347
"""测试 ControlAPI._get_gpdb_client"""
348348

349-
@patch("agentrun.utils.control_api.GPDBClient")
349+
@patch("alibabacloud_gpdb20160503.client.Client")
350350
def test_get_gpdb_client_known_region(self, mock_client_class):
351351
"""测试已知 region 使用通用 endpoint"""
352352
config = Config(
@@ -365,7 +365,7 @@ def test_get_gpdb_client_known_region(self, mock_client_class):
365365
config_arg = call_args[0][0]
366366
assert config_arg.endpoint == "gpdb.aliyuncs.com"
367367

368-
@patch("agentrun.utils.control_api.GPDBClient")
368+
@patch("alibabacloud_gpdb20160503.client.Client")
369369
def test_get_gpdb_client_unknown_region(self, mock_client_class):
370370
"""测试未知 region 使用区域级别 endpoint"""
371371
config = Config(
@@ -384,7 +384,7 @@ def test_get_gpdb_client_unknown_region(self, mock_client_class):
384384
config_arg = call_args[0][0]
385385
assert config_arg.endpoint == "gpdb.us-west-1.aliyuncs.com"
386386

387-
@patch("agentrun.utils.control_api.GPDBClient")
387+
@patch("alibabacloud_gpdb20160503.client.Client")
388388
def test_get_gpdb_client_all_known_regions(self, mock_client_class):
389389
"""测试所有已知 region 使用通用 endpoint"""
390390
known_regions = [

0 commit comments

Comments
 (0)