import pandasai as pai
from pandasai_openai.openai import OpenAI
llm = OpenAI("OPEN_AI_API_KEY")
pai.config.set({
    "llm": llm
})
# Sample DataFrame
df = pai.DataFrame({
    "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
    "revenue": [5000, 3200, 2900, 4100, 2300, 2100, 2500, 2600, 4500, 7000]
})
df.chat('Which are the top 5 countries by sales?')

df.chat(
    "What is the total sales for the top 3 countries by sales?"
)

df.chat(
    "Plot the histogram of countries showing for each one the gd. Use different colors for each bar",
)

import pandasai as pai
from pandasai_openai.openai import OpenAI
employees_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Name': ['John', 'Emma', 'Liam', 'Olivia', 'William'],
    'Department': ['HR', 'Sales', 'IT', 'Marketing', 'Finance']
}
salaries_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Salary': [5000, 6000, 4500, 7000, 5500]
}
llm = OpenAI("OPEN_AI_API_KEY")
pai.config.set({
    "llm": llm
})
employees_df = pai.DataFrame(employees_data)
salaries_df = pai.DataFrame(salaries_data)
pai.chat("Who gets paid the most?", employees_df, salaries_df)

import pandasai as pai
from pandasai_docker import DockerSandbox
from pandasai_openai.openai import OpenAI
# Initialize the sandbox
sandbox = DockerSandbox()
sandbox.start()
employees_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Name': ['John', 'Emma', 'Liam', 'Olivia', 'William'],
    'Department': ['HR', 'Sales', 'IT', 'Marketing', 'Finance']
}
salaries_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Salary': [5000, 6000, 4500, 7000, 5500]
}
llm = OpenAI("OPEN_AI_API_KEY")
pai.config.set({
    "llm": llm
})
employees_df = pai.DataFrame(employees_data)
salaries_df = pai.DataFrame(salaries_data)
pai.chat("Who gets paid the most?", employees_df, salaries_df, sandbox=sandbox)
# Don't forget to stop the sandbox when done
sandbox.stop()