在本文中，我想分享向 LlamaIndex 添加新的数据加载器的过程。首先，我们将了解 LlamaIndex 是什么，并尝试一个简单的示例，使用简单的 CSV 加载器为 LLM 查询提供额外上下文。然后，我们将看到向 LlamaIndex 添加新的图数据库加载器是多么容易。最后，我们将实践中尝试我添加的那个新加载器以及另一个用于 GraphQL API 的加载器，并看看它们的额外上下文如何帮助 LLM 更好地回答问题。

背景/上下文

我当时正在听 "This Week in ML" (twiml) 播客，LlamaIndex（前身为 GPT-Index）的 Jerry Liu 在播客中解释了这个库背后的想法，即利用来自各种来源的数据来丰富 LLM 的查询上下文。

LlamaIndex 是一个工具包，用于使用上下文学习 (in-context learning) 来增强 LLM，使其可以使用您自己的（私有）数据。它负责从大型知识库中选择要检索的正确上下文。为了实现这一点，它利用了许多连接器或加载器（来自 LlamaHub）和数据结构（索引），以有效地将预处理的数据作为 Documents 提供。

每种类型的索引都以不同的方式存储文档，例如通过嵌入向量用于向量搜索，或者作为简单的列表、图或树结构。这些索引被用作 LLM 的查询接口，透明地嵌入相关的上下文。

除了从 LLM 获得更高质量的响应外，您还会获得用于构建答案的文档。LlamaIndex 还支持思维链推理、比较/对比查询以及数据库的自然语言查询。

参见 来自 Jerry 的 这份 演示文稿

本博客文章的所有代码都可以在这个 Colab Notebook 中找到。

使用基本 CSV 加载器

这里是使用基本 CSV 加载器为 LlamaIndex 提供文档的示例。

在我们的 Notebook 中，我们通过 Countries List 项目 (MIT) (原始源文件) 下载 countries.csv。

我们的依赖项是 llama-index 和 python-dotenv。

!pip install llama-index==0.6.19 python-dotenv

我们需要提供我们的 OpenAI API 密钥，为了避免在 Notebook 中意外泄露，我上传了一个 openai.env 文件，并使用 dotenv 库将内容加载为环境变量。

下一步，我们加载 env 文件并准备 OpenAI ChatGPTLLMPredictor (默认使用 gpt-3.5-turbo)，并将其添加到 ServiceContext 中。

import os
from pathlib import Path
from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, ServiceContext, GPTListIndex
from llama_index.llm_predictor.chatgpt import ChatGPTLLMPredictor
from dotenv import load_dotenv
from llama_index import download_loader

load_dotenv("openai.env")

llm_predictor = ChatGPTLLMPredictor()
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)

现在我们可以使用加载器加载 CSV 并将其转换为文档，创建一个 GPT 索引（本例中是 VectorStoreIndex），LlamaIndex 就可以利用它来检索相关信息，并将其作为上下文传递给 LLM。

初始化 CSV 加载器和 GPTVectorStoreIndex

SimpleCSVReader = download_loader("SimpleCSVReader")
loader = SimpleCSVReader(concat_rows=False)
documents = loader.load_data(file=Path('./countries.csv'))

print(documents)
index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)

来自 CSV 加载器的文档

[Document(text='country, capital, type', doc_id='67c30c68-7d9f-4906-945b-9affc96f95d2', embedding=None, doc_hash='3a506ebea9c04655b51406d79fdf5e3a87c3d8ff5b5387aace3e5a79711a21b8', extra_info=None),
Document(text='Abkhazia, Sukhumi, countryCapital', doc_id='6e6be4b5-051f-48e0-8774-6d48e0444785', embedding=None, doc_hash='ea387d0eab94cc6c59f98c473ac1f0ee64093901673b43e1c0d163bbc203026e', extra_info=None),
...]

默认情况下，CSV 加载器并没有为每个 CSV 行创建一个文档，而是只为整个文档创建一个文档，但您可以配置它，使其将 CSV 转换为每行一个文档。

LlamaIndex 支持更复杂的设置，包括不同类型的索引，允许将它们串联起来，甚至有条件地选择一个或另一个。这里我们只做最基本的操作来演示我们的加载器。

在使用适当的加载器设置索引并连接索引后，我们现在可以使用索引作为 LLM 查询引擎来执行我们的用户查询。

为了证明 LLM 仍然能够利用其世界知识，我们可以用英语（系统）、德语（问题）和法语（要求回答）混合提问。

queryEngine = index.as_query_engine()

queryEngine.query("""
Provide the answer in French.
Question: Was ist die Hauptstadt von Albanien?
""")

正如您在下面的回复中所见，它不仅用法语正确回答了我们的问题 La capitale de l’Albanie est Tirana.（阿尔巴尼亚的首都是地拉那），还提供了它用于生成答案的文档。

Response(response="La capitale de l'Albanie est Tirana.", 
source_nodes=[NodeWithScore(node=Node(text='              <td>Albania</td>', doc_id='3decbee1-98cc-4650-a071-ed25cd3e00d5', embedding=None, doc_hash='7d9d85082095471a9663690742d2d49fc37b2ec37cc5acf4e99e006a68a17742', extra_info=None, 
node_info={'start': 0, 'end': 30, '_node_type': <NodeType.TEXT: '1'>}, 
relationships={<DocumentRelationship.SOURCE: '1'>: '7b6c861f-2c2f-4905-a047-edfc25f7df19'}), score=0.7926356007369129), 
NodeWithScore(node=Node(text='              <td>Algiers</td>', doc_id='8111b737-9f45-4855-8cd8-f958d4eb0ccd', embedding=None, doc_hash='8570a02a057a6ebbd0aff6d3f63c9f29a0ee858a81d913298d31b025101d1e44', 
extra_info=None, node_info={'start': 0, 'end': 30, '_node_type': <NodeType.TEXT: '1'>}, relationships={<DocumentRelationship.SOURCE: '1'>: '22e11ac6-8375-4d0c-91c6-4750fc63a375'}), score=0.7877589022795918)], extra_info={'3decbee1-98cc-4650-a071-ed25cd3e00d5': None, '8111b737-9f45-4855-8cd8-f958d4eb0ccd': None})

LlamaIndex 加载器

LlamaHub 中现有数据源的数量令人印象深刻，我在 仓库 中数了 100 多个集成。您可以找到从 Google Docs 到 GitHub，再到关系数据库的各种集成。

但我缺少了我最喜欢的两种技术：Facebook 开源的 API 查询语言 GraphQL，以及 Neo4j 等图数据库，它们是存储和管理大量连接数据（例如知识图谱）的最佳方式。

所以我想：“添加它们能有多难呢？:)"

添加新的加载器

添加新的加载器非常简单直接。llama-hub 仓库中有一个脚本可以帮助添加新的加载器。运行 ./add-loader.sh <folder> 会添加骨架文件。

为了熟悉现有的实现，我查看了 数据库（关系型） 和 MongoDB 集成，前者是为图数据库而看，后者是为 GraphQL 而看。

这足够容易了，我们只需要加载器所需的依赖项，用简单的 API 实现 base.py，以及包含解释和代码示例的 README.md`。

我的加载器与现有加载器的主要区别在于，它们不使用硬编码的字段名称来提取查询结果中的相关值，而是将结果转换为 YAML 格式。

我选择 YAML 并不是因为我喜欢它，而是因为它最接近于嵌套键值对树的文本表示，用户可以将其写成嵌套的无序列表。

下面是图数据库实现的示例代码（GraphQL 的类似）。

添加图数据库加载器

我添加了 neo4j 依赖项的要求，这是一个基于 Bolt 协议的 Cypher 查询语言 Python 驱动程序，它也适用于 Memgraph 和 AWS Neptune。

然后我添加了 __init__ 的代码，用于接收数据库服务器 URI、数据库名称和凭据，以连接并创建驱动程序实例。

load_data 方法接收要运行的查询和可选参数。它通过调用驱动程序的 execute_query 方法来实现。

每行结果都被映射到一个 LlamaIndex Document 中，其中文档的 text 属性是结果的 YAML 表示。

"""Graph Database Cypher Reader."""

from typing import Dict, List, Optional

from llama_index.readers.base import BaseReader
from llama_index.readers.schema.base import Document

import yaml

class GraphDBCypherReader(BaseReader):
    """Graph database Cypher reader.

    Combines all Cypher query results into the Document type used by LlamaIndex.

    Args:
        uri (str): Graph Database URI
        username (str): Username
        password (str): Password

    """

    def __init__(
        self,
        uri: str,
        username: str,
        password: str,
        database: str
    ) -> None:
        """Initialize with parameters."""
        try:
            from neo4j import GraphDatabase, basic_auth

        except ImportError:
            raise ImportError(
                "`neo4j` package not found, please run `pip install neo4j`"
            )
        if uri:
            if uri is None:
                raise ValueError("`uri` must be provided.")
            self.client = GraphDatabase.driver(uri=uri, auth=basic_auth(username, password))
            self.database = database

    def load_data(
        self, query: str, parameters: Optional[Dict] = None
    ) -> List[Document]:
        """Run the Cypher with optional parameters and turn results into documents

        Args:
            query (str): Graph Cypher query string.
            parameters (Optional[Dict]): optional query parameters.

        Returns:
            List[Document]: A list of documents.

        """
        if parameters is None:
            parameters = {}

        records, summary, keys = self.client.execute_query(query, parameters, database_ = self.database)

        documents = [Document(yaml.dump(entry.data())) for entry in records]

        return documents

您现在可以使用数据加载器了。如果您想在代码中使用它，只需从相关文件中导入 GraphDBCypherReader 并按照以下步骤操作。

如果您希望将加载器提交到 LlamaHub，过程相当简单。在 readme 中添加了一个使用始终在线演示服务器（包含 StackOverflow 数据）的示例后，我就可以创建一个 拉取请求 了。经过简短讨论，该 PR 很快被合并。

非常感谢 Jerry 带来了顺畅的体验。

现在让我们看看如何使用我们的两个加载器。

使用图数据库加载器

GraphDB Cypher 加载器连接到图数据库，图数据库是专门的数据库，它们不像表格那样存储数据，而是存储实体（节点）及其关系。由于它们是无模式的，您可以存储真实世界的知识而不会牺牲其丰富性。

关系也可以包含属性，这些属性可以表示时间、权重、成本或任何定义具体关系的属性。任何节点都可以根据需要拥有任意数量的属性或关系。

要查询图数据库，您可以使用 Cypher 查询语言，这是一种基于模式的语言，它以可视化的 ASCII 艺术模式表达这些关系。您用括号将节点括起来 ()，并用箭头 --> 绘制关系，并在方括号中添加其他约束。除此之外，Cypher 提供了许多从 SQL 中熟知的功能，还支持许多图操作，以及处理嵌套文档、列表和字典等数据结构。

让我们使用一个电影图数据库，并向 LLM 提问关于常见动作片情节的问题。

设置 ServiceContext 和 ChatGPTLLMPredictor 与之前相同。

然后我们获取 GraphDBCypherReader 并将其连接到我们的数据库（经许可使用来自 TheMovieDB 的小型电影图示例）。

GraphDBCypherReader = download_loader('GraphDBCypherReader')

reader = GraphDBCypherReader(uri = "neo4j+s://demo.neo4jlabs.com", \
    username = "recommendations", password = "recommendations", database = "recommendations")

然后我们定义对图数据库的查询，包含一个年份参数，以便我们选择更近期的电影。加载数据时，每行结果都应转换为一个 Document，其中文档的 text 属性是该行的 YAML 表示。

query = """
    MATCH (m:Movie)-[rel:ACTED_IN|DIRECTED|IN_GENRE]-(other)
    WHERE $year < m.year and m.imdbRating > $rating
    WITH m, type(rel) as relation, collect(other.name) as names
    RETURN m.title as title, m.year as year, m.plot as plot, relation, names
    ORDER BY m.year ASC
"""

documents = reader.load_data(query, parameters = {"year":1990,"rating":8})
index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)

print(len(documents))
print(documents[0:5])

输出将类似于以下内容

829
[Document(text='names:\n- Saifei He\n- Li Gong\n- Jingwu Ma\n- Cuifen Cao\nplot: A young woman becomes the fourth wife of a wealthy lord, and must learn to live\n  with the strict rules and tensions within the household.\nrelation: ACTED_IN\ntitle: Raise the Red Lantern (Da hong deng long gao gao gua)\nyear: 1991\n', doc_id='782d9a63-251b-4bb8-aa3d-5d8f6d1fb5d2', embedding=None, doc_hash='f9fd966bc5f2234e94d09efebd3be008db8c891f8666c1a364abf7812f5d7a1c', extra_info=None), Document(text='names:\n- Yimou Zhang\nplot: A young woman becomes the fourth wife of a wealthy lord, and must learn to live\n  with the strict rules and tensions within the household.\nrelation: DIRECTED\ntitle: Raise the Red Lantern (Da hong deng long gao gao gua)\nyear: 1991\n', doc_id='2e13caf6-b9cf-4263-a264-7121bc77d1ee', embedding=None, doc_hash='e1f340ed1fac2f1b8d6076cfc2c9e9cb0109d5d11e5dcdbf3a467332f5995cb1', extra_info=None), ...]

现在我们可以使用我们的 index 来运行 LLM 查询以回答我们想提出的问题。

queryEngine= index.as_query_engine()

queryEngine.query("""
What are the most common plots in action movies?
""")

答案显示，LLM 可以利用输入，理解“动作片”类型，并能概括它们的情节。以下是它的回答。

根据给定的上下文信息，动作片中最常见的情节似乎是抢劫和对抗控制势力。然而，需要注意的是，这一结论基于有限的样本量，可能无法代表所有动作片。

Response(response='Based on the given context information, it appears that the most common plots in action movies are heists and battles against controlling forces. However, it is important to note that this conclusion is based on a limited sample size and may not be representative of all action movies.',


source_nodes=[NodeWithScore(node=Node(text='names:\n- Action\n- Crime\n- Thriller\nplot: A group of professional bank robbers start to feel the heat from police when\n  they unknowingly leave a clue at their latest heist.\nrelation: IN_GENRE\ntitle: Heat\nyear: 1995\n', doc_id='bb117618-1cce-4cec-bd9b-8645ab0b50a3', embedding=None, doc_hash='4d493a9f33eb7a1c071756f61e1975ae5c313ecd42243f81a8827919a618468b', extra_info=None, node_info={'start': 0, 'end': 215, '_node_type': <NodeType.TEXT: '1'>}, relationships={<DocumentRelationship.SOURCE: '1'>: 'dbfffdae-d88c-49e2-9d6b-83dad427a3f3'}), score=0.8247381316731472), NodeWithScore(node=Node(text='names:\n- Thriller\n- Sci-Fi\n- Action\nplot: A computer hacker learns from mysterious rebels about the true nature of his\n  reality and his role in the war against its controllers.\nrelation: IN_GENRE\ntitle: Matrix, The\nyear: 1999\n', doc_id='c4893c61-32ee-4d05-b559-1f65a5197e5e', embedding=None, doc_hash='0b6a080bf712548099c5c8c1b033884a38742c73dc23d420ac2e677e7ece82f4', extra_info=None, node_info={'start': 0, 'end': 227, '_node_type': <NodeType.TEXT: '1'>}, relationships={<DocumentRelationship.SOURCE: '1'>: '6c8dea11-1371-4f5a-a1a1-7f517f027008'}), score=0.8220633045996049)], extra_info={'bb117618-1cce-4cec-bd9b-8645ab0b50a3': None, 'c4893c61-32ee-4d05-b559-1f65a5197e5e': None})

使用 GraphQL 加载器

GraphQL 加载器同样易于使用。

GraphQL 不是数据库查询语言，而是一种基于严格模式的 API 查询语言，该模式以“类型定义”表示。在其中，您可以表达您的实体、它们的属性（字段），包括标量数据类型以及指向其他实体的对象数据类型。

GraphQL 本身是一种基于树的查询语言，它表达了您想从根查询开始获取数据的嵌套结构。从该查询返回的每个实体的字段都可以被选择，对于对象字段，您可以进一步选择引用实体的字段，依此类推，几乎无限延伸（受限于 API 限制）。

有许多 GraphQL 库，最著名的是 JavaScript 参考实现，还有 Python 的 gql，以及与数据库（如 Hasura、Prisma 或 Neo4j-GraphQL-Library）的集成。现在一些大型项目提供 GraphQL API，包括 GitHub、Spotify、Twitter。

这个演示类似于我们的第一个演示。我们使用一个公共 GraphQL 端点 (https://countries.trevorblades.com/)，它提供了一个大陆→国家→首都的结构。（根据 MIT 许可）

这里是类型定义的子集。

type Query {
    continent(code: ID!): Continent
    continents(filter: ContinentFilterInput = {}): [Continent!]!
    countries(filter: CountryFilterInput = {}): [Country!]!
    country(code: ID!): Country
    language(code: ID!): Language
    languages(filter: LanguageFilterInput = {}): [Language!]!
}

type Continent {
    code: ID!
    countries: [Country!]!
    name: String!
}

type Country {
    awsRegion: String!
    capital: String
    code: ID!
    continent: Continent!
    currencies: [String!]!
    currency: String
    emoji: String!
    emojiU: String!
    languages: [Language!]!
    name(lang: String): String!
    native: String!
    phone: String!
    phones: [String!]!
    states: [State!]!
    subdivisions: [Subdivision!]!
}
...

在我们的演示中，我们再次像之前一样使用 ChatGPTLLMPredictor 定义了 ServiceContext。然后我们获取 GraphQLReader 加载器并将其指向端点的 URL。您还可以提供额外的 HTTP 头部，例如用于身份验证。

GraphQLReader = download_loader('GraphQLReader')
reader = GraphQLReader(uri = "https://countries.trevorblades.com/", headers = {})

query = """
query getContinents {
  continents {
    name
    countries {
      name
      capital
    }
  }
}
"""
documents = reader.load_data(query, variables = {})
print(len(documents))
print(documents)

我们看到它找到了 7 个大陆及其国家和首都，每个根结果（大陆）都转换为一个文档。

7
[Document(text='countries:\n- capital: Luanda\n  name: Angola\n- capital: Ouagadougou\n  name: Burkina Faso\n- capital: Bujumbura\n  name: Burundi\n- capital: Porto-Novo\n  name: Benin\n- capital: Gaborone\n  name: Botswana\n- capital: Kinshasa\n  name: Democratic Republic of the Congo\n- capital: Bangui\n  name: Central African Republic\n....',doc_id='b82fec36-5e82-4246-b7ab-f590bf6741ab', embedding=None, doc_hash='a4caa760423d6ca861b9332f386add3c449f1683168391ae10f7f73a691a2240', extra_info=None)]

我们再次用德语稍微刁难了一下 LLM，问它“北美有哪些首都”。

index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
queryEngine= index.as_query_engine()

response = queryEngine.query("""
Question: Welche Hauptstädte liegen in Nordamerika?
Answer:
""")

response.response

我很惊讶，因为我只预料到几个国家和城市。但我们得到了北美洲的 27 个国家。这表明我们的认知受到西方世界观的偏见影响。

Die Hauptstädte, die in Nordamerika liegen, sind Ottawa, San Jos\xE9, Havana, Willemstad, Roseau, Santo Domingo, St. George's, Nuuk, Guatemala City, Tegucigalpa, Port-au-Prince, Kingston, Basseterre, George Town, Castries, Marigot, Fort-de-France, Plymouth, Mexico City, Managua, Panama City, Saint-Pierre, San Juan, San Salvador, Philipsburg, Cockburn Town, Port of Spain, Washington D.C., Kingstown und Road Town.

我们也可以反转 GraphQL 查询，然后获取 250 个国家及其各自的首都和大陆。

query = """
query getCountries {
  countries {
    name
    capital
    continent {
        name
    }
  }
}
"""
documents = reader.load_data(query, variables = {})
print(len(documents))
print(documents)

这两个文档列表应该同样有效，但让我们看看。

这次 LLM 的回答更加有限。我不确定是不是因为索引提供给 LLM 选择的文档较少。

index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
queryEngine= index.as_query_engine()

response = queryEngine.query("""
Question: Which capitals are in North America?
Answer:
""")

response.response

Washington D.C. and Mexico City are in North America.

结论

添加新的数据加载器到 LlamaHub 确实非常顺畅，非常感谢 Jerry Liu 让这一切如此容易。请告诉我您在使用这些加载器时做了些什么，以及您是否有任何反馈。

如果在接下来的几周内我能找到时间，我也想研究一下 KnowledgeGraphIndex，看看我的图数据库加载器能否很好地填充它。