{"cells": [{"attachments": {}, "cell_type": "markdown", "id": "70b86d36", "metadata": {}, "source": ["<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/agent/openai_assistant_query_cookbook.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"在 Colab 中打开\"/></a>\n"]}, {"attachments": {}, "cell_type": "markdown", "id": "b50c4af8-fec3-4396-860a-1322089d76cb", "metadata": {}, "source": ["# OpenAI助手高级检索手册\n", "\n", "在这个笔记本中，我们尝试使用OpenAI助手API来进行高级检索任务，通过插入各种查询引擎工具和数据集。我们使用的包装抽象是我们的`OpenAIAssistantAgent`类，它允许我们插入自定义工具。我们探讨了`OpenAIAssistant`如何通过其代理执行+函数调用循环来补充/替代现有的由我们的检索器/查询引擎解决的工作流程。\n", "\n", "- 联合问答+摘要\n", "- 自动检索\n", "- 联合SQL和向量搜索\n"]}, {"cell_type": "code", "execution_count": null, "id": "29a3e4fa", "metadata": {}, "outputs": [], "source": ["%pip install llama-index-agent-openai\n", "%pip install llama-index-vector-stores-pinecone\n", "%pip install llama-index-readers-wikipedia\n", "%pip install llama-index-llms-openai"]}, {"cell_type": "code", "execution_count": null, "id": "5887c147-29f1-4a9b-b766-ea6fc01235bc", "metadata": {}, "outputs": [], "source": ["!pip install llama-index"]}, {"cell_type": "code", "execution_count": null, "id": "5befaedd-10f1-4bd7-be46-c235ec237976", "metadata": {}, "outputs": [], "source": ["import nest_asyncio\n", "\n", "nest_asyncio.apply()"]}, {"cell_type": "markdown", "id": "f428bbe5-d7de-49eb-857c-564a49f948ce", "metadata": {}, "source": ["## 联合问答和摘要\n", "\n", "在本节中，我们将展示如何让助手代理同时回答基于事实的问题和摘要问题。这是内部检索工具难以实现的功能。\n"]}, {"cell_type": "markdown", "id": "93f80dcf-f07b-43b6-ab4f-4826cfbcff2c", "metadata": {}, "source": ["```python\n", "# 加载数据\n", "```\n", "\n", "这里是加载数据的部分。\n"]}, {"cell_type": "code", "execution_count": null, "id": "de92789e-850a-4668-a459-96318cd74b0a", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["--2023-11-11 09:40:13--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n", "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8002::154, 2606:50c0:8003::154, 2606:50c0:8000::154, ...\n", "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8002::154|:443... connected.\n", "HTTP request sent, awaiting response... 200 OK\n", "Length: 75042 (73K) [text/plain]\n", "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n", "\n", "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.009s  \n", "\n", "2023-11-11 09:40:14 (8.24 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n", "\n"]}], "source": ["!mkdir -p 'data/paul_graham/'\n", "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"]}, {"cell_type": "code", "execution_count": null, "id": "1eec54ef-b469-46c4-973f-8e213577d063", "metadata": {}, "outputs": [], "source": ["from llama_index.core import SimpleDirectoryReader", "", "# 加载文档", "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"]}, {"cell_type": "markdown", "id": "227675b4-96ff-45f4-a794-21a831143dfa", "metadata": {}, "source": ["### 设置向量 + 摘要索引/查询引擎/工具\n"]}, {"cell_type": "code", "execution_count": null, "id": "4d4baa45-fbd0-4285-ba92-6e9eeffbd91c", "metadata": {}, "outputs": [], "source": ["from llama_index.llms.openai import OpenAI", "from llama_index.core import Settings", "from llama_index.core import StorageContext, VectorStoreIndex", "from llama_index.core import SummaryIndex", "", "# 初始化设置（设置块大小）", "Settings.llm = OpenAI()", "Settings.chunk_size = 1024", "nodes = Settings.node_parser.get_nodes_from_documents(documents)", "", "# 初始化存储上下文（默认情况下为内存）", "storage_context = StorageContext.from_defaults()", "storage_context.docstore.add_documents(nodes)", "", "# 定义摘要索引和向量索引相同的数据", "summary_index = SummaryIndex(nodes, storage_context=storage_context)", "vector_index = VectorStoreIndex(nodes, storage_context=storage_context)", "", "# 定义查询引擎", "summary_query_engine = summary_index.as_query_engine(", "    response_mode=\"tree_summarize\",", "    use_async=True,", ")", "vector_query_engine = vector_index.as_query_engine()"]}, {"cell_type": "code", "execution_count": null, "id": "f8474aef-4702-46a3-ad0a-da2d94e73e4b", "metadata": {}, "outputs": [], "source": ["from llama_index.core.tools import QueryEngineTool\n", "\n", "summary_tool = QueryEngineTool.from_defaults(\n", "    query_engine=summary_query_engine,\n", "    name=\"summary_tool\",\n", "    description=(\n", "        \"Useful for summarization questions related to the author's life\"\n", "    ),\n", ")\n", "\n", "vector_tool = QueryEngineTool.from_defaults(\n", "    query_engine=vector_query_engine,\n", "    name=\"vector_tool\",\n", "    description=(\n", "        \"Useful for retrieving specific context to answer specific questions about the author's life\"\n", "    ),\n", ")"]}, {"cell_type": "markdown", "id": "214c21d9-d530-44f0-83aa-6443354a3d70", "metadata": {}, "source": ["### 定义助理代理\n", "\n", "在这个部分，我们将定义一个助理代理，它将帮助我们执行各种任务。\n"]}, {"cell_type": "code", "execution_count": null, "id": "3f7553e1-6ca4-43b8-b485-50487f4691ae", "metadata": {}, "outputs": [], "source": ["from llama_index.agent.openai import OpenAIAssistantAgent\n", "\n", "agent = OpenAIAssistantAgent.from_new(\n", "    name=\"QA bot\",\n", "    instructions=\"You are a bot designed to answer questions about the author\",\n", "    openai_tools=[],\n", "    tools=[summary_tool, vector_tool],\n", "    verbose=True,\n", "    run_retrieve_sleep_time=1.0,\n", ")"]}, {"cell_type": "markdown", "id": "a7694116-9052-4656-a4ca-ba1d4d95b170", "metadata": {}, "source": ["#### 结果：有点不稳定\n"]}, {"cell_type": "code", "execution_count": null, "id": "afa22f7a-6cb3-4d75-a99d-3c80e3a9fb09", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: summary_tool with args: {\"input\":\"Can you give me a summary about the author's life?\"}\n", "Got output: The author, Paul Graham, had a strong interest in writing and programming from a young age. They started writing short stories and experimenting with programming in high school. In college, they initially studied philosophy but switched to studying artificial intelligence. However, they realized that the AI being practiced at the time was not going to lead to true understanding of natural language. This led them to focus on Lisp programming and eventually write a book about Lisp hacking. Despite being in a PhD program in computer science, the author also developed a passion for art and decided to pursue it further. They attended the Accademia di Belli Arti in Florence but found that it did not teach them much. They then returned to the US and got a job at a software company. Afterward, they attended the Rhode Island School of Design but dropped out due to the focus on developing a signature style rather than teaching the fundamentals of art. They then moved to New York City and became interested in the World Wide Web, eventually starting a company called Viaweb. They later founded Y Combinator, an investment firm, and created Hacker News.\n", "========================\n", "Paul Graham is an author with eclectic interests and a varied career path. He began with interests in writing and programming, engaged in philosophy and artificial intelligence during college, and authored a book on Lisp programming. With an equally strong passion for art, he studied at the Accademia di Belli Arti in Florence and briefly at the Rhode Island School of Design before immersing himself in the tech industry by starting Viaweb and later founding the influential startup accelerator Y Combinator. He also created Hacker News, a social news website focused on computer science and entrepreneurship. Graham's life reflects a blend of technology, entrepreneurship, and the arts.\n"]}], "source": ["response = agent.chat(\"Can you give me a summary about the author's life?\")\n", "print(str(response))"]}, {"cell_type": "code", "execution_count": null, "id": "f5be4399-dd5c-4111-95be-0889ab0c7f3d", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: vector_tool with args: {\"input\":\"After RICS\"}\n", "Got output: After RICS, the author moved back to Providence to continue at RISD. However, it became clear that art school, specifically the painting department, did not have the same relationship to art as medical school had to medicine. Painting students were expected to express themselves and develop a distinctive signature style.\n", "========================\n", "After the author's time at the Royal Institution of Chartered Surveyors (RICS), they moved back to Providence to continue their studies at the Rhode Island School of Design (RISD). There, the author noted a significant difference in the educational approaches between RISD and medical school, specifically in the painting department. At RISD, students were encouraged to express themselves and to develop a unique and distinctive signature style in their artwork.\n"]}], "source": ["response = agent.query(\"What did the author do after RICS?\")\n", "print(str(response))"]}, {"attachments": {}, "cell_type": "markdown", "id": "db402a8b-90d6-4e1d-8df6-347c54624f26", "metadata": {}, "source": ["## 从向量数据库自动检索\n", "\n", "我们现有的“自动检索”功能（在`VectorIndexAutoRetriever`中）允许LLM推断向量数据库的正确查询参数，包括查询字符串和元数据过滤器。\n", "\n", "由于助手API可以调用函数并推断函数参数，我们在这里探索它在执行自动检索方面的能力。\n"]}, {"attachments": {}, "cell_type": "markdown", "id": "50444bd8", "metadata": {}, "source": ["如果您在colab上打开这个笔记本，您可能需要安装LlamaIndex 🦙。\n"]}, {"cell_type": "code", "execution_count": null, "id": "efba05c6-a80e-4992-8a90-fd8fadd53587", "metadata": {}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/pinecone/index.py:4: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n", "  from tqdm.autonotebook import tqdm\n"]}], "source": ["import pinecone\n", "import os\n", "\n", "api_key = os.environ[\"PINECONE_API_KEY\"]\n", "pinecone.init(api_key=api_key, environment=\"us-west1-gcp\")"]}, {"cell_type": "code", "execution_count": null, "id": "e114fd4e-eefd-48d3-b082-1d918c15dfc9", "metadata": {}, "outputs": [], "source": ["# 文本嵌入ada-002的维度", "try:", "    pinecone.create_index(", "        \"quickstart\", dimension=1536, metric=\"euclidean\", pod_type=\"p1\"", "    )", "except Exception:", "    # 很可能索引已经存在", "    pass"]}, {"cell_type": "code", "execution_count": null, "id": "8b815ed5-d770-4ec4-8e2c-57b5a4c0941c", "metadata": {}, "outputs": [], "source": ["pinecone_index = pinecone.Index(\"quickstart\")"]}, {"cell_type": "code", "execution_count": null, "id": "e889542b-8514-4483-97a6-ff5d032fa75f", "metadata": {}, "outputs": [{"data": {"text/plain": ["{}"]}, "execution_count": null, "metadata": {}, "output_type": "execute_result"}], "source": ["# 可选：删除pinecone索引中的数据", "pinecone_index.delete(deleteAll=True, namespace=\"test\")"]}, {"cell_type": "code", "execution_count": null, "id": "f7764e84-8f0c-4890-9d6a-5d97be6ed1f9", "metadata": {}, "outputs": [], "source": ["from llama_index.core import VectorStoreIndex, StorageContext\n", "from llama_index.vector_stores.pinecone import PineconeVectorStore"]}, {"cell_type": "code", "execution_count": null, "id": "d5b32f02-e09a-4cfa-b7ba-fa937bdcec9b", "metadata": {}, "outputs": [], "source": ["from llama_index.core.schema import TextNode\n", "\n", "nodes = [\n", "    TextNode(\n", "        text=(\n", "            \"Michael Jordan is a retired professional basketball player,\"\n", "            \" widely regarded as one of the greatest basketball players of all\"\n", "            \" time.\"\n", "        ),\n", "        metadata={\n", "            \"category\": \"Sports\",\n", "            \"country\": \"United States\",\n", "        },\n", "    ),\n", "    TextNode(\n", "        text=(\n", "            \"Angelina Jolie is an American actress, filmmaker, and\"\n", "            \" humanitarian. She has received numerous awards for her acting\"\n", "            \" and is known for her philanthropic work.\"\n", "        ),\n", "        metadata={\n", "            \"category\": \"Entertainment\",\n", "            \"country\": \"United States\",\n", "        },\n", "    ),\n", "    TextNode(\n", "        text=(\n", "            \"Elon Musk is a business magnate, industrial designer, and\"\n", "            \" engineer. He is the founder, CEO, and lead designer of SpaceX,\"\n", "            \" Tesla, Inc., Neuralink, and The Boring Company.\"\n", "        ),\n", "        metadata={\n", "            \"category\": \"Business\",\n", "            \"country\": \"United States\",\n", "        },\n", "    ),\n", "    TextNode(\n", "        text=(\n", "            \"Rihanna is a Barbadian singer, actress, and businesswoman. She\"\n", "            \" has achieved significant success in the music industry and is\"\n", "            \" known for her versatile musical style.\"\n", "        ),\n", "        metadata={\n", "            \"category\": \"Music\",\n", "            \"country\": \"Barbados\",\n", "        },\n", "    ),\n", "    TextNode(\n", "        text=(\n", "            \"Cristiano Ronaldo is a Portuguese professional footballer who is\"\n", "            \" considered one of the greatest football players of all time. He\"\n", "            \" has won numerous awards and set multiple records during his\"\n", "            \" career.\"\n", "        ),\n", "        metadata={\n", "            \"category\": \"Sports\",\n", "            \"country\": \"Portugal\",\n", "        },\n", "    ),\n", "]"]}, {"cell_type": "code", "execution_count": null, "id": "c7974714-552d-4297-9cc3-6497409496e8", "metadata": {}, "outputs": [], "source": ["vector_store = PineconeVectorStore(\n", "    pinecone_index=pinecone_index, namespace=\"test\"\n", ")\n", "storage_context = StorageContext.from_defaults(vector_store=vector_store)"]}, {"cell_type": "code", "execution_count": null, "id": "f6f3fff1-974f-4a72-8f09-874e61c54230", "metadata": {}, "outputs": [{"data": {"application/json": {"ascii": false, "bar_format": null, "colour": null, "elapsed": 0.0056912899017333984, "initial": 0, "n": 0, "ncols": null, "nrows": 22, "postfix": null, "prefix": "Upserted vectors", "rate": null, "total": 5, "unit": "it", "unit_divisor": 1000, "unit_scale": false}, "application/vnd.jupyter.widget-view+json": {"model_id": "c9975a4d317443a2bf3b9f1b98b832a2", "version_major": 2, "version_minor": 0}, "text/plain": ["Upserted vectors:   0%|          | 0/5 [00:00<?, ?it/s]"]}, "metadata": {}, "output_type": "display_data"}], "source": ["index = VectorStoreIndex(nodes, storage_context=storage_context)"]}, {"attachments": {}, "cell_type": "markdown", "id": "0b557703-3587-41c6-a85b-949094344ca8", "metadata": {}, "source": ["#### 定义函数工具\n", "\n", "在这里，我们定义了函数接口，该接口将传递给OpenAI以执行自动检索。\n", "\n", "我们无法让OpenAI使用嵌套的pydantic对象或元组作为参数，因此我们将元数据过滤键和值转换为列表，以便函数API能够使用。\n"]}, {"cell_type": "code", "execution_count": null, "id": "113b467e-b924-4994-9011-86ab98167545", "metadata": {}, "outputs": [], "source": ["# 定义函数工具", "from llama_index.core.tools import FunctionTool", "from llama_index.core.vector_stores import (", "    VectorStoreInfo,", "    MetadataInfo,", "    ExactMatchFilter,", "    MetadataFilters,", ")", "from llama_index.core.retrievers import VectorIndexRetriever", "from llama_index.core.query_engine import RetrieverQueryEngine", "", "from typing import List, Tuple, Any", "from pydantic import BaseModel, Field", "", "# 暂时硬编码 top k", "top_k = 3", "", "# 定义描述向量存储架构的向量存储信息", "vector_store_info = VectorStoreInfo(", "    content_info=\"名人简介\",", "    metadata_info=[", "        MetadataInfo(", "            name=\"category\",", "            type=\"str\",", "            description=(", "                \"名人的类别，包括[体育、娱乐、商业、音乐]\"", "            ),", "        ),", "        MetadataInfo(", "            name=\"country\",", "            type=\"str\",", "            description=(", "                \"名人的国家，包括[美国、巴巴多斯、葡萄牙]\"", "            ),", "        ),", "    ],", ")", "", "", "# 为自动检索函数定义 pydantic 模型", "class AutoRetrieveModel(BaseModel):", "    query: str = Field(..., description=\"自然语言查询字符串\")", "    filter_key_list: List[str] = Field(", "        ..., description=\"元数据过滤字段名称列表\"", "    )", "    filter_value_list: List[str] = Field(", "        ...,", "        description=(", "            \"元数据过滤字段值列表（对应于 filter_key_list 中指定的名称）\"", "        ),", "    )", "", "", "def auto_retrieve_fn(", "    query: str, filter_key_list: List[str], filter_value_list: List[str]", "):", "    \"\"\"自动检索函数。", "", "    从向量数据库中执行自动检索，然后应用一组过滤器。", "", "    \"\"\"", "    query = query or \"查询\"", "", "    exact_match_filters = [", "        ExactMatchFilter(key=k, value=v)", "        for k, v in zip(filter_key_list, filter_value_list)", "    ]", "    retriever = VectorIndexRetriever(", "        index,", "        filters=MetadataFilters(filters=exact_match_filters),", "        top_k=top_k,", "    )", "    results = retriever.retrieve(query)", "    return [r.get_content() for r in results]", "", "", "description = f\"\"\"\\", "使用此工具查找有关名人的传记信息。", "以下是向量数据库架构：", "{vector_store_info.json()}", "\"\"\"", "", "auto_retrieve_tool = FunctionTool.from_defaults(", "    fn=auto_retrieve_fn,", "    name=\"celebrity_bios\",", "    description=description,", "    fn_schema=AutoRetrieveModel,", ")"]}, {"cell_type": "code", "execution_count": null, "id": "475cc850-dbc1-4073-896d-a3d811fabfd3", "metadata": {}, "outputs": [{"data": {"text/plain": ["['Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received numerous awards for her acting and is known for her philanthropic work.',\n", " 'Michael Jordan is a retired professional basketball player, widely regarded as one of the greatest basketball players of all time.']"]}, "execution_count": null, "metadata": {}, "output_type": "execute_result"}], "source": ["auto_retrieve_fn(\n", "    \"celebrity from the United States\",\n", "    filter_key_list=[\"country\"],\n", "    filter_value_list=[\"United States\"],\n", ")"]}, {"attachments": {}, "cell_type": "markdown", "id": "61f39c99-311d-4609-b13b-ca8f4b4631e9", "metadata": {}, "source": ["```python\n", "class Agent:\n", "    def __init__(self, name):\n", "        self.name = name\n", "        self.health = 100\n", "        self.mana = 100\n", "```\n", "初始化代理人\n"]}, {"cell_type": "code", "execution_count": null, "id": "49addd6e-c325-4222-b694-c19be2583650", "metadata": {}, "outputs": [], "source": ["from llama_index.agent.openai import OpenAIAssistantAgent\n", "\n", "agent = OpenAIAssistantAgent.from_new(\n", "    name=\"Celebrity bot\",\n", "    instructions=\"You are a bot designed to answer questions about celebrities.\",\n", "    tools=[auto_retrieve_tool],\n", "    verbose=True,\n", ")"]}, {"cell_type": "code", "execution_count": null, "id": "a6c3fc23-3f09-492f-b1d8-2291cd091a0c", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: celebrity_bios with args: {\"query\": \"celebrity from United States\", \"filter_key_list\": [\"country\"], \"filter_value_list\": [\"United States\"]}\n", "Got output: ['Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received numerous awards for her acting and is known for her philanthropic work.', 'Michael Jordan is a retired professional basketball player, widely regarded as one of the greatest basketball players of all time.']\n", "========================\n", "=== Calling Function ===\n", "Calling function: celebrity_bios with args: {\"query\": \"celebrity from United States\", \"filter_key_list\": [\"country\"], \"filter_value_list\": [\"United States\"]}\n", "Got output: ['Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received numerous awards for her acting and is known for her philanthropic work.', 'Michael Jordan is a retired professional basketball player, widely regarded as one of the greatest basketball players of all time.']\n", "========================\n", "Here is some information about two celebrities from the United States:\n", "\n", "1. Angelina Jolie - Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received numerous awards for her acting and is known for her philanthropic work. Over the years, Jolie has starred in several critically acclaimed and commercially successful films, and she has also been involved in various humanitarian causes, advocating for refugees and children's education, among other things.\n", "\n", "2. Michael Jordan - Michael Jordan is a retired professional basketball player, widely regarded as one of the greatest basketball players of all time. During his career, Jordan dominated the NBA with his scoring ability, athleticism, and competitiveness. He won six NBA championships with the Chicago Bulls and earned the NBA Most Valuable Player Award five times. Jordan has also been a successful businessman and the principal owner of the Charlotte Hornets basketball team.\n", "\n", "Both figures have made significant impacts in their respective fields and continue to be influential even after reaching the peaks of their careers.\n"]}], "source": ["response = agent.chat(\"Tell me about two celebrities from the United States. \")\n", "print(str(response))"]}, {"attachments": {}, "cell_type": "markdown", "id": "d641f03d-d0cb-4c11-a571-33a9cb7221d5", "metadata": {}, "source": ["## 文本到SQL和语义搜索的联合\n", "\n", "这目前由我们的`SQLAutoVectorQueryEngine`处理。\n", "\n", "让我们尝试通过让我们的`OpenAIAssistantAgent`访问两个查询工具来实现这一点：SQL和向量搜索。\n", "\n", "**注意：** 任何文本到SQL应用程序都应该意识到执行任意SQL查询可能存在安全风险。建议采取必要的预防措施，例如使用受限角色、只读数据库、沙箱等。\n"]}, {"attachments": {}, "cell_type": "markdown", "id": "550dc181-3725-4b96-8528-2669483955ba", "metadata": {}, "source": ["#### 加载和索引结构化数据\n", "\n", "我们将样本结构化数据点加载到 SQL 数据库中并进行索引。\n"]}, {"cell_type": "code", "execution_count": null, "id": "bef04b5c-0e7f-4c44-b71d-8cb26406d50c", "metadata": {}, "outputs": [], "source": ["from sqlalchemy import (\n", "    create_engine,\n", "    MetaData,\n", "    Table,\n", "    Column,\n", "    String,\n", "    Integer,\n", "    select,\n", "    column,\n", ")\n", "from llama_index.core import SQLDatabase\n", "from llama_index.core.indices import SQLStructStoreIndex\n", "\n", "engine = create_engine(\"sqlite:///:memory:\", future=True)\n", "metadata_obj = MetaData()"]}, {"cell_type": "code", "execution_count": null, "id": "1f3739b3-d1c5-4600-985d-7b6e310eb8eb", "metadata": {}, "outputs": [], "source": ["# 创建城市SQL表", "table_name = \"city_stats\"", "city_stats_table = Table(", "    table_name,", "    metadata_obj,", "    Column(\"city_name\", String(16), primary_key=True),", "    Column(\"population\", Integer),", "    Column(\"country\", String(16), nullable=False),", ")", "", "metadata_obj.create_all(engine)"]}, {"cell_type": "code", "execution_count": null, "id": "251a170a-e1cd-4651-8689-a01d9c82002b", "metadata": {}, "outputs": [{"data": {"text/plain": ["dict_keys(['city_stats'])"]}, "execution_count": null, "metadata": {}, "output_type": "execute_result"}], "source": ["# 打印表格", "metadata_obj.tables.keys()"]}, {"cell_type": "code", "execution_count": null, "id": "e3ea511f-5a01-493e-a3f7-ceb144823996", "metadata": {}, "outputs": [], "source": ["from sqlalchemy import insert\n", "\n", "rows = [\n", "    {\"city_name\": \"Toronto\", \"population\": 2930000, \"country\": \"Canada\"},\n", "    {\"city_name\": \"Tokyo\", \"population\": 13960000, \"country\": \"Japan\"},\n", "    {\"city_name\": \"Berlin\", \"population\": 3645000, \"country\": \"Germany\"},\n", "]\n", "for row in rows:\n", "    stmt = insert(city_stats_table).values(**row)\n", "    with engine.begin() as connection:\n", "        cursor = connection.execute(stmt)"]}, {"cell_type": "code", "execution_count": null, "id": "26dfe2b8-4831-4603-a88f-ae3e571f9436", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["[('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Berlin', 3645000, 'Germany')]\n"]}], "source": ["with engine.connect() as connection:\n", "    cursor = connection.exec_driver_sql(\"SELECT * FROM city_stats\")\n", "    print(cursor.fetchall())"]}, {"cell_type": "code", "execution_count": null, "id": "ff4e85e7-364b-480b-8a7b-a1db956b7011", "metadata": {}, "outputs": [], "source": ["sql_database = SQLDatabase(engine, include_tables=[\"city_stats\"])"]}, {"cell_type": "code", "execution_count": null, "id": "faadf01c", "metadata": {}, "outputs": [], "source": ["from llama_index.core.query_engine import NLSQLTableQueryEngine"]}, {"cell_type": "code", "execution_count": null, "id": "299ad520-1ceb-4191-b843-31b3c2ac3cdc", "metadata": {}, "outputs": [], "source": ["query_engine = NLSQLTableQueryEngine(\n", "    sql_database=sql_database,\n", "    tables=[\"city_stats\"],\n", ")"]}, {"attachments": {}, "cell_type": "markdown", "id": "a2f1fc66-6627-4a36-9a5b-bbbd1d89dc05", "metadata": {}, "source": ["#### 加载和索引非结构化数据\n", "\n", "我们将非结构化数据加载到由Pinecone支持的向量索引中。\n"]}, {"cell_type": "code", "execution_count": null, "id": "0b7fdd9d-365a-4978-9352-67c2d6e84599", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Requirement already satisfied: wikipedia in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (1.4.0)\n", "Requirement already satisfied: requests<3.0.0,>=2.0.0 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (2.28.2)\n", "Requirement already satisfied: beautifulsoup4 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (4.12.2)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.1.0)\n", "Requirement already satisfied: idna<4,>=2.5 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.4)\n", "Requirement already satisfied: certifi>=2017.4.17 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (2022.12.7)\n", "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (1.26.15)\n", "Requirement already satisfied: soupsieve>1.2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from beautifulsoup4->wikipedia) (2.4.1)\n", "\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"]}], "source": ["# 安装维基百科的Python包", "!pip install wikipedia"]}, {"cell_type": "code", "execution_count": null, "id": "e38dd18c-e4e6-402d-93b9-8bb46423e944", "metadata": {}, "outputs": [], "source": ["from llama_index.readers.wikipedia import WikipediaReader\n", "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex"]}, {"cell_type": "code", "execution_count": null, "id": "0bf7fc15-b699-4414-8968-5ae67d59ab7c", "metadata": {}, "outputs": [], "source": ["cities = [\"Toronto\", \"Berlin\", \"Tokyo\"]\n", "wiki_docs = WikipediaReader().load_data(pages=cities)"]}, {"cell_type": "code", "execution_count": null, "id": "38f97c34-0f9d-4030-a6e5-5526a4f015c9", "metadata": {}, "outputs": [], "source": ["from llama_index.core import Settings", "from llama_index.core import StorageContext", "from llama_index.core.node_parser import TokenTextSplitter", "from llama_index.llms.openai import OpenAI", "", "# 定义节点解析器和LLM", "Settings.chunk_size = 1024", "Settings.llm = OpenAI(temperature=0, model=\"gpt-4\")", "", "text_splitter = TokenTextSplitter(chunk_size=1024)", "", "# 使用默认的内存存储", "storage_context = StorageContext.from_defaults()", "vector_index = VectorStoreIndex([], storage_context=storage_context)"]}, {"cell_type": "code", "execution_count": null, "id": "4dc76bc8-6265-4a0e-957c-1f3d6ac0b802", "metadata": {}, "outputs": [], "source": ["# 将文档插入向量索引", "# 每个文档都附带城市的元数据", "for city, wiki_doc in zip(cities, wiki_docs):", "    nodes = text_splitter.get_nodes_from_documents([wiki_doc])", "    # 为每个节点添加元数据", "    for node in nodes:", "        node.metadata = {\"title\": city}", "    vector_index.insert_nodes(nodes)"]}, {"attachments": {}, "cell_type": "markdown", "id": "c6e73f5a-5c75-4963-8642-dc5060c0a735", "metadata": {}, "source": ["#### 定义查询引擎/工具\n"]}, {"cell_type": "code", "execution_count": null, "id": "21086866-120d-45fe-9233-20b070643e6a", "metadata": {}, "outputs": [], "source": ["from llama_index.core.tools import QueryEngineTool"]}, {"cell_type": "code", "execution_count": null, "id": "13a86ca7-1240-4ea2-9f5f-6d93cca00978", "metadata": {}, "outputs": [], "source": ["sql_tool = QueryEngineTool.from_defaults(\n", "    query_engine=query_engine,\n", "    name=\"sql_tool\",\n", "    description=(\n", "        \"Useful for translating a natural language query into a SQL query over\"\n", "        \" a table containing: city_stats, containing the population/country of\"\n", "        \" each city\"\n", "    ),\n", ")\n", "vector_tool = QueryEngineTool.from_defaults(\n", "    query_engine=vector_index.as_query_engine(similarity_top_k=2),\n", "    name=\"vector_tool\",\n", "    description=(\n", "        f\"Useful for answering semantic questions about different cities\"\n", "    ),\n", ")"]}, {"attachments": {}, "cell_type": "markdown", "id": "73780d81-150a-46a7-a935-ab5af22d52ce", "metadata": {}, "source": ["```python\n", "class Agent:\n", "    def __init__(self, name):\n", "        self.name = name\n", "```\n", "初始化代理\n"]}, {"cell_type": "code", "execution_count": null, "id": "651570b6-e2d8-45f6-9c0c-bf5871af49eb", "metadata": {}, "outputs": [], "source": ["from llama_index.agent.openai import OpenAIAssistantAgent\n", "\n", "agent = OpenAIAssistantAgent.from_new(\n", "    name=\"City bot\",\n", "    instructions=\"You are a bot designed to answer questions about cities (both unstructured and structured data)\",\n", "    tools=[sql_tool, vector_tool],\n", "    verbose=True,\n", ")"]}, {"cell_type": "code", "execution_count": null, "id": "f467f7d3-1cb1-47e9-8b67-a356f8be55f2", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: sql_tool with args: {\"input\":\"SELECT name, country FROM city_stats ORDER BY population DESC LIMIT 1\"}\n", "Got output: The city with the highest population is Tokyo, Japan.\n", "========================\n", "=== Calling Function ===\n", "Calling function: vector_tool with args: {\"input\":\"What are the arts and culture like in Tokyo, Japan?\"}\n", "Got output: Tokyo has a vibrant arts and culture scene. The city is home to many museums, including the Tokyo National Museum, which specializes in traditional Japanese art, the National Museum of Western Art, and the Edo-Tokyo Museum. There are also theaters for traditional forms of Japanese drama, such as the National Noh Theatre and the Kabuki-za. Tokyo hosts modern Japanese and international pop and rock music concerts, and the New National Theater Tokyo is a hub for opera, ballet, contemporary dance, and drama. The city also celebrates various festivals throughout the year, including the Sannō, Sanja, and Kanda Festivals. Additionally, Tokyo is known for its youth style, fashion, and cosplay in the Harajuku neighborhood.\n", "========================\n", "Tokyo, Japan, which has the highest population of any city, boasts a rich and diverse arts and culture landscape. The city is a hub for traditional Japanese art as showcased in prominent institutions like the Tokyo National Museum, and it also features artwork from different parts of the world at the National Museum of Western Art. Tokyo has a deep appreciation for its historical roots, with the Edo-Tokyo Museum presenting the past in a detailed and engaging manner.\n", "\n", "The traditional performing arts have a significant presence in Tokyo, with theaters such as the National Noh Theatre presenting classical Noh dramas and the iconic Kabuki-za offering enchanting Kabuki performances. For enthusiasts of modern entertainment, Tokyo is a prime spot for contemporary music, including both Japanese pop and rock as well as international acts.\n", "\n", "Opera, ballet, contemporary dance, and drama find a prestigious platform at the New National Theater Tokyo. Tokyo's calendar is filled with a variety of festivals that reflect the city's vibrant cultural heritage, including the Sannō, Sanja, and Kanda Festivals. Additionally, Tokyo is at the forefront of fashion and youth culture, particularly in the Harajuku district, which is famous for its unique fashion, style, and cosplay.\n", "\n", "This mix of traditional and modern, local and international arts and culture makes Tokyo a dynamic and culturally rich city.\n"]}], "source": ["response = agent.chat(\n", "    \"Tell me about the arts and culture of the city with the highest\"\n", "    \" population\"\n", ")\n", "print(str(response))"]}, {"cell_type": "code", "execution_count": null, "id": "9ffc35ba-3162-4fa9-81a9-174aa0f1dad3", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: vector_tool with args: {\"input\":\"What is the history of Berlin, Germany?\"}\n", "Got output: Berlin has a rich and diverse history. It was first documented in the 13th century and has served as the capital of various entities throughout history, including the Margraviate of Brandenburg, the Kingdom of Prussia, the German Empire, the Weimar Republic, and Nazi Germany. After World War II, the city was divided, with West Berlin becoming a part of West Germany and East Berlin becoming the capital of East Germany. Following German reunification in 1990, Berlin once again became the capital of all of Germany. Throughout its history, Berlin has been a center of scientific, artistic, and philosophical activity, and has experienced periods of economic growth and cultural flourishing. Today, it is a world city of culture, politics, media, and science, known for its vibrant arts scene, diverse architecture, and high quality of life.\n", "========================\n", "Berlin, the capital city of Germany, has a rich and complex history that stretches back to its first documentation in the 13th century. Throughout the centuries, Berlin has been at the heart of numerous important historical movements and events.\n", "\n", "Initially a small town, Berlin grew in significance as the capital of the Margraviate of Brandenburg. Later on, it ascended in prominence as the capital of the Kingdom of Prussia. With the unification of Germany, Berlin became the imperial capital of the German Empire, a position it retained until the end of World War I.\n", "\n", "The interwar period saw Berlin as the capital of the Weimar Republic, and it was during this time that the city became known for its vibrant cultural scene. However, the rise of the Nazi regime in the 1930s led to a dark period in Berlin's history, and the city was heavily damaged during World War II.\n", "\n", "Following the war's end, Berlin became a divided city. The division was physical, represented by the Berlin Wall, and ideological, with West Berlin aligning with democratic West Germany while East Berlin became the capital of the socialist East Germany.\n", "\n", "The fall of the Berlin Wall in November 1989 was a historic moment, leading to German reunification in 1990. Berlin was once again chosen as the capital of a united Germany. Since reunification, Berlin has undergone massive reconstruction and has become a hub of contemporary culture, politics, media, and science.\n", "\n", "Today, Berlin celebrates its diverse heritage, from its grand historical landmarks like the Brandenburg Gate and the Reichstag, to its remembrance of the past with monuments such as the Berlin Wall Memorial and the Holocaust Memorial. It is a city known for its cultural dynamism, thriving arts and music scenes, and a high quality of life. Berlin's history has shaped it into a unique world city that continues to play a significant role on the global stage.\n"]}], "source": ["response = agent.chat(\"Tell me about the history of Berlin\")\n", "print(str(response))"]}, {"cell_type": "code", "execution_count": null, "id": "0873463d-4790-4c17-bfe9-2ece610fe4b3", "metadata": {}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["=== Calling Function ===\n", "Calling function: sql_tool with args: {\"input\":\"SELECT name, country FROM city_stats\"}\n", "Got output: The cities in the city_stats table are Toronto from Canada, Tokyo from Japan, and Berlin from Germany.\n", "========================\n", "Here are the countries corresponding to each city:\n", "\n", "- Toronto: Canada\n", "- Tokyo: Japan\n", "- Berlin: Germany\n"]}], "source": ["response = agent.chat(\n", "    \"Can you give me the country corresponding to each city?\"\n", ")\n", "print(str(response))"]}], "metadata": {"kernelspec": {"display_name": "llama_index_v2", "language": "python", "name": "llama_index_v2"}, "language_info": {"codemirror_mode": {"name": "ipython", "version": 3}, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3"}}, "nbformat": 4, "nbformat_minor": 5}