#!/usr/bin/env python
# coding: utf-8

# <a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/llm/ai21.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="在 Colab 中打开"/></a>
# 

# AI21
# 
# AI21是一个人工智能平台，旨在为用户提供各种人工智能相关的工具和服务。它包括自然语言处理、机器学习、数据分析和其他人工智能技术。AI21的目标是帮助用户利用人工智能技术解决各种问题，并提供高质量的人工智能解决方案。
# 

# ## 基本用法
# 

# #### 使用提示调用`complete`
# 

# 如果您在colab上打开这个笔记本，您可能需要安装LlamaIndex 🦙。
# 

# In[ ]:


get_ipython().run_line_magic('pip', 'install llama-index-llms-ai21')


# In[ ]:


get_ipython().system('pip install llama-index')


# In[ ]:


from llama_index.llms.ai21 import AI21

api_key = "Your api key"
resp = AI21(api_key=api_key).complete("Paul Graham is ")


# In[ ]:


print(resp)


# #### 使用消息列表调用`chat`
# 

# In[ ]:


from llama_index.core.llms import ChatMessage
from llama_index.llms.ai21 import AI21

messages = [
    ChatMessage(role="user", content="hello there"),
    ChatMessage(
        role="assistant", content="Arrrr, matey! How can I help ye today?"
    ),
    ChatMessage(role="user", content="What is your name"),
]

resp = AI21(api_key=api_key).chat(
    messages, preamble_override="You are a pirate with a colorful personality"
)


# In[ ]:


print(resp)


# ## 配置模型
# 

# In[ ]:


from llama_index.llms.ai21 import AI21

llm = AI21(model="j2-mid", api_key=api_key)


# In[ ]:


resp = llm.complete("Paul Graham is ")


# In[ ]:


print(resp)


# ## 在每个实例级别设置API密钥
# 如果需要的话，您可以让不同的LLM实例使用不同的API密钥。
# 

# In[ ]:


from llama_index.llms.ai21 import AI21

llm_good = AI21(api_key=api_key)
llm_bad = AI21(model="j2-mid", api_key="BAD_KEY")

resp = llm_good.complete("Paul Graham is ")
print(resp)

resp = llm_bad.complete("Paul Graham is ")
print(resp)