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

# # LLM编译器代理食谱
# 
# <a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/agent/llm_compiler.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="在Colab中打开"/></a>
# 
# **注意**: 感谢[LLMCompiler源代码库](https://github.com/SqueezeAILab/LLMCompiler)。我们的许多实现都是从这个代码库中提取的（并且使用了LlamaIndex模块进行了适应）。
# 
# 在这本食谱中，我们将展示如何使用我们的LLMCompiler代理实现来处理各种设置。这包括使用一些简单的函数工具进行数学运算，还可以回答关于多个文档的更高级RAG用例的多部分查询。
# 
# 我们看到LLMCompilerAgent能够并行调用函数，比通过ReAct顺序执行能够更快地得到结果。
# 

# In[ ]:


get_ipython().run_line_magic('pip', 'install llama-index-agent-llm-compiler')
get_ipython().run_line_magic('pip', 'install llama-index-readers-wikipedia')
get_ipython().run_line_magic('pip', 'install llama-index-llms-openai')


# In[ ]:


import nest_asyncio

nest_asyncio.apply()


# In[ ]:


from llama_index.agent.llm_compiler import LLMCompilerAgentWorker


# ## 使用简单函数测试LLMCompiler代理
# 
# 在这里，我们将使用简单的数学函数（加法、乘法）来测试LLMCompiler代理，以说明它的工作原理。
# 

# In[ ]:


import json
from typing import Sequence, List

from llama_index.llms.openai import OpenAI
from llama_index.core.llms import ChatMessage
from llama_index.core.tools import BaseTool, FunctionTool


# ### 定义函数
# 

# In[ ]:


def multiply(a: int, b: int) -> int:    """将两个整数相乘并返回结果整数"""    return a * bmultiply_tool = FunctionTool.from_defaults(fn=multiply)def add(a: int, b: int) -> int:    """将两个整数相加并返回结果整数"""    return a + badd_tool = FunctionTool.from_defaults(fn=add)tools = [multiply_tool, add_tool]


# In[ ]:


multiply_tool.metadata.fn_schema_str


# ### 设置LLMCompiler代理
# 
# 我们导入`LLMCompilerAgentWorker`并将其与`AgentRunner`结合在一起。
# 

# In[ ]:


from llama_index.core.agent import AgentRunner


# In[ ]:


llm = OpenAI(model="gpt-4")


# In[ ]:


callback_manager = llm.callback_manager


# In[ ]:


agent_worker = LLMCompilerAgentWorker.from_tools(
    tools, llm=llm, verbose=True, callback_manager=callback_manager
)
agent = AgentRunner(agent_worker, callback_manager=callback_manager)


# ### 测试一些查询
# 

# In[ ]:


response = agent.chat("What is (121 * 3) + 42?")


# In[ ]:


response


# In[ ]:


agent.memory.get_all()


# ## 尝试使用LLMCompiler进行RAG
# 
# 现在让我们尝试使用LLMCompiler进行RAG的用例。具体来说，我们将加载一个包含各种城市维基百科文章的数据集，并对其提出问题。
# 

# ### 设置数据
# 
# 我们使用我们的 `WikipediaReader` 来加载各个城市的数据。
# 

# In[ ]:


get_ipython().run_line_magic('pip', 'install llama-index-readers-wikipedia wikipedia')


# In[ ]:


from llama_index.readers.wikipedia import WikipediaReader


# In[ ]:


wiki_titles = ["Toronto", "Seattle", "Chicago", "Boston", "Miami"]


# In[ ]:


city_docs = {}
reader = WikipediaReader()
for wiki_title in wiki_titles:
    docs = reader.load_data(pages=[wiki_title])
    city_docs[wiki_title] = docs


# ### 设置LLM + 服务上下文
# 

# In[ ]:


from llama_index.core import ServiceContext
from llama_index.llms.openai import OpenAI
from llama_index.core.callbacks import CallbackManager

llm = OpenAI(temperature=0, model="gpt-4")
service_context = ServiceContext.from_defaults(llm=llm)
callback_manager = CallbackManager([])


# ### 定义工具集
# 
# 在这个项目中，我们将使用以下工具集：
# 
# - Python 3：一种流行的编程语言，用于开发机器学习模型和处理数据。
# - NumPy：用于进行科学计算的Python库，特别适用于处理多维数组和矩阵运算。
# - Pandas：提供数据分析和处理工具的Python库，特别适用于处理结构化数据。
# - Matplotlib：用于创建可视化图表和图形的Python库。
# - Scikit-learn：用于机器学习建模和数据分析的Python库，包含了许多常用的机器学习算法和工具。
# 
# 确保在运行代码之前安装了这些工具集。
# 

# In[ ]:


from llama_index.core import load_index_from_storage, StorageContextfrom llama_index.core.node_parser import SentenceSplitterfrom llama_index.core.tools import QueryEngineTool, ToolMetadatafrom llama_index.core import VectorStoreIndeximport osnode_parser = SentenceSplitter()# 构建代理字典query_engine_tools = []for idx, wiki_title in enumerate(wiki_titles):    nodes = node_parser.get_nodes_from_documents(city_docs[wiki_title])    if not os.path.exists(f"./data/{wiki_title}"):        # 构建向量索引        vector_index = VectorStoreIndex(            nodes,            service_context=service_context,            callback_manager=callback_manager,        )        vector_index.storage_context.persist(            persist_dir=f"./data/{wiki_title}"        )    else:        vector_index = load_index_from_storage(            StorageContext.from_defaults(persist_dir=f"./data/{wiki_title}"),            service_context=service_context,            callback_manager=callback_manager,        )    # 定义查询引擎    vector_query_engine = vector_index.as_query_engine()    # 定义工具    query_engine_tools.append(        QueryEngineTool(            query_engine=vector_query_engine,            metadata=ToolMetadata(                name=f"vector_tool_{wiki_title}",                description=(                    "用于与特定方面相关的问题（例如历史、艺术与文化、体育、人口统计等）"                    f" {wiki_title} 相关的问题。"                ),            ),        )    )


# ### 设置LLMCompilerAgent
# 
# ```python
# from llmcompileragent import LLMCompilerAgent
# 
# # 创建LLMCompilerAgent对象
# agent = LLMCompilerAgent()
# ```
# 

# In[ ]:


from llama_index.core.agent import AgentRunner
from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-4")
agent_worker = LLMCompilerAgentWorker.from_tools(
    query_engine_tools,
    llm=llm,
    verbose=True,
    callback_manager=callback_manager,
)
agent = AgentRunner(agent_worker, callback_manager=callback_manager)


# ### 测试查询
# 
# 这里我们将测试一些查询。
# 

# In[ ]:


response = agent.chat(
    "Tell me about the demographics of Miami, and compare that with the demographics of Chicago?"
)
print(str(response))


# In[ ]:


response = agent.chat(
    "Is the climate of Chicago or Seattle better during the wintertime?"
)
print(str(response))