Spring Boot 使用 Spring AI 构建知识库服务

目录

前言

环境准备

JDK17

Spring Boot 3.2.4

Ollama

PostgreSQL16

下载向量化模型

pom

yml

 EmbeddingController

向量化示例

向量化文本

向量化检索

ChatController

知识库示例

---

前言

做 AI 大模型技术调研时，参考的开源项目 Maxkb，它基于大模型做了一个知识库的应用，用户可构建自己的知识库，创建自己的应用然后关联知识库，这样可以基于知识库里的内容让大模型的回答更加符合我们地预期。

虽然 Maxkb 是使用 Python 写的，不过参考它用到的相关模型和数据库，可以用 Spring Boot 来构建一个自己的知识库服务。Spring Boot 也提供了 AI 相关的库，这使得接入 AI 大模型能力也十分方便。

以下是使用Spring Boot 构建一个知识库的样例。

环境准备

JDK17

Spring Boot 3.2.4

Ollama

如果没有安装，可参考以下博客安装并下载 qwen2 大模型

Ollama 部署本地大模型-CSDN博客文章浏览阅读456次，点赞15次，收藏13次。工作需要，被指派对大模型做技术调研，又不想花钱买各大产商的key，所以就考虑部署本地的大模型，了解到 Ollama 开源框架，专为在本地机器上便捷部署和运行大型语言模型（LLM）而设计。便在自己电脑上安装了，记录下安装以及使用过程。https://blog.csdn.net/typeracer/article/details/140707448

PostgreSQL16

如果没有安装，可参考以下博客安装并安装 vector 扩展

Windows 安装 PostgreSQL 并安装 vector 扩展-CSDN博客文章浏览阅读502次，点赞5次，收藏5次。调研大模型时，了解到一些大模型的应用，其中一个就是知识库，用户可以上传文档到知识库中，系统解析文档并将内容向量化保存起来，以便在和模型交互时使用。在和大模型对话时，可以先向量化检索自己的知识库，如果命中，则返回文档内容，然后将文档内容作为大模型的输入，以便让大模型回答得更加符合我们的预期。这里就涉及到了向量化数据的存储，可以为 PostgreSQL 安装 vector 扩展来存储向量化数据。记录 Windows 安装 PostgreSQL 以及 vector 扩展的步骤。https://blog.csdn.net/typeracer/article/details/140711057选择一个数据库实例，安装扩展，然后创建 vector_store 表

CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS hstore;
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";

CREATE TABLE IF NOT EXISTS vector_store (
	id uuid DEFAULT uuid_generate_v4() PRIMARY KEY,
	content text,
	metadata json,
	embedding vector(768)
);

CREATE INDEX ON vector_store USING HNSW (embedding vector_cosine_ops);

下载向量化模型

shibing624_text2vec-base-chinese

下载地址：

https://huggingface.co/shibing624/text2vec-base-chinesehttps://huggingface.co/shibing624/text2vec-base-chinese将下载的模型放到项目工程 resources 文件夹下

pom

    <dependencyManagement>
        <dependencies>
            <dependency>
                <groupId>org.springframework.boot</groupId>
                <artifactId>spring-boot-dependencies</artifactId>
                <version>3.2.4</version>
                <type>pom</type>
                <scope>import</scope>
            </dependency>
            <dependency>
                <groupId>org.springframework.ai</groupId>
                <artifactId>spring-ai-bom</artifactId>
                <version>1.0.0-SNAPSHOT</version>
                <type>pom</type>
                <scope>import</scope>
            </dependency>
        </dependencies>
    </dependencyManagement>

    <dependencies>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
        </dependency>

        <dependency>
            <groupId>org.postgresql</groupId>
            <artifactId>postgresql</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-tika-document-reader</artifactId>
        </dependency>

    </dependencies>

    <repositories>
        <!-- 里程碑（Milestone）-->
        <repository>
            <id>spring-milestones</id>
            <name>Spring Milestones</name>
            <url>https://repo.spring.io/milestone</url>
            <snapshots>
                <enabled>true</enabled>
            </snapshots>
        </repository>
        <!-- 快照（Snapshot）-->
        <repository>
            <id>spring-snapshots</id>
            <name>Spring Snapshots</name>
            <url>https://repo.spring.io/snapshot</url>
            <releases>
                <enabled>false</enabled>
            </releases>
        </repository>
    </repositories>

yml

spring:
  datasource:
    url: jdbc:postgresql://localhost:5432/postgres
    username: postgres
    password: postgres
    driver-class-name: org.postgresql.Driver
  ai:
    ollama:
      base-url: http://localhost:11434
      chat:
        options:
          model: qwen2

    embedding:
      transformer:
        onnx:
          modelUri: classpath:/shibing624_text2vec-base-chinese/onnx/model.onnx
        tokenizer:
          uri: classpath:/shibing624_text2vec-base-chinese/onnx/tokenizer.json
    vectorstore:
      pgvector:
        index-type: HNSW
        distance-type: COSINE_DISTANCE
        dimensions: 768

server:
  port: 8888

 更多ai模型和向量模型的详细配置可参照官方文档

ETL Pipeline :: Spring AI Referencehttps://docs.spring.io/spring-ai/reference/api/etl-pipeline.html

 EmbeddingController

@RestController
public class EmbeddingController {

    @Autowired
    @Qualifier("embeddingModel")
    EmbeddingModel embeddingModel;

    @Autowired
    VectorStore vectorStore;

    @GetMapping("/ai/embedding")
    public void embed(@RequestParam(value = "message") String message) {
        System.out.println(embeddingModel.embed(message));
        System.out.println("ok");
    }

    @PostMapping("/ai/vectorStore")
    public List<String> vectorStore(@RequestParam(name = "file") MultipartFile file) throws Exception {
        // 从IO流中读取文件
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new InputStreamResource(file.getInputStream()));
        // 将文本内容划分成更小的块
        List<Document> splitDocuments = new TokenTextSplitter()
                .apply(tikaDocumentReader.read());
        // 存入向量数据库，这个过程会自动调用embeddingModel,将文本变成向量再存入。
        vectorStore.add(splitDocuments);

        return splitDocuments.stream().map(Document::getContent).collect(toList());
    }

    @GetMapping("/ai/vectorSearch")
    public List<String> vectorSearch(@RequestParam(name = "text") String text) {

        List<Document> documents = vectorStore.similaritySearch(SearchRequest.query(text).withTopK(1));
        
        return documents.stream().map(Document::getContent).collect(toList());
    }
}

向量化示例

使用 AI 随便生成一段文本内容，这里让 AI 生成了 ollama 相关的文本

向量化文本

上传文档到知识库，其实就是调用向量化接口 /ai/vectorStore

查询向量数据库，可以看到插入了两条数据

代码中使用了 TokenTextSplitter 进行了分段处理，文档内容过多会被切割成多个再进行向量化

向量化检索

调用 /ai/vectorSearch 接口进行向量化检索

查询 vector_store 表，返回最匹配的一行数据

ChatController

@RestController
public class ChatController {

    @Autowired
    OllamaChatModel ollamaChatModel;
    @Autowired
    OllamaApi ollamaApi;
    @Autowired
    VectorStore vectorStore;

    @GetMapping("/ai/generate")
    public String generate(@RequestParam(value = "message") String message) {
        return ollamaChatModel.call(message);
    }

    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return ollamaChatModel.stream(prompt);
    }

    @GetMapping("/ai/ollamaApi")
    public OllamaApi.ChatResponse ollamaApi(@RequestParam(value = "message") String message) {
        //从知识库检索相关信息，再将检索得到的信息同用户的输入一起构建一个prompt，最后调用ollama api
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.query(message).withTopK(1));
        String targetMessage = String.format("已知信息：%s\n 用户提问：%s\n",
                documents.get(0).getContent(), message);

        OllamaApi.ChatRequest request = OllamaApi.ChatRequest.builder("qwen2")
                .withStream(false) // not streaming
                .withMessages(List.of(
                        OllamaApi.Message.builder(OllamaApi.Message.Role.USER)
                                .withContent(targetMessage)
                                .build()))
                .withOptions(OllamaOptions.create().withTemperature(0.9f))
                .build();

        return ollamaApi.chat(request);
    }
}

知识库示例

调用 /ai/ollamaApi 接口

先从知识库检索相关信息

再将检索得到的信息同用户的输入一起构建一个 prompt

最后调用 ollama api