add cua cost to task step (#2131)

This commit is contained in:
Shuchang Zheng
2025-04-11 17:25:43 -07:00
committed by GitHub
parent 92c7862a77
commit 438a49f470
4 changed files with 62 additions and 5 deletions

View File

@@ -1247,6 +1247,22 @@ class ForgeAgent:
truncation="auto",
)
previous_response = first_response
input_tokens = first_response.usage.input_tokens or 0
output_tokens = first_response.usage.output_tokens or 0
first_response.usage.total_tokens or 0
cached_tokens = first_response.usage.input_tokens_details.cached_tokens or 0
reasoning_tokens = first_response.usage.output_tokens_details.reasoning_tokens or 0
llm_cost = (3.0 / 1000000) * input_tokens + (12.0 / 1000000) * output_tokens
await app.DATABASE.update_step(
task_id=task.task_id,
step_id=step.step_id,
organization_id=task.organization_id,
incremental_cost=llm_cost,
incremental_input_tokens=input_tokens if input_tokens > 0 else None,
incremental_output_tokens=output_tokens if output_tokens > 0 else None,
incremental_reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else None,
incremental_cached_tokens=cached_tokens if cached_tokens > 0 else None,
)
computer_calls = [item for item in previous_response.output if item.type == "computer_call"]
if not computer_calls:
@@ -1284,6 +1300,22 @@ class ForgeAgent:
},
truncation="auto",
)
input_tokens = current_response.usage.input_tokens or 0
output_tokens = current_response.usage.output_tokens or 0
current_response.usage.total_tokens or 0
cached_tokens = current_response.usage.input_tokens_details.cached_tokens or 0
reasoning_tokens = current_response.usage.output_tokens_details.reasoning_tokens or 0
llm_cost = (3.0 / 1000000) * input_tokens + (12.0 / 1000000) * output_tokens
await app.DATABASE.update_step(
task_id=task.task_id,
step_id=step.step_id,
organization_id=task.organization_id,
incremental_cost=llm_cost,
incremental_input_tokens=input_tokens if input_tokens > 0 else None,
incremental_output_tokens=output_tokens if output_tokens > 0 else None,
incremental_reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else None,
incremental_cached_tokens=cached_tokens if cached_tokens > 0 else None,
)
return parse_cua_actions(task, step, current_response), current_response

View File

@@ -1503,17 +1503,20 @@ async def run_task(
request=request,
background_tasks=background_tasks,
)
run_type = RunType.task_v1
if run_request.engine == RunEngine.openai_cua:
run_type = RunType.openai_cua
# build the task run response
return TaskRunResponse(
run_id=task_v1_response.task_id,
run_type=RunType.task_v1,
run_type=run_type,
status=str(task_v1_response.status),
output=task_v1_response.extracted_information,
failure_reason=task_v1_response.failure_reason,
created_at=task_v1_response.created_at,
modified_at=task_v1_response.modified_at,
run_request=TaskRunRequest(
engine=RunEngine.skyvern_v1,
engine=run_request.engine,
prompt=task_v1_response.navigation_goal,
url=task_v1_response.url,
webhook_url=task_v1_response.webhook_callback_url,