介绍

🤖️ 一种利用 langchain 思想实现的基于本地知识库的问答应用，目标期望建立一套对中文场景与开源模型支持友好、可离线运行的知识库问答解决方案。

💡 受 GanymedeNil 的项目 document.ai 和 AlexZhangji 创建的 ChatGLM-6B Pull Request 启发，建立了全流程可使用开源模型实现的本地知识库问答应用。本项目的最新版本中通过使用 FastChat 接入 Vicuna, Alpaca, LLaMA, Koala, RWKV 等模型，依托于 langchain 框架支持通过基于 FastAPI 提供的 API 调用服务，或使用基于 Streamlit 的 WebUI 进行操作。

✅ 依托于本项目支持的开源 LLM 与 Embedding 模型，本项目可实现全部使用开源模型离线私有部署。与此同时，本项目也支持 OpenAI GPT API 的调用，并将在后续持续扩充对各类模型及模型 API 的接入。

⛓️ 本项目实现原理如下图所示，过程包括加载文件 -> 读取文本 -> 文本分割 -> 文本向量化 -> 问句向量化 -> 在文本向量中匹配出与问句向量最相似的 top k个 -> 匹配出的文本作为上下文和问题一起添加到 prompt中 -> 提交给 LLM生成回答。

实现原理：

SSE(Server Sent Event)，直译为服务器发送事件，也就是服务器主动发送事件，客户端可以获取到服务器发送的事件。

1、SseEmitter类简介

SpringBoot 利用 SseEmitter 来支持SSE，并对SSE规范做了一些封装，使用起来非常简单。我们操作SseEmitter对象，关注消息文本即可。

SseEmitter类的几个方法：

send()：发送数据，如果传入的是一个非SseEventBuilder对象，那么传递参数会被封装到 data 中。

complete()：表示执行完毕，会断开连接。

onTimeout()：连接超时时回调触发。

onCompletion()：结束之后的回调触发。

onError()：报错时的回调触发。

2.1 创建 SseServer

我们创建一个 SseServer来简单封装一下业务操作SSE的方法。

import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

/**

* SseServer业务封装类来操作SEE

*/

@Slf4j

public class SseServer {

    /**

    * 当前连接总数

    */

    private static AtomicInteger currentConnectTotal = new AtomicInteger(0);

    /**

    * messageId的 SseEmitter对象映射集

    */

    private static Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    /**

    * 创建sse连接

    *

    * @param messageId - 消息id（唯一）

    * @return

    */

    public static SseEmitter createConnect(String messageId) {

        /**

        * 设置连接超时时间。0表示不过期，默认是30秒，超过时间未完成会抛出异常

        */

        SseEmitter sseEmitter = new SseEmitter(0L);

        /*

        // 超时时间设置为3s，设置前端的重试时间为1s。重连时，注意总数的统计

        SseEmitter sseEmitter = new SseEmitter(3_000L);

        try {

            sseEmitter.send(

                    SseEmitter.event()

                    .reconnectTime(1000L)

                    //.data("前端重连成功") // 重连成功的提示信息

            );

        } catch (IOException e) {

            log.error("前端重连异常 ==> messageId={}, 异常信息：", messageId, e.getMessage());

            e.printStackTrace();

        }*/

        // 注册回调

        sseEmitter.onCompletion(completionCallBack(messageId));

        sseEmitter.onTimeout(timeOutCallBack(messageId));

        sseEmitter.onError(errorCallBack(messageId));

        sseEmitterMap.put(messageId, sseEmitter);

        //记录一下连接总数。数量+1

        int count = currentConnectTotal.incrementAndGet();

        log.info("创建sse连接成功 ==> 当前连接总数={}， messageId={}", count, messageId);

        return sseEmitter;

    }

    /**

    * 给指定 messageId发消息

    *

    * @param messageId - 消息id（唯一）

    * @param message  - 消息文本

    */

    public static void sendMessage(String messageId, String message) {

        if (sseEmitterMap.containsKey(messageId)) {

            try {

                sseEmitterMap.get(messageId).send(message);

            } catch (IOException e) {

                log.error("发送消息异常 ==> messageId={}, 异常信息：", messageId, e.getMessage());

                e.printStackTrace();

            }

        } else {

            throw new RuntimeException("连接不存在或者超时， messageId=" + messageId);

        }

    }

    /**

    * 给所有 messageId广播发送消息

    *

    * @param message

    */

    public static void batchAllSendMessage(String message) {

        sseEmitterMap.forEach((messageId, sseEmitter) -> {

            try {

                sseEmitter.send(message, MediaType.APPLICATION_JSON);

            } catch (IOException e) {

                log.error("广播发送消息异常 ==> messageId={}, 异常信息：", messageId, e.getMessage());

                removeMessageId(messageId);

            }

        });

    }

    /**

    * 给指定 messageId集合群发消息

    *

    * @param messageIds

    * @param message

    */

    public static void batchSendMessage(List<String> messageIds, String message) {

        if (CollectionUtils.isEmpty(messageIds)) {

            return;

        }

        // 去重

        messageIds = messageIds.stream().distinct().collect(Collectors.toList());

        messageIds.forEach(userId -> sendMessage(userId, message));

    }

    /**

    * 给指定组群发消息（即组播，我们让 messageId满足我们的组命名确定即可）

    *

    * @param groupId

    * @param message

    */

    public static void groupSendMessage(String groupId, String message) {

        if (MapUtils.isEmpty(sseEmitterMap)) {

            return;

        }

        sseEmitterMap.forEach((messageId, sseEmitter) -> {

            try {

                // 这里 groupId作为前缀

                if (messageId.startsWith(groupId)) {

                    sseEmitter.send(message, MediaType.APPLICATION_JSON);

                }

            } catch (IOException e) {

                log.error("组播发送消息异常 ==> groupId={}, 异常信息：", groupId, e.getMessage());

                removeMessageId(messageId);

            }

        });

    }

    /**

    * 移除 MessageId

    *

    * @param messageId

    */

    public static void removeMessageId(String messageId) {

        sseEmitterMap.remove(messageId);

        //数量-1

        currentConnectTotal.getAndDecrement();

        log.info("remove messageId={}", messageId);

    }

    /**

    * 获取所有的 MessageId集合

    *

    * @return

    */

    public static List<String> getMessageIds() {

        return new ArrayList<>(sseEmitterMap.keySet());

    }

    /**

    * 获取当前连接总数

    *

    * @return

    */

    public static int getConnectTotal() {

        return currentConnectTotal.intValue();

    }

    /**

    * 断开SSE连接时的回调

    *

    * @param messageId

    * @return

    */

    private static Runnable completionCallBack(String messageId) {

        return () -> {

            log.info("结束连接 ==> messageId={}", messageId);

            removeMessageId(messageId);

        };

    }

    /**

    * 连接超时时回调触发

    *

    * @param messageId

    * @return

    */

    private static Runnable timeOutCallBack(String messageId) {

        return () -> {

            log.info("连接超时 ==> messageId={}", messageId);

            removeMessageId(messageId);

        };

    }

    /**

    * 连接报错时回调触发。

    *

    * @param messageId

    * @return

    */

    private static Consumer<Throwable> errorCallBack(String messageId) {

        return throwable -> {

            log.error("连接异常 ==> messageId={}", messageId);

            removeMessageId(messageId);

        };

    }

}

@RestController

@CrossOrigin

@RequestMapping("/sse")

public class SseDemoController {

    /**

    * 用户SSE连接

    * 它返回一个SseEmitter实例，这时候连接就已经创建了.

    *

    * @return

    */

    @GetMapping("/userConnect")

    public SseEmitter connect() {

        /**

        * 一般取登录用户账号作为 messageId。分组的话需要约定 messageId的格式。

        * 这里模拟创建一个用户连接

        */

        String userId = "userId-" + RandomUtils.nextInt(1, 10);

        return SseServer.createConnect(userId);

    }

    /**

    * 模拟实例：下载进度条显示。 前端访问下载接口之前，先建立用户SSE连接，然后访问下载接口，服务端推送消息。

    * http://localhost:8080/sse/downLoad/userId-1

    *

    * @throws InterruptedException

    */

    @GetMapping("/downLoad/{userId}")

    public void pushOne(@PathVariable("userId") String userId) throws InterruptedException {

        for (int i = 0; i <= 100; i++) {

            if (i > 50 && i < 70) {

                Thread.sleep(500L);

            } else {

                Thread.sleep(100L);

            }

            System.out.println("sendMessage --> 消息=" + i);

            SseServer.sendMessage(userId, String.valueOf(i));

        }

        System.out.println("下载成功");

    }

    /**

    * 广播发送。http://localhost:8080/sse/pushAllUser

    *

    * @throws InterruptedException

    */

    @GetMapping("/pushAllUser")

    public void pushAllUser() throws InterruptedException {

        for (int i = 0; i <= 100; i++) {

            if (i > 50 && i < 70) {

                Thread.sleep(500L);

            } else {

                Thread.sleep(100L);

            }

            System.out.println("batchAllSendMessage --> 消息=" + i);

            SseServer.batchAllSendMessage(String.valueOf(i));

        }

    }

}

技术选型：

使用 FastChat 提供开源 LLM 模型的 API，以 OpenAI API 接口形式接入，提升 LLM 模型加载效果；

使用 langchain 中已有 Chain 的实现，便于后续接入不同类型 Chain，并将对 Agent 接入开展测试；

使用 FastAPI 提供 API 服务，全部接口可在 FastAPI 自动生成的 docs 中开展测试，且所有对话接口支持通过参数设置流式或非流式输出；

使用 Streamlit 提供 WebUI 服务，可选是否基于 API 服务启动 WebUI，增加会话管理，可以自定义会话主题并切换，且后续可支持不同形式输出内容的显示；

项目中默认 LLM 模型改为 THUDM/chatglm2-6b，默认 Embedding 模型改为 moka-ai/m3e-base，文件加载方式与文段划分方式也有调整，后续将重新实现上下文扩充，并增加可选设置；

项目中扩充了对不同类型向量库的支持，除支持 FAISS 向量库外，还提供 Milvus, PGVector 向量库的接入；

项目中搜索引擎对话，除 Bing 搜索外，增加 DuckDuckGo 搜索选项，DuckDuckGo 搜索无需配置 API Key，在可访问国外服务环境下可直接使用。