Skip to content

Commit 0e1d3f6

Browse files
SodawyxSodawyx
committed
fix(sandbox): optimize connect sandbox template type
Signed-off-by: Sodawyx <sodawyx@126.com>
1 parent 54bea1d commit 0e1d3f6

File tree

2 files changed

+41
-62
lines changed

2 files changed

+41
-62
lines changed

‎agentrun/sandbox/__sandbox_async_template.py‎

Lines changed: 14 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -308,7 +308,9 @@ async def connect_async(
308308
309309
Args:
310310
sandbox_id: Sandbox ID
311-
type: 可选的类型参数，用于类型提示和运行时验证
311+
template_type: 可选的类型参数，用于类型提示和运行时类型决定。
312+
提供时直接使用该类型决定返回的子类，不调用 get_template（无需 AKSK）。
313+
未提供时通过 get_template 获取类型（需要 AKSK）。
312314
config: 配置对象
313315
314316
Returns:
@@ -325,47 +327,38 @@ async def connect_async(
325327
sandbox_id, config=config
326328
)
327329

328-
# 根据 template_name 获取 template 类型
329-
if sandbox.template_name is None:
330-
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
330+
resolved_type = template_type
331+
if resolved_type is None:
332+
if sandbox.template_name is None:
333+
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
331334

332-
template = await cls.get_template_async(
333-
sandbox.template_name, config=config
334-
)
335-
336-
# 如果提供了 type 参数，验证类型是否匹配
337-
if (
338-
template_type is not None
339-
and template.template_type != template_type
340-
):
341-
raise ValueError(
342-
f"Sandbox {sandbox_id} has template type"
343-
f" {template.template_type}, but expected {template_type}"
335+
template = await cls.get_template_async(
336+
sandbox.template_name, config=config
344337
)
338+
resolved_type = template.template_type
345339

346-
# 根据 template 类型创建相应的 Sandbox 子类
347340
from agentrun.sandbox.aio_sandbox import AioSandbox
348341
from agentrun.sandbox.browser_sandbox import BrowserSandbox
349342
from agentrun.sandbox.code_interpreter_sandbox import (
350343
CodeInterpreterSandbox,
351344
)
352345

353346
result = None
354-
if template.template_type == TemplateType.CODE_INTERPRETER:
347+
if resolved_type == TemplateType.CODE_INTERPRETER:
355348
result = CodeInterpreterSandbox.model_validate(
356349
sandbox.model_dump(by_alias=False)
357350
)
358-
elif template.template_type == TemplateType.BROWSER:
351+
elif resolved_type == TemplateType.BROWSER:
359352
result = BrowserSandbox.model_validate(
360353
sandbox.model_dump(by_alias=False)
361354
)
362-
elif template.template_type == TemplateType.AIO:
355+
elif resolved_type == TemplateType.AIO:
363356
result = AioSandbox.model_validate(
364357
sandbox.model_dump(by_alias=False)
365358
)
366359
else:
367360
raise ValueError(
368-
f"Unsupported template type: {template.template_type}. "
361+
f"Unsupported template type: {resolved_type}. "
369362
"Expected 'code-interpreter', 'browser' or 'aio'"
370363
)
371364

‎agentrun/sandbox/sandbox.py‎

Lines changed: 27 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -535,7 +535,9 @@ async def connect_async(
535535
536536
Args:
537537
sandbox_id: Sandbox ID
538-
type: 可选的类型参数，用于类型提示和运行时验证
538+
template_type: 可选的类型参数，用于类型提示和运行时类型决定。
539+
提供时直接使用该类型决定返回的子类，不调用 get_template（无需 AKSK）。
540+
未提供时通过 get_template 获取类型（需要 AKSK）。
539541
config: 配置对象
540542
541543
Returns:
@@ -552,47 +554,38 @@ async def connect_async(
552554
sandbox_id, config=config
553555
)
554556

555-
# 根据 template_name 获取 template 类型
556-
if sandbox.template_name is None:
557-
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
557+
resolved_type = template_type
558+
if resolved_type is None:
559+
if sandbox.template_name is None:
560+
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
558561

559-
template = await cls.get_template_async(
560-
sandbox.template_name, config=config
561-
)
562-
563-
# 如果提供了 type 参数，验证类型是否匹配
564-
if (
565-
template_type is not None
566-
and template.template_type != template_type
567-
):
568-
raise ValueError(
569-
f"Sandbox {sandbox_id} has template type"
570-
f" {template.template_type}, but expected {template_type}"
562+
template = await cls.get_template_async(
563+
sandbox.template_name, config=config
571564
)
565+
resolved_type = template.template_type
572566

573-
# 根据 template 类型创建相应的 Sandbox 子类
574567
from agentrun.sandbox.aio_sandbox import AioSandbox
575568
from agentrun.sandbox.browser_sandbox import BrowserSandbox
576569
from agentrun.sandbox.code_interpreter_sandbox import (
577570
CodeInterpreterSandbox,
578571
)
579572

580573
result = None
581-
if template.template_type == TemplateType.CODE_INTERPRETER:
574+
if resolved_type == TemplateType.CODE_INTERPRETER:
582575
result = CodeInterpreterSandbox.model_validate(
583576
sandbox.model_dump(by_alias=False)
584577
)
585-
elif template.template_type == TemplateType.BROWSER:
578+
elif resolved_type == TemplateType.BROWSER:
586579
result = BrowserSandbox.model_validate(
587580
sandbox.model_dump(by_alias=False)
588581
)
589-
elif template.template_type == TemplateType.AIO:
582+
elif resolved_type == TemplateType.AIO:
590583
result = AioSandbox.model_validate(
591584
sandbox.model_dump(by_alias=False)
592585
)
593586
else:
594587
raise ValueError(
595-
f"Unsupported template type: {template.template_type}. "
588+
f"Unsupported template type: {resolved_type}. "
596589
"Expected 'code-interpreter', 'browser' or 'aio'"
597590
)
598591

@@ -612,7 +605,9 @@ def connect(
612605
613606
Args:
614607
sandbox_id: Sandbox ID
615-
type: 可选的类型参数，用于类型提示和运行时验证
608+
template_type: 可选的类型参数，用于类型提示和运行时类型决定。
609+
提供时直接使用该类型决定返回的子类，不调用 get_template（无需 AKSK）。
610+
未提供时通过 get_template 获取类型（需要 AKSK）。
616611
config: 配置对象
617612
618613
Returns:
@@ -627,45 +622,36 @@ def connect(
627622
# 先获取 sandbox 信息
628623
sandbox = cls.__get_client().get_sandbox(sandbox_id, config=config)
629624

630-
# 根据 template_name 获取 template 类型
631-
if sandbox.template_name is None:
632-
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
625+
resolved_type = template_type
626+
if resolved_type is None:
627+
if sandbox.template_name is None:
628+
raise ValueError(f"Sandbox {sandbox_id} has no template_name")
633629

634-
template = cls.get_template(sandbox.template_name, config=config)
630+
template = cls.get_template(sandbox.template_name, config=config)
631+
resolved_type = template.template_type
635632

636-
# 如果提供了 type 参数，验证类型是否匹配
637-
if (
638-
template_type is not None
639-
and template.template_type != template_type
640-
):
641-
raise ValueError(
642-
f"Sandbox {sandbox_id} has template type"
643-
f" {template.template_type}, but expected {template_type}"
644-
)
645-
646-
# 根据 template 类型创建相应的 Sandbox 子类
647633
from agentrun.sandbox.aio_sandbox import AioSandbox
648634
from agentrun.sandbox.browser_sandbox import BrowserSandbox
649635
from agentrun.sandbox.code_interpreter_sandbox import (
650636
CodeInterpreterSandbox,
651637
)
652638

653639
result = None
654-
if template.template_type == TemplateType.CODE_INTERPRETER:
640+
if resolved_type == TemplateType.CODE_INTERPRETER:
655641
result = CodeInterpreterSandbox.model_validate(
656642
sandbox.model_dump(by_alias=False)
657643
)
658-
elif template.template_type == TemplateType.BROWSER:
644+
elif resolved_type == TemplateType.BROWSER:
659645
result = BrowserSandbox.model_validate(
660646
sandbox.model_dump(by_alias=False)
661647
)
662-
elif template.template_type == TemplateType.AIO:
648+
elif resolved_type == TemplateType.AIO:
663649
result = AioSandbox.model_validate(
664650
sandbox.model_dump(by_alias=False)
665651
)
666652
else:
667653
raise ValueError(
668-
f"Unsupported template type: {template.template_type}. "
654+
f"Unsupported template type: {resolved_type}. "
669655
"Expected 'code-interpreter', 'browser' or 'aio'"
670656
)
671657

0 commit comments

Comments
 (0)