import erniebot

# List supported models
models = erniebot.Model.list()

print(models)
# ernie-bot             文心一言模型（ernie-bot）
# ernie-bot-turbo       文心一言模型（ernie-bot-turbo）
# ernie-bot-4           文心一言模型（ernie-bot-4）
# ernie-bot-8k          文心一言模型（ernie-bot-8k）
# ernie-text-embedding  文心百中语义模型
# ernie-vilg-v2         文心一格模型

# Set authentication params
erniebot.api_type = "aistudio"
erniebot.access_token = "<access-token-for-aistudio>"

# Create a chat completion
response = erniebot.ChatCompletion.create(model="ernie-bot", messages=[{"role": "user", "content": "你好，请介绍下你自己"}])

print(response.get_result())

import erniebot

erniebot.api_type = "aistudio"
erniebot.access_token = "<access-token-for-aistudio>"

response = erniebot.ChatCompletion.create(
    model="ernie-bot",
    messages=[{
        "role": "user",
        "content": "请问你是谁？"
    }, {
        "role": "assistant",
        "content":
        "我是百度公司开发的人工智能语言模型，我的中文名是文心一言，英文名是ERNIE-Bot，可以协助您完成范围广泛的任务并提供有关各种主题的信息，比如回答问题，提供定义和解释及建议。如果您有任何问题，请随时向我提问。"
    }, {
        "role": "user",
        "content": "我在深圳，周末可以去哪里玩？"
    }])

print(response.get_result())

import erniebot

erniebot.api_type = "aistudio"
erniebot.access_token = "<access-token-for-aistudio>"

response = erniebot.Embedding.create(
    model="ernie-text-embedding",
    input=[
        "我是百度公司开发的人工智能语言模型，我的中文名是文心一言，英文名是ERNIE-Bot，可以协助您完成范围广泛的任务并提供有关各种主题的信息，比如回答问题，提供定义和解释及建议。如果您有任何问题，请随时向我提问。",
        "2018年深圳市各区GDP"
        ])

print(response.get_result())

import erniebot

erniebot.api_type = "yinian"
erniebot.access_token = "<access-token-for-yinian>"

response = erniebot.Image.create(
    model="ernie-vilg-v2",
    prompt="雨后的桃花，8k，辛烷值渲染",
    width=512,
    height=512
)

print(response.get_result())

import erniebot

erniebot.api_type = "aistudio"
erniebot.access_token = "<access-token-for-aistudio>"

response = erniebot.ChatCompletion.create(
    model="ernie-bot",
    messages=[{
        "role": "user",
        "content": "深圳市今天气温多少摄氏度？",
    }, ],
    functions = [
        {
            "name": "get_current_temperature",
            "description": "获取指定城市的气温",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "城市名称",
                    },
                    "unit": {
                        "type": "string",
                        "enum": [
                            "摄氏度",
                            "华氏度",
                        ],
                    },
                },
                "required": [
                    "location",
                    "unit",
                ],
            },
            "responses": {
                "type": "object",
                "properties": {
                    "temperature": {
                        "type": "integer",
                        "description": "城市气温",
                    },
                    "unit": {
                        "type": "string",
                        "enum": [
                            "摄氏度",
                            "华氏度",
                        ],
                    },
                },
            },
        },
    ],
    stream=False)
print(response.get_result())