import pandas as pd
import pandasai as pai
from pandasai import SmartDataframe
from pandasai_litellm import LiteLLM

llm = LiteLLM(
    api_key="sk-", 
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", 
    // OpenAI-Compatible Endpoints | liteLLM： https://docs.litellm.ai/docs/providers/openai_compatible#usage---completion
    // LiteLLm的规范是必须加上 "openai/"的前缀才能用。
    model="openai/qwen-max-latest"
)

pai.config.set({"llm": llm})

df = pd.read_csv("data/supermarket_sales.csv")

smart_df = SmartDataframe(df, config={"llm": llm})

response = smart_df.chat("How are the sales in each city? Show me the pictures. ")

print(response)

import pandasai as pai
from pandasai_litellm import LiteLLM

llm = LiteLLM(
    api_key="sk-", 
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", 
    model="openai/qwen-max-latest"
)

pai.config.set({"llm": llm})

file = pai.read_csv("data/supermarket_sales.csv")

# Create a dataset
pai.create(
    # 数据集的名称
    path="company/supermarket-sales",
    # 数据集的描述
    description="supermarket sales data",
    # 数据集的文件
    df = file,
)

# Load the dataset
supermarket_sales_data = pai.load("company/supermarket-sales")

# Query the data
response = supermarket_sales_data.chat("How are the sales in each city? Show me the pictures")

print(response)