vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2026-05-01 21:20:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

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
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 14 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
skyvern/forge/sdk/api/llm/api_handler_factory.py
View file

@ -115,11 +115,15 @@ class LLMAPIHandlerFactory:
data=response.model_dump_json(indent=2).encode("utf-8"),
)
llm_cost = litellm.completion_cost(completion_response=response)
prompt_tokens = response.get("usage", {}).get("prompt_tokens", 0)
completion_tokens = response.get("usage", {}).get("completion_tokens", 0)
await app.DATABASE.update_step(
task_id=step.task_id,
step_id=step.step_id,
organization_id=step.organization_id,
incremental_cost=llm_cost,
incremental_input_tokens=prompt_tokens if prompt_tokens > 0 else None,
incremental_output_tokens=completion_tokens if completion_tokens > 0 else None,
)
parsed_response = parse_api_response(response, llm_config.add_assistant_prefix)
if step:
@ -206,11 +210,15 @@ class LLMAPIHandlerFactory:
data=response.model_dump_json(indent=2).encode("utf-8"),
)
llm_cost = litellm.completion_cost(completion_response=response)
prompt_tokens = response.get("usage", {}).get("prompt_tokens", 0)
completion_tokens = response.get("usage", {}).get("completion_tokens", 0)
await app.DATABASE.update_step(
task_id=step.task_id,
step_id=step.step_id,
organization_id=step.organization_id,
incremental_cost=llm_cost,
incremental_input_tokens=prompt_tokens if prompt_tokens > 0 else None,
incremental_output_tokens=completion_tokens if completion_tokens > 0 else None,
)
parsed_response = parse_api_response(response, llm_config.add_assistant_prefix)
if step: