cache task run uses block level model override when ai fallback happens (#4073)

This commit is contained in:
Shuchang Zheng
2025-11-21 22:48:20 -08:00
committed by GitHub
parent b52982d3c8
commit 7729d7cffe
3 changed files with 130 additions and 19 deletions

View File

@@ -673,16 +673,30 @@ def _build_action_statement(
last_line=cst.SimpleWhitespace(INDENT),
),
),
cst.Arg(
keyword=cst.Name("label"),
value=_value(block_title),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
),
comma=cst.Comma(),
),
]
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("label"):
args.append(
cst.Arg(
keyword=cst.Name("label"),
value=_value(block.get("label")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
),
comma=cst.Comma(),
)
)
_mark_last_arg_as_comma(args)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("action")),
args=args,
@@ -733,15 +747,30 @@ def _build_extract_statement(
last_line=cst.SimpleWhitespace(INDENT),
),
),
cst.Arg(
keyword=cst.Name("label"),
value=_value(block_title),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
),
comma=cst.Comma(),
),
]
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("label"):
args.append(
cst.Arg(
keyword=cst.Name("label"),
value=_value(block_title),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
),
comma=cst.Comma(),
)
)
_mark_last_arg_as_comma(args)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("extract")),
@@ -882,6 +911,18 @@ def _build_validate_statement(
)
)
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
# Add label if it exists
if block.get("label") is not None:
args.append(
@@ -1114,6 +1155,18 @@ def _build_pdf_parser_statement(block: dict[str, Any]) -> cst.SimpleStatementLin
)
)
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("label") is not None:
args.append(
cst.Arg(
@@ -1172,6 +1225,18 @@ def _build_file_url_parser_statement(block: dict[str, Any]) -> cst.SimpleStateme
)
)
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("label") is not None:
args.append(
cst.Arg(
@@ -1319,6 +1384,18 @@ def _build_prompt_statement(block: dict[str, Any]) -> cst.SimpleStatementLine:
)
)
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("label") is not None:
args.append(
cst.Arg(
@@ -1331,7 +1408,7 @@ def _build_prompt_statement(block: dict[str, Any]) -> cst.SimpleStatementLine:
)
)
if block.get("parameters") is not None:
if block.get("parameters"):
parameters = block.get("parameters", [])
parameter_list = [parameter["key"] for parameter in parameters]
args.append(
@@ -1340,10 +1417,12 @@ def _build_prompt_statement(block: dict[str, Any]) -> cst.SimpleStatementLine:
value=_value(parameter_list),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
_mark_last_arg_as_comma(args)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("prompt")),
args=args,
@@ -1568,6 +1647,17 @@ def __build_base_task_statement(
),
)
)
if block.get("model"):
args.append(
cst.Arg(
keyword=cst.Name("model"),
value=_value(block.get("model")),
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
),
)
)
if block.get("block_type") == "task_v2":
args.append(
cst.Arg(