From aa8ceed2faa6e7303cce04c748e4183addc09cf2 Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Fri, 28 Mar 2025 12:24:19 +0800
Subject: [PATCH] feat: track and log accumulated completion cost in
 LiteLLMAPIBackend (#727)

* feat: Track and log accumulated completion cost in LiteLLMAPIBackend

* refactor: Use variable for retry count and update import formatting

* lint

* fix: Allow chat_max_tokens to be None in LLMSettings

* feat: Add logging for LiteLLM settings and finish reason in cost calculation
---
 rdagent/oai/backend/base.py    |  5 +++--
 rdagent/oai/backend/litellm.py | 16 +++++++++++++++-
 rdagent/oai/llm_conf.py        |  2 +-
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py
index cb81fca0..139f67ba 100644
--- a/rdagent/oai/backend/base.py
+++ b/rdagent/oai/backend/base.py
@@ -391,7 +391,8 @@ class APIBackend(ABC):
 
         all_response = ""
         new_messages = deepcopy(messages)
-        for _ in range(6):  # for some long code, 3 times may not enough for reasoning models
+        try_n = 6
+        for _ in range(try_n):  # for some long code, 3 times may not enough for reasoning models
             if "json_mode" in kwargs:
                 del kwargs["json_mode"]
             response, finish_reason = self._create_chat_completion_add_json_in_prompt(
@@ -412,7 +413,7 @@ class APIBackend(ABC):
                     self.cache.chat_set(input_content_json, all_response)
                 return all_response
             new_messages.append({"role": "assistant", "content": response})
-        raise RuntimeError("Failed to continue the conversation after 3 retries.")
+        raise RuntimeError(f"Failed to continue the conversation after {try_n} retries.")
 
     def _create_embedding_with_cache(
         self, input_content_list: list[str], *args: Any, **kwargs: Any
diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py
index 1689ed73..893a654b 100644
--- a/rdagent/oai/backend/litellm.py
+++ b/rdagent/oai/backend/litellm.py
@@ -1,6 +1,12 @@
 from typing import Any
 
-from litellm import completion, embedding, supports_response_schema, token_counter
+from litellm import (
+    completion,
+    completion_cost,
+    embedding,
+    supports_response_schema,
+    token_counter,
+)
 
 from rdagent.log import LogColors
 from rdagent.log import rdagent_logger as logger
@@ -18,6 +24,8 @@ class LiteLLMSettings(LLMSettings):
 
 
 LITELLM_SETTINGS = LiteLLMSettings()
+logger.info(f"{LITELLM_SETTINGS}")
+ACC_COST = 0.0
 
 
 class LiteLLMAPIBackend(APIBackend):
@@ -111,4 +119,10 @@ class LiteLLMAPIBackend(APIBackend):
             )
             logger.info(f"{LogColors.BLUE}assistant:{LogColors.END} {finish_reason_str}\n{content}", tag="llm_messages")
 
+        global ACC_COST
+        cost = completion_cost(model=LITELLM_SETTINGS.chat_model, messages=messages, completion=content)
+        ACC_COST += cost
+        logger.info(
+            f"Current Cost: ${float(cost):.10f}; Accumulated Cost: ${float(ACC_COST):.10f}; {finish_reason=}",
+        )
         return content, finish_reason
diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py
index bbb13665..fcd6e350 100644
--- a/rdagent/oai/llm_conf.py
+++ b/rdagent/oai/llm_conf.py
@@ -51,7 +51,7 @@ class LLMSettings(ExtendedBaseSettings):
     chat_openai_base_url: str | None = None  #
     chat_azure_api_base: str = ""
     chat_azure_api_version: str = ""
-    chat_max_tokens: int = 3000
+    chat_max_tokens: int | None = None
     chat_temperature: float = 0.5
     chat_stream: bool = True
     chat_seed: int | None = None