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

track reasoning token and cached token (#1985)

Browse Source
This commit is contained in:
Shuchang Zheng
2025-03-20 16:42:57 -07:00
committed by GitHub
parent 185464f8ec
commit eb3eb4eede
9 changed files with 112 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -492,6 +492,8 @@ class AgentDB:
incremental_cost: float | None = None,
incremental_input_tokens: int | None = None,
incremental_output_tokens: int | None = None,
incremental_reasoning_tokens: int | None = None,
incremental_cached_tokens: int | None = None,
) -> Step:
try:
async with self.Session() as session:
@@ -517,6 +519,10 @@ class AgentDB:
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)
if incremental_reasoning_tokens is not None:
step.reasoning_token_count = incremental_reasoning_tokens + (step.reasoning_token_count or 0)
if incremental_cached_tokens is not None:
step.cached_token_count = incremental_cached_tokens + (step.cached_token_count or 0)
await session.commit()
updated_step = await self.get_step(task_id, step_id, organization_id)
@@ -2290,6 +2296,8 @@ class AgentDB:
output: dict[str, Any] | None = None,
input_token_count: int | None = None,
output_token_count: int | None = None,
reasoning_token_count: int | None = None,
cached_token_count: int | None = None,
thought_cost: float | None = None,
organization_id: str | None = None,
) -> Thought:
@@ -2309,6 +2317,8 @@ class AgentDB:
output=output,
input_token_count=input_token_count,
output_token_count=output_token_count,
reasoning_token_count=reasoning_token_count,
cached_token_count=cached_token_count,
thought_cost=thought_cost,
organization_id=organization_id,
)
@@ -2330,6 +2340,8 @@ class AgentDB:
output: dict[str, Any] | None = None,
input_token_count: int | None = None,
output_token_count: int | None = None,
reasoning_token_count: int | None = None,
cached_token_count: int | None = None,
thought_cost: float | None = None,
organization_id: str | None = None,
) -> Thought:
@@ -2362,6 +2374,10 @@ class AgentDB:
thought_obj.input_token_count = input_token_count
if output_token_count:
thought_obj.output_token_count = output_token_count
if reasoning_token_count:
thought_obj.reasoning_token_count = reasoning_token_count
if cached_token_count:
thought_obj.cached_token_count = cached_token_count
if thought_cost:
thought_obj.thought_cost = thought_cost
await session.commit()

4
skyvern/forge/sdk/db/models.py
View File

@@ -117,6 +117,8 @@ class StepModel(Base):
)
input_token_count = Column(Integer, default=0)
output_token_count = Column(Integer, default=0)
reasoning_token_count = Column(Integer, default=0)
cached_token_count = Column(Integer, default=0)
step_cost = Column(Numeric, default=0)
@@ -612,6 +614,8 @@ class ThoughtModel(Base):
answer = Column(String, nullable=True)
input_token_count = Column(Integer, nullable=True)
output_token_count = Column(Integer, nullable=True)
reasoning_token_count = Column(Integer, nullable=True)
cached_token_count = Column(Integer, nullable=True)
thought_cost = Column(Numeric, nullable=True)
observer_thought_type = Column(String, nullable=True, default=ThoughtType.plan)

2
skyvern/forge/sdk/db/utils.py
View File

@@ -108,6 +108,8 @@ def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step:
organization_id=step_model.organization_id,
input_token_count=step_model.input_token_count,
output_token_count=step_model.output_token_count,
reasoning_token_count=step_model.reasoning_token_count,
cached_token_count=step_model.cached_token_count,
step_cost=step_model.step_cost,
)