使用云 AI 提供商创建嵌入

Cypher 函数 ai.text.embed 和过程 ai.text.embedBatch 允许您通过外部 AI 提供商为一个或多个文本片段生成嵌入。您需要拥有受支持提供商(OpenAI、Vertex AI、Azure OpenAI、Amazon Bedrock)中的任意一个的 API 令牌。

本页假设您已 导入了推荐数据集配置好环境,展示如何基于电影的标题和情节生成并存储 Movie 节点的嵌入。

嵌入始终在 Neo4j 之外 生成,但 存储在 Neo4j 数据库中

设置环境

编码函数是 Neo4j GenAI 插件的一部分。

  • 在 Aura 实例上,插件默认已启用,因此如果您使用的是 Aura 上的 Neo4j,无需执行其他操作。

  • 对于自托管实例,需要先安装插件。您可以将 neo4j‑genai.jar 文件从 /products 移动到 Neo4j 主目录下的 /plugins,或在启动 Docker 容器时添加额外参数 --env NEO4J_PLUGINS='["genai"]'
    更多信息,请参见 配置 → 插件

为电影创建嵌入

下面的示例从数据库中获取所有 Movie 节点,生成电影标题与情节拼接后的嵌入,并将其作为额外的 embedding 属性添加到每个节点上。

import neo4j


URI = '<database-uri>'
AUTH = ('<username>', '<password>')
DB_NAME = '<database-name>'  # examples: 'recommendations-5.26', 'neo4j'

openAI_token = '<OpenAI API token>'


def main():
    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:  (1)
        driver.verify_connectivity()

        batch_size = 100
        batch_n = 1
        movies_batch = []
        with driver.session(database=DB_NAME) as session:
            # Fetch `Movie` nodes
            result = session.run('MATCH (m:Movie) RETURN m.plot AS plot, m.title AS title')
            for record in result:
                title = record.get('title')
                plot = record.get('plot')

                if title is not None and plot is not None:
                    movies_batch.append({
                        'title': title,
                        'plot': plot,
                        'to_encode': f'Title: {title}\nPlot: {plot}'  (2)
                    })

                # Import a batch; flush buffer
                if len(movies_batch) == batch_size:  (3)
                    import_batch(driver, movies_batch, batch_n)
                    movies_batch = []
                    batch_n += 1

            # Flush last batch
            import_batch(driver, movies_batch, batch_n)

        # Import complete, show counters
        records, _, _ = driver.execute_query('''
        MATCH (m:Movie WHERE m.embedding IS NOT NULL)
        RETURN count(*) AS countMoviesWithEmbeddings, size(m.embedding) AS embeddingSize
        ''', database_=DB_NAME)
        print(f"""
    Embeddings generated and attached to nodes.
    Movie nodes with embeddings: {records[0].get('countMoviesWithEmbeddings')}.
    Embedding size: {records[0].get('embeddingSize')}.
        """)


def import_batch(driver, nodes, batch_n):
    # Generate and store embeddings for Movie nodes
    driver.execute_query('''
    CALL ai.text.embedBatch($listToEncode, 'OpenAI', { (4)
      token: $token, model: 'text-embedding-3-small'
    }) YIELD index, vector  (5)
    MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
    CALL db.create.setNodeVectorProperty(m, 'embedding', toFloatList(vector))  (6)
    ''', movies=nodes, listToEncode=[movie['to_encode'] for movie in nodes], token=openAI_token,
    database_=DB_NAME)
    print(f'Processed batch {batch_n}')


if __name__ == '__main__':
    main()

'''
Movie nodes with embeddings: 9083.
Embedding size: 1536.
'''
1 driver 对象是与您的 Neo4j 实例交互的接口。更多信息,请参见 使用 Neo4j 与 Python 构建应用程序
2 需要让 OpenAI 编码为嵌入的字符串。
3 在提交完整批次到数据库之前,会先收集一定数量的嵌入。这可以避免一次性将整个数据集加载到内存中,从而防止潜在的超时(对大数据集尤为重要)。
4 过程 ai.text.embedBatch() 将批次提交给 OpenAI 进行编码。OpenAI 的模型 text-embedding-3-small 将文本嵌入为 1536 维向量(即包含 1536 个数字的列表)。有关受支持提供商和选项的列表,请参见 GenAI 提供商始终对同一数据集使用相同的模型生成嵌入:选择一个模型并在整个项目中坚持使用。
5 ai.text.embedBatch 返回的 index 用于将嵌入与电影对应,从而可以检索每个电影节点并为其附加相应的嵌入;vector 的类型为 VECTOR,该存储方式仅在 Enterprise Edition 中可用。
6 导入查询在每个节点 m 上设置一个新的 embedding 属性,值为嵌入 vector。Cypher 过程 db.create.setNodeVectorProperty 比将向量存为列表更高效。由于 VECTOR 类型的属性仅在 Enterprise Edition 中可用,额外调用 toFloatList 可确保在所有版本间兼容。若需在关系上设置向量属性,请使用 db.create.setRelationshipVectorProperty

使用 Enterprise Edition 时,您可以省略对 db.create.setNodeVectorProperty 的调用,直接通过 Cypher 子句 SET 将嵌入作为属性写入。相应的查询将会是如下形式:

CALL ai.text.embedBatch($listToEncode, 'OpenAI', {
  token: $token, model: 'text-embedding-3-small'
}) YIELD index, vector
MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
SET m.embedding = vector

嵌入写入数据库后,您即可利用它们 比较电影之间的相似度