chroe(2585): keep polishing the code of codex token usage (#2689)
This commit is contained in:
@@ -21,13 +21,40 @@ from langchain_core.callbacks import CallbackManagerForLLMRun
|
|||||||
from langchain_core.language_models.chat_models import BaseChatModel
|
from langchain_core.language_models.chat_models import BaseChatModel
|
||||||
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
|
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
|
||||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||||
from langchain_openai.chat_models.base import _create_usage_metadata_responses
|
|
||||||
|
|
||||||
from deerflow.models.credential_loader import CodexCliCredential, load_codex_cli_credential
|
from deerflow.models.credential_loader import CodexCliCredential, load_codex_cli_credential
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
||||||
|
|
||||||
|
|
||||||
|
def _build_usage_metadata(oai_usage: dict) -> dict:
|
||||||
|
"""Convert Codex/Responses API usage dict to LangChain usage_metadata format.
|
||||||
|
|
||||||
|
Maps OpenAI Responses API token usage fields to the dict structure that
|
||||||
|
LangChain AIMessage.usage_metadata expects. This avoids depending on
|
||||||
|
langchain_openai private helpers like ``_create_usage_metadata_responses``.
|
||||||
|
"""
|
||||||
|
input_tokens = oai_usage.get("input_tokens", 0)
|
||||||
|
output_tokens = oai_usage.get("output_tokens", 0)
|
||||||
|
total_tokens = oai_usage.get("total_tokens", input_tokens + output_tokens)
|
||||||
|
metadata: dict = {
|
||||||
|
"input_tokens": input_tokens,
|
||||||
|
"output_tokens": output_tokens,
|
||||||
|
"total_tokens": total_tokens,
|
||||||
|
}
|
||||||
|
input_details = oai_usage.get("input_tokens_details") or {}
|
||||||
|
output_details = oai_usage.get("output_tokens_details") or {}
|
||||||
|
cache_read = input_details.get("cached_tokens")
|
||||||
|
if cache_read is not None:
|
||||||
|
metadata["input_token_details"] = {"cache_read": cache_read}
|
||||||
|
reasoning = output_details.get("reasoning_tokens")
|
||||||
|
if reasoning is not None:
|
||||||
|
metadata["output_token_details"] = {"reasoning": reasoning}
|
||||||
|
return metadata
|
||||||
|
|
||||||
|
|
||||||
MAX_RETRIES = 3
|
MAX_RETRIES = 3
|
||||||
|
|
||||||
|
|
||||||
@@ -347,7 +374,7 @@ class CodexChatModel(BaseChatModel):
|
|||||||
)
|
)
|
||||||
|
|
||||||
usage = response.get("usage", {})
|
usage = response.get("usage", {})
|
||||||
usage_metadata = _create_usage_metadata_responses(usage) if usage else None
|
usage_metadata = _build_usage_metadata(usage) if usage else None
|
||||||
additional_kwargs = {}
|
additional_kwargs = {}
|
||||||
if reasoning_content:
|
if reasoning_content:
|
||||||
additional_kwargs["reasoning_content"] = reasoning_content
|
additional_kwargs["reasoning_content"] = reasoning_content
|
||||||
|
|||||||
@@ -103,13 +103,13 @@ def test_parse_response_populates_usage_metadata():
|
|||||||
|
|
||||||
result = model._parse_response(response)
|
result = model._parse_response(response)
|
||||||
|
|
||||||
assert result.generations[0].message.usage_metadata == {
|
meta = result.generations[0].message.usage_metadata
|
||||||
"input_tokens": 10,
|
assert meta is not None
|
||||||
"output_tokens": 5,
|
assert meta["input_tokens"] == 10
|
||||||
"total_tokens": 15,
|
assert meta["output_tokens"] == 5
|
||||||
"input_token_details": {"cache_read": 3},
|
assert meta["total_tokens"] == 15
|
||||||
"output_token_details": {"reasoning": 2},
|
assert meta["input_token_details"]["cache_read"] == 3
|
||||||
}
|
assert meta["output_token_details"]["reasoning"] == 2
|
||||||
|
|
||||||
|
|
||||||
def test_parse_response_reasoning_content():
|
def test_parse_response_reasoning_content():
|
||||||
|
|||||||
Reference in New Issue
Block a user