From 6c179442bfe59b465ba2c2fb4c4aee1df1967c2b Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Fri, 25 Apr 2025 19:26:08 +0800
Subject: [PATCH] chore: log cost object (#829)

---
 rdagent/oai/backend/litellm.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py
index 97be2207..66273af7 100644
--- a/rdagent/oai/backend/litellm.py
+++ b/rdagent/oai/backend/litellm.py
@@ -139,4 +139,16 @@ class LiteLLMAPIBackend(APIBackend):
         logger.info(
             f"Current Cost: ${float(cost):.10f}; Accumulated Cost: ${float(ACC_COST):.10f}; {finish_reason=}",
         )
+        prompt_tokens = token_counter(model=model, messages=messages)
+        completion_tokens = token_counter(model=model, text=content)
+        logger.log_object(
+            {
+                "model": model,
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "cost": cost,
+                "accumulated_cost": ACC_COST,
+            },
+            tag="token_cost",
+        )
         return content, finish_reason