mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-17 23:41:46 +00:00
Community : Add OpenAI prompt caching and reasoning tokens tracking (#27135)
Added Token tracking for OpenAI's prompt caching and reasoning tokens Costs updated from https://openai.com/api/pricing/ usage example ```python from langchain_community.callbacks import get_openai_callback from langchain_openai import ChatOpenAI llm = ChatOpenAI(model_name="o1-mini",temperature=1) with get_openai_callback() as cb: response = llm.invoke("hi "*1500) print(cb) ``` Output ``` Tokens Used: 1720 Prompt Tokens: 1508 Prompt Tokens Cached: 1408 Completion Tokens: 212 Reasoning Tokens: 192 Successful Requests: 1 Total Cost (USD): $0.0049559999999999995 ``` --------- Co-authored-by: Chester Curme <chester.curme@gmail.com>
This commit is contained in:
@@ -3,7 +3,8 @@ from uuid import uuid4
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from langchain_core.outputs import LLMResult
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.outputs import ChatGeneration, LLMResult
|
||||
from langchain_core.utils.pydantic import get_fields
|
||||
|
||||
from langchain_community.callbacks import OpenAICallbackHandler
|
||||
@@ -35,6 +36,43 @@ def test_on_llm_end(handler: OpenAICallbackHandler) -> None:
|
||||
assert handler.total_cost > 0
|
||||
|
||||
|
||||
def test_on_llm_end_with_chat_generation(handler: OpenAICallbackHandler) -> None:
|
||||
response = LLMResult(
|
||||
generations=[
|
||||
[
|
||||
ChatGeneration(
|
||||
text="Hello, world!",
|
||||
message=AIMessage(
|
||||
content="Hello, world!",
|
||||
usage_metadata={
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 4,
|
||||
"input_token_details": {
|
||||
"cache_read": 1,
|
||||
},
|
||||
"output_token_details": {
|
||||
"reasoning": 1,
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
]
|
||||
],
|
||||
llm_output={
|
||||
"model_name": get_fields(BaseOpenAI)["model_name"].default,
|
||||
},
|
||||
)
|
||||
handler.on_llm_end(response)
|
||||
assert handler.successful_requests == 1
|
||||
assert handler.total_tokens == 4
|
||||
assert handler.prompt_tokens == 2
|
||||
assert handler.prompt_tokens_cached == 1
|
||||
assert handler.completion_tokens == 2
|
||||
assert handler.reasoning_tokens == 1
|
||||
assert handler.total_cost > 0
|
||||
|
||||
|
||||
def test_on_llm_end_custom_model(handler: OpenAICallbackHandler) -> None:
|
||||
response = LLMResult(
|
||||
generations=[],
|
||||
|
Reference in New Issue
Block a user