From b441123518d96708b135dc4f0a80545427322e4a Mon Sep 17 00:00:00 2001 From: ltn18 Date: Wed, 9 Oct 2024 23:19:30 -0400 Subject: [PATCH] [feature] add gpt-4o, fix openai API version compatibility Signed-off-by: ltn18 --- src/tot/models.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/tot/models.py b/src/tot/models.py index b3c4fe0..709085f 100644 --- a/src/tot/models.py +++ b/src/tot/models.py @@ -30,10 +30,10 @@ def chatgpt(messages, model="gpt-4", temperature=0.7, max_tokens=1000, n=1, stop cnt = min(n, 20) n -= cnt res = completions_with_backoff(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, n=cnt, stop=stop) - outputs.extend([choice["message"]["content"] for choice in res["choices"]]) + outputs.extend([choice.message.content for choice in res.choices]) # log completion tokens - completion_tokens += res["usage"]["completion_tokens"] - prompt_tokens += res["usage"]["prompt_tokens"] + completion_tokens += res.usage.completion_tokens + prompt_tokens += res.usage.prompt_tokens return outputs def gpt_usage(backend="gpt-4"): @@ -42,4 +42,6 @@ def gpt_usage(backend="gpt-4"): cost = completion_tokens / 1000 * 0.06 + prompt_tokens / 1000 * 0.03 elif backend == "gpt-3.5-turbo": cost = completion_tokens / 1000 * 0.002 + prompt_tokens / 1000 * 0.0015 + elif backend == "gpt-4o": + cost = completion_tokens / 1000 * 0.00250 + prompt_tokens / 1000 * 0.01 return {"completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, "cost": cost}