From fca945c77f5afa14fe66ce982e158efdba3d3820 Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Thu, 29 May 2025 22:58:29 +0800
Subject: [PATCH] fix: default cost to NaN when calculation fails in LiteLLM
 backend (#912)

* fix: default cost to NaN when calculation fails in LiteLLM backend

* lint

* lint
---
 rdagent/oai/backend/deprec.py  | 2 ++
 rdagent/oai/backend/litellm.py | 2 ++
 rdagent/utils/env.py           | 2 +-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/rdagent/oai/backend/deprec.py b/rdagent/oai/backend/deprec.py
index 96bcf496..79f5f22e 100644
--- a/rdagent/oai/backend/deprec.py
+++ b/rdagent/oai/backend/deprec.py
@@ -17,6 +17,7 @@ from typing import Any, Optional, cast
 import numpy as np
 import openai
 import tiktoken
+from openai.types.chat import ChatCompletion
 
 from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass, import_class
 from rdagent.log import LogColors
@@ -385,6 +386,7 @@ class DeprecBackend(APIBackend):
                     if len(chunk.choices) > 0 and chunk.choices[0].finish_reason is not None:
                         finish_reason = chunk.choices[0].finish_reason
             else:
+                response = cast(ChatCompletion, response)
                 resp = response.choices[0].message.content
                 finish_reason = response.choices[0].finish_reason
                 if LLM_SETTINGS.log_llm_chat_content:
diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py
index 58a6d0c7..0a65bed0 100644
--- a/rdagent/oai/backend/litellm.py
+++ b/rdagent/oai/backend/litellm.py
@@ -1,5 +1,6 @@
 from typing import Any, Literal, cast
 
+import numpy as np
 from litellm import (
     completion,
     completion_cost,
@@ -139,6 +140,7 @@ class LiteLLMAPIBackend(APIBackend):
             cost = completion_cost(model=model, messages=messages, completion=content)
         except Exception as e:
             logger.warning(f"Cost calculation failed for model {model}: {e}. Skip cost statistics.")
+            cost = np.nan
         else:
             ACC_COST += cost
             logger.info(
diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index 1f3192c8..6e6e8dba 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -826,7 +826,7 @@ class QTDockerEnv(DockerEnv):
     def __init__(self, conf: DockerConf = QlibDockerConf()):
         super().__init__(conf)
 
-    def prepare(self, *args, **kwargs) -> None:  # type: ignore[explicit-override, no-untyped-def]
+    def prepare(self, *args, **kwargs) -> None:  # type: ignore[no-untyped-def]
         """
         Download image & data if it doesn't exist
         """