Leon/docutranslate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复token计数失败导致中断的问题

Browse Source
This commit is contained in:
xunbu
2025-09-23 23:21:47 +08:00
parent 612d4de26a
commit 75e85e9639
2 changed files with 36 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
docutranslate/agents/agent.py
View File

@@ -107,39 +107,41 @@ def extract_token_info(response_data: dict) -> tuple[int, int, int, int]:
# 初始化token详细统计 # 初始化token详细统计
cached_tokens = 0 cached_tokens = 0
reasoning_tokens = 0 reasoning_tokens = 0
try:
# 尝试从不同格式获取cached_tokens
# 格式1: input_tokens_details.cached_tokens
if (
"input_tokens_details" in usage
and "cached_tokens" in usage["input_tokens_details"]
):
cached_tokens = usage["input_tokens_details"]["cached_tokens"]
# 格式2: prompt_tokens_details.cached_tokens
elif (
"prompt_tokens_details" in usage
and "cached_tokens" in usage["prompt_tokens_details"]
):
cached_tokens = usage["prompt_tokens_details"]["cached_tokens"]
# 格式3: prompt_cache_hit_tokens (直接在usage下)
elif "prompt_cache_hit_tokens" in usage:
cached_tokens = usage["prompt_cache_hit_tokens"]
# 尝试从不同格式获取cached_tokens # 尝试从不同格式获取reasoning_tokens
# 格式1: input_tokens_details.cached_tokens # 格式1: output_tokens_details.reasoning_tokens
if ( if (
"input_tokens_details" in usage "output_tokens_details" in usage
and "cached_tokens" in usage["input_tokens_details"] and "reasoning_tokens" in usage["output_tokens_details"]
): ):
cached_tokens = usage["input_tokens_details"]["cached_tokens"] reasoning_tokens = usage["output_tokens_details"]["reasoning_tokens"]
# 格式2: prompt_tokens_details.cached_tokens # 格式2: completion_tokens_details.reasoning_tokens
elif ( elif (
"prompt_tokens_details" in usage "completion_tokens_details" in usage
and "cached_tokens" in usage["prompt_tokens_details"] and "reasoning_tokens" in usage["completion_tokens_details"]
): ):
cached_tokens = usage["prompt_tokens_details"]["cached_tokens"] reasoning_tokens = usage["completion_tokens_details"]["reasoning_tokens"]
# 格式3: prompt_cache_hit_tokens (直接在usage下) return input_tokens, cached_tokens, output_tokens, reasoning_tokens
elif "prompt_cache_hit_tokens" in usage: except TypeError as e:
cached_tokens = usage["prompt_cache_hit_tokens"] print(f"获取token发生错误:{e.__repr__()}")
return -1,-1,-1,-1
# 尝试从不同格式获取reasoning_tokens
# 格式1: output_tokens_details.reasoning_tokens
if (
"output_tokens_details" in usage
and "reasoning_tokens" in usage["output_tokens_details"]
):
reasoning_tokens = usage["output_tokens_details"]["reasoning_tokens"]
# 格式2: completion_tokens_details.reasoning_tokens
elif (
"completion_tokens_details" in usage
and "reasoning_tokens" in usage["completion_tokens_details"]
):
reasoning_tokens = usage["completion_tokens_details"]["reasoning_tokens"]
return input_tokens, cached_tokens, output_tokens, reasoning_tokens
class TokenCounter: class TokenCounter:

2
pyproject.toml
View File

@@ -26,12 +26,14 @@ dynamic = ["version"]
docling = [ docling = [
"opencv-python>=4.11.0.86", "opencv-python>=4.11.0.86",
"docling>=2.40.0", "docling>=2.40.0",
"hf-xet>=1.1.10"
] ]
[dependency-groups] [dependency-groups]
dev = [ dev = [
"opencv-python>=4.11.0.86", "opencv-python>=4.11.0.86",
"docling>=2.40.0", "docling>=2.40.0",
"hf-xet>=1.1.10"
] ]
[project.scripts] [project.scripts]
docutranslate = "docutranslate.cli:main" docutranslate = "docutranslate.cli:main"