simplify code

This commit is contained in:
better629 2025-03-01 18:00:30 +08:00
parent 119385c5ce
commit cc774db7b3
4 changed files with 19 additions and 17 deletions

View file

@ -79,7 +79,7 @@ class OpenAILLM(BaseLLM):
def _get_proxy_params(self) -> dict:
params = {}
if self.config.proxy:
params = {"proxies": self.config.proxy}
params = {"proxy": self.config.proxy}
if self.config.base_url:
params["base_url"] = self.config.base_url
@ -94,13 +94,16 @@ class OpenAILLM(BaseLLM):
collected_reasoning_messages = []
has_finished = False
async for chunk in response:
if hasattr(chunk.choices[0].delta, "reasoning_content"):
collected_reasoning_messages.append(chunk.choices[0].delta.reasoning_content) # for deepseek
if not chunk.choices:
continue
chunk_message = chunk.choices[0].delta.content or "" if chunk.choices else "" # extract the message
finish_reason = (
chunk.choices[0].finish_reason if chunk.choices and hasattr(chunk.choices[0], "finish_reason") else None
)
choice0 = chunk.choices[0]
choice_delta = choice0.delta
if hasattr(choice_delta, "reasoning_content") and choice_delta.reasoning_content:
collected_reasoning_messages.append(choice_delta.reasoning_content) # for deepseek
continue
chunk_message = choice_delta.content or "" # extract the message
finish_reason = choice0.finish_reason if hasattr(choice0, "finish_reason") else None
log_llm_stream(chunk_message)
collected_messages.append(chunk_message)
chunk_has_usage = hasattr(chunk, "usage") and chunk.usage
@ -111,13 +114,10 @@ class OpenAILLM(BaseLLM):
if finish_reason:
if chunk_has_usage:
# Some services have usage as an attribute of the chunk, such as Fireworks
if isinstance(chunk.usage, CompletionUsage):
usage = chunk.usage
else:
usage = CompletionUsage(**chunk.usage)
elif hasattr(chunk.choices[0], "usage"):
usage = CompletionUsage(**chunk.usage) if isinstance(chunk.usage, dict) else chunk.usage
elif hasattr(choice0, "usage"):
# The usage of some services is an attribute of chunk.choices[0], such as Moonshot
usage = CompletionUsage(**chunk.choices[0].usage)
usage = CompletionUsage(**choice0.usage)
has_finished = True
log_llm_stream("\n")

View file

@ -144,6 +144,6 @@ class FireworksCostManager(CostManager):
cost = (prompt_tokens * token_costs["prompt"] + completion_tokens * token_costs["completion"]) / 1000000
self.total_cost += cost
logger.info(
f"Total running cost: ${self.total_cost:.4f}"
f"Total running cost: ${self.total_cost:.4f}, "
f"Current cost: ${cost:.4f}, prompt_tokens: {prompt_tokens}, completion_tokens: {completion_tokens}"
)