mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
52714cedd4
# OpanAI finetuned model giving zero tokens cost Very simple fix to the previously committed solution to allowing finetuned Openai models. Improves #5127 --------- Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
import pytest
|
|
|
|
from langchain.callbacks import OpenAICallbackHandler
|
|
from langchain.llms.openai import BaseOpenAI
|
|
from langchain.schema import LLMResult
|
|
|
|
|
|
@pytest.fixture
|
|
def handler() -> OpenAICallbackHandler:
|
|
return OpenAICallbackHandler()
|
|
|
|
|
|
def test_on_llm_end(handler: OpenAICallbackHandler) -> None:
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"token_usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_name": BaseOpenAI.__fields__["model_name"].default,
|
|
},
|
|
)
|
|
handler.on_llm_end(response)
|
|
assert handler.successful_requests == 1
|
|
assert handler.total_tokens == 3
|
|
assert handler.prompt_tokens == 2
|
|
assert handler.completion_tokens == 1
|
|
assert handler.total_cost > 0
|
|
|
|
|
|
def test_on_llm_end_custom_model(handler: OpenAICallbackHandler) -> None:
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"token_usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_name": "foo-bar",
|
|
},
|
|
)
|
|
handler.on_llm_end(response)
|
|
assert handler.total_cost == 0
|
|
|
|
|
|
def test_on_llm_end_finetuned_model(handler: OpenAICallbackHandler) -> None:
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"token_usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_name": "ada:ft-your-org:custom-model-name-2022-02-15-04-21-04",
|
|
},
|
|
)
|
|
handler.on_llm_end(response)
|
|
assert handler.total_cost > 0
|