Skyvern/skyvern/forge/sdk/api/chat_completion_price.py

27 lines
986 B
Python

from typing import Callable
from pydantic import BaseModel
openai_model_to_price_lambdas = {
"gpt-4-vision-preview": (0.01, 0.03),
"gpt-4-1106-preview": (0.01, 0.03),
"gpt-4-0125-preview": (0.01, 0.03),
"gpt-3.5-turbo": (0.001, 0.002),
"gpt-3.5-turbo-1106": (0.001, 0.002),
"gpt-3.5-turbo-0125": (0.0005, 0.0015),
}
class ChatCompletionPrice(BaseModel):
input_token_count: int
output_token_count: int
openai_model_to_price_lambda: Callable[[int, int], float]
def __init__(self, input_token_count: int, output_token_count: int, model_name: str):
input_token_price, output_token_price = openai_model_to_price_lambdas[model_name]
super().__init__(
input_token_count=input_token_count,
output_token_count=output_token_count,
openai_model_to_price_lambda=lambda input_token, output_token: input_token_price * input_token / 1000
+ output_token_price * output_token / 1000,
)