vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Keep track of token counts in steps table (#412)

Browse Source
This commit is contained in:
Kerem Yilmaz
2024-06-03 15:55:34 -07:00
committed by GitHub
parent 4603894119
commit 04c6e55848
2 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
skyvern/forge/sdk/db/client.py
View File

@@ -293,6 +293,8 @@ class AgentDB:
retry_index: int | None = None,
organization_id: str | None = None,
incremental_cost: float | None = None,
incremental_input_tokens: int | None = None,
incremental_output_tokens: int | None = None,
) -> Step:
try:
async with self.Session() as session:
@@ -314,6 +316,10 @@ class AgentDB:
step.retry_index = retry_index
if incremental_cost is not None:
step.step_cost = incremental_cost + float(step.step_cost or 0)
if incremental_input_tokens is not None:
step.input_token_count = incremental_input_tokens + (step.input_token_count or 0)
if incremental_output_tokens is not None:
step.output_token_count = incremental_output_tokens + (step.output_token_count or 0)
await session.commit()
updated_step = await self.get_step(task_id, step_id, organization_id)