get_openai_callback bug when model is gpt-4o

[nehcneb]

System Info

macOS m3 14.4.1
Python 3.11.7
pandasai 2.0.40

🐛 Describe the bug

Hi Gabriele,
Thank you so much for incorporating gpt-4o so quickly and for making the amazing Pandas-AI! I've just updated to version 2.0.40, and am getting an error from running get_openai_callback.

Consider the following example:

import pandas as pd
from pandasai import Agent
from pandasai.llm.openai import OpenAI
from pandasai.helpers.openai_info import get_openai_callback

df = pd.DataFrame({
    "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
    "gdp": [19294482071552, 2891615567872, 2411255037952, 3435817336832, 1745433788416, 1181205135360, 1607402389504, 1490967855104, 4380756541440, 14631844184064],
    "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12]
})


llm = OpenAI(api_token=some_token, model = "gpt-4o")

df = Agent(df, config={"llm": llm, },)

with get_openai_callback() as cb:
    response = df.chat('Which are the 5 happiest countries?')
    print(response)
    print(cb)

This example yields the following error:

Unfortunately, I was not able to answer your question, because of the following error:

Unknown model: gpt-4o-2024-05-13-completion. Please provide a valid OpenAI model name.Known models are: gpt-4, gpt-4-0613, gpt-4-turbo-preview, gpt-4-0125-preview, gpt-4-1106-preview, gpt-4o, gpt-4o-2024-05-13, gpt-4-32k, gpt-4-32k-0613, gpt-4-completion, gpt-4-0613-completion, gpt-4-turbo-preview-completion, gpt-4-0125-preview-completion, gpt-4-1106-preview-completion, gpt-4-32k-completion, gpt-4-32k-0613-completion, gpt-3.5-turbo, gpt-3.5-turbo-0125, gpt-3.5-turbo-1106, gpt-3.5-turbo-0613, gpt-3.5-turbo-16k, gpt-3.5-turbo-16k-0613, gpt-3.5-turbo-instruct, gpt-3.5-turbo-completion, gpt-3.5-turbo-0125-completion, gpt-3.5-turbo-1106-completion, gpt-3.5-turbo-0613-completion, gpt-3.5-turbo-16k-completion, gpt-3.5-turbo-16k-0613-completion, gpt-3.5-turbo-instruct-completion, gpt-35-turbo, gpt-35-turbo-0613, gpt-35-turbo-instruct, gpt-35-turbo-16k, gpt-35-turbo-16k-0613, gpt-35-turbo-completion, gpt-35-turbo-0613-completion, gpt-35-turbo-instruct-completion, gpt-35-turbo-16k-completion, gpt-35-turbo-16k-0613-completion, gpt-3.5-turbo-0613-finetuned, gpt-3.5-turbo-1106-finetuned, gpt-3.5-turbo-0613-finetuned-completion, gpt-3.5-turbo-1106-finetuned-completion, gpt-35-turbo-0613-azure-finetuned, gpt-35-turbo-0613-azure-finetuned-completion

I wonder if this error is caused by the v2.0.38 amendments to openai_info.py? This error disappears if I remove with get_openai_callback() as cb:

Perhaps adding 'gpt-4o-2024-05-13-completion' under the list commencing # GPT-4 output will resolve the error.

Kind regards,
Ben

[gventuri]

Hi @nehcneb I confirm this should be the issue! Would you mind opening a PR for that?

Thanks a lot!

[nehcneb]

Just did, thank you so much!