Neo4j实现向量检索

最近因为Dify、RagFlow这样的智能体的镜像拉取的速度实在太麻烦，一狠心想实现自己的最简单的RAG。
因为之前图数据库使用到了neo4j，查阅资料才发现​​Neo4j从5.11版本开始支持向量索引，提供一个真实可用的单元测试案例。

Neo4j建向量索引表

-- 创建索引语句
CREATE VECTOR INDEX doc_embedding_index
FOR (n:Doc) ON (n.embedding)
OPTIONS {indexConfig: {`vector.dimensions`: 1024,`vector.similarity_function`: 'cosine'
}}

这里是1024是根据不同的Embedding模型设置的，bge-large-zh-v1.5中使用1024

Java调用Embedding模型插入数据

// 连接 Neo4j
Driver driver = GraphDatabase.driver("bolt://127.0.0.1:7687", AuthTokens.basic("neo4j", "123456"));
Session session = driver.session()// 初始化 OpenAI 模型
// 这里使用bge-large-zh-v1.5
OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder().baseUrl("base url").apiKey("api key").modelName("bge-large-zh-v1.5").build();
// 创建语料库
String[] contents = {"性别","年龄","男性","女性","高收入","低收入",
};
for (int i = 0; i < contents.length; i++) {String text = contents[i];Embedding embedding = embeddingModel.embed(text).content();List<Float> vector = embedding.vectorAsList();session.run("CREATE (d:Doc {id: $id, content: $content})", Values.parameters("id", "doc" + i, "content", text));session.run("""MATCH (d:Doc {id: $id})CALL db.create.setNodeVectorProperty(d, 'embedding', $vector)""", Values.parameters("id", "doc" + i, "vector", vector));
}String question = " 帮我圈个高价值年轻男性客群";// 向量化问题
Embedding questionEmbedding = embeddingModel.embed(question).content();
List<Float> queryVector = questionEmbedding.vectorAsList();// 检索数据
List<String> docs;
try (Session session = driver.session()) {Result result = session.run("""CALL db.index.vector.queryNodes('doc_embedding_index', 3, $vector)YIELD node, scoreRETURN node.content AS content""", Values.parameters("vector", queryVector));// 这里的3可以换成其他数字，指取多少个Stream<String> content = result.stream().map(record -> record.get("content").asString());docs = content.collect(Collectors.toList());System.out.println("docs: " + docs);
}

CALL db.index.vector.queryNodes('doc_embedding_index', 3, $vector)YIELD node, scoreRETURN node.content AS content

这里的3是检索的条数，可以换成自己需要检索的数量。

相关依赖

基于JDK17

<dependency><groupId>dev.langchain4j</groupId><artifactId>langchain4j-open-ai</artifactId><version>1.0.0-rc1</version>
</dependency>
<dependency><groupId>dev.langchain4j</groupId><artifactId>langchain4j</artifactId><version>1.0.0-rc1</version>
</dependency>
<dependency><groupId>org.neo4j.driver</groupId><artifactId>neo4j-java-driver</artifactId><version>5.28.5</version>
</dependency>