Skip to content

Commit 54bb23b

Browse files
CopilotOhYee
andcommitted
refactor(core): use generic package mapping for optional dependency checks
Refactor the import error handling to use a more maintainable and scalable approach with a package mapping structure instead of hardcoded string checks. Changes: - Added _OPTIONAL_PACKAGES list with (import_name, install_name) tuples - Refactored __getattr__ to iterate through the package mapping - Now easier to add new optional dependencies in the future - All three server optional dependencies (fastapi, uvicorn, ag_ui) registered This makes the codebase more maintainable when adding new optional dependencies in the future. Addresses feedback from @OhYee on PR #23. Co-authored-by: OhYee <[email protected]>
1 parent 5c01863 commit 54bb23b

File tree

1 file changed

+16
-6
lines changed

1 file changed

+16
-6
lines changed

agentrun/__init__.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,14 @@
289289
"MergeOptions",
290290
}
291291

292+
# 可选依赖包映射:导入错误的包名 -> 安装使用的包名
293+
# Optional dependency mapping: import error package name -> installation package name
294+
_OPTIONAL_PACKAGES = [
295+
("fastapi", "agentrun-sdk[server]"),
296+
("uvicorn", "agentrun-sdk[server]"),
297+
("ag_ui", "agentrun-sdk[server]"),
298+
]
299+
292300

293301
def __getattr__(name: str):
294302
"""延迟加载 server 模块的导出,避免可选依赖导致导入失败
@@ -303,12 +311,14 @@ def __getattr__(name: str):
303311
return getattr(server, name)
304312
except ImportError as e:
305313
# 检查是否是缺少可选依赖导致的错误
306-
if "fastapi" in str(e) or "uvicorn" in str(e) or "ag_ui" in str(e):
307-
raise ImportError(
308-
f"'{name}' requires the 'server' optional dependencies. "
309-
"Install with: pip install agentrun-sdk[server]\n"
310-
f"Original error: {e}"
311-
) from e
314+
error_str = str(e)
315+
for package_name, install_name in _OPTIONAL_PACKAGES:
316+
if package_name in error_str:
317+
raise ImportError(
318+
f"'{name}' requires optional dependencies. "
319+
f"Install with: pip install {install_name}\n"
320+
f"Original error: {e}"
321+
) from e
312322
# 其他导入错误继续抛出
313323
raise
314324

0 commit comments

Comments
 (0)