使用云 AI 提供商创建嵌入
Cypher 函数 ai.text.embed 和过程 ai.text.embedBatch 允许您通过外部 AI 提供商为一个或多个文本片段生成嵌入。您需要拥有受支持提供商(OpenAI、Vertex AI、Azure OpenAI、Amazon Bedrock)中的任意一个的 API 令牌。
| 嵌入始终在 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 时,您可以省略对
|
嵌入写入数据库后,您即可利用它们 比较电影之间的相似度。