From 20d3e1d297195b995af7082b60ef56100c88caeb Mon Sep 17 00:00:00 2001
From: Pedro Gomes <113145167+PedroGomes02@users.noreply.github.com>
Date: Sun, 4 Aug 2024 07:42:22 +0100
Subject: [PATCH] Fix increase_usage of total_price in agent_runner (#6688)

---
 api/core/agent/cot_agent_runner.py | 1 +
 api/core/agent/fc_agent_runner.py  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py
index 9bd8f37d85..06492bb12f 100644
--- a/api/core/agent/cot_agent_runner.py
+++ b/api/core/agent/cot_agent_runner.py
@@ -79,6 +79,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
                 llm_usage.completion_tokens += usage.completion_tokens
                 llm_usage.prompt_price += usage.prompt_price
                 llm_usage.completion_price += usage.completion_price
+                llm_usage.total_price += usage.total_price
 
         model_instance = self.model_instance
 
diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py
index 7019b5e39f..3ee6e47742 100644
--- a/api/core/agent/fc_agent_runner.py
+++ b/api/core/agent/fc_agent_runner.py
@@ -62,6 +62,7 @@ class FunctionCallAgentRunner(BaseAgentRunner):
                 llm_usage.completion_tokens += usage.completion_tokens
                 llm_usage.prompt_price += usage.prompt_price
                 llm_usage.completion_price += usage.completion_price
+                llm_usage.total_price += usage.total_price
 
         model_instance = self.model_instance